From e5bfe4c0a9e3fa9ae2e60b05def53462c6ded89e Mon Sep 17 00:00:00 2001
From: Viraj Bhartiya <vlbhartiya@gmail.com>
Date: Thu, 4 Sep 2025 17:40:46 +0530
Subject: [PATCH] feat: add RLP encoding size checks for embedded nodes in
 utils.cairo and utils.py

---
 python/mpt/src/mpt/utils.cairo | 17 ++++++++++++++++-
 python/mpt/src/mpt/utils.py    | 15 +++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/python/mpt/src/mpt/utils.cairo b/python/mpt/src/mpt/utils.cairo
index ca8f92c3c..9dffead42 100644
--- a/python/mpt/src/mpt/utils.cairo
+++ b/python/mpt/src/mpt/utils.cairo
@@ -14,7 +14,7 @@ from ethereum.prague.trie import (
     SubnodesStruct,
     bytes_to_nibble_list,
 )
-from ethereum_rlp.rlp import Extended, ExtendedEnum, ExtendedImpl, decode
+from ethereum_rlp.rlp import Extended, ExtendedEnum, ExtendedImpl, decode, encode
 from ethereum_types.bytes import Bytes, BytesStruct
 from starkware.cairo.common.alloc import alloc
 from starkware.cairo.common.math_cmp import is_le_felt
@@ -184,6 +184,11 @@ func check_branch_node(node: BranchNode) {
         if (x.value.sequence.value.len == 0) {
             raise_ValueError('LTTwoNonNullSubnodes');
         }
+        // Embedded nodes must have RLP encoding size < 32
+        let encoded = encode(x);
+        if (encoded.value.len >= 32) {
+            raise_ValueError('EmbeddedNodeTooLarge');
+        }
     }
 
     // Check that the second subnode is not None and not empty
@@ -201,6 +206,11 @@ func check_branch_node(node: BranchNode) {
         if (y.value.sequence.value.len == 0) {
             raise_ValueError('LTTwoNonNullSubnodes');
         }
+        // Embedded nodes must have RLP encoding size < 32
+        let encoded = encode(y);
+        if (encoded.value.len >= 32) {
+            raise_ValueError('EmbeddedNodeTooLarge');
+        }
     }
 
     return ();
@@ -273,6 +283,11 @@ func check_extension_node(node: ExtensionNode, parent_node: OptionalInternalNode
         if (sequence_variant.len == 0) {
             raise_ValueError('EmptySubnode');
         }
+        // Embedded node must have RLP encoding size < 32
+        let encoded = encode(subnode);
+        if (encoded.value.len >= 32) {
+            raise_ValueError('EmbeddedNodeTooLarge');
+        }
         return ();
     }
 
diff --git a/python/mpt/src/mpt/utils.py b/python/mpt/src/mpt/utils.py
index 5bcdfbee7..089a07ade 100644
--- a/python/mpt/src/mpt/utils.py
+++ b/python/mpt/src/mpt/utils.py
@@ -91,6 +91,15 @@ def check_branch_node(node: BranchNode) -> None:
     if len(non_null_subnodes) < 2:
         raise ValueError("LTTwoNonNullSubnodes")
 
+    # For each embedded subnode (represented as a list), ensure its RLP encoding is < 32 bytes
+    for subnode in node.subnodes:
+        if subnode in (None, b""):
+            continue
+        if isinstance(subnode, list):
+            encoded = rlp.encode(subnode)
+            if len(encoded) >= 32:
+                raise ValueError("EmbeddedNodeTooLarge")
+
 
 def check_leaf_node(path: Bytes, node: LeafNode) -> None:
     """
@@ -121,3 +130,9 @@ def check_extension_node(node: ExtensionNode, parent: Optional[InternalNode]) ->
 
     if parent is not None and isinstance(parent, ExtensionNode):
         raise ValueError("InvalidParent")
+
+    # If subnode is embedded (represented as a list), ensure its RLP encoding is < 32 bytes
+    if isinstance(node.subnode, list):
+        encoded = rlp.encode(node.subnode)
+        if len(encoded) >= 32:
+            raise ValueError("EmbeddedNodeTooLarge")