Skip to content

Commit

Permalink
5.1 Missing Length Validation in MPT Verify
Browse files Browse the repository at this point in the history
  • Loading branch information
@tiagofneto
tiagofneto committed Oct 26, 2023
1 parent 99b1c0e commit 287c3ec
Show file tree
Hide file tree
Showing 4 changed files with 131 additions and 133 deletions.
121 changes: 36 additions & 85 deletions src/data_structures/eth_mpt.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,19 +19,21 @@ impl MPTDefault of Default<MPT> {
}

// @notice Represents a node in the MPT
#[derive(Drop)]
#[derive(Drop, PartialEq)]
enum MPTNode {
// @param 16 hashes of children
// @param Value of the node
// @param hashes 16 hashes of children
// @param value value of the node
Branch: (Span<Words64>, Words64),
// @param shared_nibbles
// @param next_node
// @param nibbles_skip Number of nibbles to skip in shared nibbles
Extension: (Words64, u256, usize),
// @param nibbles_skip number of nibbles to skip in shared nibbles
// @param n_nibbles number of shared nibbles
Extension: (Words64, u256, usize, usize),
// @param key_end
// @param value of the node
// @param nibbles_skip Number of nibbles to skip in the key end
Leaf: (Words64, Words64, usize)
// @param n_nibbles number of nibbles in key_end
Leaf: (Words64, Words64, usize, usize)
}

