tests.rs

use std::collections::BTreeMap;

use miden_node_proto::generated::{self as proto};
use miden_node_proto::server::validator_api;
use miden_node_store::{BlockStore, GenesisState};
use miden_node_utils::fee::test_fee_params;
use miden_protocol::Word;
use miden_protocol::block::{BlockHeader, BlockInputs, ProposedBlock};
use miden_protocol::crypto::dsa::ecdsa_k256_keccak::SigningKey;
use miden_protocol::testing::random_secret_key::random_secret_key;
use miden_protocol::transaction::PartialBlockchain;
use miden_tx::utils::serde::Serializable;

use super::{ValidatorError, ValidatorService};
use crate::ValidatorSigner;
use crate::db::{load_chain_tip, setup, upsert_block_header};

// TEST HELPERS
// ================================================================================================

/// Test harness that wraps a [`Validator`] and tracks the chain MMR state needed to construct valid
/// [`ProposedBlock`]s.
struct TestValidator {
    server: ValidatorService,
    chain: PartialBlockchain,
    chain_tip: BlockHeader,
}

impl TestValidator {
    /// Creates a correctly configured [`ValidatorService`]: the validator signs blocks with the
    /// same key that is designated as the `validator_key` in the genesis block.
    async fn new() -> Self {
        let key = random_secret_key();
        let signer = ValidatorSigner::new_local(key.clone());
        let (db, block_store, genesis_header) = setup_db_with_genesis(&key).await;

        Self {
            server: ValidatorService::new(signer, db, block_store, 0, 0, 0).await.unwrap(),
            chain: PartialBlockchain::default(),
            chain_tip: genesis_header,
        }
    }

    /// Builds an empty [`ProposedBlock`] extending the current chain tip.
    fn propose_empty_block(&self) -> ProposedBlock {
        empty_block(&self.chain_tip, &self.chain)
    }

    /// Calls `sign_block` on the validator server.
    async fn call_sign_block(
        &self,
        proposed_block: &ProposedBlock,
    ) -> Result<proto::blockchain::SignBlockResponse, tonic::Status> {
        let request = tonic::Request::new(proto::blockchain::ProposedBlock {
            proposed_block: proposed_block.to_bytes(),
        });
        validator_api::SignBlock::full(&self.server, request).await
    }

    /// Opens a block subscription starting from `block_from`.
    async fn call_block_subscription(
        &self,
        block_from: u32,
    ) -> <ValidatorService as proto::server::validator_api::BlockSubscription>::ItemStream {
        let request =
            tonic::Request::new(proto::validator::BlockSubscriptionRequest { block_from });
        validator_api::BlockSubscription::full(&self.server, request)
            .await
            .expect("subscription should open")
    }

    /// Loads the current chain tip from the validator's database.
    async fn load_chain_tip(&self) -> BlockHeader {
        self.server
            .db
            .query("load_chain_tip", load_chain_tip)
            .await
            .unwrap()
            .expect("chain tip should exist")
    }

    /// Builds, submits, and applies an empty block, advancing the chain tip.
    ///
    /// Panics if the block is rejected.
    async fn apply_empty_block(&mut self) {
        let proposed = self.propose_empty_block();
        self.call_sign_block(&proposed).await.unwrap();
        let (header, _) = proposed.into_header_and_body().unwrap();
        // Advance our local chain state to match what the server now has.
        self.chain.add_block(&self.chain_tip, false);
        self.chain_tip = header;
    }
}

/// Creates a validator database seeded with a genesis block whose `validator_key` is the public key
/// of `key`. Returns the database handle and the genesis block header.
async fn setup_db_with_genesis(key: &SigningKey) -> (miden_node_db::Db, BlockStore, BlockHeader) {
    let genesis_state = GenesisState::new(vec![], test_fee_params(), 1, 0, key.public_key());
    let genesis_block = genesis_state.into_block(key).unwrap();
    let genesis_header = genesis_block.inner().header().clone();

    let dir = tempfile::tempdir().unwrap();
    let db = setup(dir.path().join("validator.sqlite3")).await.unwrap();
    let block_store =
        BlockStore::bootstrap(dir.path().join("blocks").clone(), &genesis_block).unwrap();

    db.transact("upsert_genesis", {
        let h = genesis_header.clone();
        move |conn| upsert_block_header(conn, &h)
    })
    .await
    .unwrap();

    (db, block_store, genesis_header)
}

