From f2d406ad5d285489f56a96eeef69df3ed07588ef Mon Sep 17 00:00:00 2001
From: steviez <steven@solana.com>
Date: Mon, 7 Feb 2022 16:21:43 -0600
Subject: [PATCH] Fix panic in ProgramTestContext::warp_to_slot() when warping
 one slot (#22977)

warp_to_slot() first warps to one slot before desired slot in order to
freeze the bank at warp slot. This would cause issues when warping by
one slot as that would attempt to warp to the same slot and hit a sanity
check assertion.
---
 program-test/src/lib.rs    | 23 +++++++++++++++--------
 program-test/tests/warp.rs | 22 +++++++++++++++++++++-
 2 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/program-test/src/lib.rs b/program-test/src/lib.rs
index b497f2bf85..5e5ec05b79 100644
--- a/program-test/src/lib.rs
+++ b/program-test/src/lib.rs
@@ -1075,26 +1075,33 @@ impl ProgramTestContext {
             bank.register_tick(&Hash::new_unique());
         }
 
-        // warp ahead to one slot *before* the desired slot because the warped
-        // bank is frozen
+        // Ensure that we are actually progressing forward
         let working_slot = bank.slot();
         if warp_slot <= working_slot {
             return Err(ProgramTestError::InvalidWarpSlot);
         }
 
+        // Warp ahead to one slot *before* the desired slot because the bank
+        // from Bank::warp_from_parent() is frozen. If the desired slot is one
+        // slot *after* the working_slot, no need to warp at all.
         let pre_warp_slot = warp_slot - 1;
-        let warp_bank = bank_forks.insert(Bank::warp_from_parent(
-            &bank,
-            &Pubkey::default(),
-            pre_warp_slot,
-        ));
+        let warp_bank = if pre_warp_slot == working_slot {
+            bank.freeze();
+            bank
+        } else {
+            bank_forks.insert(Bank::warp_from_parent(
+                &bank,
+                &Pubkey::default(),
+                pre_warp_slot,
+            ))
+        };
         bank_forks.set_root(
             pre_warp_slot,
             &solana_runtime::accounts_background_service::AbsRequestSender::default(),
             Some(pre_warp_slot),
         );
 
-        // warp bank is frozen, so go forward one slot from it
+        // warp_bank is frozen so go forward to get unfrozen bank at warp_slot
         bank_forks.insert(Bank::new_from_parent(
             &warp_bank,
             &Pubkey::default(),
diff --git a/program-test/tests/warp.rs b/program-test/tests/warp.rs
index e3c392788f..6e29fa3ad7 100644
--- a/program-test/tests/warp.rs
+++ b/program-test/tests/warp.rs
@@ -132,7 +132,7 @@ async fn clock_sysvar_updated_from_warp() {
     );
 
     let mut context = program_test.start_with_context().await;
-    let expected_slot = 100_000;
+    let mut expected_slot = 100_000;
     let instruction = Instruction::new_with_bincode(
         program_id,
         &expected_slot,
@@ -175,6 +175,26 @@ async fn clock_sysvar_updated_from_warp() {
         .await
         .unwrap();
 
+    // Try warping ahead one slot (corner case in warp logic)
+    expected_slot += 1;
+    assert!(context.warp_to_slot(expected_slot).is_ok());
+    let instruction = Instruction::new_with_bincode(
+        program_id,
+        &expected_slot,
+        vec![AccountMeta::new_readonly(clock::id(), false)],
+    );
+    let transaction = Transaction::new_signed_with_payer(
+        &[instruction],
+        Some(&context.payer.pubkey()),
+        &[&context.payer],
+        context.last_blockhash,
+    );
+    context
+        .banks_client
+        .process_transaction(transaction)
+        .await
+        .unwrap();
+
     // Try warping again to the same slot
     assert_eq!(
         context.warp_to_slot(expected_slot).unwrap_err(),