From 0d764bf7ffbbf81b4d51d1bfd5e51929e7fb0460 Mon Sep 17 00:00:00 2001
From: Csongor Kiss <kiss.csongor.kiss@gmail.com>
Date: Fri, 1 Mar 2024 10:25:13 +0000
Subject: [PATCH] solana: enforce no fees have been charged

we could support this case instead, but the calculation would be a
little more difficult on the client side
---
 .../src/error.rs                              |  4 +++
 .../src/instructions/transfer.rs              | 26 +++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/solana/programs/example-native-token-transfers/src/error.rs b/solana/programs/example-native-token-transfers/src/error.rs
index 1ab9f508a..7db832da7 100644
--- a/solana/programs/example-native-token-transfers/src/error.rs
+++ b/solana/programs/example-native-token-transfers/src/error.rs
@@ -37,4 +37,8 @@ pub enum NTTError {
     DisabledTransceiver,
     #[msg("InvalidDeployer")]
     InvalidDeployer,
+    #[msg("BadAmountAfterTransfer")]
+    BadAmountAfterTransfer,
+    #[msg("BadAmountAfterBurn")]
+    BadAmountAfterBurn,
 }
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 31cba2789..7ece7d2be 100644
--- a/solana/programs/example-native-token-transfers/src/instructions/transfer.rs
+++ b/solana/programs/example-native-token-transfers/src/instructions/transfer.rs
@@ -134,6 +134,8 @@ pub fn transfer_burn(ctx: Context<TransferBurn>, args: TransferArgs) -> Result<(
         accs.peer.token_decimals,
     );
 
+    let before = accs.common.from.amount;
+
     token_interface::burn(
         CpiContext::new_with_signer(
             accs.common.token_program.to_account_info(),
@@ -152,6 +154,13 @@ pub fn transfer_burn(ctx: Context<TransferBurn>, args: TransferArgs) -> Result<(
         amount,
     )?;
 
+    accs.common.from.reload()?;
+    let after = accs.common.from.amount;
+
+    if after != before - amount {
+        return Err(NTTError::BadAmountAfterBurn.into());
+    }
+
     let recipient_ntt_manager = accs.peer.address;
 
     insert_into_outbox(
@@ -227,6 +236,8 @@ pub fn transfer_lock(ctx: Context<TransferLock>, args: TransferArgs) -> Result<(
         accs.peer.token_decimals,
     );
 
+    let before = accs.custody.amount;
+
     token_interface::transfer_checked(
         CpiContext::new_with_signer(
             accs.common.token_program.to_account_info(),
@@ -247,6 +258,21 @@ pub fn transfer_lock(ctx: Context<TransferLock>, args: TransferArgs) -> Result<(
         accs.common.mint.decimals,
     )?;
 
+    accs.custody.reload()?;
+    let after = accs.custody.amount;
+
+    // NOTE: we currently do not support tokens with fees. Support could be
+    // added, but it would require the client to calculate the amount _before_
+    // paying fees that results in an amount that can safely be trimmed.
+    // Otherwise, if the amount after paying fees has dust, then that amount
+    // would be lost.
+    // To support fee tokens, we would first transfer the amount, _then_ assert
+    // that the resulting amount has no dust (instead of removing dust before
+    // the transfer like we do now).
+    if after != before + amount {
+        return Err(NTTError::BadAmountAfterTransfer.into());
+    }
+
     let recipient_ntt_manager = accs.peer.address;
 
     insert_into_outbox(