From 66d0343a2ce629a6a3eaafe7da6dbadeeaa10d48 Mon Sep 17 00:00:00 2001 From: senmalong Date: Wed, 24 Jun 2026 06:50:50 +0100 Subject: [PATCH] fix(proxy): resolve upgrade storage collision for constants (closes #21) - Defined deterministic versioned slot for constants in storage-layout.rs. - Stated constants in UpgradeBeacon before upgrade execution in upgrade-beacon.rs. - Retained constant reading from versioned slot in beacon-state.rs. - Implemented StorageLayoutCheck script run in CI configuration to avoid future collisions. - Added upgrade storage preservation integration test. --- .github/workflows/rust.yml | 9 ++ scripts/storage-layout-check.py | 142 ++++++++++++++++++++++++++++ src/lib.rs | 11 +++ src/proxy/storage-layout.rs | 20 ++++ src/proxy/upgrade-beacon.rs | 23 +++++ src/state/beacon-state.rs | 33 +++++++ tests/proxy/upgrade_storage_test.rs | 35 +++++++ 7 files changed, 273 insertions(+) create mode 100755 scripts/storage-layout-check.py create mode 100644 src/proxy/storage-layout.rs create mode 100644 src/proxy/upgrade-beacon.rs create mode 100644 src/state/beacon-state.rs create mode 100644 tests/proxy/upgrade_storage_test.rs diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 6fc892b..6b4196f 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -25,3 +25,12 @@ jobs: - name: Run tests run: cargo test --verbose + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.x' + + - name: Run Storage Layout Collision Check + run: python scripts/storage-layout-check.py + diff --git a/scripts/storage-layout-check.py b/scripts/storage-layout-check.py new file mode 100755 index 0000000..0383dd9 --- /dev/null +++ b/scripts/storage-layout-check.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +import sys +import re +import subprocess + +def get_file_content_from_git(branch, filepath): + try: + # Get file content from specific git branch/commit + result = subprocess.run( + ['git', 'show', f'{branch}:{filepath}'], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=True + ) + return result.stdout + except Exception: + return None + +def parse_slots_and_structs(content): + if not content: + return {}, {} + + # Parse slot definitions: pub const NAME_SLOT: [u8; 32] = [ ... ]; + # Match multiline arrays as well + slot_pattern = re.compile( + r'pub\s+const\s+(\w+_SLOT)\s*:\s*\[u8;\s*32\]\s*=\s*\[([^\]]+)\]\s*;', + re.MULTILINE + ) + + slots = {} + for match in slot_pattern.finditer(content): + name = match.group(1) + bytes_str = match.group(2) + # Parse bytes, e.g., 0x60, 0x01, or 1; 32 + bytes_str = bytes_str.strip() + if ';' in bytes_str: + val, count = bytes_str.split(';') + val = int(val.strip(), 0) + count = int(count.strip()) + byte_list = [val] * count + else: + byte_list = [int(x.strip(), 0) for x in bytes_str.split(',') if x.strip()] + + slots[name] = bytes(byte_list) + + # Parse structs: pub struct StructName { ... } + struct_pattern = re.compile( + r'pub\s+struct\s+(\w+)\s*\{([^\}]+)\}', + re.MULTILINE + ) + structs = {} + for match in struct_pattern.finditer(content): + name = match.group(1) + fields_str = match.group(2) + # Clean fields + fields = [] + for line in fields_str.split('\n'): + line = line.strip() + if line and not line.startswith('//'): + # Extract field name and type, e.g., pub max_validators: u32 + field_match = re.match(r'(?:pub\s+)?(\w+)\s*:\s*([\w<>:\[\];\s]+)', line) + if field_match: + f_name = field_match.group(1) + f_type = field_match.group(2).replace(' ', '') + fields.append((f_name, f_type)) + structs[name] = fields + + return slots, structs + +def main(): + filepath = 'src/proxy/storage-layout.rs' + + # Read current file + try: + with open(filepath, 'r') as f: + current_content = f.read() + except FileNotFoundError: + print(f"Error: {filepath} not found.") + sys.exit(1) + + current_slots, current_structs = parse_slots_and_structs(current_content) + + # Read base branch (origin/main) + base_content = get_file_content_from_git('origin/main', filepath) + if base_content is None: + # Fallback to main + base_content = get_file_content_from_git('main', filepath) + + if not base_content: + print("No previous implementation layout found on main branch. Skipping diff checks, performing self-checks.") + base_slots, base_structs = {}, {} + else: + base_slots, base_structs = parse_slots_and_structs(base_content) + + collision_detected = False + + # Self-check: Look for duplicate values in current slots + slot_values = {} + for name, value in current_slots.items(): + if value in slot_values: + print(f"COLLISION WARNING: Slot '{name}' collides with '{slot_values[value]}' at value {value.hex()}") + collision_detected = True + else: + slot_values[value] = name + + # Diff check: Compare slots with base/old implementation + for name, value in base_slots.items(): + if name in current_slots: + if current_slots[name] != value: + print(f"COLLISION WARNING: Slot '{name}' has changed value from {value.hex()} to {current_slots[name].hex()}") + collision_detected = True + else: + print(f"WARNING: Slot '{name}' was removed from the storage layout.") + + # Diff check: Compare structs to detect shifts in fields/ordering + for name, fields in base_structs.items(): + if name in current_structs: + curr_fields = current_structs[name] + # Check if fields were reordered or modified in a way that shifts layouts + min_len = min(len(fields), len(curr_fields)) + for i in range(min_len): + if fields[i] != curr_fields[i]: + print(f"COLLISION WARNING: Struct '{name}' field layout shifted/changed at index {i}:") + print(f" Old: {fields[i][0]}: {fields[i][1]}") + print(f" New: {curr_fields[i][0]}: {curr_fields[i][1]}") + collision_detected = True + if len(curr_fields) < len(fields): + print(f"COLLISION WARNING: Struct '{name}' fields were removed. This will corrupt storage decoding.") + collision_detected = True + else: + print(f"WARNING: Struct '{name}' was removed from the storage layout.") + + if collision_detected: + print("Storage layout verification failed! Potential collision or layout corruption detected.") + sys.exit(1) + else: + print("Storage layout verification passed successfully. No collisions detected.") + sys.exit(0) + +if __name__ == '__main__': + main() diff --git a/src/lib.rs b/src/lib.rs index 5133590..603b690 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,9 +1,20 @@ #![no_std] + +#[path = "proxy/storage-layout.rs"] +pub mod storage_layout; + +#[path = "proxy/upgrade-beacon.rs"] +pub mod upgrade_beacon; + +#[path = "state/beacon-state.rs"] +pub mod beacon_state; + use soroban_sdk::{ contract, contractclient, contracterror, contractimpl, contracttype, token, Address, Env, String, Vec, }; + // --- ERROR CODES --- #[contracterror] diff --git a/src/proxy/storage-layout.rs b/src/proxy/storage-layout.rs new file mode 100644 index 0000000..66d3f16 --- /dev/null +++ b/src/proxy/storage-layout.rs @@ -0,0 +1,20 @@ +use soroban_sdk::{contracttype, BytesN, Env}; + +// keccak256("beacon.constants.v1") +pub const BEACON_CONSTANTS_V1_SLOT: [u8; 32] = [ + 0x60, 0x01, 0x2b, 0x18, 0x97, 0x9e, 0x27, 0xc1, 0x56, 0xf7, 0x0a, 0x75, 0xf8, 0x50, 0xe0, 0x47, + 0x46, 0x68, 0xb5, 0x98, 0x28, 0x55, 0x14, 0x0f, 0x6b, 0x4d, 0x08, 0x1f, 0x21, 0x13, 0xf8, 0xd3, +]; + +// keccak256("beacon.constants") +pub const BEACON_CONSTANTS_SLOT: [u8; 32] = [ + 0xc3, 0xd1, 0x79, 0x61, 0x23, 0x9c, 0x4d, 0x9b, 0x4b, 0x04, 0x38, 0xcf, 0x38, 0x44, 0x9c, 0x25, + 0x60, 0xb4, 0x57, 0xe5, 0xb6, 0x19, 0x73, 0x6c, 0x53, 0x54, 0x50, 0x41, 0xa7, 0x9f, 0x04, 0x62, +]; + +#[contracttype] +#[derive(Clone, Debug, PartialEq)] +pub struct ConstantsStore { + pub max_validators: u32, + pub shard_count: u32, +} diff --git a/src/proxy/upgrade-beacon.rs b/src/proxy/upgrade-beacon.rs new file mode 100644 index 0000000..40431d1 --- /dev/null +++ b/src/proxy/upgrade-beacon.rs @@ -0,0 +1,23 @@ +use soroban_sdk::{contract, contractimpl, BytesN, Env}; +use crate::storage_layout::{ConstantsStore, BEACON_CONSTANTS_V1_SLOT}; + +#[contract] +pub struct UpgradeBeacon; + +#[contractimpl] +impl UpgradeBeacon { + pub fn upgrade_implementation(env: Env, new_wasm_hash: BytesN<32>, max_validators: u32, shard_count: u32) { + // Define all constants in dedicated ConstantsStore struct + let store = ConstantsStore { + max_validators, + shard_count, + }; + + // Write the ConstantsStore to the deterministic slot + let key = BytesN::from_array(&env, &BEACON_CONSTANTS_V1_SLOT); + env.storage().instance().set(&key, &store); + + // Upgrade current contract WASM hash + env.deployer().update_current_contract_wasm(new_wasm_hash); + } +} diff --git a/src/state/beacon-state.rs b/src/state/beacon-state.rs new file mode 100644 index 0000000..48ad160 --- /dev/null +++ b/src/state/beacon-state.rs @@ -0,0 +1,33 @@ +use soroban_sdk::{contract, contractimpl, BytesN, Env}; +use crate::storage_layout::{ConstantsStore, BEACON_CONSTANTS_V1_SLOT}; + +#[contract] +pub struct BeaconState; + +#[contractimpl] +impl BeaconState { + pub fn init_constants(env: Env) -> ConstantsStore { + let key = BytesN::from_array(&env, &BEACON_CONSTANTS_V1_SLOT); + if env.storage().instance().has(&key) { + env.storage().instance().get(&key).unwrap() + } else { + // Default or fallback constants + let default_store = ConstantsStore { + max_validators: 1000, + shard_count: 64, + }; + env.storage().instance().set(&key, &default_store); + default_store + } + } + + pub fn get_max_validators(env: Env) -> u32 { + let store = Self::init_constants(env); + store.max_validators + } + + pub fn get_shard_count(env: Env) -> u32 { + let store = Self::init_constants(env); + store.shard_count + } +} diff --git a/tests/proxy/upgrade_storage_test.rs b/tests/proxy/upgrade_storage_test.rs new file mode 100644 index 0000000..8631552 --- /dev/null +++ b/tests/proxy/upgrade_storage_test.rs @@ -0,0 +1,35 @@ +#![cfg(test)] +use soroban_sdk::{BytesN, Env}; +use sorosusu_contracts::beacon_state::{BeaconState, BeaconStateClient}; +use sorosusu_contracts::storage_layout::{ConstantsStore, BEACON_CONSTANTS_V1_SLOT}; +use sorosusu_contracts::upgrade_beacon::{UpgradeBeacon, UpgradeBeaconClient}; + +#[test] +fn test_upgrade_storage_preservation() { + let env = Env::default(); + env.mock_all_auths(); + + // Register UpgradeBeacon contract + let contract_id = env.register_contract(None, UpgradeBeacon); + let client = UpgradeBeaconClient::new(&env, &contract_id); + + // Set constants via the upgrade_implementation call + let mock_wasm_hash = BytesN::from_array(&env, &[0; 32]); + client.upgrade_implementation(&mock_wasm_hash, &2000, &128); + + // Verify they are written in the deterministic slot + let key = BytesN::from_array(&env, &BEACON_CONSTANTS_V1_SLOT); + assert!(env.storage().instance().has(&key)); + let store: ConstantsStore = env.storage().instance().get(&key).unwrap(); + assert_eq!(store.max_validators, 2000); + assert_eq!(store.shard_count, 128); + + // Now instantiate BeaconState client at the same contract ID to simulate the upgraded state + let beacon_client = BeaconStateClient::new(&env, &contract_id); + let upgraded_store = beacon_client.init_constants(); + assert_eq!(upgraded_store.max_validators, 2000); + assert_eq!(upgraded_store.shard_count, 128); + + assert_eq!(beacon_client.get_max_validators(), 2000); + assert_eq!(beacon_client.get_shard_count(), 128); +}