/// Builds an empty [`ProposedBlock`] that extends the given parent block header using the provided
/// partial blockchain state.
fn empty_block(parent_header: &BlockHeader, chain: &PartialBlockchain) -> ProposedBlock {
    let block_inputs = BlockInputs::new(
        parent_header.clone(),
        chain.clone(),
        BTreeMap::new(),
        BTreeMap::new(),
        BTreeMap::new(),
    );
    ProposedBlock::new(block_inputs, vec![]).unwrap()
}

// TESTS
// ================================================================================================

/// A validator whose signing key does not match the `validator_key` designated by the chain
/// (carried forward from genesis) must fail to start, rather than coming up and silently producing
/// signatures that the block producer cannot verify.
#[tokio::test]
async fn signing_key_mismatch_rejected() {
    // Seed a database whose genesis designates `genesis_key` as the validator key.
    let genesis_key = random_secret_key();
    let (db, block_store, genesis_header) = setup_db_with_genesis(&genesis_key).await;

    // Start a validator with a different key, modelling a validator configured with the wrong key.
    let rogue_signer = ValidatorSigner::new_local(random_secret_key());
    assert_ne!(
        rogue_signer.public_key(),
        *genesis_header.validator_key(),
        "test requires a signing key that differs from the genesis validator key",
    );

    let result = ValidatorService::new(rogue_signer, db, block_store, 0, 0, 0).await;
    assert!(
        matches!(result, Err(ValidatorError::ValidatorKeyMismatch { .. })),
        "expected ValidatorKeyMismatch error",
    );
}

/// The `SignBlock` response reports the commitment of the block the validator signed, and it
/// matches the commitment the caller derives from the same proposed block. This lets the block
/// producer detect a block-hash mismatch between itself and the validator.
#[tokio::test]
async fn sign_block_returns_signed_commitment() {
    let tv = TestValidator::new().await;

    let proposed = tv.propose_empty_block();
    let response = tv.call_sign_block(&proposed).await.expect("block should be signed");

    let (header, _) = proposed.into_header_and_body().unwrap();
    let returned: Word = response
        .block_commitment
        .expect("response should carry the signed commitment")
        .try_into()
        .unwrap();
    assert_eq!(
        returned,
        header.commitment(),
        "returned commitment must match the proposed block's commitment",
    );
}

/// An empty block at chain tip + 1 with the correct previous block commitment should be accepted.
#[tokio::test]
async fn chain_tip_plus_one_succeeds() {
    let tv = TestValidator::new().await;

    let proposed = tv.propose_empty_block();
    let result = tv.call_sign_block(&proposed).await;

    assert!(result.is_ok(), "chain tip + 1 should succeed, got: {:?}", result.err());
}

/// A replacement block at the same height as the current chain tip should be accepted.
#[tokio::test]
async fn chain_tip_replacement_succeeds() {
    let mut tv = TestValidator::new().await;

    // The genesis block can never be replaced, so we advance the chain to block 1, which we can
    // then replace.
    let genesis_header = tv.chain_tip.clone();
    let chain_at_genesis = tv.chain.clone();
    tv.apply_empty_block().await;
    let original_header = tv.chain_tip.clone();

    // Submit a different block at the same height (block 1), which is a replacement. Use an
    // explicit timestamp far in the future to ensure the replacement block differs.
    let block_inputs = BlockInputs::new(
        genesis_header.clone(),
        chain_at_genesis.clone(),
        BTreeMap::new(),
        BTreeMap::new(),
        BTreeMap::new(),
    );
    let far_future_timestamp = genesis_header.timestamp() + 1_000_000;
    let replacement = ProposedBlock::new_at(block_inputs, vec![], far_future_timestamp).unwrap();
    let (replacement_header, _) = replacement.clone().into_header_and_body().unwrap();

    assert_eq!(replacement_header.block_num(), original_header.block_num());
    assert_ne!(
        replacement_header.commitment(),
        original_header.commitment(),
        "replacement block should differ from the original"
    );

    let result = tv.call_sign_block(&replacement).await;
    assert!(result.is_ok(), "chain tip replacement should succeed, got: {:?}", result.err());

    // Verify that the chain tip in the database is now the replacement block, not the original.
    let new_chain_tip = tv.load_chain_tip().await;
    assert_eq!(
        new_chain_tip.commitment(),
        replacement_header.commitment(),
        "chain tip should be the replacement block"
    );
    assert_ne!(
        new_chain_tip.commitment(),
        original_header.commitment(),
        "chain tip should no longer be the original block"
    );
}