#[generate_trait]
Expand Down Expand Up @@ -97,7 +99,7 @@ impl MPTImpl of MPTTrait {
key_pow2 = key_pow2 / 16;
},
MPTNode::Extension((
shared_nibbles, next_node, nibbles_skip
shared_nibbles, next_node, nibbles_skip, n_nibbles
)) => {
let mut shared_nibbles_pow2 = pow(2, nibbles_skip.into() * 4);

Expand All @@ -113,7 +115,11 @@ impl MPTImpl of MPTTrait {

let mut shared_nibbles_word_idx = nibbles_skip / 16;
let mut shared_nibbles_word = *shared_nibbles.at(shared_nibbles_word_idx);
let mut i_nibbles = 0;
let next_hash = loop {
if i_nibbles == n_nibbles {
break Result::Ok(next_node);
}
if key_pow2 == 0 {
break Result::Err('Key reached end');
}
Expand All @@ -123,7 +129,7 @@ impl MPTImpl of MPTTrait {
& 0xf;
let current_nibble_key = (key / key_pow2) & 0xf;
if current_nibble_shared_nibbles.into() != current_nibble_key {
break Result::Ok(next_node);
break Result::Err('Extension nibbles not matching');
}

if shared_nibbles_pow2 == 0x100000000000000 {
Expand All @@ -140,6 +146,7 @@ impl MPTImpl of MPTTrait {

in_byte = !in_byte;
key_pow2 = key_pow2 / 16;
i_nibbles += 1;
};

match next_hash {
Expand All @@ -152,7 +159,7 @@ impl MPTImpl of MPTTrait {
}
},
MPTNode::Leaf((
key_end, value, nibbles_skip
key_end, value, nibbles_skip, n_nibbles
)) => {
let mut key_end_pow2 = pow(2, nibbles_skip.into() * 4);

Expand All @@ -168,8 +175,9 @@ impl MPTImpl of MPTTrait {

let mut key_end_word_idx = nibbles_skip / 16;
let mut key_end_word = *key_end.at(key_end_word_idx);
let mut i_nibbles = 0;
break loop {
if key_pow2 == 0 {
if key_pow2 == 0 && i_nibbles == n_nibbles {
break Result::Ok(value);
}

Expand All @@ -193,6 +201,7 @@ impl MPTImpl of MPTTrait {

in_byte = !in_byte;
key_pow2 = key_pow2 / 16;
i_nibbles += 1;
};
}
};
Expand All @@ -215,33 +224,41 @@ impl MPTImpl of MPTTrait {
let mut i: usize = 0;
loop {
if i == 16 {
let value = *l.at(16);
let (value, _) = *l.at(16);
break Result::Ok(MPTNode::Branch((nibble_hashes.span(), value)));
}

nibble_hashes.append(*l.at(i));
let (current_hash, _) = *l.at(i);
nibble_hashes.append(current_hash);
i += 1;
}
} else if len == 2 {
let first = *l.at(0);
let (first, first_len) = *l.at(0);
let (second, _) = *l.at(1);
// Unwrap impossible to fail, as we are making with 0xff, meaning the result always fits in a byte
let prefix_byte: Byte = (*first.at(0) & 0xff).try_into().unwrap();
let (prefix, _) = prefix_byte.extract_nibbles();

let n_nibbles = (first_len * 2) - 1;

if prefix == 0 {
match (*l.at(1)).try_into() {
Option::Some(n) => Result::Ok(MPTNode::Extension((first, n, 2))),
match second.try_into() {
Option::Some(n) => Result::Ok(
MPTNode::Extension((first, n, 2, n_nibbles - 1))
),
Option::None(_) => Result::Err('Invalid next node')
}
} else if prefix == 1 {
match (*l.at(1)).try_into() {
Option::Some(n) => Result::Ok(MPTNode::Extension((first, n, 1))),
match second.try_into() {
Option::Some(n) => Result::Ok(
MPTNode::Extension((first, n, 1, n_nibbles))
),
Option::None(_) => Result::Err('Invalid next node')
}
} else if prefix == 2 {
Result::Ok(MPTNode::Leaf((first, *l.at(1), 2)))
Result::Ok(MPTNode::Leaf((first, second, 2, n_nibbles - 1)))
} else if prefix == 3 {
Result::Ok(MPTNode::Leaf((first, *l.at(1), 1)))
Result::Ok(MPTNode::Leaf((first, second, 1, n_nibbles)))
} else {
Result::Err('Invalid RLP prefix')
}
Expand All @@ -259,69 +276,3 @@ impl MPTImpl of MPTTrait {
keccak_cairo_words64(rlp)
}
}

impl MPTNodePartialEq of PartialEq<MPTNode> {
fn eq(lhs: @MPTNode, rhs: @MPTNode) -> bool {
match lhs {
MPTNode::Branch((
lhs_nibbles, lhs_value
)) => {
match rhs {
MPTNode::Branch((
rhs_nibbles, rhs_value
)) => {
if (*lhs_nibbles).len() != (*rhs_nibbles).len() {
return false;
}
let mut i: usize = 0;
loop {
if i >= (*lhs_nibbles).len() {
break lhs_value == rhs_value;
}
if (*lhs_nibbles).at(i) != (*rhs_nibbles).at(i) {
break false;
}
i += 1;
}
},
MPTNode::Extension(_) => false,
MPTNode::Leaf(_) => false
}
},
MPTNode::Extension((
lhs_shared_nibbles, lhs_next_node, lhs_nibbles_skip
)) => {
match rhs {
MPTNode::Branch(_) => false,
MPTNode::Extension((
rhs_shared_nibbles, rhs_next_node, rhs_nibbles_skip
)) => {
lhs_shared_nibbles == rhs_shared_nibbles
&& lhs_next_node == rhs_next_node
&& lhs_nibbles_skip == rhs_nibbles_skip
},
MPTNode::Leaf(_) => false
}
},
MPTNode::Leaf((
lhs_key_end, lhs_value, lhs_nibbles_skip
)) => {
match rhs {
MPTNode::Branch(_) => false,
MPTNode::Extension(_) => false,
MPTNode::Leaf((
rhs_key_end, rhs_value, rhs_nibbles_skip
)) => {
lhs_key_end == rhs_key_end
&& lhs_value == rhs_value
&& lhs_nibbles_skip == rhs_nibbles_skip
}
}
}
}
}

fn ne(lhs: @MPTNode, rhs: @MPTNode) -> bool {
!(lhs == rhs)
}
}
4 changes: 2 additions & 2 deletions src/data_structures/tests/test_eth_mpt.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ fn test_decode_rlp_node_leaf_odd() {
];

let decoded = MPTTrait::decode_rlp_node(rlp_node.span()).unwrap();
let expected_node = MPTNode::Leaf((expected_key_end.span(), expected_value.span(), 1));
let expected_node = MPTNode::Leaf((expected_key_end.span(), expected_value.span(), 1, 57));
assert(decoded == expected_node, 'Even leaf node differs');
}

Expand Down Expand Up @@ -193,7 +193,7 @@ fn test_decode_rlp_node_leaf_even() {
];

let decoded = MPTTrait::decode_rlp_node(rlp_node.span()).unwrap();
let expected_node = MPTNode::Leaf((expected_key_end.span(), expected_value.span(), 2));
let expected_node = MPTNode::Leaf((expected_key_end.span(), expected_value.span(), 2, 56));
assert(decoded == expected_node, 'Even leaf node differs');
}

Expand Down
14 changes: 7 additions & 7 deletions src/encoding/rlp.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,15 +38,15 @@ impl RLPTypeImpl of RLPTypeTrait {
// @notice Represent a RLP item
#[derive(Drop)]
enum RLPItem {
Bytes: Words64,
Bytes: (Words64, usize),
// Should be Span<RLPItem> to allow for any depth/recursion, not yet supported by the compiler
List: Span<Words64>
List: Span<(Words64, usize)>
}

// @notice RLP decodes a rlp encoded byte array
// For more info: https://ethereum.org/en/developers/docs/data-structures-and-encoding/rlp/
// @param input RLP encoded input, in little endian 64 bits words
// @return Result with RLPItem and size of the decoded item
// @return Result with RLPItem and size of the encoded item
fn rlp_decode(input: Words64) -> Result<(RLPItem, usize), felt252> {
// It's guaranteed to fid in 32 bits, as we are masking with 0xff
let prefix: u32 = (*input.at(0) & 0xff).try_into().unwrap();
Expand All @@ -56,13 +56,13 @@ fn rlp_decode(input: Words64) -> Result<(RLPItem, usize), felt252> {
match rlp_type {
RLPType::String(()) => {
let mut arr = array![prefix.into()];
Result::Ok((RLPItem::Bytes(arr.span()), 1))
Result::Ok((RLPItem::Bytes((arr.span(), 1)), 1))
},
RLPType::StringShort(()) => {
let len = prefix.into() - 0x80;
let res = input.slice_le(6, len);

Result::Ok((RLPItem::Bytes(res), 1 + len))
Result::Ok((RLPItem::Bytes((res, len)), 1 + len))
},
RLPType::StringLong(()) => {
let len_len = prefix - 0xb7;
Expand All @@ -76,7 +76,7 @@ fn rlp_decode(input: Words64) -> Result<(RLPItem, usize), felt252> {
.unwrap();
let res = input.slice_le(6 - len_len, len);

Result::Ok((RLPItem::Bytes(res), 1 + len_len + len))
Result::Ok((RLPItem::Bytes((res, len)), 1 + len_len + len))
},
RLPType::ListShort(()) => {
let mut len = prefix - 0xc0;
Expand Down Expand Up @@ -106,7 +106,7 @@ fn rlp_decode(input: Words64) -> Result<(RLPItem, usize), felt252> {
// @param input RLP encoded input, in little endian 64 bits words
// @param len Length of the input
// @return Result with RLPItem::List
fn rlp_decode_list(ref input: Words64, len: usize) -> Result<Span<Words64>, felt252> {
fn rlp_decode_list(ref input: Words64, len: usize) -> Result<Span<(Words64, usize)>, felt252> {
let mut i = 0;
let mut output = ArrayTrait::new();
let mut total_len = len;
Expand Down
Loading

0 comments on commit 287c3ec

Please sign in to comment.