fix(proxy): resolve upgrade storage collision for constants

.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
+

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()

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,
+}

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);
+    }
+}

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
+    }
+}

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);
+}