/// A block at chain tip + 2 (skipping a block number) should be rejected.
#[tokio::test]
async fn chain_tip_plus_two_rejected() {
    let mut tv = TestValidator::new().await;

    // Apply block 1.
    tv.apply_empty_block().await;

    // Build block 2 locally without applying it, then build block 3 on top.
    let block_2 = tv.propose_empty_block();
    let (block_2_header, _) = block_2.into_header_and_body().unwrap();
    let mut chain_after_1 = tv.chain.clone();
    chain_after_1.add_block(&tv.chain_tip, false);
    let block_3 = empty_block(&block_2_header, &chain_after_1);

    let result = tv.call_sign_block(&block_3).await;
    assert!(result.is_err(), "chain tip + 2 should be rejected");
    let status = result.unwrap_err();
    assert!(
        status.message().contains("block number mismatch"),
        "expected block number mismatch error, got: {}",
        status.message()
    );
}

/// A block at chain tip - 1 (behind the tip) should be rejected.
#[tokio::test]
async fn chain_tip_minus_one_rejected() {
    let mut tv = TestValidator::new().await;

    // Save genesis state.
    let genesis_header = tv.chain_tip.clone();
    let chain_at_genesis = tv.chain.clone();

    // Advance the chain to block 2.
    tv.apply_empty_block().await;
    tv.apply_empty_block().await;

    // Try to submit a block at height 1 (chain tip - 1). This is neither a replacement (which would
    // need to match tip height 2) nor the next block (which would be 3).
    let stale_block = empty_block(&genesis_header, &chain_at_genesis);

    let result = tv.call_sign_block(&stale_block).await;
    assert!(result.is_err(), "chain tip - 1 should be rejected");
    let status = result.unwrap_err();
    assert!(
        status.message().contains("block number mismatch"),
        "expected block number mismatch error, got: {}",
        status.message()
    );
}

/// A block with the wrong previous block commitment should be rejected.
#[tokio::test]
async fn commitment_mismatch_rejected() {
    let tv = TestValidator::new().await;

    // Build a valid ProposedBlock on a *different* genesis so its prev_block_commitment won't match
    // the validator's actual chain tip.
    let other_genesis_signer = random_secret_key();
    let other_genesis_state =
        GenesisState::new(vec![], test_fee_params(), 1, 1, other_genesis_signer.public_key());
    let other_genesis_block = other_genesis_state.into_block(&other_genesis_signer).unwrap();
    let other_genesis_header = other_genesis_block.inner().header().clone();
    let mismatched_block = empty_block(&other_genesis_header, &PartialBlockchain::default());

    let result = tv.call_sign_block(&mismatched_block).await;
    assert!(result.is_err(), "commitment mismatch should be rejected");
    let status = result.unwrap_err();
    assert!(
        status.message().contains("previous block commitment"),
        "expected commitment mismatch error, got: {}",
        status.message()
    );
}

/// A replacement block (same height as chain tip) with the wrong parent commitment should be
/// rejected.
#[tokio::test]
async fn replacement_commitment_mismatch_rejected() {
    let mut tv = TestValidator::new().await;

    // Advance past genesis so we have a replaceable block.
    tv.apply_empty_block().await;

    // Build a replacement block at the same height but using a *different* genesis so its
    // prev_block_commitment won't match the validator's actual parent of the chain tip.
    let other_genesis_signer = random_secret_key();
    let other_genesis_state =
        GenesisState::new(vec![], test_fee_params(), 1, 1, other_genesis_signer.public_key());
    let other_genesis_block = other_genesis_state.into_block(&other_genesis_signer).unwrap();
    let other_genesis_header = other_genesis_block.inner().header().clone();
    let mismatched_replacement = empty_block(&other_genesis_header, &PartialBlockchain::default());

    let result = tv.call_sign_block(&mismatched_replacement).await;
    assert!(result.is_err(), "replacement with mismatched commitment should be rejected");
    let status = result.unwrap_err();
    assert!(
        status.message().contains("previous block commitment"),
        "expected commitment mismatch error, got: {}",
        status.message()
    );
}

/// An empty block (no transactions, no batches) should be accepted.
#[tokio::test]
async fn empty_block_succeeds() {
    let tv = TestValidator::new().await;

    let proposed = tv.propose_empty_block();
    assert_eq!(proposed.transactions().count(), 0, "block should have no transactions");

    let result = tv.call_sign_block(&proposed).await;
    assert!(result.is_ok(), "empty block should succeed, got: {:?}", result.err());
}

