From e59a0dc18626de2fca2a2acebffe3fa7b4819171 Mon Sep 17 00:00:00 2001
From: Jon Cinque <jon.cinque@gmail.com>
Date: Wed, 27 Apr 2022 17:48:08 +0200
Subject: [PATCH] token-2022: Set transfer fee two epochs ahead (#3120)

---
 token/program-2022-test/tests/transfer_fee.rs | 50 ++++++++++++++++++-
 .../src/extension/transfer_fee/processor.rs   |  7 +--
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/token/program-2022-test/tests/transfer_fee.rs b/token/program-2022-test/tests/transfer_fee.rs
index 2542661e26a..4ee4a7a87ba 100644
--- a/token/program-2022-test/tests/transfer_fee.rs
+++ b/token/program-2022-test/tests/transfer_fee.rs
@@ -267,6 +267,18 @@ async fn set_fee() {
         }])
         .await
         .unwrap();
+
+    // warp to first normal slot to easily calculate epochs
+    let epoch_schedule = context.context.lock().await.genesis_config().epoch_schedule;
+    let first_normal_slot = epoch_schedule.first_normal_slot;
+    let slots_per_epoch = epoch_schedule.slots_per_epoch;
+    context
+        .context
+        .lock()
+        .await
+        .warp_to_slot(first_normal_slot)
+        .unwrap();
+
     let token = context.token_context.unwrap().token;
 
     // set to something new, old fee not touched
@@ -315,9 +327,43 @@ async fn set_fee() {
     );
     assert_eq!(extension.older_transfer_fee, newer_transfer_fee);
 
-    // warp forward one epoch, new fee becomes old fee during set
+    // warp forward one epoch, old fee still not touched when set
+    let new_transfer_fee_basis_points = 10;
+    let new_maximum_fee = 10;
+    context
+        .context
+        .lock()
+        .await
+        .warp_to_slot(first_normal_slot + slots_per_epoch)
+        .unwrap();
+    token
+        .set_transfer_fee(
+            &transfer_fee_config_authority,
+            new_transfer_fee_basis_points,
+            new_maximum_fee,
+        )
+        .await
+        .unwrap();
+    let state = token.get_mint_info().await.unwrap();
+    let extension = state.get_extension::<TransferFeeConfig>().unwrap();
+    assert_eq!(
+        extension.newer_transfer_fee.transfer_fee_basis_points,
+        new_transfer_fee_basis_points.into()
+    );
+    assert_eq!(
+        extension.newer_transfer_fee.maximum_fee,
+        new_maximum_fee.into()
+    );
+    assert_eq!(extension.older_transfer_fee, newer_transfer_fee);
+
+    // warp forward two epochs, old fee is replaced on set
     let newer_transfer_fee = extension.newer_transfer_fee;
-    context.context.lock().await.warp_to_slot(10_000).unwrap();
+    context
+        .context
+        .lock()
+        .await
+        .warp_to_slot(first_normal_slot + 3 * slots_per_epoch)
+        .unwrap();
     let new_transfer_fee_basis_points = MAX_FEE_BASIS_POINTS;
     let new_maximum_fee = u64::MAX;
     token
diff --git a/token/program-2022/src/extension/transfer_fee/processor.rs b/token/program-2022/src/extension/transfer_fee/processor.rs
index adce6e8306a..60ef5944a51 100644
--- a/token/program-2022/src/extension/transfer_fee/processor.rs
+++ b/token/program-2022/src/extension/transfer_fee/processor.rs
@@ -91,15 +91,16 @@ fn process_set_transfer_fee(
     // When setting the transfer fee, we have two situations:
     // * newer transfer fee epoch <= current epoch:
     //     newer transfer fee is the active one, so overwrite older transfer fee with newer, then overwrite newer transfer fee
-    // * newer transfer fee epoch == next epoch:
+    // * newer transfer fee epoch >= next epoch:
     //     it was never used, so just overwrite next transfer fee
     let epoch = Clock::get()?.epoch;
-    let next_epoch = epoch.saturating_add(1);
     if u64::from(extension.newer_transfer_fee.epoch) <= epoch {
         extension.older_transfer_fee = extension.newer_transfer_fee;
     }
+    // set two epochs ahead to avoid rug pulls at the end of an epoch
+    let newer_fee_start_epoch = epoch.saturating_add(2);
     let transfer_fee = TransferFee {
-        epoch: next_epoch.into(),
+        epoch: newer_fee_start_epoch.into(),
         transfer_fee_basis_points: transfer_fee_basis_points.into(),
         maximum_fee: maximum_fee.into(),
     };