From cea766bb6452049c47b4ccb4e0e10ffc0c26b1d0 Mon Sep 17 00:00:00 2001 From: Evan Batsell Date: Fri, 29 Nov 2024 13:27:38 -0500 Subject: [PATCH] Some more tests, fixes CI? --- .github/workflows/ci.yaml | 25 +- Cargo.lock | 1 + Cargo.toml | 2 +- core/src/ballot_box.rs | 237 ++++++++++++++- core/src/constants.rs | 16 +- core/src/error.rs | 4 +- .../tests/fixtures/test_builder.rs | 12 +- integration_tests/tests/tests.rs | 1 - .../tests/{ => tip_router}/bpf/mod.rs | 0 .../{ => tip_router}/bpf/set_merkle_root.rs | 0 integration_tests/tests/tip_router/mod.rs | 1 + meta_merkle_tree/Cargo.toml | 2 + meta_merkle_tree/src/meta_merkle_tree.rs | 286 ++++++++---------- meta_merkle_tree/src/tree_node.rs | 27 +- meta_merkle_tree/src/utils.rs | 4 +- program/src/cast_vote.rs | 2 +- 16 files changed, 408 insertions(+), 212 deletions(-) rename integration_tests/tests/{ => tip_router}/bpf/mod.rs (100%) rename integration_tests/tests/{ => tip_router}/bpf/set_merkle_root.rs (100%) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index a50f657..cd97962 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -25,6 +25,8 @@ jobs: with: crate: cargo-audit - run: cargo audit --ignore RUSTSEC-2022-0093 --ignore RUSTSEC-2023-0065 --ignore RUSTSEC-2024-0344 + env: + ANCHOR_IDL_BUILD_SKIP_LINT: TRUE code_gen: name: code generation @@ -39,6 +41,8 @@ jobs: toolchain: nightly-2024-07-25 - name: Regenerate Shank IDL files run: cargo b --release -p jito-tip-router-shank-cli && ./target/release/jito-tip-router-shank-cli + env: + ANCHOR_IDL_BUILD_SKIP_LINT: TRUE - name: Verify no changed files uses: tj-actions/verify-changed-files@v20 with: @@ -84,8 +88,14 @@ jobs: with: crate: cargo-sort - run: cargo sort --workspace --check + env: + ANCHOR_IDL_BUILD_SKIP_LINT: TRUE - run: cargo fmt --all --check + env: + ANCHOR_IDL_BUILD_SKIP_LINT: TRUE - run: cargo clippy --all-features -- -D warnings -D clippy::all -D clippy::nursery -D clippy::integer_division -D clippy::arithmetic_side_effects -D clippy::style -D clippy::perf + env: + ANCHOR_IDL_BUILD_SKIP_LINT: TRUE build: name: build @@ -105,11 +115,13 @@ jobs: run: cargo-build-sbf env: TIP_ROUTER_PROGRAM_ID: ${{ env.TIP_ROUTER_PROGRAM_ID }} + SBF_OUT_DIR: ${{ github.workspace }}/target/sbf-solana-solana/release + ANCHOR_IDL_BUILD_SKIP_LINT: TRUE - name: Upload MEV Tip Distribution NCN program uses: actions/upload-artifact@v4 with: name: jito_tip_router_program.so - path: target/sbf-solana-solana/release/jito_tip_router_program.so + path: target/sbf-solana-solana/release/ if-no-files-found: error # coverage: @@ -150,11 +162,16 @@ jobs: uses: actions/download-artifact@v4 with: name: jito_tip_router_program.so - path: target/sbf-solana-solana/release/ + path: integration_tests/tests/fixtures/ - uses: taiki-e/install-action@nextest - - run: cargo nextest run --all-features + # Test the non-BPF tests and the BPF tests separately + - run: cargo nextest run --all-features -E 'not test(bpf)' env: - SBF_OUT_DIR: ${{ github.workspace }}/target/sbf-solana-solana/release + ANCHOR_IDL_BUILD_SKIP_LINT: TRUE + - run: cargo nextest run --all-features -E 'test(bpf)' + env: + SBF_OUT_DIR: ${{ github.workspace }}/integration_tests/tests/fixtures + ANCHOR_IDL_BUILD_SKIP_LINT: TRUE # create_release: # name: Create Release diff --git a/Cargo.lock b/Cargo.lock index b943619..8c58407 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2940,6 +2940,7 @@ dependencies = [ "serde_json", "shank", "solana-program 1.18.26", + "solana-sdk", "spl-associated-token-account", "spl-math", "spl-token", diff --git a/Cargo.toml b/Cargo.toml index 3ba83dd..856aa07 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,7 +6,7 @@ members = [ "integration_tests", "meta_merkle_tree", "program", - "shank_cli" + "shank_cli", ] resolver = "2" diff --git a/core/src/ballot_box.rs b/core/src/ballot_box.rs index caf6f22..b493302 100644 --- a/core/src/ballot_box.rs +++ b/core/src/ballot_box.rs @@ -10,7 +10,7 @@ use solana_program::{ }; use spl_math::precise_number::PreciseNumber; -use crate::{constants::PRECISE_CONSENSUS, discriminators::Discriminators, error::TipRouterError}; +use crate::{constants::precise_consensus, discriminators::Discriminators, error::TipRouterError}; #[derive(Debug, Clone, PartialEq, Eq, Copy, Zeroable, ShankType, Pod, ShankType)] #[repr(C)] @@ -345,12 +345,26 @@ impl BallotBox { ballot: Ballot, stake_weight: u128, current_slot: u64, + valid_slots_after_consensus: u64, ) -> Result<(), TipRouterError> { + if !self.is_voting_valid(current_slot, valid_slots_after_consensus) { + return Err(TipRouterError::VotingNotValid); + } + let ballot_index = self.increment_or_create_ballot_tally(&ballot, stake_weight)?; + let consensus_reached = self.is_consensus_reached(); + for vote in self.operator_votes.iter_mut() { if vote.operator().eq(&operator) { - return Err(TipRouterError::DuplicateVoteCast); + if consensus_reached { + return Err(TipRouterError::ConsensusAlreadyReached); + } + + let operator_vote = + OperatorVote::new(ballot_index, operator, current_slot, stake_weight); + *vote = operator_vote; + return Ok(()); } if vote.is_empty() { @@ -397,8 +411,7 @@ impl BallotBox { .checked_div(&precise_total_stake_weight) .ok_or(TipRouterError::DenominatorIsZero)?; - let target_precise_percentage = - PreciseNumber::new(PRECISE_CONSENSUS).ok_or(TipRouterError::NewPreciseNumberError)?; + let target_precise_percentage = precise_consensus()?; let consensus_reached = ballot_percentage_of_total.greater_than_or_equal(&target_precise_percentage); @@ -450,8 +463,6 @@ impl BallotBox { } } -// merkle tree of merkle trees struct - #[cfg(test)] mod tests { use super::*; @@ -468,4 +479,218 @@ mod tests { .verify_merkle_root(Pubkey::default(), vec![], [0u8; 32], 0, 0) .unwrap(); } + + #[test] + fn test_cast_vote() { + let ncn = Pubkey::new_unique(); + let operator = Pubkey::new_unique(); + let current_slot = 100; + let epoch = 1; + let stake_weight: u128 = 1000; + let valid_slots_after_consensus = 10; + let mut ballot_box = BallotBox::new(ncn, epoch, 0, current_slot); + let ballot = Ballot::new([1; 32]); + + // Test initial cast vote + ballot_box + .cast_vote( + operator, + ballot, + stake_weight, + current_slot, + valid_slots_after_consensus, + ) + .unwrap(); + + // Verify vote was recorded correctly + let operator_vote = ballot_box + .operator_votes + .iter() + .find(|v| v.operator() == operator) + .unwrap(); + assert_eq!(operator_vote.stake_weight(), stake_weight); + assert_eq!(operator_vote.slot_voted(), current_slot); + + // Verify ballot tally + let tally = ballot_box + .ballot_tallies + .iter() + .find(|t| t.ballot() == ballot) + .unwrap(); + assert_eq!(tally.stake_weight(), stake_weight); + + // Test re-vote with different ballot + let new_ballot = Ballot::new([2u8; 32]); + let new_slot = current_slot + 1; + ballot_box + .cast_vote( + operator, + new_ballot, + stake_weight, + new_slot, + valid_slots_after_consensus, + ) + .unwrap(); + + // Verify new ballot tally increased + let new_tally = ballot_box + .ballot_tallies + .iter() + .find(|t| t.ballot() == new_ballot) + .unwrap(); + assert_eq!(new_tally.stake_weight(), stake_weight); + + // Test error on changing vote after consensus + ballot_box.set_winning_ballot(new_ballot); + ballot_box.slot_consensus_reached = PodU64::from(new_slot); + let result = ballot_box.cast_vote( + operator, + ballot, + stake_weight, + new_slot + 1, + valid_slots_after_consensus, + ); + assert!(matches!( + result, + Err(TipRouterError::ConsensusAlreadyReached) + )); + + // Test voting window expired after consensus + let result = ballot_box.cast_vote( + operator, + ballot, + stake_weight, + new_slot + valid_slots_after_consensus + 1, + valid_slots_after_consensus, + ); + assert!(matches!(result, Err(TipRouterError::VotingNotValid))); + } + + #[test] + fn test_increment_or_create_ballot_tally() { + let mut ballot_box = BallotBox::new(Pubkey::new_unique(), 1, 1, 1); + let ballot = Ballot::new([1u8; 32]); + let stake_weight = 100; + + // Test creating new ballot tally + let tally_index = ballot_box + .increment_or_create_ballot_tally(&ballot, stake_weight) + .unwrap(); + assert_eq!(tally_index, 0); + assert_eq!(ballot_box.unique_ballots(), 1); + assert_eq!(ballot_box.ballot_tallies[0].stake_weight(), stake_weight); + assert_eq!(ballot_box.ballot_tallies[0].ballot(), ballot); + + // Test incrementing existing ballot tally + let tally_index = ballot_box + .increment_or_create_ballot_tally(&ballot, stake_weight) + .unwrap(); + assert_eq!(tally_index, 0); + assert_eq!(ballot_box.unique_ballots(), 1); + assert_eq!( + ballot_box.ballot_tallies[0].stake_weight(), + stake_weight * 2 + ); + assert_eq!(ballot_box.ballot_tallies[0].ballot(), ballot); + + // Test creating second ballot tally + let ballot2 = Ballot::new([2u8; 32]); + let tally_index = ballot_box + .increment_or_create_ballot_tally(&ballot2, stake_weight) + .unwrap(); + assert_eq!(tally_index, 1); + assert_eq!(ballot_box.unique_ballots(), 2); + assert_eq!(ballot_box.ballot_tallies[1].stake_weight(), stake_weight); + assert_eq!(ballot_box.ballot_tallies[1].ballot(), ballot2); + + // Test error when ballot tallies are full + for i in 3..=32 { + let ballot = Ballot::new([i as u8; 32]); + ballot_box + .increment_or_create_ballot_tally(&ballot, stake_weight) + .unwrap(); + } + let ballot_full = Ballot::new([33u8; 32]); + let result = ballot_box.increment_or_create_ballot_tally(&ballot_full, stake_weight); + assert!(matches!(result, Err(TipRouterError::BallotTallyFull))); + } + + #[test] + fn test_tally_votes() { + let ncn = Pubkey::new_unique(); + let current_slot = 100; + let epoch = 1; + let stake_weight: u128 = 1000; + let total_stake_weight: u128 = 1000; + let mut ballot_box = BallotBox::new(ncn, epoch, 0, current_slot); + let ballot = Ballot::new([1; 32]); + + // Test no consensus when below threshold + ballot_box + .increment_or_create_ballot_tally(&ballot, stake_weight / 2) + .unwrap(); + ballot_box + .tally_votes(total_stake_weight, current_slot) + .unwrap(); + assert!(!ballot_box.is_consensus_reached()); + assert_eq!(ballot_box.slot_consensus_reached(), 0); + assert!(matches!( + ballot_box.get_winning_ballot(), + Err(TipRouterError::ConsensusNotReached) + )); + + // Test consensus reached when above threshold + ballot_box + .increment_or_create_ballot_tally(&ballot, stake_weight / 2) + .unwrap(); + ballot_box + .tally_votes(total_stake_weight, current_slot) + .unwrap(); + assert!(ballot_box.is_consensus_reached()); + assert_eq!(ballot_box.slot_consensus_reached(), current_slot); + assert_eq!(ballot_box.get_winning_ballot().unwrap(), ballot); + + // Consensus remains after additional votes + let ballot2 = Ballot::new([2; 32]); + ballot_box + .increment_or_create_ballot_tally(&ballot2, stake_weight) + .unwrap(); + ballot_box + .tally_votes(total_stake_weight, current_slot + 1) + .unwrap(); + assert!(ballot_box.is_consensus_reached()); + assert_eq!(ballot_box.slot_consensus_reached(), current_slot); + assert_eq!(ballot_box.get_winning_ballot().unwrap(), ballot); + + // Test with multiple competing ballots + let mut ballot_box = BallotBox::new(ncn, epoch, 0, current_slot); + let ballot1 = Ballot::new([1; 32]); + let ballot2 = Ballot::new([2; 32]); + let ballot3 = Ballot::new([3; 32]); + + ballot_box + .increment_or_create_ballot_tally(&ballot1, stake_weight / 4) + .unwrap(); + ballot_box + .increment_or_create_ballot_tally(&ballot2, stake_weight / 4) + .unwrap(); + ballot_box + .increment_or_create_ballot_tally(&ballot3, stake_weight / 2) + .unwrap(); + + ballot_box + .tally_votes(total_stake_weight, current_slot) + .unwrap(); + assert!(!ballot_box.is_consensus_reached()); + + // Add more votes to reach consensus + ballot_box + .increment_or_create_ballot_tally(&ballot3, stake_weight / 2) + .unwrap(); + ballot_box + .tally_votes(total_stake_weight, current_slot) + .unwrap(); + assert!(ballot_box.is_consensus_reached()); + assert_eq!(ballot_box.get_winning_ballot().unwrap(), ballot3); + } } diff --git a/core/src/constants.rs b/core/src/constants.rs index df8fbe0..7b02431 100644 --- a/core/src/constants.rs +++ b/core/src/constants.rs @@ -1,4 +1,18 @@ +use spl_math::precise_number::PreciseNumber; + +use crate::error::TipRouterError; + pub const MAX_FEE_BPS: u64 = 10_000; pub const MAX_OPERATORS: usize = 256; pub const MAX_VAULT_OPERATOR_DELEGATIONS: usize = 64; -pub const PRECISE_CONSENSUS: u128 = 666_666_666_666; +const PRECISE_CONSENSUS_NUMERATOR: u128 = 2; +const PRECISE_CONSENSUS_DENOMINATOR: u128 = 3; +pub fn precise_consensus() -> Result { + PreciseNumber::new(PRECISE_CONSENSUS_NUMERATOR) + .ok_or(TipRouterError::NewPreciseNumberError)? + .checked_div( + &PreciseNumber::new(PRECISE_CONSENSUS_DENOMINATOR) + .ok_or(TipRouterError::NewPreciseNumberError)?, + ) + .ok_or(TipRouterError::DenominatorIsZero) +} diff --git a/core/src/error.rs b/core/src/error.rs index 1f0ff6c..b343bb6 100644 --- a/core/src/error.rs +++ b/core/src/error.rs @@ -70,14 +70,14 @@ pub enum TipRouterError { OperatorVotesFull, #[error("Merkle root tally full")] BallotTallyFull, - #[error("Consensus already reached")] + #[error("Consensus already reached, cannot change vote")] ConsensusAlreadyReached, #[error("Consensus not reached")] ConsensusNotReached, #[error("Epoch snapshot not finalized")] EpochSnapshotNotFinalized, - #[error("Voting not valid")] + #[error("Voting not valid, too many slots after consensus reached")] VotingNotValid, #[error("Tie breaker admin invalid")] TieBreakerAdminInvalid, diff --git a/integration_tests/tests/fixtures/test_builder.rs b/integration_tests/tests/fixtures/test_builder.rs index 7d2dd02..d80aca3 100644 --- a/integration_tests/tests/fixtures/test_builder.rs +++ b/integration_tests/tests/fixtures/test_builder.rs @@ -55,8 +55,9 @@ impl Debug for TestBuilder { impl TestBuilder { pub async fn new() -> Self { - // TODO explain difference - let program_test = if std::env::vars().any(|(key, value)| key.eq("SBF_OUT_DIR")) { + let run_as_bpf = std::env::vars().any(|(key, value)| key.eq("SBF_OUT_DIR")); + + let program_test = if run_as_bpf { let mut program_test = ProgramTest::new( "jito_tip_router_program", jito_tip_router_program::id(), @@ -69,7 +70,6 @@ impl TestBuilder { // Anchor programs do not expose a compatible entrypoint for solana_program_test::processor! program_test.add_program("jito_tip_distribution", jito_tip_distribution::id(), None); - // program_test.prefer_bpf(true); program_test } else { let mut program_test = ProgramTest::new( @@ -88,12 +88,6 @@ impl TestBuilder { processor!(jito_restaking_program::process_instruction), ); - // program_test.add_program( - // "jito_tip_distribution", - // jito_tip_distribution::id(), - // processor!(jito_tip_router_program::process_instruction), - // ); - program_test }; diff --git a/integration_tests/tests/tests.rs b/integration_tests/tests/tests.rs index febdbf7..ffd4b7d 100644 --- a/integration_tests/tests/tests.rs +++ b/integration_tests/tests/tests.rs @@ -1,4 +1,3 @@ -mod bpf; mod fixtures; mod helpers; mod tip_router; diff --git a/integration_tests/tests/bpf/mod.rs b/integration_tests/tests/tip_router/bpf/mod.rs similarity index 100% rename from integration_tests/tests/bpf/mod.rs rename to integration_tests/tests/tip_router/bpf/mod.rs diff --git a/integration_tests/tests/bpf/set_merkle_root.rs b/integration_tests/tests/tip_router/bpf/set_merkle_root.rs similarity index 100% rename from integration_tests/tests/bpf/set_merkle_root.rs rename to integration_tests/tests/tip_router/bpf/set_merkle_root.rs diff --git a/integration_tests/tests/tip_router/mod.rs b/integration_tests/tests/tip_router/mod.rs index 43d86a4..e486042 100644 --- a/integration_tests/tests/tip_router/mod.rs +++ b/integration_tests/tests/tip_router/mod.rs @@ -1,4 +1,5 @@ mod admin_update_weight_table; +mod bpf; mod initialize_epoch_snapshot; mod initialize_ncn_config; mod initialize_operator_snapshot; diff --git a/meta_merkle_tree/Cargo.toml b/meta_merkle_tree/Cargo.toml index b3c3d0c..f8d94be 100644 --- a/meta_merkle_tree/Cargo.toml +++ b/meta_merkle_tree/Cargo.toml @@ -32,3 +32,5 @@ spl-math = { workspace = true } spl-token = { workspace = true } thiserror = { workspace = true } +[dev-dependencies] +solana-sdk = { workspace = true } diff --git a/meta_merkle_tree/src/meta_merkle_tree.rs b/meta_merkle_tree/src/meta_merkle_tree.rs index 0d4c03d..4bb00ed 100644 --- a/meta_merkle_tree/src/meta_merkle_tree.rs +++ b/meta_merkle_tree/src/meta_merkle_tree.rs @@ -39,7 +39,7 @@ pub type Result = result::Result; impl MetaMerkleTree { pub fn new(mut tree_nodes: Vec) -> Result { - // TODO Consider correctness of a sorting step here + // Sort by hash to ensure consistent trees tree_nodes.sort_by_key(|node| node.hash()); let hashed_nodes = tree_nodes @@ -82,7 +82,7 @@ impl MetaMerkleTree { Self::new(tree_nodes) } - // TODO uncomment if we need to load this from a file (for operator?) + // TODO if we need to load this from a file (for operator?) /// Load a merkle tree from a csv path // pub fn new_from_csv(path: &PathBuf) -> Result { // let csv_entries = CsvEntry::new_from_file(path)?; @@ -200,169 +200,119 @@ impl MetaMerkleTree { } } -// TODO rewrite tests for MetaMerkleTree - -// #[cfg(test)] -// mod tests { -// use std::path::PathBuf; - -// use solana_program::{pubkey, pubkey::Pubkey}; -// use solana_sdk::{ -// signature::{EncodableKey, Keypair}, -// signer::Signer, -// }; - -// use super::*; - -// pub fn new_test_key() -> Pubkey { -// let kp = Keypair::new(); -// let out_path = format!("./test_keys/{}.json", kp.pubkey()); - -// kp.write_to_file(out_path) -// .expect("Failed to write to signer"); - -// kp.pubkey() -// } - -// fn new_test_merkle_tree(num_nodes: u64, path: &PathBuf) { -// let mut tree_nodes = vec![]; - -// fn rand_balance() -> u64 { -// rand::random::() % 100 * u64::pow(10, 9) -// } - -// for _ in 0..num_nodes { -// // choose amount unlocked and amount locked as a random u64 between 0 and 100 -// tree_nodes.push(TreeNode { -// vote_account: new_test_key(), -// proof: None, -// total_unlocked_staker: rand_balance(), -// total_locked_staker: rand_balance(), -// total_unlocked_searcher: rand_balance(), -// total_locked_searcher: rand_balance(), -// total_unlocked_validator: rand_balance(), -// total_locked_validator: rand_balance(), -// }); -// } - -// let merkle_tree = MetaMerkleTree::new(tree_nodes).unwrap(); - -// merkle_tree.write_to_file(path); -// } - -// #[test] -// fn test_verify_new_merkle_tree() { -// let tree_nodes = vec![TreeNode { -// vote_account: Pubkey::default(), -// proof: None, -// total_unlocked_staker: 2, -// total_locked_staker: 3, -// total_unlocked_searcher: 4, -// total_locked_searcher: 5, -// total_unlocked_validator: 6, -// total_locked_validator: 7, -// }]; -// let merkle_tree = MetaMerkleTree::new(tree_nodes).unwrap(); -// assert!(merkle_tree.verify_proof().is_ok(), "verify failed"); -// } - -// #[test] -// fn test_write_merkle_distributor_to_file() { -// // create a merkle root from 3 tree nodes and write it to file, then read it -// let tree_nodes = vec![ -// TreeNode { -// vote_account: pubkey!("FLYqJsmJ5AGMxMxK3Qy1rSen4ES2dqqo6h51W3C1tYS"), -// proof: None, -// total_unlocked_staker: (100 * u64::pow(10, 9)), -// total_locked_staker: (100 * u64::pow(10, 9)), -// total_unlocked_searcher: 0, -// total_locked_searcher: 0, -// total_unlocked_validator: 0, -// total_locked_validator: 0, -// }, -// TreeNode { -// vote_account: pubkey!("EDGARWktv3nDxRYjufjdbZmryqGXceaFPoPpbUzdpqED"), -// proof: None, -// total_unlocked_staker: 100 * u64::pow(10, 9), -// total_locked_staker: (100 * u64::pow(10, 9)), -// total_unlocked_searcher: 0, -// total_locked_searcher: 0, -// total_unlocked_validator: 0, -// total_locked_validator: 0, -// }, -// TreeNode { -// vote_account: pubkey!("EDGARWktv3nDxRYjufjdbZmryqGXceaFPoPpbUzdpqEH"), -// proof: None, -// total_locked_staker: (100 * u64::pow(10, 9)), -// total_unlocked_staker: (100 * u64::pow(10, 9)), -// total_unlocked_searcher: 0, -// total_locked_searcher: 0, -// total_unlocked_validator: 0, -// total_locked_validator: 0, -// }, -// ]; - -// let merkle_distributor_info = MetaMerkleTree::new(tree_nodes).unwrap(); -// let path = PathBuf::from("merkle_tree.json"); - -// // serialize merkle distributor to file -// merkle_distributor_info.write_to_file(&path); -// // now test we can successfully read from file -// let merkle_distributor_read: MetaMerkleTree = MetaMerkleTree::new_from_file(&path).unwrap(); - -// assert_eq!(merkle_distributor_read.tree_nodes.len(), 3); -// } - -// #[test] -// fn test_new_test_merkle_tree() { -// new_test_merkle_tree(100, &PathBuf::from("merkle_tree_test_csv.json")); -// } - -// // Test creating a merkle tree from Tree Nodes, where claimants are not unique -// #[test] -// fn test_new_merkle_tree_duplicate_claimants() { -// let duplicate_pubkey = Pubkey::new_unique(); -// let tree_nodes = vec![ -// TreeNode { -// vote_account: duplicate_pubkey, -// proof: None, -// total_unlocked_staker: 10, -// total_locked_staker: 20, -// total_unlocked_searcher: 30, -// total_locked_searcher: 40, -// total_unlocked_validator: 50, -// total_locked_validator: 60, -// }, -// TreeNode { -// vote_account: duplicate_pubkey, -// proof: None, -// total_unlocked_staker: 1, -// total_locked_staker: 2, -// total_unlocked_searcher: 3, -// total_locked_searcher: 4, -// total_unlocked_validator: 5, -// total_locked_validator: 6, -// }, -// TreeNode { -// vote_account: Pubkey::new_unique(), -// proof: None, -// total_unlocked_staker: 0, -// total_locked_staker: 0, -// total_unlocked_searcher: 0, -// total_locked_searcher: 0, -// total_unlocked_validator: 0, -// total_locked_validator: 0, -// }, -// ]; - -// let tree = MetaMerkleTree::new(tree_nodes).unwrap(); -// // Assert that the merkle distributor correctly combines the two tree nodes -// assert_eq!(tree.tree_nodes.len(), 2); -// assert_eq!(tree.tree_nodes[0].total_unlocked_staker, 11); -// assert_eq!(tree.tree_nodes[0].total_locked_staker, 22); -// assert_eq!(tree.tree_nodes[0].total_unlocked_searcher, 33); -// assert_eq!(tree.tree_nodes[0].total_locked_searcher, 44); -// assert_eq!(tree.tree_nodes[0].total_unlocked_validator, 55); -// assert_eq!(tree.tree_nodes[0].total_locked_validator, 66); -// } -// } +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use solana_program::pubkey::Pubkey; + use solana_sdk::{ + signature::{EncodableKey, Keypair}, + signer::Signer, + }; + + use super::*; + + pub fn new_test_key() -> Pubkey { + let kp = Keypair::new(); + let out_path = format!("./test_keys/{}.json", kp.pubkey()); + + kp.write_to_file(out_path) + .expect("Failed to write to signer"); + + kp.pubkey() + } + + fn new_test_merkle_tree(num_nodes: u64, path: &PathBuf) { + let mut tree_nodes = vec![]; + + fn rand_balance() -> u64 { + rand::random::() % 100 * u64::pow(10, 9) + } + + for _ in 0..num_nodes { + tree_nodes.push(TreeNode::new( + new_test_key(), + [0; 32], + rand_balance(), + rand_balance(), + )); + } + + let merkle_tree = MetaMerkleTree::new(tree_nodes).unwrap(); + + merkle_tree.write_to_file(path); + } + + #[test] + fn test_verify_new_merkle_tree() { + let tree_nodes = vec![TreeNode::new(Pubkey::default(), [0; 32], 100, 10)]; + let merkle_tree = MetaMerkleTree::new(tree_nodes).unwrap(); + assert!(merkle_tree.verify_proof().is_ok(), "verify failed"); + } + + #[test] + fn test_write_merkle_distributor_to_file() { + // create a merkle root from 3 tree nodes and write it to file, then read it + let tree_nodes = vec![ + TreeNode::new( + new_test_key(), + [0; 32], + 100 * u64::pow(10, 9), + 100 * u64::pow(10, 9), + ), + TreeNode::new( + new_test_key(), + [0; 32], + 100 * u64::pow(10, 9), + 100 * u64::pow(10, 9), + ), + TreeNode::new( + new_test_key(), + [0; 32], + 100 * u64::pow(10, 9), + 100 * u64::pow(10, 9), + ), + ]; + + let merkle_distributor_info = MetaMerkleTree::new(tree_nodes).unwrap(); + let path = PathBuf::from("merkle_tree.json"); + + // serialize merkle distributor to file + merkle_distributor_info.write_to_file(&path); + // now test we can successfully read from file + let merkle_distributor_read: MetaMerkleTree = MetaMerkleTree::new_from_file(&path).unwrap(); + + assert_eq!(merkle_distributor_read.tree_nodes.len(), 3); + } + + #[test] + fn test_new_test_merkle_tree() { + new_test_merkle_tree(100, &PathBuf::from("merkle_tree_test_csv.json")); + } + + // Test creating a merkle tree from Tree Nodes + #[test] + fn test_new_merkle_tree() { + let pubkey1 = Pubkey::new_unique(); + let pubkey2 = Pubkey::new_unique(); + let pubkey3 = Pubkey::new_unique(); + + let mut tree_nodes = vec![ + TreeNode::new(pubkey1, [0; 32], 10, 20), + TreeNode::new(pubkey2, [0; 32], 1, 2), + TreeNode::new(pubkey3, [0; 32], 3, 4), + ]; + + // Sort by hash + tree_nodes.sort_by_key(|node| node.hash()); + + let tree = MetaMerkleTree::new(tree_nodes).unwrap(); + + assert_eq!(tree.tree_nodes.len(), 3); + assert_eq!(tree.tree_nodes[0].max_total_claim, 10); + assert_eq!(tree.tree_nodes[0].max_num_nodes, 20); + assert_eq!(tree.tree_nodes[0].validator_merkle_root, [0; 32]); + assert_eq!(tree.tree_nodes[0].tip_distribution_account, pubkey1); + assert!(tree.tree_nodes[0].proof.is_some()); + } +} diff --git a/meta_merkle_tree/src/tree_node.rs b/meta_merkle_tree/src/tree_node.rs index d946ef2..8d14697 100644 --- a/meta_merkle_tree/src/tree_node.rs +++ b/meta_merkle_tree/src/tree_node.rs @@ -1,5 +1,3 @@ -use std::str::FromStr; - use serde::{Deserialize, Serialize}; use solana_program::{ hash::{hashv, Hash}, @@ -41,6 +39,7 @@ impl TreeNode { pub fn hash(&self) -> Hash { hashv(&[ + &self.tip_distribution_account.to_bytes(), &self.validator_merkle_root, &self.max_total_claim.to_le_bytes(), &self.max_num_nodes.to_le_bytes(), @@ -61,25 +60,21 @@ impl From for TreeNode { } } -// TODO rewrite tests for MetaMerkleTree TreeNode #[cfg(test)] mod tests { use super::*; #[test] fn test_serialize_tree_node() { - // let tree_node = TreeNode { - // claimant: Pubkey::default(), - // proof: None, - // total_unlocked_staker: 0, - // total_locked_staker: 0, - // total_unlocked_searcher: 0, - // total_locked_searcher: 0, - // total_unlocked_validator: 0, - // total_locked_validator: 0, - // }; - // let serialized = serde_json::to_string(&tree_node).unwrap(); - // let deserialized: TreeNode = serde_json::from_str(&serialized).unwrap(); - // assert_eq!(tree_node, deserialized); + let tree_node = TreeNode { + tip_distribution_account: Pubkey::default(), + proof: None, + validator_merkle_root: [0; 32], + max_total_claim: 0, + max_num_nodes: 0, + }; + let serialized = serde_json::to_string(&tree_node).unwrap(); + let deserialized: TreeNode = serde_json::from_str(&serialized).unwrap(); + assert_eq!(tree_node, deserialized); } } diff --git a/meta_merkle_tree/src/utils.rs b/meta_merkle_tree/src/utils.rs index 30383ed..713bae2 100644 --- a/meta_merkle_tree/src/utils.rs +++ b/meta_merkle_tree/src/utils.rs @@ -1,6 +1,4 @@ -use solana_program::pubkey::Pubkey; - -use crate::{merkle_tree::MerkleTree, tree_node::TreeNode}; +use crate::merkle_tree::MerkleTree; pub fn get_proof(merkle_tree: &MerkleTree, index: usize) -> Vec<[u8; 32]> { let mut proof = Vec::new(); diff --git a/program/src/cast_vote.rs b/program/src/cast_vote.rs index f6537c3..dfa69dd 100644 --- a/program/src/cast_vote.rs +++ b/program/src/cast_vote.rs @@ -84,7 +84,7 @@ pub fn process_cast_vote( let ballot = Ballot::new(meta_merkle_root); - ballot_box.cast_vote(*operator.key, ballot, operator_stake_weight, slot)?; + ballot_box.cast_vote(*operator.key, ballot, operator_stake_weight, slot, valid_slots_after_consensus)?; ballot_box.tally_votes(total_stake_weight, slot)?;