/// A block containing transactions that were not previously validated should be rejected.
#[tokio::test]
async fn unknown_transactions_rejected() {
    use miden_protocol::Word;
    use miden_protocol::asset::FungibleAsset;
    use miden_protocol::batch::{BatchAccountUpdate, BatchId, ProvenBatch};
    use miden_protocol::block::BlockNumber;
    use miden_protocol::testing::account_id::ACCOUNT_ID_SENDER;
    use miden_protocol::transaction::{
        InputNoteCommitment,
        InputNotes,
        OrderedTransactionHeaders,
        TransactionHeader,
    };

    let tv = TestValidator::new().await;
    let genesis_header = tv.chain_tip.clone();

    // Build a dummy transaction header with a transaction ID that has NOT been submitted through
    // `submit_proven_transaction`.
    let account_id = ACCOUNT_ID_SENDER.try_into().unwrap();
    let fee = FungibleAsset::new(test_fee_params().fee_faucet_id(), 0).unwrap();
    let tx_header = TransactionHeader::new(
        account_id,
        Word::default(),
        Word::default(),
        InputNotes::<InputNoteCommitment>::default(),
        vec![],
        fee,
    );
    let tx_id = tx_header.id();

    // Build a ProvenBatch containing this transaction.
    let batch = ProvenBatch::new_unchecked(
        BatchId::from_ids(std::iter::once((tx_id, account_id))),
        genesis_header.commitment(),
        BlockNumber::GENESIS,
        BTreeMap::from([(
            account_id,
            BatchAccountUpdate::new_unchecked(
                account_id,
                Word::default(),
                Word::default(),
                miden_protocol::account::delta::AccountUpdateDetails::Private,
            ),
        )]),
        InputNotes::default(),
        vec![],
        BlockNumber::MAX,
        OrderedTransactionHeaders::new_unchecked(vec![tx_header]),
    )
    .unwrap();

    // Build a ProposedBlock containing the batch with the unknown transaction.
    let block_inputs = BlockInputs::new(
        genesis_header.clone(),
        PartialBlockchain::default(),
        BTreeMap::new(),
        BTreeMap::new(),
        BTreeMap::new(),
    );
    let proposed = ProposedBlock::new(block_inputs, vec![batch]).unwrap();

    let result = tv.server.validate_block(proposed, genesis_header).await;
    assert!(result.is_err(), "block with unknown transactions should be rejected");
    match result.unwrap_err() {
        ValidatorError::UnvalidatedTransactions(ids) => {
            assert_eq!(ids, vec![tx_id], "should report the unknown transaction ID");
        },
        other => panic!("expected UnvalidatedTransactions error, got: {other}"),
    }
}

/// After replacing the chain tip, a new block built against the pre-replacement tip should be
/// rejected because its previous block commitment no longer matches.
#[tokio::test]
async fn new_block_after_replacement_with_stale_commitment_rejected() {
    let mut tv = TestValidator::new().await;

    // Advance to block 1 and save the state needed to build on top of it.
    let genesis_header = tv.chain_tip.clone();
    let chain_at_genesis = tv.chain.clone();
    tv.apply_empty_block().await;
    let original_block_1_header = tv.chain_tip.clone();
    let chain_after_block_1 = tv.chain.clone();

    // Replace block 1 with a different block at the same height.
    let block_inputs = BlockInputs::new(
        genesis_header.clone(),
        chain_at_genesis.clone(),
        BTreeMap::new(),
        BTreeMap::new(),
        BTreeMap::new(),
    );
    let far_future_timestamp = genesis_header.timestamp() + 1_000_000;
    let replacement = ProposedBlock::new_at(block_inputs, vec![], far_future_timestamp).unwrap();
    let (replacement_header, _) = replacement.clone().into_header_and_body().unwrap();
    assert_ne!(
        replacement_header.commitment(),
        original_block_1_header.commitment(),
        "replacement block should differ from the original"
    );
    tv.call_sign_block(&replacement).await.unwrap();

    // Now try to submit block 2 built on top of the *original* block 1. Its prev_block_commitment
    // points to the old block 1, not the replacement.
    let stale_block_2 = empty_block(&original_block_1_header, &chain_after_block_1);

    let result = tv.call_sign_block(&stale_block_2).await;
    assert!(
        result.is_err(),
        "block with stale commitment after replacement should be rejected"
    );
    let status = result.unwrap_err();
    assert!(
        status.message().contains("previous block commitment"),
        "expected commitment mismatch error, got: {}",
        status.message()
    );
}

