From 0a1441b007db48510fca9dec2a1fc7d46b162603 Mon Sep 17 00:00:00 2001 From: mdulin2 Date: Fri, 26 Jun 2026 10:37:25 -0700 Subject: [PATCH 1/2] Fix missing timestamp update, and scaling bugs --- sui/packages/ntt/sources/inbox.move | 33 + sui/packages/ntt/sources/ntt.move | 2 +- sui/packages/ntt/sources/outbox.move | 5 + sui/packages/ntt/sources/rate_limit.move | 12 + sui/packages/ntt/tests/ntt_tests.move | 1115 ++++++++++++++++++ sui/packages/ntt/tests/rate_limit_tests.move | 213 ++++ 6 files changed, 1379 insertions(+), 1 deletion(-) create mode 100644 sui/packages/ntt/tests/rate_limit_tests.move diff --git a/sui/packages/ntt/sources/inbox.move b/sui/packages/ntt/sources/inbox.move index d4699788a..46adf9e9f 100644 --- a/sui/packages/ntt/sources/inbox.move +++ b/sui/packages/ntt/sources/inbox.move @@ -126,4 +126,37 @@ module ntt::inbox { public fun borrow_inbox_item(inbox: &Inbox, key: InboxKey): &InboxItem { inbox.entries.borrow(key) } + + #[test_only] + /// True if the item is queued for delayed release (`ReleaseAfter`). + public fun is_release_after(self: &InboxItem): bool { + match (&self.release_status) { + ReleaseStatus::ReleaseAfter(_) => true, + _ => false, + } + } + + #[test_only] + /// The release timestamp of a queued item. Aborts unless the item is in the + /// `ReleaseAfter` state. + public fun release_after_timestamp(self: &InboxItem): u64 { + match (&self.release_status) { + ReleaseStatus::ReleaseAfter(release_timestamp) => *release_timestamp, + _ => abort ETransferCannotBeRedeemed, + } + } + + #[test_only] + public fun is_released(self: &InboxItem): bool { + match (&self.release_status) { + ReleaseStatus::Released => true, + _ => false, + } + } + + #[test_only] + /// Number of votes recorded on the item, regardless of which transceivers are enabled. + public fun count_votes(self: &InboxItem): u8 { + self.votes.count_ones() + } } diff --git a/sui/packages/ntt/sources/ntt.move b/sui/packages/ntt/sources/ntt.move index 58ba80522..0e2aba8d5 100644 --- a/sui/packages/ntt/sources/ntt.move +++ b/sui/packages/ntt/sources/ntt.move @@ -160,7 +160,7 @@ module ntt::ntt { // consumed. refill inbox rate limit state.borrow_peer_mut(recipient_chain) .borrow_inbound_rate_limit_mut() - .refill(clock, trimmed_amount.amount()); + .refill(clock, trimmed_amount.untrim(coin_meta.get_decimals())); clock.timestamp_ms() }; diff --git a/sui/packages/ntt/sources/outbox.move b/sui/packages/ntt/sources/outbox.move index 9509ee484..93eeade9f 100644 --- a/sui/packages/ntt/sources/outbox.move +++ b/sui/packages/ntt/sources/outbox.move @@ -79,6 +79,11 @@ module ntt::outbox { &outbox_item.recipient_ntt_manager } + #[test_only] + public fun borrow_release_timestamp(outbox_item: &OutboxItem): u64 { + outbox_item.release_timestamp + } + public(package) fun try_release( outbox: &mut Outbox, key: OutboxKey, diff --git a/sui/packages/ntt/sources/rate_limit.move b/sui/packages/ntt/sources/rate_limit.move index 7d91b45bc..3d576171c 100644 --- a/sui/packages/ntt/sources/rate_limit.move +++ b/sui/packages/ntt/sources/rate_limit.move @@ -111,6 +111,8 @@ module ntt::rate_limit { let new_amount = min!(new_amount, 0xFFFF_FFFF_FFFF_FFFF) as u64; self.capacity_at_last_tx = min!(new_amount, self.limit); + let now = clock.timestamp_ms(); + self.last_tx_timestamp = now; } public fun set_limit(self: &mut RateLimitState, limit: u64, clock: &Clock) { @@ -139,4 +141,14 @@ module ntt::rate_limit { self.capacity_at_last_tx = new_capacity.min(limit); self.last_tx_timestamp = now; } + + #[test_only] + public fun limit(self: &RateLimitState): u64 { self.limit } + + #[test_only] + public fun capacity_at_last_tx(self: &RateLimitState): u64 { self.capacity_at_last_tx } + + #[test_only] + public fun last_tx_timestamp(self: &RateLimitState): u64 { self.last_tx_timestamp } + } diff --git a/sui/packages/ntt/tests/ntt_tests.move b/sui/packages/ntt/tests/ntt_tests.move index 4060d40b5..e1948be51 100644 --- a/sui/packages/ntt/tests/ntt_tests.move +++ b/sui/packages/ntt/tests/ntt_tests.move @@ -904,4 +904,1119 @@ module ntt::ntt_tests { test_scenario::return_to_address(admin, admin_cap); scenario.end(); } + + + #[test] + fun test_transfer_with_refill() { + let (_, user_a, _, _) = ntt_scenario::test_addresses(); + let mut scenario = test_scenario::begin(user_a); + ntt_scenario::setup(&mut scenario); + + scenario.next_tx(user_a); + + // Take state and clock + let mut state = ntt_scenario::take_state(&scenario); + let mut clock = ntt_scenario::take_clock(&mut scenario); + clock.increment_for_testing(1_000); + + // Drain limit for the rate limit refill + let drain: u64 = 2_000_000_000; // + let result = state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut() + .consume_or_delay(&clock, drain); + assert!(result.is_consumed()); // sanity: drain < limit (5e9), so consumed + + let coin_meta = ntt_scenario::take_coin_metadata(&scenario); + + let coins = state.mint_for_test(TEST_AMOUNT, scenario.ctx()); + + // Create transfer ticket + let recipient = x"000000000000000000000000000000000000000000000000000000000000dead"; + let (ticket, dust) = ntt::prepare_transfer( + &state, + coins, + &coin_meta, + ntt_scenario::peer_chain_id(), // recipient_chain + recipient, + option::none(), + false // should_queue + ); + + assert!(dust.value() == TEST_DUST); + + // Initial balance check + let initial_balance = if (state.borrow_mode().is_locking()) { + state.borrow_balance().value() + } else { + state.borrow_treasury_cap().total_supply() + }; + + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_limit_before = outbound_capacity.limit(); + let outbound_capacity_before = outbound_capacity.capacity_at_last_tx(); + let outbound_last_tx_timestamp_before = outbound_capacity.last_tx_timestamp(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_limit_before = inbound.limit(); + let inbound_capacity_before = inbound.capacity_at_last_tx(); + + // Execute transfer + let outbox_key = ntt::transfer_tx_sender( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + ticket, + &clock, + scenario.ctx() + ); + + // Rate limit data after the transfer + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_limit_after = outbound_capacity.limit(); + let outbound_capacity_after = outbound_capacity.capacity_at_last_tx(); + let outbound_last_tx_timestamp_after = outbound_capacity.last_tx_timestamp(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_limit_after = inbound.limit(); + let inbound_capacity_after = inbound.capacity_at_last_tx(); + let inbound_last_tx_after = inbound.last_tx_timestamp(); + + // Outbound and inbound rate limit comparisons + assert!(outbound_limit_after == outbound_limit_before); + assert!(outbound_last_tx_timestamp_after > outbound_last_tx_timestamp_before); + assert!(outbound_last_tx_timestamp_after == clock.timestamp_ms()); + assert!(outbound_capacity_after == (outbound_capacity_before) - (TEST_AMOUNT - TEST_DUST)); + assert!(inbound_limit_after == inbound_limit_before); + assert!(inbound_last_tx_after == clock.timestamp_ms()); + assert!(inbound_capacity_after == (inbound_capacity_before) + (TEST_AMOUNT - TEST_DUST)); + + // Verify state after transfer + if (state.borrow_mode().is_locking()) { + // In locking mode, tokens should be in the state's balance + assert!(state.borrow_balance().value() == initial_balance + (TEST_AMOUNT - TEST_DUST)) + } else { + assert!(state.borrow_treasury_cap().total_supply() == initial_balance - (TEST_AMOUNT - TEST_DUST)) + }; + + // Verify outbox item + let message = *state.borrow_outbox().borrow(outbox_key).borrow_data(); + + // Verify message contents + let (message_id, _, transfer) = message.destruct(); + let (trimmed_amount, _, recipient_addr, to_chain, _payload) = transfer.destruct(); + assert!(trimmed_amount.untrim(ntt_scenario::decimals()) == TEST_AMOUNT - TEST_DUST); + assert!(to_chain == ntt_scenario::peer_chain_id()); + assert!(recipient_addr.to_bytes() == recipient); + + // Advance past the 24h queue delay so the item becomes releasable. + clock.increment_for_testing(86_400_000); + + let transceiver_a_message = state.create_transceiver_message( + message_id, + &clock + ); + + let transceiver_b_message = state.create_transceiver_message( + message_id, + &clock + ); + + let (manager_message_a, source_manager_a, recipient_manager_a) = + transceiver_a_message.unwrap_outbound_message(&test_transceiver_a::auth()); + + let (manager_message_b, source_manager_b, recipient_manager_b) = + transceiver_b_message.unwrap_outbound_message(&test_transceiver_b::auth()); + + assert!(manager_message_a == manager_message_b); + assert!(source_manager_a == source_manager_b); + assert!(recipient_manager_a == recipient_manager_b); + + assert!(source_manager_a == external_address::from_address(object::id_address(&state))); + + let manager_message = ntt_manager_message::map!(manager_message_a, |x| native_token_transfer::parse(x)); + + assert!(manager_message == ntt_manager_message::new( + message_id, + external_address::from_address(user_a), + native_token_transfer::new( + ntt_common::trimmed_amount::new( + TEST_AMOUNT / 10, // token has 9 decimals + 8 + ), + external_address::from_id(object::id(&coin_meta)), + recipient_addr, + ntt_scenario::peer_chain_id(), + option::none() + ) + )); + + // Clean up + ntt_scenario::return_state(state); + ntt_scenario::return_clock(clock); + ntt_scenario::return_coin_metadata(coin_meta); + std::unit_test::destroy(dust); + test_scenario::end(scenario); + } + + #[test] + fun test_transfer_queued_released_after_delay() { + let (_, user_a, _, _) = ntt_scenario::test_addresses(); + let mut scenario = test_scenario::begin(user_a); + ntt_scenario::setup(&mut scenario); + + scenario.next_tx(user_a); + + // Take state, adminCap, and clock + let mut state = ntt_scenario::take_state(&scenario); + let admin_cap = ntt_scenario::take_admin_cap(&scenario); + let mut clock = ntt_scenario::take_clock(&mut scenario); + clock.increment_for_testing(1_000); + let coin_meta = ntt_scenario::take_coin_metadata(&scenario); + + // Set the rate limit + let limit: u64 = 1; + state::set_outbound_rate_limit(&admin_cap, &mut state, limit, &clock); + + let coins = state.mint_for_test(TEST_AMOUNT, scenario.ctx()); + + // Create transfer ticket + let recipient = x"000000000000000000000000000000000000000000000000000000000000dead"; + let (ticket, dust) = ntt::prepare_transfer( + &state, + coins, + &coin_meta, + ntt_scenario::peer_chain_id(), // recipient_chain + recipient, + option::none(), + true // should_queue + ); + + // Rate limit information before the transfer + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_limit_before = outbound_capacity.limit(); + let outbound_capacity_before = outbound_capacity.capacity_at_last_tx(); + let outbound_last_tx_timestamp_before = outbound_capacity.last_tx_timestamp(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_limit_before = inbound.limit(); + let inbound_capacity_before = inbound.capacity_at_last_tx(); + let inbound_last_tx_before = inbound.last_tx_timestamp(); + + // Execute transfer but with a later timestamp + let outbox_key = ntt::transfer_tx_sender( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + ticket, + &clock, + scenario.ctx() + ); + + // Rate limit data after the transfer + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_limit_after = outbound_capacity.limit(); + let outbound_capacity_after = outbound_capacity.capacity_at_last_tx(); + let outbound_last_tx_timestamp_after = outbound_capacity.last_tx_timestamp(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_limit_after = inbound.limit(); + let inbound_capacity_after = inbound.capacity_at_last_tx(); + let inbound_last_tx_after = inbound.last_tx_timestamp(); + + // Outbound and inbound rate limit comparisons + assert!(outbound_limit_after == outbound_limit_before); + assert!(outbound_last_tx_timestamp_after == outbound_last_tx_timestamp_before); + assert!(outbound_capacity_after == outbound_capacity_before); + assert!(inbound_limit_after == inbound_limit_before); + assert!(inbound_last_tx_after == inbound_last_tx_before); + assert!(inbound_capacity_after == inbound_capacity_before); + + // queued transfer at clock t: release is now + RATE_LIMIT_DURATION (24h)= + let release_timestamp = state.borrow_outbox().borrow(outbox_key).borrow_release_timestamp(); + assert!(release_timestamp == clock.timestamp_ms() + 86_400_000); + + let message = *state.borrow_outbox().borrow(outbox_key).borrow_data(); + + // Verify it's the transfer we expect. + let (message_id, _sender, transfer) = message.destruct(); + let (trimmed_amount, _src_token, recipient_addr, to_chain, _payload) = transfer.destruct(); + assert!(trimmed_amount.untrim(ntt_scenario::decimals()) == TEST_AMOUNT - TEST_DUST); + assert!(to_chain == ntt_scenario::peer_chain_id()); + assert!(recipient_addr.to_bytes() == recipient); + + // Advance past the 24h queue delay so the item becomes releasable. + clock.increment_for_testing(86_400_000); + + let transceiver_a_message = state.create_transceiver_message( + message_id, + &clock + ); + + let transceiver_b_message = state.create_transceiver_message( + message_id, + &clock + ); + + let (manager_message_a, source_manager_a, recipient_manager_a) = + transceiver_a_message.unwrap_outbound_message(&test_transceiver_a::auth()); + + let (manager_message_b, source_manager_b, recipient_manager_b) = + transceiver_b_message.unwrap_outbound_message(&test_transceiver_b::auth()); + + assert!(manager_message_a == manager_message_b); + assert!(source_manager_a == source_manager_b); + assert!(recipient_manager_a == recipient_manager_b); + + assert!(source_manager_a == external_address::from_address(object::id_address(&state))); + + let manager_message = ntt_manager_message::map!(manager_message_a, |x| native_token_transfer::parse(x)); + + assert!(manager_message == ntt_manager_message::new( + message_id, + external_address::from_address(user_a), + native_token_transfer::new( + ntt_common::trimmed_amount::new( + TEST_AMOUNT / 10, // token has 9 decimals + 8 + ), + external_address::from_id(object::id(&coin_meta)), + recipient_addr, + ntt_scenario::peer_chain_id(), + option::none() + ) + )); + + // Clean up + ntt_scenario::return_state(state); + ntt_scenario::return_clock(clock); + ntt_scenario::return_coin_metadata(coin_meta); + ntt_scenario::return_admin_cap(admin_cap); + std::unit_test::destroy(dust); + test_scenario::end(scenario); + } + + #[test, expected_failure(abort_code = ::ntt::ntt::ETransferExceedsRateLimit)] + fun test_transfer_exceeds_rate_limit_no_queue() { + let (_, user_a, _, _) = ntt_scenario::test_addresses(); + let mut scenario = test_scenario::begin(user_a); + ntt_scenario::setup(&mut scenario); + + scenario.next_tx(user_a); + + // Take state, adminCap, and clock + let mut state = ntt_scenario::take_state(&scenario); + let admin_cap = ntt_scenario::take_admin_cap(&scenario); + let mut clock = ntt_scenario::take_clock(&mut scenario); + clock.increment_for_testing(1_000); + let coin_meta = ntt_scenario::take_coin_metadata(&scenario); + + // Set the rate limit + let limit: u64 = 1; + state::set_outbound_rate_limit(&admin_cap, &mut state, limit, &clock); + + let coins = state.mint_for_test(TEST_AMOUNT, scenario.ctx()); + + // Create transfer ticket + let recipient = x"000000000000000000000000000000000000000000000000000000000000dead"; + let (ticket, dust) = ntt::prepare_transfer( + &state, + coins, + &coin_meta, + ntt_scenario::peer_chain_id(), // recipient_chain + recipient, + option::none(), + false // should_queue + ); + + // Execute transfer but with a later timestamp + let _outbox_key = ntt::transfer_tx_sender( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + ticket, + &clock, + scenario.ctx() + ); + + // Clean up + ntt_scenario::return_state(state); + ntt_scenario::return_clock(clock); + ntt_scenario::return_coin_metadata(coin_meta); + ntt_scenario::return_admin_cap(admin_cap); + std::unit_test::destroy(dust); + test_scenario::end(scenario); + + } + #[test, expected_failure] + fun test_transfer_queued_release_too_early() { + let (_, user_a, _, _) = ntt_scenario::test_addresses(); + let mut scenario = test_scenario::begin(user_a); + ntt_scenario::setup(&mut scenario); + + scenario.next_tx(user_a); + + // Take state, adminCap, and clock + let mut state = ntt_scenario::take_state(&scenario); + let admin_cap = ntt_scenario::take_admin_cap(&scenario); + let mut clock = ntt_scenario::take_clock(&mut scenario); + clock.increment_for_testing(1_000); + let coin_meta = ntt_scenario::take_coin_metadata(&scenario); + + // Set the rate limit + let limit: u64 = 1; + state::set_outbound_rate_limit(&admin_cap, &mut state, limit, &clock); + + let coins = state.mint_for_test(TEST_AMOUNT, scenario.ctx()); + + // Create transfer ticket + let recipient = x"000000000000000000000000000000000000000000000000000000000000dead"; + let (ticket, dust) = ntt::prepare_transfer( + &state, + coins, + &coin_meta, + ntt_scenario::peer_chain_id(), // recipient_chain + recipient, + option::none(), + true // should_queue + ); + + // Rate limit information before the transfer + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_limit_before = outbound_capacity.limit(); + let outbound_capacity_before = outbound_capacity.capacity_at_last_tx(); + let outbound_last_tx_timestamp_before = outbound_capacity.last_tx_timestamp(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_limit_before = inbound.limit(); + let inbound_capacity_before = inbound.capacity_at_last_tx(); + let inbound_last_tx_before = inbound.last_tx_timestamp(); + + // Execute transfer but with a later timestamp + let outbox_key = ntt::transfer_tx_sender( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + ticket, + &clock, + scenario.ctx() + ); + + // Rate limit data after the transfer + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_limit_after = outbound_capacity.limit(); + let outbound_capacity_after = outbound_capacity.capacity_at_last_tx(); + let outbound_last_tx_timestamp_after = outbound_capacity.last_tx_timestamp(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_limit_after = inbound.limit(); + let inbound_capacity_after = inbound.capacity_at_last_tx(); + let inbound_last_tx_after = inbound.last_tx_timestamp(); + + // Outbound and inbound rate limit comparisons + assert!(outbound_limit_after == outbound_limit_before); + assert!(outbound_last_tx_timestamp_after == outbound_last_tx_timestamp_before); + assert!(outbound_capacity_after == outbound_capacity_before); + assert!(inbound_limit_after == inbound_limit_before); + assert!(inbound_last_tx_after == inbound_last_tx_before); + assert!(inbound_capacity_after == inbound_capacity_before); + + // queued transfer at clock t: release is now + RATE_LIMIT_DURATION (24h)= + let release_timestamp = state.borrow_outbox().borrow(outbox_key).borrow_release_timestamp(); + assert!(release_timestamp == clock.timestamp_ms() + 86_400_000); + + let message = *state.borrow_outbox().borrow(outbox_key).borrow_data(); + + // Verify it's the transfer we expect. + let (message_id, _sender, transfer) = message.destruct(); + let (trimmed_amount, _src_token, recipient_addr, to_chain, _payload) = transfer.destruct(); + assert!(trimmed_amount.untrim(ntt_scenario::decimals()) == TEST_AMOUNT - TEST_DUST); + assert!(to_chain == ntt_scenario::peer_chain_id()); + assert!(recipient_addr.to_bytes() == recipient); + + // DO NOT advance past the 24h queue delay. Don't want this to be releasable. + // Fails here + let transceiver_a_message = state.create_transceiver_message( + message_id, + &clock + ); + + // Clean up + std::unit_test::destroy(transceiver_a_message); + ntt_scenario::return_state(state); + ntt_scenario::return_clock(clock); + ntt_scenario::return_coin_metadata(coin_meta); + ntt_scenario::return_admin_cap(admin_cap); + std::unit_test::destroy(dust); + test_scenario::end(scenario); + } + + #[test] + fun test_transfer_queued_rate_limit_unchanged() { + let (_, user_a, _, _) = ntt_scenario::test_addresses(); + let mut scenario = test_scenario::begin(user_a); + ntt_scenario::setup(&mut scenario); + + scenario.next_tx(user_a); + + // Take state, adminCap, and clock + let mut state = ntt_scenario::take_state(&scenario); + let admin_cap = ntt_scenario::take_admin_cap(&scenario); + let mut clock = ntt_scenario::take_clock(&mut scenario); + clock.increment_for_testing(1_000); + let coin_meta = ntt_scenario::take_coin_metadata(&scenario); + + // Set the rate limit + let limit: u64 = 1; + state::set_outbound_rate_limit(&admin_cap, &mut state, limit, &clock); + + let coins = state.mint_for_test(TEST_AMOUNT, scenario.ctx()); + + // Create transfer ticket + let recipient = x"000000000000000000000000000000000000000000000000000000000000dead"; + let (ticket, dust) = ntt::prepare_transfer( + &state, + coins, + &coin_meta, + ntt_scenario::peer_chain_id(), // recipient_chain + recipient, + option::none(), + true // should_queue + ); + + // Rate limit information before the transfer + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_limit_before = outbound_capacity.limit(); + let outbound_capacity_before = outbound_capacity.capacity_at_last_tx(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_limit_before = inbound.limit(); + let inbound_capacity_before = inbound.capacity_at_last_tx(); + let inbound_last_tx_before = inbound.last_tx_timestamp(); + + // Execute transfer but with a later timestamp + let outbox_key = ntt::transfer_tx_sender( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + ticket, + &clock, + scenario.ctx() + ); + + // Rate limit data after the transfer + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_limit_after = outbound_capacity.limit(); + let outbound_capacity_after = outbound_capacity.capacity_at_last_tx(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_limit_after = inbound.limit(); + let inbound_capacity_after = inbound.capacity_at_last_tx(); + let inbound_last_tx_after = inbound.last_tx_timestamp(); + + // Outbound and inbound rate limit comparisons + assert!(outbound_limit_after == outbound_limit_before); + assert!(outbound_capacity_after == outbound_capacity_before); + assert!(inbound_limit_after == inbound_limit_before); + assert!(inbound_last_tx_after == inbound_last_tx_before); + assert!(inbound_capacity_after == inbound_capacity_before); + + // queued transfer at clock t: release is now + RATE_LIMIT_DURATION (24h)= + let release_timestamp = state.borrow_outbox().borrow(outbox_key).borrow_release_timestamp(); + assert!(release_timestamp == clock.timestamp_ms() + 86_400_000); + + let message = *state.borrow_outbox().borrow(outbox_key).borrow_data(); + + // Verify it's the transfer we expect. + let (_message_id, _sender, transfer) = message.destruct(); + let (trimmed_amount, _src_token, recipient_addr, to_chain, _payload) = transfer.destruct(); + assert!(trimmed_amount.untrim(ntt_scenario::decimals()) == TEST_AMOUNT - TEST_DUST); + assert!(to_chain == ntt_scenario::peer_chain_id()); + assert!(recipient_addr.to_bytes() == recipient); + + + // Clean up + ntt_scenario::return_state(state); + ntt_scenario::return_clock(clock); + ntt_scenario::return_coin_metadata(coin_meta); + ntt_scenario::return_admin_cap(admin_cap); + std::unit_test::destroy(dust); + test_scenario::end(scenario); + } + + #[test] + fun test_transfer_rate_limit_timestamps() { + let (_, user_a, _, _) = ntt_scenario::test_addresses(); + let mut scenario = test_scenario::begin(user_a); + ntt_scenario::setup(&mut scenario); + + scenario.next_tx(user_a); + + // Take state, adminCap, and clock + let mut state = ntt_scenario::take_state(&scenario); + let admin_cap = ntt_scenario::take_admin_cap(&scenario); + let mut clock = ntt_scenario::take_clock(&mut scenario); + clock.increment_for_testing(1_000); + let coin_meta = ntt_scenario::take_coin_metadata(&scenario); + + let coins = state.mint_for_test(TEST_AMOUNT, scenario.ctx()); + + // Increase time + clock.increment_for_testing(1_000); + + // Create transfer ticket + let recipient = x"000000000000000000000000000000000000000000000000000000000000dead"; + let (ticket, dust) = ntt::prepare_transfer( + &state, + coins, + &coin_meta, + ntt_scenario::peer_chain_id(), // recipient_chain + recipient, + option::none(), + true // should_queue + ); + + // Rate limit information before the transfer + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_last_tx_before = outbound_capacity.last_tx_timestamp(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_last_tx_before = inbound.last_tx_timestamp(); + + // Execute transfer but with a later timestamp + let outbox_key = ntt::transfer_tx_sender( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + ticket, + &clock, + scenario.ctx() + ); + + // Rate limit data after the transfer + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_last_tx_after = outbound_capacity.last_tx_timestamp(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_last_tx_after = inbound.last_tx_timestamp(); + + // Outbound and inbound rate limit comparisons + assert!(inbound_last_tx_after != inbound_last_tx_before); + assert!(inbound_last_tx_after > inbound_last_tx_before); + assert!(outbound_last_tx_after != outbound_last_tx_before); + assert!(outbound_last_tx_after > outbound_last_tx_before); + assert!(outbound_last_tx_after == inbound_last_tx_after); + + let message = *state.borrow_outbox().borrow(outbox_key).borrow_data(); + + // Verify it's the transfer we expect. + let (_message_id, _sender, transfer) = message.destruct(); + let (trimmed_amount, _src_token, recipient_addr, to_chain, _payload) = transfer.destruct(); + assert!(trimmed_amount.untrim(ntt_scenario::decimals()) == TEST_AMOUNT - TEST_DUST); + assert!(to_chain == ntt_scenario::peer_chain_id()); + assert!(recipient_addr.to_bytes() == recipient); + + + // Clean up + ntt_scenario::return_state(state); + ntt_scenario::return_clock(clock); + ntt_scenario::return_coin_metadata(coin_meta); + ntt_scenario::return_admin_cap(admin_cap); + std::unit_test::destroy(dust); + test_scenario::end(scenario); + } + + #[test] + fun test_redeem_queued() { + let (_, user_a, user_b, _) = ntt_scenario::test_addresses(); + let mut scenario = test_scenario::begin(user_a); + ntt_scenario::setup(&mut scenario); + + // Take state, adminCap, and clock + let mut state = ntt_scenario::take_state(&scenario); + let admin_cap = ntt_scenario::take_admin_cap(&scenario); + let mut clock = ntt_scenario::take_clock(&mut scenario); + clock.increment_for_testing(1_000); + let coin_meta = ntt_scenario::take_coin_metadata(&scenario); + + // Set the rate limit + // set_peer performs a rate limit update if called in this way + let low_inbound_limit: u64 = 1; + state::set_peer( + &admin_cap, + &mut state, + ntt_scenario::peer_chain_id(), + ntt_scenario::peer_manager_address(), + ntt_scenario::decimals(), + low_inbound_limit, + &clock, + ); + + let message_id = wormhole::bytes32::from_u256_be(100); + let manager_message = ntt_manager_message::new( + message_id, + external_address::from_address(user_a), + native_token_transfer::new( + ntt_common::trimmed_amount::new( + TEST_AMOUNT / 10, // token has 9 decimals + 8 + ), + external_address::from_id(object::id(&coin_meta)), + external_address::from_address(user_b), + ntt_scenario::chain_id(), + option::none() + ) + ); + + let manager_message_encoded = ntt_manager_message::map!(manager_message, |x| x.to_bytes()); + + let validated_transceiver_message_a = ntt_common::validated_transceiver_message::new( + &test_transceiver_a::auth(), + ntt_scenario::peer_chain_id(), + ntt_common::transceiver_message_data::new( + ntt_scenario::peer_manager_address(), + external_address::from_address(object::id_address(&state)), + manager_message_encoded + ) + ); + + let validated_transceiver_message_b = ntt_common::validated_transceiver_message::new( + &test_transceiver_b::auth(), + ntt_scenario::peer_chain_id(), + ntt_common::transceiver_message_data::new( + ntt_scenario::peer_manager_address(), + external_address::from_address(object::id_address(&state)), + manager_message_encoded + ) + ); + + ntt::redeem( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + validated_transceiver_message_a, + &clock + ); + + ntt::redeem( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + validated_transceiver_message_b, + &clock + ); + + let inbox_item = state::borrow_inbox_item(&state, ntt_scenario::peer_chain_id(), manager_message); + + // Both transceivers voted, and the low inbound limit forced the transfer + // into the delayed-release queue (24h from the redeem timestamp). + assert!(inbox_item.count_votes() == 2); + assert!(inbox_item.is_release_after()); + assert!(!inbox_item.is_released()); + assert!(inbox_item.release_after_timestamp() == clock.timestamp_ms() + 86_400_000); + + // Wait out the queue delay, then complete it: + clock.increment_for_testing(86_400_000); + ntt::release( + &mut state, + upgrades::new_version_gated(), + ntt_scenario::peer_chain_id(), + manager_message, + &coin_meta, + &clock, + scenario.ctx() + ); + + let inbox_item = state::borrow_inbox_item(&state, ntt_scenario::peer_chain_id(), manager_message); + assert!(inbox_item.is_released()); + + scenario.next_tx(user_a); + + let coins = scenario.take_from_address>(user_b); + + assert!(coins.value() == TEST_AMOUNT - TEST_DUST); + + ntt_scenario::return_state(state); + ntt_scenario::return_clock(clock); + ntt_scenario::return_coin_metadata(coin_meta); + ntt_scenario::return_admin_cap(admin_cap); + std::unit_test::destroy(coins); + scenario.end(); + } + + #[test] + fun test_redeem_rate_limit_values() { + let (_, user_a, user_b, _) = ntt_scenario::test_addresses(); + let mut scenario = test_scenario::begin(user_a); + ntt_scenario::setup(&mut scenario); + + // Take state, adminCap, and clock + let mut state = ntt_scenario::take_state(&scenario); + let admin_cap = ntt_scenario::take_admin_cap(&scenario); + let mut clock = ntt_scenario::take_clock(&mut scenario); + clock.increment_for_testing(1_000); + let coin_meta = ntt_scenario::take_coin_metadata(&scenario); + + // Drain outbox limit to test backflow + let drain: u64 = 2_000_000_000; + let r = state.borrow_outbox_mut().borrow_rate_limit_mut().consume_or_delay(&clock, drain); + assert!(r.is_consumed()); + + let message_id = wormhole::bytes32::from_u256_be(100); + let manager_message = ntt_manager_message::new( + message_id, + external_address::from_address(user_a), + native_token_transfer::new( + ntt_common::trimmed_amount::new( + TEST_AMOUNT / 10, // token has 9 decimals + 8 + ), + external_address::from_id(object::id(&coin_meta)), + external_address::from_address(user_b), + ntt_scenario::chain_id(), + option::none() + ) + ); + + let manager_message_encoded = ntt_manager_message::map!(manager_message, |x| x.to_bytes()); + + let validated_transceiver_message_a = ntt_common::validated_transceiver_message::new( + &test_transceiver_a::auth(), + ntt_scenario::peer_chain_id(), + ntt_common::transceiver_message_data::new( + ntt_scenario::peer_manager_address(), + external_address::from_address(object::id_address(&state)), + manager_message_encoded + ) + ); + + let validated_transceiver_message_b = ntt_common::validated_transceiver_message::new( + &test_transceiver_b::auth(), + ntt_scenario::peer_chain_id(), + ntt_common::transceiver_message_data::new( + ntt_scenario::peer_manager_address(), + external_address::from_address(object::id_address(&state)), + manager_message_encoded + ) + ); + + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_limit_before = outbound_capacity.limit(); + let outbound_capacity_before = outbound_capacity.capacity_at_last_tx(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_limit_before = inbound.limit(); + let inbound_capacity_before = inbound.capacity_at_last_tx(); + + ntt::redeem( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + validated_transceiver_message_a, + &clock + ); + + ntt::redeem( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + validated_transceiver_message_b, + &clock + ); + + ntt::release( + &mut state, + upgrades::new_version_gated(), + ntt_scenario::peer_chain_id(), + manager_message, + &coin_meta, + &clock, + scenario.ctx() + ); + + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_limit_after = outbound_capacity.limit(); + let outbound_capacity_after = outbound_capacity.capacity_at_last_tx(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_limit_after = inbound.limit(); + let inbound_capacity_after = inbound.capacity_at_last_tx(); + + // Rate limit validations + assert!(inbound_limit_after == inbound_limit_before); + assert!(outbound_limit_after == outbound_limit_before); + assert!(outbound_capacity_after == (outbound_capacity_before + (TEST_AMOUNT - TEST_DUST))); + assert!(inbound_capacity_after == (inbound_capacity_before - (TEST_AMOUNT - TEST_DUST))); + + scenario.next_tx(user_a); + + let coins = scenario.take_from_address>(user_b); + + assert!(coins.value() == TEST_AMOUNT - TEST_DUST); + + ntt_scenario::return_state(state); + ntt_scenario::return_clock(clock); + ntt_scenario::return_coin_metadata(coin_meta); + ntt_scenario::return_admin_cap(admin_cap); + std::unit_test::destroy(coins); + scenario.end(); + } + + #[test] + fun test_redeem_rate_limit_timestamps() { + let (_, user_a, user_b, _) = ntt_scenario::test_addresses(); + let mut scenario = test_scenario::begin(user_a); + ntt_scenario::setup(&mut scenario); + + // Take state, adminCap, and clock + let mut state = ntt_scenario::take_state(&scenario); + let admin_cap = ntt_scenario::take_admin_cap(&scenario); + let mut clock = ntt_scenario::take_clock(&mut scenario); + clock.increment_for_testing(1_000); + let coin_meta = ntt_scenario::take_coin_metadata(&scenario); + + // Drain outbox limit to test backflow + let drain: u64 = 2_000_000_000; + let r = state.borrow_outbox_mut().borrow_rate_limit_mut().consume_or_delay(&clock, drain); + assert!(r.is_consumed()); + + let message_id = wormhole::bytes32::from_u256_be(100); + let manager_message = ntt_manager_message::new( + message_id, + external_address::from_address(user_a), + native_token_transfer::new( + ntt_common::trimmed_amount::new( + TEST_AMOUNT / 10, // token has 9 decimals + 8 + ), + external_address::from_id(object::id(&coin_meta)), + external_address::from_address(user_b), + ntt_scenario::chain_id(), + option::none() + ) + ); + + let manager_message_encoded = ntt_manager_message::map!(manager_message, |x| x.to_bytes()); + + let validated_transceiver_message_a = ntt_common::validated_transceiver_message::new( + &test_transceiver_a::auth(), + ntt_scenario::peer_chain_id(), + ntt_common::transceiver_message_data::new( + ntt_scenario::peer_manager_address(), + external_address::from_address(object::id_address(&state)), + manager_message_encoded + ) + ); + + let validated_transceiver_message_b = ntt_common::validated_transceiver_message::new( + &test_transceiver_b::auth(), + ntt_scenario::peer_chain_id(), + ntt_common::transceiver_message_data::new( + ntt_scenario::peer_manager_address(), + external_address::from_address(object::id_address(&state)), + manager_message_encoded + ) + ); + + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_last_tx_timestamp_before = outbound_capacity.last_tx_timestamp(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_last_tx_timestamp_before = inbound.last_tx_timestamp(); + + // Increment the clock for the transfer + let time_increase = 500; + clock.increment_for_testing(time_increase); + + ntt::redeem( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + validated_transceiver_message_a, + &clock + ); + + ntt::redeem( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + validated_transceiver_message_b, + &clock + ); + + ntt::release( + &mut state, + upgrades::new_version_gated(), + ntt_scenario::peer_chain_id(), + manager_message, + &coin_meta, + &clock, + scenario.ctx() + ); + + let outbound_capacity = state.borrow_outbox_mut().borrow_rate_limit_mut(); + let outbound_last_tx_timestamp_after = outbound_capacity.last_tx_timestamp(); + + let inbound = + state.borrow_peer_mut(ntt_scenario::peer_chain_id()) + .borrow_inbound_rate_limit_mut(); + + let inbound_last_tx_timestamp_after = inbound.last_tx_timestamp(); + + // Ensure that the inbound, and outbound last_tx_timestamp value was updated + assert!(inbound_last_tx_timestamp_after == clock.timestamp_ms()); + assert!(inbound_last_tx_timestamp_after > inbound_last_tx_timestamp_before); + assert!(outbound_last_tx_timestamp_after == clock.timestamp_ms()); + assert!(outbound_last_tx_timestamp_after > outbound_last_tx_timestamp_before); + + scenario.next_tx(user_a); + + let coins = scenario.take_from_address>(user_b); + + assert!(coins.value() == TEST_AMOUNT - TEST_DUST); + + ntt_scenario::return_state(state); + ntt_scenario::return_clock(clock); + ntt_scenario::return_coin_metadata(coin_meta); + ntt_scenario::return_admin_cap(admin_cap); + std::unit_test::destroy(coins); + scenario.end(); + } + + #[test, expected_failure(abort_code = ::ntt::ntt::ECantReleaseYet)] + fun test_redeem_queued_release_too_early() { + let (_, user_a, user_b, _) = ntt_scenario::test_addresses(); + let mut scenario = test_scenario::begin(user_a); + ntt_scenario::setup(&mut scenario); + + // Take state, adminCap, and clock + let mut state = ntt_scenario::take_state(&scenario); + let admin_cap = ntt_scenario::take_admin_cap(&scenario); + let mut clock = ntt_scenario::take_clock(&mut scenario); + clock.increment_for_testing(1_000); + let coin_meta = ntt_scenario::take_coin_metadata(&scenario); + + // Set the rate limit + // set_peer performs a rate limit update if called in this way + let low_inbound_limit: u64 = 1; + state::set_peer( + &admin_cap, + &mut state, + ntt_scenario::peer_chain_id(), + ntt_scenario::peer_manager_address(), + ntt_scenario::decimals(), + low_inbound_limit, + &clock, + ); + + let message_id = wormhole::bytes32::from_u256_be(100); + let manager_message = ntt_manager_message::new( + message_id, + external_address::from_address(user_a), + native_token_transfer::new( + ntt_common::trimmed_amount::new( + TEST_AMOUNT / 10, // token has 9 decimals + 8 + ), + external_address::from_id(object::id(&coin_meta)), + external_address::from_address(user_b), + ntt_scenario::chain_id(), + option::none() + ) + ); + + let manager_message_encoded = ntt_manager_message::map!(manager_message, |x| x.to_bytes()); + + let validated_transceiver_message_a = ntt_common::validated_transceiver_message::new( + &test_transceiver_a::auth(), + ntt_scenario::peer_chain_id(), + ntt_common::transceiver_message_data::new( + ntt_scenario::peer_manager_address(), + external_address::from_address(object::id_address(&state)), + manager_message_encoded + ) + ); + + let validated_transceiver_message_b = ntt_common::validated_transceiver_message::new( + &test_transceiver_b::auth(), + ntt_scenario::peer_chain_id(), + ntt_common::transceiver_message_data::new( + ntt_scenario::peer_manager_address(), + external_address::from_address(object::id_address(&state)), + manager_message_encoded + ) + ); + + ntt::redeem( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + validated_transceiver_message_a, + &clock + ); + + ntt::redeem( + &mut state, + upgrades::new_version_gated(), + &coin_meta, + validated_transceiver_message_b, + &clock + ); + + // DON'T wait out the queue delay. Fail if withdrawn too early. + ntt::release( + &mut state, + upgrades::new_version_gated(), + ntt_scenario::peer_chain_id(), + manager_message, + &coin_meta, + &clock, + scenario.ctx() + ); + + ntt_scenario::return_state(state); + ntt_scenario::return_clock(clock); + ntt_scenario::return_coin_metadata(coin_meta); + ntt_scenario::return_admin_cap(admin_cap); + scenario.end(); + } + } diff --git a/sui/packages/ntt/tests/rate_limit_tests.move b/sui/packages/ntt/tests/rate_limit_tests.move new file mode 100644 index 000000000..d24f8439d --- /dev/null +++ b/sui/packages/ntt/tests/rate_limit_tests.move @@ -0,0 +1,213 @@ +#[test_only] +/// Unit tests for the rate limiter's underlying properties +module ntt::rate_limit_tests { + use sui::clock::{Self, Clock}; + use ntt::rate_limit::{Self, RateLimitState}; + + // Mirrors the private RATE_LIMIT_DURATION in rate_limit.move (24h in ms). + const DURATION: u64 = 24 * 60 * 60 * 1000; + + fun new_clock(): Clock { + let mut ctx = sui::tx_context::dummy(); + clock::create_for_testing(&mut ctx) + } + + // RateLimitState has `store` but not `drop`, so it must be explicitly consumed. + fun teardown(state: RateLimitState, clock: Clock) { + std::unit_test::destroy(state); + clock::destroy_for_testing(clock); + } + + #[test] + fun test_new_starts_full() { + let clock = new_clock(); + let state = rate_limit::new(1_000); + assert!(state.limit() == 1_000); + assert!(state.capacity_at_last_tx() == 1_000); + assert!(state.last_tx_timestamp() == 0); + assert!(state.capacity_at(0) == 1_000); + teardown(state, clock); + } + + #[test] + fun test_consume_reduces_capacity_and_stamps_time() { + let mut clock = new_clock(); + clock.increment_for_testing(5_000); + let mut state = rate_limit::new(1_000); + + let r = state.consume_or_delay(&clock, 400); + assert!(r.is_consumed()); + assert!(state.capacity_at_last_tx() == 600); + assert!(state.last_tx_timestamp() == 5_000); + teardown(state, clock); + } + + #[test] + fun test_consume_exact_capacity_ok() { + let clock = new_clock(); + let mut state = rate_limit::new(1_000); + + let r = state.consume_or_delay(&clock, 1_000); + assert!(r.is_consumed()); + assert!(state.capacity_at_last_tx() == 0); + teardown(state, clock); + } + + #[test] + /// A delayed (over-capacity) consume must NOT mutate the limiter counters: + /// "Transactions that exceeded the capacity do not count, they are just delayed." + fun test_delay_leaves_counters_untouched() { + let mut clock = new_clock(); + clock.increment_for_testing(7_000); + let mut state = rate_limit::new(1_000); + + // Consume some first, so we have a non-default state to detect mutation against. + // capacity 800, last_tx 7_000 + let r0 = state.consume_or_delay(&clock, 200); + assert!(r0.is_consumed()); + + // t = 8_000 + clock.increment_for_testing(1_000); + let cap_before = state.capacity_at_last_tx(); + let ts_before = state.last_tx_timestamp(); + + // Far over capacity. + let r = state.consume_or_delay(&clock, 100_000); + assert!(r.is_delayed()); + assert!(r.delayed_until() == 8_000 + DURATION); + + // Counters untouched by the delayed path. + assert!(state.capacity_at_last_tx() == cap_before); + assert!(state.last_tx_timestamp() == ts_before); + teardown(state, clock); + } + + #[test] + /// Capacity recovers proportionally to elapsed time. Choosing limit == DURATION + /// makes the refill rate exactly 1 unit/ms, so after half the window exactly + /// half the limit is back. + fun test_partial_time_refill_is_proportional() { + let mut clock = new_clock(); + let mut state = rate_limit::new(DURATION); + + // Drain fully at t=0. + let r = state.consume_or_delay(&clock, DURATION); + assert!(r.is_consumed()); + assert!(state.capacity_at(0) == 0); + + clock.increment_for_testing(DURATION / 2); + assert!(state.capacity_at(DURATION / 2) == DURATION / 2); + teardown(state, clock); + } + + #[test] + /// Time-based refill saturates at `limit` — it never exceeds it no matter how + /// much time passes. + fun test_time_refill_caps_at_limit() { + let mut clock = new_clock(); + let mut state = rate_limit::new(DURATION); + + let r = state.consume_or_delay(&clock, DURATION); + assert!(r.is_consumed()); + + // Exactly back to limit after one full window. + clock.increment_for_testing(DURATION); + assert!(state.capacity_at(DURATION) == DURATION); + + // Still capped after additional time. + clock.increment_for_testing(2 * DURATION); + assert!(state.capacity_at(3 * DURATION) == DURATION); + teardown(state, clock); + } + + #[test] + fun test_refill_adds_within_headroom() { + let clock = new_clock(); + let mut state = rate_limit::new(1_000); + + // capacity 400 at t=0 + let r = state.consume_or_delay(&clock, 600); + assert!(r.is_consumed()); + + // 400 + 300 = 700, under limit + state.refill(&clock, 300); + assert!(state.capacity_at_last_tx() == 700); + assert!(state.last_tx_timestamp() == 0); + teardown(state, clock); + } + + #[test] + fun test_refill_saturates_at_limit() { + let clock = new_clock(); + let mut state = rate_limit::new(1_000); + + // capacity 400 + let r = state.consume_or_delay(&clock, 600); + assert!(r.is_consumed()); + + // Would overshoot the limit; must be capped. + state.refill(&clock, 10_000); + assert!(state.capacity_at_last_tx() == 1_000); + teardown(state, clock); + } + + #[test] + /// Lowering the limit reduces current capacity by the same delta. + fun test_set_limit_decrease_reduces_capacity() { + let clock = new_clock(); + // full: capacity 1_000 + let mut state = rate_limit::new(1_000); + + state.set_limit(600, &clock); + assert!(state.limit() == 600); + // 1_000 - 400 + assert!(state.capacity_at_last_tx() == 600); + teardown(state, clock); + } + + #[test] + /// If the limit decrease exceeds current capacity, capacity floors at 0. + fun test_set_limit_decrease_floors_at_zero() { + let clock = new_clock(); + let mut state = rate_limit::new(1_000); + + // capacity 50 + let r = state.consume_or_delay(&clock, 950); + assert!(r.is_consumed()); + + // Decrease of 200 > capacity 50. + state.set_limit(800, &clock); + assert!(state.capacity_at_last_tx() == 0); + teardown(state, clock); + } + + #[test] + /// Raising the limit increases current capacity by the same delta. + fun test_set_limit_increase_raises_capacity() { + let clock = new_clock(); + // full: capacity 1_000 + let mut state = rate_limit::new(1_000); + + state.set_limit(1_500, &clock); + assert!(state.limit() == 1_500); + // 1_000 + 500 + assert!(state.capacity_at_last_tx() == 1_500); + teardown(state, clock); + } + + #[test, expected_failure] + /// `capacity_at` requires `last_tx_timestamp <= now`; querying the past aborts. + fun test_capacity_at_before_last_tx_aborts() { + let mut clock = new_clock(); + clock.increment_for_testing(5_000); + let mut state = rate_limit::new(1_000); + + // last_tx = 5_000 + let r = state.consume_or_delay(&clock, 100); + assert!(r.is_consumed()); + + // now < last_tx -> abort + let _ = state.capacity_at(4_000); + teardown(state, clock); + } +} From a87eaf42a4cd18bf9bd26ea9071c0b63defe37c2 Mon Sep 17 00:00:00 2001 From: mdulin2 Date: Tue, 30 Jun 2026 11:35:59 -0700 Subject: [PATCH 2/2] Precise value comparison on last_tx_timestamp tests --- sui/packages/ntt/tests/ntt_tests.move | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/sui/packages/ntt/tests/ntt_tests.move b/sui/packages/ntt/tests/ntt_tests.move index e1948be51..f0cbda2da 100644 --- a/sui/packages/ntt/tests/ntt_tests.move +++ b/sui/packages/ntt/tests/ntt_tests.move @@ -1529,12 +1529,21 @@ module ntt::ntt_tests { let inbound_last_tx_after = inbound.last_tx_timestamp(); - // Outbound and inbound rate limit comparisons - assert!(inbound_last_tx_after != inbound_last_tx_before); - assert!(inbound_last_tx_after > inbound_last_tx_before); - assert!(outbound_last_tx_after != outbound_last_tx_before); - assert!(outbound_last_tx_after > outbound_last_tx_before); - assert!(outbound_last_tx_after == inbound_last_tx_after); + // The transfer records the current clock time as last_tx_timestamp for + // both the inbound and outbound rate limiters. The clock is not advanced + // between the snapshots, so `before + time == after` holds exactly. + let time = clock.timestamp_ms(); + + // Both limiters started at 0 (set during setup at clock time 0). + assert!(inbound_last_tx_before == 0); + assert!(outbound_last_tx_before == 0); + + // Inbound and outbound timestamps advance to exactly the transfer time. + assert!(inbound_last_tx_after == inbound_last_tx_before + time); + assert!(outbound_last_tx_after == outbound_last_tx_before + time); + + // Both limiters are updated to the same timestamp. + assert!(outbound_last_tx_after == inbound_last_tx_after); let message = *state.borrow_outbox().borrow(outbox_key).borrow_data();