solana: fail on wrong mode earlier; one less timestamp call

wormhole-foundation · Feb 22, 2024 · ad8e4a3 · ad8e4a3
1 parent 79eefa0
commit ad8e4a3
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 41 deletions.
diff --git 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
@@ -95,6 +95,12 @@ pub struct TransferBurn<'info> {
 
 // TODO: fees for relaying?
 pub fn transfer_burn(ctx: Context<TransferBurn>, 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<TransferBurn>, 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<TransferLock>, 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<TransferLock>, 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
@@ -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;
     }
 }