/// Verify that `validate_block` rejects blocks with a non-sequential block number.
#[tokio::test]
async fn validate_block_number_mismatch() {
    let mut tv = TestValidator::new().await;

    // Advance to block 1.
    tv.apply_empty_block().await;
    let block_1_header = tv.chain_tip.clone();

    // Build block 2 and 3 locally, then try to submit block 3 with chain_tip = block 1.
    let mut chain = tv.chain.clone();
    let block_2 = empty_block(&block_1_header, &chain);
    let (block_2_header, _) = block_2.into_header_and_body().unwrap();

    chain.add_block(&block_1_header, false);
    let block_3 = empty_block(&block_2_header, &chain);

    let result = tv.server.validate_block(block_3, block_1_header).await;
    assert!(result.is_err());
    assert!(
        matches!(result.unwrap_err(), ValidatorError::BlockNumberMismatch { .. }),
        "expected BlockNumberMismatch error"
    );
}

/// A block subscription replays the finalized backed-up blocks from the requested height and then
/// streams blocks live as they are finalized (i.e. once a child block is signed on top of them).
#[tokio::test]
async fn block_subscription_replays_then_follows() {
    use std::time::Duration;

    use miden_protocol::block::SignedBlock;
    use miden_tx::utils::serde::Deserializable;
    use tokio_stream::StreamExt;

    let mut tv = TestValidator::new().await;

    // Sign blocks 1, 2 and 3. The signed tip is 3, so blocks 1 and 2 are finalized while block 3
    // remains the provisional, replaceable tip and is withheld from subscribers.
    tv.apply_empty_block().await;
    tv.apply_empty_block().await;
    tv.apply_empty_block().await;

    // Subscribe from the first signed block and confirm only the finalized blocks (1 and 2) are
    // replayed; the provisional tip (block 3) is withheld.
    let mut stream = tv.call_block_subscription(1).await;
    for expected in 1..=2 {
        let response = tokio::time::timeout(Duration::from_secs(5), stream.next())
            .await
            .expect("replayed block should arrive promptly")
            .expect("stream should not end")
            .expect("stream item should not be an error");
        let block = SignedBlock::read_from_bytes(&response.block).expect("valid signed block");
        assert_eq!(block.header().block_num().as_u32(), expected);
        assert_eq!(response.committed_chain_tip, 2);
    }

    // Sign a new block (4), finalizing block 3, and confirm block 3 is now streamed live.
    tv.apply_empty_block().await;
    let response = tokio::time::timeout(Duration::from_secs(5), stream.next())
        .await
        .expect("live block should arrive promptly")
        .expect("stream should not end")
        .expect("stream item should not be an error");
    let block = SignedBlock::read_from_bytes(&response.block).expect("valid signed block");
    assert_eq!(block.header().block_num().as_u32(), 3);
    assert_eq!(response.committed_chain_tip, 3);
}

/// The provisional chain tip (the most recently signed block) must be withheld from subscribers
/// until a child block is signed on top of it, since it can still be replaced.
#[tokio::test]
async fn provisional_tip_is_withheld_from_subscribers() {
    use std::time::Duration;

    use miden_protocol::block::SignedBlock;
    use miden_tx::utils::serde::Deserializable;
    use tokio_stream::StreamExt;

    let mut tv = TestValidator::new().await;

    // Sign blocks 1 and 2. Block 1 is finalized (it has a child); block 2 is the provisional tip.
    tv.apply_empty_block().await;
    tv.apply_empty_block().await;

    let mut stream = tv.call_block_subscription(1).await;

    // Block 1 is finalized and streamed.
    let response = tokio::time::timeout(Duration::from_secs(5), stream.next())
        .await
        .expect("finalized block should arrive promptly")
        .expect("stream should not end")
        .expect("stream item should not be an error");
    let block = SignedBlock::read_from_bytes(&response.block).expect("valid signed block");
    assert_eq!(block.header().block_num().as_u32(), 1);
    assert_eq!(response.committed_chain_tip, 1);

    // Block 2 is the provisional tip and must not be streamed while it remains replaceable.
    let next = tokio::time::timeout(Duration::from_millis(500), stream.next()).await;
    assert!(next.is_err(), "provisional tip (block 2) should be withheld from subscribers");
}