[Governance] Formal Upgrade Framework with Version Tracking, Migration Hooks, and Rollback (Issue #16)

contracts/utility_contracts/src/lib.rs

@@ -333,6 +333,7 @@ pub mod nonce_sync;
 pub mod secure_call_interface;
 pub mod tariff_oracle;
 pub mod temporary_storage;
+pub mod upgrade_framework;
 pub mod velocity_limit;
 
 #[cfg(test)]
@@ -1037,6 +1038,15 @@ pub enum DataKey {
     UpgradeApproval(u64, Address),
     UpgradeProposalCounter,
     ActiveUpgradeProposalId,
+    // Issue #16 - Upgrade Framework
+    ContractVersion,
+    LastUpgradeLedger,
+    VersionInfo(u32),
+    VersionCount,
+    StorageSchemaVersion(BytesN<32>),
+    RollbackPoint(u32),
+    RollbackCount,
+    PreviousWasmHash(u32),
 }
 
 #[contracterror(export = false)]
@@ -1174,6 +1184,11 @@ pub enum ContractError {
     // Issue #23 - Token Security
     UnapprovedToken = 117,
     TokenBalanceMismatch = 118,
+    // Issue #16 - Upgrade Framework
+    UpgradeDelayNotElapsed = 119,
+    RollbackPointNotFound = 120,
+    RollbackAlreadyConsumed = 121,
+    MigrationFailed = 122,
 }
 
 #[contracttype]
@@ -1419,9 +1434,7 @@ fn require_approved_token(env: &Env, token: &Address) {
     // Skip whitelist enforcement in test mode
     #[cfg(not(test))]
     {
-        let approved: Option<Vec<Address>> = env.storage()
-            .instance()
-            .get(&DataKey::ApprovedTokens);
+        let approved: Option<Vec<Address>> = env.storage().instance().get(&DataKey::ApprovedTokens);
         if let Some(tokens) = approved {
             if tokens.len() > 0 && !tokens.contains(token) {
                 panic_with_error!(env, ContractError::UnapprovedToken);
@@ -3056,13 +3069,16 @@ impl UtilityContract {
     /// Only callable by the contract admin.
     pub fn approve_token(env: Env, token: Address, decimals: u32) {
         require_admin_auth(&env);
-        let mut approved: Vec<Address> = env.storage()
+        let mut approved: Vec<Address> = env
+            .storage()
             .instance()
             .get(&DataKey::ApprovedTokens)
             .unwrap_or(Vec::new(&env));
         if !approved.contains(&token) {
             approved.push_back(token.clone());
-            env.storage().instance().set(&DataKey::ApprovedTokens, &approved);
+            env.storage()
+                .instance()
+                .set(&DataKey::ApprovedTokens, &approved);
         }
         let info = TokenInfo {
             token: token.clone(),
@@ -3071,20 +3087,25 @@ impl UtilityContract {
             approved_at: env.ledger().timestamp(),
             approved_by: get_admin_or_panic(&env),
         };
-        env.storage().instance().set(&DataKey::TokenInfo(token), &info);
+        env.storage()
+            .instance()
+            .set(&DataKey::TokenInfo(token), &info);
     }
 
     /// Revoke a token from the protocol whitelist.
     /// Only callable by the contract admin.
     pub fn revoke_token(env: Env, token: Address) {
         require_admin_auth(&env);
-        let mut approved: Vec<Address> = env.storage()
+        let mut approved: Vec<Address> = env
+            .storage()
             .instance()
             .get(&DataKey::ApprovedTokens)
             .unwrap_or(Vec::new(&env));
         if let Some(pos) = approved.first_index_of(&token) {
             approved.remove(pos);
-            env.storage().instance().set(&DataKey::ApprovedTokens, &approved);
+            env.storage()
+                .instance()
+                .set(&DataKey::ApprovedTokens, &approved);
             env.storage().instance().remove(&DataKey::TokenInfo(token));
         }
     }
@@ -3099,9 +3120,7 @@ impl UtilityContract {
 
     /// Get token info for a specific token.
     pub fn get_token_info(env: Env, token: Address) -> Option<TokenInfo> {
-        env.storage()
-            .instance()
-            .get(&DataKey::TokenInfo(token))
+        env.storage().instance().get(&DataKey::TokenInfo(token))
     }
 
     pub fn set_admin(env: Env, admin_address: Address) {
@@ -4854,14 +4873,12 @@ impl UtilityContract {
         let mut cost = signed_data.units_consumed.saturating_mul(discounted_rate);
 
         // Apply SLA Penalty if active
-        if let Some(config) = &meter.sla_config {
-            if meter.sla_state.is_penalty_active
-                || meter.sla_state.accumulated_downtime >= config.threshold_seconds
-            {
-                cost = cost
-                    .saturating_mul(config.penalty_multiplier_bps)
-                    .saturating_div(10000);
-            }
+        if meter.sla_state.is_penalty_active
+            || meter.sla_state.accumulated_downtime >= meter.sla_config.threshold_seconds
+        {
+            cost = cost
+                .saturating_mul(meter.sla_config.penalty_multiplier_bps)
+                .saturating_div(10000);
         }
 
         // Apply provider withdrawal limits

contracts/utility_contracts/src/nonce_sync_tests.rs

@@ -2,9 +2,10 @@ use crate::nonce_sync::{
     DeviceNonceState, NonceAlertType, NonceDesyncAlert, NonceResetRequest, NonceSyncManager,
     SignedHeartbeat, NONCE_WINDOW_SIZE,
 };
-use crate::std::string::ToString;
 use crate::{ContractError, DataKey};
-use soroban_sdk::{testutils::Address as _, testutils::BytesN as _, Address, BytesN, Env};
+use soroban_sdk::{
+    testutils::Address as _, testutils::BytesN as _, Address, BytesN, Env, String, Vec,
+};
 
 #[cfg(test)]
 pub mod nonce_sync_fuzz_tests {
@@ -329,6 +330,7 @@ pub mod nonce_sync_fuzz_tests {
 /// and prevents any form of nonce reuse or manipulation.
 #[cfg(test)]
 mod property_tests {
+    use super::nonce_sync_fuzz_tests::create_test_heartbeat;
     use super::*;
     use proptest::prelude::*;
 

contracts/utility_contracts/src/tariff_oracle_tests.rs

@@ -460,6 +460,7 @@ pub mod tariff_oracle_tests {
 /// Property-based tests for tariff calculations
 #[cfg(test)]
 mod tariff_property_tests {
+    use super::tariff_oracle_tests::create_test_schedule;
     use super::*;
     use proptest::prelude::*;
 

contracts/utility_contracts/src/temporary_storage_tests.rs

@@ -9,7 +9,9 @@ mod tests {
         temporary_storage::{OptimizedFlowCalculator, OptimizedUsageTracker, TempStorageManager},
         BillingType, ContinuousFlow, DataKey, Meter, StreamStatus, UsageData,
     };
-    use soroban_sdk::{testutils::Address as _, testutils::Ledger, Address, BytesN, Env, Symbol};
+    use soroban_sdk::{
+        testutils::Address as _, testutils::Ledger, Address, BytesN, Env, Symbol, Vec,
+    };
 
     fn create_test_env() -> Env {
         let env = Env::default();