From 8681cd713121a7c7ea7607d39e2a98502e8e144b Mon Sep 17 00:00:00 2001
From: 0xaatif <169152398+0xaatif@users.noreply.github.com>
Date: Mon, 21 Oct 2024 15:45:00 +0100
Subject: [PATCH] fix: more robust SMT parsing (#733)

* mark: 0xaatif/smt-parse-tweaking

* robustness
---
 trace_decoder/src/type2.rs | 3 +++
 trace_decoder/src/wire.rs  | 4 ----
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/trace_decoder/src/type2.rs b/trace_decoder/src/type2.rs
index a71761533..44d13e89a 100644
--- a/trace_decoder/src/type2.rs
+++ b/trace_decoder/src/type2.rs
@@ -173,8 +173,10 @@ fn visit(
             address,
             value,
         }) => {
+            ensure!(address.len() == Address::len_bytes());
             let address = Address::from_slice(&address);
             let collated = leaves.entry(address).or_default();
+            ensure!(value.len() <= 32);
             let value = U256::from_big_endian(&value);
             macro_rules! ensure {
                 ($expr:expr) => {
@@ -195,6 +197,7 @@ fn visit(
                     collated.code = Some(value)
                 }
                 SmtLeafType::Storage(slot) => {
+                    ensure!(slot.len() <= 32);
                     let clobbered = collated.storage.insert(U256::from_big_endian(&slot), value);
                     ensure!(clobbered.is_none())
                 }
diff --git a/trace_decoder/src/wire.rs b/trace_decoder/src/wire.rs
index 63dee6040..359043057 100644
--- a/trace_decoder/src/wire.rs
+++ b/trace_decoder/src/wire.rs
@@ -80,8 +80,6 @@ pub enum Instruction {
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
-// TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/705
-//                `address` and `value` should be fixed length fields
 pub struct SmtLeaf {
     pub node_type: SmtLeafType,
     pub address: NonEmpty<Vec<u8>>,
@@ -89,8 +87,6 @@ pub struct SmtLeaf {
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
-// TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/705
-//                `Storage` should contain a fixed length field
 pub enum SmtLeafType {
     Balance,
     Nonce,