chore(security): stage follow-up atomic module fixes

Author: jamesatomc

scripts/apply_followup_security_fixes.py

@@ -0,0 +1,198 @@
+#!/usr/bin/env python3
+from pathlib import Path
+import re
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def load(path: str) -> str:
+    return (ROOT / path).read_text()
+
+
+def save(path: str, text: str) -> None:
+    (ROOT / path).write_text(text)
+
+
+def replace_once(text: str, old: str, new: str, label: str) -> str:
+    count = text.count(old)
+    if count != 1:
+        raise RuntimeError(f"{label}: expected one match, found {count}")
+    return text.replace(old, new, 1)
+
+
+# Use the protocol-selected gas configuration in all feature modes.
+path = "crates/kanari-core/src/engine/produce_dag_vertex.rs"
+s = load(path)
+s = replace_once(
+    s,
+    "kanari_types::gas_v2::GasConfig::default().max_gas_per_block",
+    "kanari_types::GasConfig::default().max_gas_per_block",
+    "protocol-selected checkpoint gas config",
+)
+save(path, s)
+
+# Restore gas_v2 as an internal implementation detail.
+path = "crates/kanari-types/src/lib.rs"
+s = load(path)
+s = replace_once(
+    s,
+    '#[cfg(feature = "zero-gas")]\npub mod gas_v2;',
+    '#[cfg(feature = "zero-gas")]\nmod gas_v2;',
+    "restore gas module visibility",
+)
+save(path, s)
+
+# Core needs the bytecode parser to derive canonical module IDs for the atomic batch.
+path = "crates/kanari-core/Cargo.toml"
+s = load(path)
+if "move-binary-format" not in s:
+    s = replace_once(
+        s,
+        "move-core-types = { workspace = true }\n",
+        "move-core-types = { workspace = true }\nmove-binary-format = { workspace = true }\n",
+        "add module parser dependency",
+    )
+save(path, s)
+
+# Put published module blobs and the module index in the same RocksDB WriteBatch
+# as state, checkpoint metadata, transaction indexes, and receipts.
+path = "crates/kanari-core/src/engine.rs"
+s = load(path)
+s = replace_once(
+    s,
+    '''        let mut updates = Vec::new();
+
+        let mut slim_chain = chain.clone();''',
+    '''        let mut updates = Vec::new();
+        let mut module_index = store
+            .load::<Vec<String>>(b"module_index")?
+            .unwrap_or_default();
+        let mut module_index_changed = false;
+
+        let mut slim_chain = chain.clone();''',
+    "load module index for atomic checkpoint",
+)
+s = replace_once(
+    s,
+    '''            for tx in checkpoint.transactions.iter() {
+                let tx_hash = tx.transaction_hash().to_vec();''',
+    '''            for tx in checkpoint.transactions.iter() {
+                if let Transaction::PublishModule { module_bytes, .. } = &tx.transaction {
+                    let compiled =
+                        move_binary_format::CompiledModule::deserialize_with_defaults(module_bytes)
+                            .map_err(|error| {
+                                anyhow::anyhow!(
+                                    "Failed to decode published module during checkpoint commit: {:?}",
+                                    error
+                                )
+                            })?;
+                    let module_id = compiled.self_id();
+                    let module_key = format!(
+                        "module:{}:{}",
+                        module_id.address().to_hex_literal(),
+                        module_id.name().as_str()
+                    );
+                    updates.push((
+                        module_key.as_bytes().to_vec(),
+                        bcs::to_bytes(module_bytes)?,
+                    ));
+                    if !module_index.iter().any(|entry| entry == &module_key) {
+                        module_index.push(module_key);
+                        module_index_changed = true;
+                    }
+                }
+
+                let tx_hash = tx.transaction_hash().to_vec();''',
+    "atomically persist published modules",
+)
+s = replace_once(
+    s,
+    '''            updates.push((
+                Self::recent_transaction_hashes_key().to_vec(),
+                bcs::to_bytes(&recent_hashes)?,
+            ));
+        }
+
+        for receipt in receipts {''',
+    '''            updates.push((
+                Self::recent_transaction_hashes_key().to_vec(),
+                bcs::to_bytes(&recent_hashes)?,
+            ));
+            if module_index_changed {
+                module_index.sort();
+                module_index.dedup();
+                updates.push((b"module_index".to_vec(), bcs::to_bytes(&module_index)?));
+            }
+        }
+
+        for receipt in receipts {''',
+    "commit module index atomically",
+)
+save(path, s)
+
+# Do not re-execute and persist modules before the atomic checkpoint batch.
+path = "crates/kanari-core/src/engine/apply_checkpoint.rs"
+s = load(path)
+s, count = re.subn(
+    r'''impl BlockchainEngine \{\n    fn requires_runtime_side_effect_persistence\(transactions: &\[SignedTransaction\]\) -> bool \{.*?\n    \}\n\n''',
+    "impl BlockchainEngine {\n",
+    s,
+    count=1,
+    flags=re.S,
+)
+if count != 1:
+    raise RuntimeError(f"remove side-effect predicate: expected one match, found {count}")
+s, count = re.subn(
+    r'''        if Self::requires_runtime_side_effect_persistence\(&to_execute\) \{.*?\n        \}\n\n        self\.finalize_checkpoint''',
+    '''        // Canonical state and published module bytes are committed together by
+        // finalize_checkpoint. Re-executing here would create pre-commit side effects.
+        let _ = to_execute;
+
+        self.finalize_checkpoint''',
+    s,
+    count=1,
+    flags=re.S,
+)
+if count != 1:
+    raise RuntimeError(f"remove side-effect replay: expected one match, found {count}")
+save(path, s)
+
+# A cache reload after an atomic module commit must rebuild the in-memory module set
+# from the canonical store rather than retaining a stale pre-commit set.
+path = "move-execution/v1/kanari-move-runtime-v1/src/move_runtime/mod.rs"
+s = load(path)
+s = replace_once(
+    s,
+    '''    pub fn reload_vm_cache(&self) -> Result<()> {
+        let new_vm = MoveVM::new(self.all_natives.as_ref().clone())
+            .map_err(|e| anyhow::anyhow!("Failed to reload MoveVM: {:?}", e))?;
+
+        // Replace the VM instance to clear internal caches.
+        *self.write_vm() = new_vm;
+
+        // Preload published modules so follow-up executions can resolve dependencies immediately.
+        self.preload_system_modules_into_vm()?;''',
+    '''    pub fn reload_vm_cache(&self) -> Result<()> {
+        let new_vm = MoveVM::new(self.all_natives.as_ref().clone())
+            .map_err(|e| anyhow::anyhow!("Failed to reload MoveVM: {:?}", e))?;
+
+        let refreshed_modules = self
+            .state
+            .get_all_module_ids()?
+            .into_iter()
+            .collect::<HashSet<_>>();
+        *self
+            .published_modules
+            .write()
+            .unwrap_or_else(|poisoned| poisoned.into_inner()) = refreshed_modules;
+
+        // Replace the VM instance to clear internal caches.
+        *self.write_vm() = new_vm;
+
+        // Preload published modules so follow-up executions can resolve dependencies immediately.
+        self.preload_system_modules_into_vm()?;''',
+    "refresh module index during cache reload",
+)
+save(path, s)
+
+print("Applied follow-up atomic module and feature-mode fixes.")