From aa29c76807dd288d0005322f0e63cc82f306b632 Mon Sep 17 00:00:00 2001
From: Awosdot <awosdot@gmail.com>
Date: Sun, 18 Oct 2026 22:02:07 +0100
Subject: [PATCH] fix(vault): bind admin in storage migration and align
 ForceUpgrade checkpoints (#1633)

---
 contracts/vault/src/lib.rs         | 29 +++++++++++----
 contracts/vault/src/proxy_tests.rs | 58 ++++++++++++++++++++++++++++--
 2 files changed, 79 insertions(+), 8 deletions(-)

diff --git a/contracts/vault/src/lib.rs b/contracts/vault/src/lib.rs
index b5bd3f42d..c084bf5f2 100644
--- a/contracts/vault/src/lib.rs
+++ b/contracts/vault/src/lib.rs
@@ -597,10 +597,23 @@ impl YieldVault {
         let pending = get_pending_upgrade(&env).ok_or(VaultError::NoPendingWithdrawal)?;
         Self::assert_timelock_ready(&env, pending.eta)?;
 
-        let pre_version = get_storage_version(&env);
-        Self::run_storage_migration(&env, STORAGE_VERSION)?;
+        set_pending_upgrade(&env, &None);
+        Self::upgrade_with_checkpoint(&env, &admin, pending.wasm_hash)
+    }
+
+    /// Shared by `execute_upgrade` and the `ForceUpgrade` emergency action:
+    /// migrates storage to `STORAGE_VERSION`, verifies the version checkpoint,
+    /// swaps the WASM, and emits `upg_chkpt` with
+    /// `(old_wasm_hash, new_wasm_hash, old_storage_version, new_storage_version)`.
+    fn upgrade_with_checkpoint(
+        env: &Env,
+        caller: &Address,
+        new_wasm_hash: BytesN<32>,
+    ) -> Result<(), VaultError> {
+        let pre_version = get_storage_version(env);
+        Self::run_storage_migration(env, STORAGE_VERSION)?;
         // Checkpoint: storage version must equal STORAGE_VERSION after migration.
-        let post_version = get_storage_version(&env);
+        let post_version = get_storage_version(env);
         assert_eq!(
             post_version, STORAGE_VERSION,
             "storage version checkpoint failed: expected {}, got {}",
@@ -614,8 +627,12 @@ impl YieldVault {
             post_version
         );
 
-        set_pending_upgrade(&env, &None);
-        Self::apply_upgrade(&env, &admin, pending.wasm_hash);
+        let previous = get_implementation(env);
+        Self::apply_upgrade(env, caller, new_wasm_hash.clone());
+        env.events().publish(
+            (symbol_short!("upg_chkpt"), caller.clone()),
+            (previous, new_wasm_hash, pre_version, post_version),
+        );
         Ok(())
     }
 
@@ -1203,7 +1220,7 @@ impl YieldVault {
 
             emergency::EmergencyActionKind::ForceUpgrade => {
                 let hash = proposal.wasm_hash.clone().expect("wasm hash required");
-                Self::apply_upgrade(&env, &confirmer, hash);
+                Self::upgrade_with_checkpoint(&env, &confirmer, hash)?;
             }
         }
 
diff --git a/contracts/vault/src/proxy_tests.rs b/contracts/vault/src/proxy_tests.rs
index 4ab464735..e5fedeb7b 100644
--- a/contracts/vault/src/proxy_tests.rs
+++ b/contracts/vault/src/proxy_tests.rs
@@ -274,8 +274,8 @@ fn test_migrate_storage_version_checkpoint() {
 
     let pre = vault.storage_version();
     // Idempotent migration to current version must pass the checkpoint.
-    vault.migrate_storage(&2);
-    assert_eq!(vault.storage_version(), 2);
+    vault.migrate_storage(&STORAGE_VERSION);
+    assert_eq!(vault.storage_version(), STORAGE_VERSION);
     assert!(vault.storage_version() >= pre);
 }
 
@@ -448,3 +448,57 @@ fn generate_storage_fingerprint(env: &Env) -> &str {
     // (Rust Vec of strings is hard to return here, so we just use it for internal assertion)
     "Admin TokenAsset Initialized"
 }
+
+#[test]
+fn test_execute_upgrade_emits_checkpoint_event_with_versions() {
+    let env = Env::default();
+    let (vault, admin, hash) = setup_upgrade(&env);
+
+    let eta = vault.propose_upgrade(&hash);
+    env.ledger().set_timestamp(eta);
+    vault.execute_upgrade();
+
+    let events = env.events().all();
+    assert!(events.iter().any(|(contract, topics, data)| {
+        contract == vault.address
+            && topics == (symbol_short!("upg_chkpt"), admin.clone()).into_val(&env)
+            && <(Option<BytesN<32>>, BytesN<32>, u32, u32)>::try_from_val(&env, &data)
+                == Ok((None, hash.clone(), STORAGE_VERSION, STORAGE_VERSION))
+    }));
+}
+
+#[test]
+fn test_force_upgrade_runs_migration_checkpoint_and_emits_event() {
+    let env = Env::default();
+    let (vault, _admin, hash) = setup_upgrade(&env);
+    let primary = Address::generate(&env);
+    let secondary = Address::generate(&env);
+    vault.set_emergency_approvers(&primary, &secondary);
+
+    // Simulate a vault whose storage predates the current schema.
+    env.as_contract(&vault.address, || set_storage_version(&env, 1));
+    assert_eq!(vault.storage_version(), 1);
+
+    let id = vault.propose_emergency_action(
+        &primary,
+        &emergency::EmergencyActionKind::ForceUpgrade,
+        &0u32,
+        &None,
+        &Some(hash.clone()),
+    );
+    let window = vault.emergency_dispute_window();
+    env.ledger()
+        .set_timestamp(env.ledger().timestamp() + window + 1);
+    vault.confirm_emergency_action(&secondary, &id);
+    let events = env.events().all();
+
+    assert_eq!(vault.storage_version(), STORAGE_VERSION);
+    assert_eq!(vault.implementation(), Some(hash.clone()));
+
+    assert!(events.iter().any(|(contract, topics, data)| {
+        contract == vault.address
+            && topics == (symbol_short!("upg_chkpt"), secondary.clone()).into_val(&env)
+            && <(Option<BytesN<32>>, BytesN<32>, u32, u32)>::try_from_val(&env, &data)
+                == Ok((None, hash.clone(), 1, STORAGE_VERSION))
+    }));
+}