diff --git a/solana/programs/example-native-token-transfers/src/config.rs b/solana/programs/example-native-token-transfers/src/config.rs index 7d04ece31..6a1668017 100644 --- a/solana/programs/example-native-token-transfers/src/config.rs +++ b/solana/programs/example-native-token-transfers/src/config.rs @@ -68,3 +68,12 @@ pub enum Mode { Burning, Locking, } + +impl std::fmt::Display for Mode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Mode::Burning => write!(f, "Burning"), + Mode::Locking => write!(f, "Locking"), + } + } +} diff --git a/solana/programs/example-native-token-transfers/src/instructions/transfer.rs b/solana/programs/example-native-token-transfers/src/instructions/transfer.rs index 5c586e1d8..abb10058a 100644 --- a/solana/programs/example-native-token-transfers/src/instructions/transfer.rs +++ b/solana/programs/example-native-token-transfers/src/instructions/transfer.rs @@ -95,6 +95,12 @@ pub struct TransferBurn<'info> { // TODO: fees for relaying? pub fn transfer_burn(ctx: Context, args: TransferArgs) -> Result<()> { + require_eq!( + ctx.accounts.common.config.mode, + Mode::Burning, + NTTError::InvalidMode + ); + let accs = ctx.accounts; let TransferArgs { amount, @@ -106,24 +112,21 @@ pub fn transfer_burn(ctx: Context, args: TransferArgs) -> Result<( // TODO: should we revert if we have dust? let amount = NormalizedAmount::remove_dust(amount, accs.common.mint.decimals); - match accs.common.config.mode { - Mode::Burning => token_interface::burn( - CpiContext::new_with_signer( - accs.common.token_program.to_account_info(), - token_interface::Burn { - mint: accs.common.mint.to_account_info(), - from: accs.common.from.to_account_info(), - authority: accs.common.token_authority.to_account_info(), - }, - &[&[ - crate::TOKEN_AUTHORITY_SEED, - &[ctx.bumps.common.token_authority], - ]], - ), - amount, - )?, - Mode::Locking => return Err(NTTError::InvalidMode.into()), - } + token_interface::burn( + CpiContext::new_with_signer( + accs.common.token_program.to_account_info(), + token_interface::Burn { + mint: accs.common.mint.to_account_info(), + from: accs.common.from.to_account_info(), + authority: accs.common.token_authority.to_account_info(), + }, + &[&[ + crate::TOKEN_AUTHORITY_SEED, + &[ctx.bumps.common.token_authority], + ]], + ), + amount, + )?; insert_into_outbox( &mut accs.common, @@ -162,6 +165,12 @@ pub struct TransferLock<'info> { // TODO: fees for relaying? // TODO: factor out common bits pub fn transfer_lock(ctx: Context, args: TransferArgs) -> Result<()> { + require_eq!( + ctx.accounts.common.config.mode, + Mode::Locking, + NTTError::InvalidMode + ); + let accs = ctx.accounts; let TransferArgs { amount, @@ -173,26 +182,23 @@ pub fn transfer_lock(ctx: Context, args: TransferArgs) -> Result<( // TODO: should we revert if we have dust? let amount = NormalizedAmount::remove_dust(amount, accs.common.mint.decimals); - match accs.common.config.mode { - Mode::Burning => return Err(NTTError::InvalidMode.into()), - Mode::Locking => token_interface::transfer_checked( - CpiContext::new_with_signer( - accs.common.token_program.to_account_info(), - token_interface::TransferChecked { - from: accs.common.from.to_account_info(), - to: accs.custody.to_account_info(), - authority: accs.common.token_authority.to_account_info(), - mint: accs.common.mint.to_account_info(), - }, - &[&[ - crate::TOKEN_AUTHORITY_SEED, - &[ctx.bumps.common.token_authority], - ]], - ), - amount, - accs.common.mint.decimals, - )?, - } + token_interface::transfer_checked( + CpiContext::new_with_signer( + accs.common.token_program.to_account_info(), + token_interface::TransferChecked { + from: accs.common.from.to_account_info(), + to: accs.custody.to_account_info(), + authority: accs.common.token_authority.to_account_info(), + mint: accs.common.mint.to_account_info(), + }, + &[&[ + crate::TOKEN_AUTHORITY_SEED, + &[ctx.bumps.common.token_authority], + ]], + ), + amount, + accs.common.mint.decimals, + )?; insert_into_outbox( &mut accs.common, diff --git a/solana/programs/example-native-token-transfers/src/queue/rate_limit.rs b/solana/programs/example-native-token-transfers/src/queue/rate_limit.rs index 8cda07f08..a0b36ea83 100644 --- a/solana/programs/example-native-token-transfers/src/queue/rate_limit.rs +++ b/solana/programs/example-native-token-transfers/src/queue/rate_limit.rs @@ -94,13 +94,14 @@ impl RateLimitState { /// Refills the capacity by the given amount. /// This is used to replenish the capacity via backflows. pub fn refill(&mut self, now: UnixTimestamp, amount: u64) { - self.capacity_at_last_tx = self.capacity().saturating_add(amount).min(self.limit); + self.capacity_at_last_tx = self.capacity_at(now).saturating_add(amount).min(self.limit); self.last_tx_timestamp = now; } pub fn set_limit(&mut self, limit: u64) { let old_limit = self.limit; - let current_capacity = self.capacity(); + let now = current_timestamp(); + let current_capacity = self.capacity_at(now); self.limit = limit; @@ -115,7 +116,7 @@ impl RateLimitState { }; self.capacity_at_last_tx = new_capacity.min(limit); - self.last_tx_timestamp = current_timestamp(); + self.last_tx_timestamp = now; } }