From f9a508bb6eff80014e3fb0339db40d53bcbd03ab Mon Sep 17 00:00:00 2001 From: Liam Phung Date: Fri, 16 Oct 2026 15:15:08 +0700 Subject: [PATCH] fix(revenue_pool): feed real layout hashes and version progression into upgrade validation (#1164) --- contracts/revenue_pool/src/lib.rs | 100 +++++++++++--- .../src/test_storage_migration.rs | 125 ++++++++++++++++++ 2 files changed, 207 insertions(+), 18 deletions(-) diff --git a/contracts/revenue_pool/src/lib.rs b/contracts/revenue_pool/src/lib.rs index a59e9eb6..381156b1 100644 --- a/contracts/revenue_pool/src/lib.rs +++ b/contracts/revenue_pool/src/lib.rs @@ -11,10 +11,15 @@ use soroban_sdk::{ contract, contractimpl, contracttype, token, Address, BytesN, Env, Map, String, Symbol, Vec, }; -/// Storage-layout version recorded by [`StorageMigrationValidator`] for this -/// contract. Bumped whenever the on-ledger schema changes so that the -/// pre-upgrade validation gate can enforce ordered, single-step migrations. -const STORAGE_MIGRATION_VERSION: u32 = 1; +/// Initial storage-layout version recorded by [`StorageMigrationValidator`] +/// when the revenue pool is initialized. +pub const STORAGE_MIGRATION_VERSION: u32 = 1; + +/// Canonical storage-layout descriptor hashed by +/// [`callora_storage_migration::layout_hash`] to detect on-ledger schema drift +/// across upgrades. +pub const STORAGE_LAYOUT_DESCRIPTOR: &str = + "revenue_pool:v1:admin,usdc,paused,emergency_paused,pending_admin,pause_guardian,max_distribute,cumulative_yield,version,emergency_drain"; // --------------------------------------------------------------------------- // Storage keys @@ -148,6 +153,15 @@ impl RevenuePool { inst.set(&Symbol::new(&env, USDC_KEY), &usdc_token); inst.set(&Symbol::new(&env, PAUSED_KEY), &false); inst.set(&Symbol::new(&env, EMERGENCY_PAUSED_KEY), &false); + let initial_layout = + callora_storage_migration::layout_hash(&env, STORAGE_LAYOUT_DESCRIPTOR); + if let Err(e) = StorageMigrationValidator::record_deployed_schema( + &env, + STORAGE_MIGRATION_VERSION, + &initial_layout, + ) { + env.panic_with_error(e); + } inst.extend_ttl(LIFETIME_THRESHOLD, BUMP_AMOUNT); env.events() .publish((events::event_init(&env), admin), usdc_token); @@ -859,7 +873,25 @@ impl RevenuePool { // Upgrade // ----------------------------------------------------------------------- - /// Admin-gated contract upgrade. Replaces the WASM and persists the version. + /// Compute the canonical storage-layout hash for this contract binary. + pub fn current_layout_hash(env: Env) -> BytesN<32> { + callora_storage_migration::layout_hash(&env, STORAGE_LAYOUT_DESCRIPTOR) + } + + /// Return the currently recorded storage-migration schema version (`0` if unrecorded). + pub fn get_storage_migration_version(env: Env) -> u32 { + Self::bump_instance_ttl(&env); + StorageMigrationValidator::current_version(&env) + } + + /// Return the currently recorded storage-layout hash, or `None` if unrecorded. + pub fn get_storage_layout_hash(env: Env) -> Option> { + Self::bump_instance_ttl(&env); + StorageMigrationValidator::current_layout_hash(&env) + } + + /// Admin-gated contract upgrade using the current contract storage layout + /// and advancing the storage-migration version by one step. /// /// # Errors /// * [`RevenuePoolError::Unauthorized`] - caller is not the current admin. @@ -867,31 +899,63 @@ impl RevenuePool { /// # Events /// Emits `upgraded` with `admin` as topic and `new_wasm_hash` as data. pub fn upgrade(env: Env, caller: Address, new_wasm_hash: BytesN<32>) { + let current_ver = StorageMigrationValidator::current_version(&env); + let target_version = current_ver + .checked_add(1) + .unwrap_or_else(|| env.panic_with_error(RevenuePoolError::Overflow)); + let layout = StorageMigrationValidator::current_layout_hash(&env).unwrap_or_else(|| { + callora_storage_migration::layout_hash(&env, STORAGE_LAYOUT_DESCRIPTOR) + }); + Self::upgrade_with_layout( + env, + caller, + new_wasm_hash, + target_version, + layout.clone(), + layout, + ); + } + + /// Admin-gated contract upgrade with explicit target version and storage-layout + /// hashes validated by [`StorageMigrationValidator`]. + /// + /// Rejects mismatched `expected_source_layout` with `SchemaMismatch`, + /// enforces single-step ordered version progression (`v -> v + 1`), and + /// records the new schema state upon upgrade. + /// + /// # Errors + /// * [`RevenuePoolError::Unauthorized`] - caller is not the current admin. + /// * [`callora_storage_migration::StorageMigrationError::SchemaMismatch`] - `expected_source_layout` does not match the deployed layout hash. + /// + /// # Events + /// Emits `upgraded` with `admin` as topic and `new_wasm_hash` as data. + pub fn upgrade_with_layout( + env: Env, + caller: Address, + new_wasm_hash: BytesN<32>, + target_version: u32, + expected_source_layout: BytesN<32>, + new_layout_hash: BytesN<32>, + ) { caller.require_auth(); Self::require_not_emergency_paused(&env); Self::require_admin(&env, &caller); - // ── Pre-upgrade storage-migration validation ────────────────────── - // Runs in the *current* (old) code, before the WASM is swapped, and - // never mutates business state. It enforces ordered, single-step - // upgrades, rejects all-zero WASM hashes, and prevents unsanctioned - // rollbacks — ensuring existing deployed data stays readable and no - // implicit destructive transformation occurs. - let placeholder_layout = callora_storage_migration::zero_layout_hash(&env); + let migration_performed = expected_source_layout != new_layout_hash; if let Err(e) = StorageMigrationValidator::validate_before_upgrade( &env, - STORAGE_MIGRATION_VERSION, - &placeholder_layout, - &placeholder_layout, + target_version, + &expected_source_layout, + &new_layout_hash, &new_wasm_hash, - false, + migration_performed, ) { env.panic_with_error(e); } if let Err(e) = StorageMigrationValidator::finalize_migration( &env, - STORAGE_MIGRATION_VERSION, - &placeholder_layout, + target_version, + &new_layout_hash, &new_wasm_hash, ) { env.panic_with_error(e); diff --git a/contracts/revenue_pool/src/test_storage_migration.rs b/contracts/revenue_pool/src/test_storage_migration.rs index d076fb5e..f93b5549 100644 --- a/contracts/revenue_pool/src/test_storage_migration.rs +++ b/contracts/revenue_pool/src/test_storage_migration.rs @@ -1,6 +1,7 @@ extern crate std; use super::*; +use callora_storage_migration::StorageMigrationError; use soroban_sdk::testutils::Address as _; use soroban_sdk::token; use soroban_sdk::{Address, BytesN, Env}; @@ -22,6 +23,27 @@ fn create_pool(env: &Env) -> (Address, RevenuePoolClient<'_>) { (address, client) } +/// Initialization records version 1 (v1) and the real non-zero storage layout hash. +#[test] +fn init_records_real_layout_hash_and_v1() { + let env = Env::default(); + env.mock_all_auths(); + let admin = Address::generate(&env); + let (_, client) = create_pool(&env); + let (usdc_address, _, _) = create_usdc(&env, &admin); + + client.init(&admin, &usdc_address); + + let expected_layout = callora_storage_migration::layout_hash(&env, STORAGE_LAYOUT_DESCRIPTOR); + let zero_layout = callora_storage_migration::zero_layout_hash(&env); + assert_ne!(expected_layout, zero_layout); + assert_eq!( + client.get_storage_migration_version(), + STORAGE_MIGRATION_VERSION + ); + assert_eq!(client.get_storage_layout_hash(), Some(expected_layout)); +} + /// A fresh, well-formed upgrade (non-zero hash, single forward step) must pass /// the pre-upgrade migration guard and preserve all previously deployed state. #[test] @@ -43,6 +65,109 @@ fn upgrade_valid_wasm_hash_passes_migration_guard() { let readback: Option> = client.get_version(); assert_eq!(readback, Some(new_hash)); + assert_eq!(client.get_storage_migration_version(), 2); +} + +/// Two sequential upgrades (v1 -> v2 -> v3) succeed both via default `upgrade` +/// and via explicit `upgrade_with_layout` across evolving layout descriptors. +#[test] +fn two_sequential_upgrades_v1_to_v2_to_v3_succeed() { + let env = Env::default(); + env.mock_all_auths(); + let admin = Address::generate(&env); + let (_, client) = create_pool(&env); + let (usdc_address, _, _) = create_usdc(&env, &admin); + + client.init(&admin, &usdc_address); + assert_eq!(client.get_storage_migration_version(), 1); + + let layout_v1 = client.current_layout_hash(); + let layout_v2 = callora_storage_migration::layout_hash( + &env, + "revenue_pool:v2:admin,usdc,paused,emergency_paused,pending_admin,pause_guardian,max_distribute,cumulative_yield,version,emergency_drain", + ); + let layout_v3 = callora_storage_migration::layout_hash( + &env, + "revenue_pool:v3:admin,usdc,paused,emergency_paused,pending_admin,pause_guardian,max_distribute,cumulative_yield,version,emergency_drain,extra_field", + ); + + let wasm_v2 = env + .deployer() + .upload_contract_wasm(soroban_sdk::Bytes::from_slice( + &env, + &[0, 97, 115, 109, 1, 0, 0, 0], + )); + client.upgrade_with_layout(&admin, &wasm_v2, &2, &layout_v1, &layout_v2); + assert_eq!(client.get_storage_migration_version(), 2); + assert_eq!(client.get_storage_layout_hash(), Some(layout_v2.clone())); + assert_eq!(client.get_version(), Some(wasm_v2)); + + let wasm_v3 = env + .deployer() + .upload_contract_wasm(soroban_sdk::Bytes::new(&env)); + client.upgrade_with_layout(&admin, &wasm_v3, &3, &layout_v2, &layout_v3); + assert_eq!(client.get_storage_migration_version(), 3); + assert_eq!(client.get_storage_layout_hash(), Some(layout_v3)); + assert_eq!(client.get_version(), Some(wasm_v3)); +} + +/// Calling `upgrade` twice sequentially also advances v1 -> v2 -> v3 cleanly. +#[test] +fn sequential_default_upgrades_advance_v1_to_v2_to_v3() { + let env = Env::default(); + env.mock_all_auths(); + let admin = Address::generate(&env); + let (_, client) = create_pool(&env); + let (usdc_address, _, _) = create_usdc(&env, &admin); + + client.init(&admin, &usdc_address); + assert_eq!(client.get_storage_migration_version(), 1); + + let wasm_v2 = env + .deployer() + .upload_contract_wasm(soroban_sdk::Bytes::new(&env)); + client.upgrade(&admin, &wasm_v2); + assert_eq!(client.get_storage_migration_version(), 2); + + let wasm_v3 = env + .deployer() + .upload_contract_wasm(soroban_sdk::Bytes::from_slice( + &env, + &[0, 97, 115, 109, 1, 0, 0, 0], + )); + client.upgrade(&admin, &wasm_v3); + assert_eq!(client.get_storage_migration_version(), 3); + assert_eq!(client.get_version(), Some(wasm_v3)); +} + +/// An upgrade with a mismatched `expected_source_layout` fails with `SchemaMismatch`. +#[test] +fn upgrade_with_mismatched_source_layout_fails_with_schema_mismatch() { + let env = Env::default(); + env.mock_all_auths(); + let admin = Address::generate(&env); + let (_, client) = create_pool(&env); + let (usdc_address, _, _) = create_usdc(&env, &admin); + + client.init(&admin, &usdc_address); + + let wrong_source_layout = + callora_storage_migration::layout_hash(&env, "revenue_pool:wrong_source_schema"); + let target_layout = client.current_layout_hash(); + let new_hash = env + .deployer() + .upload_contract_wasm(soroban_sdk::Bytes::new(&env)); + + let res = + client.try_upgrade_with_layout(&admin, &new_hash, &2, &wrong_source_layout, &target_layout); + assert_eq!( + res, + Err(Ok(soroban_sdk::Error::from_contract_error( + StorageMigrationError::SchemaMismatch as u32 + ))) + ); + assert_eq!(client.get_storage_migration_version(), 1); + assert_eq!(client.get_version(), None); } /// The pre-upgrade storage-migration gate must reject an all-zero WASM hash