From 3c6849cedd3860bfdf8b550ea4ff63008c6524a3 Mon Sep 17 00:00:00 2001
From: Ofuree <iamdrew404@gmail.com>
Date: Tue, 28 Jul 2026 03:25:17 +0200
Subject: [PATCH] Feat: Timelocked-Upgrades and Mandatory 48-Hour Delay Window
 on WASM Replacement

---
 src/governance.rs        |  2 ++
 src/lib.rs               | 12 +++++++-----
 src/upgrades/mod.rs      |  1 +
 src/upgrades/timelock.rs | 32 ++++++++++++++++++++++++++++++++
 4 files changed, 42 insertions(+), 5 deletions(-)
 create mode 100644 src/upgrades/timelock.rs

diff --git a/src/governance.rs b/src/governance.rs
index 4a8c727..4f302cc 100644
--- a/src/governance.rs
+++ b/src/governance.rs
@@ -158,6 +158,8 @@ pub struct StagedUpgrade {
     pub new_wasm_hash: BytesN<32>,
     pub proposer: Address,
     pub staged_at: u64,
+    /// Earliest ledger timestamp at which the replacement may execute.
+    pub execute_at: u64,
 }
 
 #[contracttype]
diff --git a/src/lib.rs b/src/lib.rs
index c41f385..d5523cb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -207,7 +207,7 @@ pub enum ContractError {
 pub(crate) const DATA_KEY: Symbol = symbol_short!("DATA");
 pub(crate) const SIGNERS_KEY: Symbol = symbol_short!("SIGNERS");
 const PENDING_UPGRADE_KEY: Symbol = symbol_short!("PENDING");
-pub(crate) const UPGRADE_DELAY_SECONDS: u64 = 48 * 60 * 60;
+pub(crate) use crate::upgrades::timelock::WASM_UPGRADE_DELAY_SECONDS as UPGRADE_DELAY_SECONDS;
 pub(crate) const STAKE_REGISTRY_KEY: Symbol = symbol_short!("STAKES");
 pub(crate) const TOTAL_STAKED_KEY: Symbol = symbol_short!("TOTAL");
 const HEARTBEAT_KEY: Symbol = symbol_short!("HBEAT");
@@ -443,7 +443,10 @@ impl TimeLockedUpgradeContract {
         if data.admin != proposer { return Err(ContractError::NotAdmin); }
         proposer.require_auth();
         consume_nonce(&env, &proposer, nonce, salt, salt_signature)?;
-        let staged = StagedUpgrade { new_wasm_hash, proposer, staged_at: env.ledger().timestamp() };
+        let staged_at = env.ledger().timestamp();
+        let execute_at = crate::upgrades::timelock::execution_timestamp(staged_at)
+            .ok_or(ContractError::MathOverflow)?;
+        let staged = StagedUpgrade { new_wasm_hash, proposer, staged_at, execute_at };
         env.storage().instance().set(&PENDING_UPGRADE_KEY, &staged);
         Ok(())
     }
@@ -464,7 +467,7 @@ impl TimeLockedUpgradeContract {
             .instance()
             .get(&PENDING_UPGRADE_KEY)
             .ok_or(ContractError::NoPendingUpgrade)?;
-        if !verify_staged_delay(pending.staged_at, env.ledger().timestamp(), UPGRADE_DELAY_SECONDS) {
+        if !crate::upgrades::timelock::is_ready(pending.execute_at, env.ledger().timestamp()) {
             return Err(ContractError::UpgradeTimelockNotSatisfied);
         }
         // Store pre-upgrade contract data snapshot for health check validation
@@ -518,8 +521,7 @@ impl TimeLockedUpgradeContract {
 
     pub fn get_upgrade_timelock_remaining(env: Env) -> Option<u64> {
         env.storage().instance().get(&PENDING_UPGRADE_KEY).map(|staged: StagedUpgrade| {
-            let elapsed = env.ledger().timestamp().saturating_sub(staged.staged_at);
-            UPGRADE_DELAY_SECONDS.saturating_sub(elapsed)
+            staged.execute_at.saturating_sub(env.ledger().timestamp())
         })
     }
 
diff --git a/src/upgrades/mod.rs b/src/upgrades/mod.rs
index a571eaa..adb7209 100644
--- a/src/upgrades/mod.rs
+++ b/src/upgrades/mod.rs
@@ -1 +1,2 @@
 pub mod migration;
+pub mod timelock;
diff --git a/src/upgrades/timelock.rs b/src/upgrades/timelock.rs
new file mode 100644
index 0000000..f0ab189
--- /dev/null
+++ b/src/upgrades/timelock.rs
@@ -0,0 +1,32 @@
+//! Timelock policy for WASM replacements.
+
+/// Minimum delay between registering a WASM hash and replacing the code.
+pub const WASM_UPGRADE_DELAY_SECONDS: u64 = 48 * 60 * 60;
+
+/// Return the earliest ledger timestamp at which an upgrade may execute.
+pub fn execution_timestamp(proposed_at: u64) -> Option<u64> {
+    proposed_at.checked_add(WASM_UPGRADE_DELAY_SECONDS)
+}
+
+/// Check whether the mandatory delay has elapsed.
+pub fn is_ready(execute_at: u64, current_timestamp: u64) -> bool {
+    current_timestamp >= execute_at
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn requires_the_full_delay() {
+        let execute_at = execution_timestamp(1_000).unwrap();
+
+        assert!(!is_ready(execute_at, execute_at - 1));
+        assert!(is_ready(execute_at, execute_at));
+    }
+
+    #[test]
+    fn rejects_deadline_overflow() {
+        assert_eq!(execution_timestamp(u64::MAX), None);
+    }
+}