Address code review feedback: add validation, edge cases, and docs

Co-authored-by: Steake <530040+Steake@users.noreply.github.com>

Author: Copilot

crates/bitcell-state/src/storage.rs

@@ -337,6 +337,12 @@ impl StorageManager {
         // Store snapshot data with metadata: height(8) | root_len(4) | state_root | accounts_data
         let mut snapshot_data = Vec::new();
         snapshot_data.extend_from_slice(&height.to_be_bytes());
+        
+        // Validate state_root length to prevent integer overflow
+        if state_root.len() > u32::MAX as usize {
+            return Err("State root too large (exceeds u32::MAX)".to_string());
+        }
+        
         snapshot_data.extend_from_slice(&(state_root.len() as u32).to_be_bytes());
         snapshot_data.extend_from_slice(state_root);
         snapshot_data.extend_from_slice(accounts_data);
@@ -390,6 +396,14 @@ impl StorageManager {
                 .map_err(|_| "Invalid snapshot height in data".to_string())?
         );
 
+        // Validate stored height matches expected height from index
+        if stored_height != height {
+            return Err(format!(
+                "Snapshot height mismatch: index says {}, data says {}",
+                height, stored_height
+            ));
+        }
+        
         let root_len = u32::from_be_bytes(
             snapshot_data[8..12].try_into()
                 .map_err(|_| "Invalid root length in data".to_string())?
@@ -434,6 +448,14 @@ impl StorageManager {
                 .map_err(|_| "Invalid snapshot height in data".to_string())?
         );
 
+        // Validate stored height matches requested height
+        if stored_height != height {
+            return Err(format!(
+                "Snapshot height mismatch: expected {}, got {}",
+                height, stored_height
+            ));
+        }
+        
         let root_len = u32::from_be_bytes(
             snapshot_data[8..12].try_into()
                 .map_err(|_| "Invalid root length in data".to_string())?
@@ -811,6 +833,33 @@ mod tests {
         assert_eq!(snap2.2.as_slice(), b"data2");
     }
 
+    #[test]
+    fn test_snapshot_edge_cases() {
+        let temp_dir = TempDir::new().unwrap();
+        let storage = StorageManager::new(temp_dir.path()).unwrap();
+        
+        // Test empty state_root
+        storage.create_snapshot(100, &[], b"data").unwrap();
+        let snap = storage.get_snapshot(100).unwrap().unwrap();
+        assert_eq!(snap.0, 100);
+        assert_eq!(snap.1.len(), 0);
+        assert_eq!(snap.2.as_slice(), b"data");
+        
+        // Test empty accounts_data  
+        storage.create_snapshot(101, b"root", &[]).unwrap();
+        let snap = storage.get_snapshot(101).unwrap().unwrap();
+        assert_eq!(snap.0, 101);
+        assert_eq!(snap.1.as_slice(), b"root");
+        assert_eq!(snap.2.len(), 0);
+        
+        // Test both empty
+        storage.create_snapshot(102, &[], &[]).unwrap();
+        let snap = storage.get_snapshot(102).unwrap().unwrap();
+        assert_eq!(snap.0, 102);
+        assert_eq!(snap.1.len(), 0);
+        assert_eq!(snap.2.len(), 0);
+    }
+
     #[test]
     fn test_account_persistence() {
         let temp_dir = TempDir::new().unwrap();

crates/bitcell-state/tests/storage_persistence_test.rs

@@ -58,11 +58,11 @@ fn test_multi_block_persistence() {
         assert!(header.is_some(), "Block {} not found", height);
     }
 
-    // Verify transactions
-    for height in 0..100 {
-        let sender = format!("sender_{:033}", height % 10);
+    // Verify transactions - check each unique sender once
+    for sender_id in 0..10 {
+        let sender = format!("sender_{:033}", sender_id);
         let txs = storage.get_transactions_by_sender(sender.as_bytes(), 0).unwrap();
-        assert!(txs.len() >= 10, "Expected at least 10 transactions for sender at height {}", height);
+        assert_eq!(txs.len(), 100, "Expected 100 transactions for sender {}", sender_id);
     }
 
     // Verify state roots

docs/STORAGE.md

@@ -30,7 +30,8 @@ The storage layer uses separate RocksDB column families for different data types
 
 1. **Multiple Indexes**: Blocks and headers are indexed by both height and hash for O(1) lookups
 2. **Sender Index**: Transactions use a composite key (sender||height||tx_hash) for efficient range queries
-3. **Atomic Writes**: All multi-key operations use `WriteBatch` for atomicity
+3. **Atomic Writes**: Multi-key operations use `WriteBatch` for atomicity (blocks, headers, state roots, snapshots)
+   > **Note:** Transaction deletion in production pruning is not fully implemented yet. Transaction storage uses atomic batches.
 4. **Snapshots**: Variable-length snapshot format with length prefix for flexibility
 5. **Separation of Concerns**: Block data, state data, and indexes are in separate column families