Skip to content

Commit 03f42b9

Browse files
james-toussaintjames-toussaint
committed
Traverse after extension node
1 parent f587b88 commit 03f42b9

File tree

2 files changed

+78
-50
lines changed

2 files changed

+78
-50
lines changed

‎contracts/utils/cryptography/TrieProof.sol‎

Lines changed: 29 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -125,13 +125,27 @@ library TrieProof {
125125

126126
if (nodeLength == branchNodeLength) {
127127
// If we've consumed the entire key, the value must be in the last slot
128-
if (keyIndex == key.length) return _validateLastItem(node.decoded[radix], trieProof, i);
128+
if (keyIndex == key.length) return _validateLastItem(node.decoded[radix], trieProof.length, i);
129129

130130
// Otherwise, continue down the branch specified by the next nibble in the key
131131
uint8 branchKey = uint8(key[keyIndex]);
132132
(nodeId, keyIndex) = (_id(node.decoded[branchKey]), keyIndex + 1);
133133
} else if (nodeLength == LEAF_OR_EXTENSION_NODE_LENGTH) {
134-
return _processLeafOrExtension(node, trieProof, key, nodeId, keyIndex, i);
134+
(uint8 prefix, bytes memory pathRemainder, bytes memory keyRemainder) = _extract(node, key, keyIndex);
135+
if (prefix == type(uint8).max) return ("", ProofError.UNKNOWN_NODE_PREFIX);
136+
// Leaf node (terminal) - return its value if key matches completely
137+
if (Prefix(prefix) == Prefix.LEAF_EVEN || Prefix(prefix) == Prefix.LEAF_ODD) {
138+
if (!string(pathRemainder).equal(string(keyRemainder)))
139+
return ("", ProofError.MISMATCH_LEAF_PATH_KEY_REMAINDERS);
140+
if (keyRemainder.length == 0) return ("", ProofError.INVALID_KEY_REMAINDER);
141+
return _validateLastItem(node.decoded[1], trieProof.length, i);
142+
}
143+
// Extension node (non-terminal) - validate shared path & continue to next node
144+
uint256 sharedNibbleLength = _sharedNibbleLength(pathRemainder, keyRemainder);
145+
// Path must match at least partially with our key
146+
if (sharedNibbleLength == 0) return ("", ProofError.INVALID_PATH_REMAINDER);
147+
// Increment keyIndex by the number of nibbles consumed and continue traversal
148+
(nodeId, keyIndex) = (_id(node.decoded[1]), keyIndex + sharedNibbleLength);
135149
}
136150
}
137151

@@ -157,43 +171,19 @@ library TrieProof {
157171
return ProofError.NO_ERROR; // No error
158172
}
159173

160-
/**
161-
* @dev Processes a leaf or extension node in the trie proof.
162-
*
163-
* For leaf nodes, validates that the key matches completely and returns the value.
164-
* For extension nodes, continues traversal by updating the node ID and key index.
165-
*/
166-
function _processLeafOrExtension(
174+
/// @dev Extract prefix, path remainder, and key remainder from a leaf or extension node.
175+
function _extract(
167176
Node memory node,
168-
Node[] memory trieProof,
169177
bytes memory key,
170-
bytes memory nodeId,
171-
uint256 keyIndex,
172-
uint256 i
173-
) private pure returns (bytes memory value, ProofError err) {
178+
uint256 keyIndex
179+
) private pure returns (uint8 prefix, bytes memory pathRemainder, bytes memory keyRemainder) {
174180
bytes memory path = _path(node);
175-
uint8 prefix = uint8(path[0]);
181+
prefix = uint8(path[0]);
182+
if (prefix > uint8(type(Prefix).max)) return (type(uint8).max, "", "");
176183
uint8 offset = 2 - (prefix % 2); // Calculate offset based on even/odd path length
177-
bytes memory pathRemainder = Bytes.slice(path, offset); // Path after the prefix
178-
bytes memory keyRemainder = Bytes.slice(key, keyIndex); // Remaining key to match
179-
if (prefix > uint8(type(Prefix).max)) return ("", ProofError.UNKNOWN_NODE_PREFIX);
180-
181-
// Leaf node (terminal) - return its value if key matches completely
182-
if (Prefix(prefix) == Prefix.LEAF_EVEN || Prefix(prefix) == Prefix.LEAF_ODD) {
183-
if (!string(pathRemainder).equal(string(keyRemainder)))
184-
return ("", ProofError.MISMATCH_LEAF_PATH_KEY_REMAINDERS);
185-
if (keyRemainder.length == 0) return ("", ProofError.INVALID_KEY_REMAINDER);
186-
return _validateLastItem(node.decoded[1], trieProof, i);
187-
}
188-
189-
// Extension node (non-terminal) - validate shared path & continue to next node
190-
uint256 sharedNibbleLength = _sharedNibbleLength(pathRemainder, keyRemainder);
191-
if (Prefix(prefix) == Prefix.EXTENSION_EVEN || Prefix(prefix) == Prefix.EXTENSION_ODD) {
192-
// Path must match at least partially with our key
193-
if (sharedNibbleLength == 0) return ("", ProofError.INVALID_PATH_REMAINDER);
194-
}
195-
// Increment keyIndex by the number of nibbles consumed
196-
(nodeId, keyIndex) = (_id(node.decoded[1]), keyIndex + sharedNibbleLength);
184+
pathRemainder = Bytes.slice(path, offset); // Path after the prefix
185+
keyRemainder = Bytes.slice(key, keyIndex); // Remaining key to match
186+
return (prefix, pathRemainder, keyRemainder);
197187
}
198188

199189
/**
@@ -202,12 +192,12 @@ library TrieProof {
202192
*/
203193
function _validateLastItem(
204194
Memory.Slice item,
205-
Node[] memory trieProof,
195+
uint256 trieProofLength,
206196
uint256 i
207197
) private pure returns (bytes memory value, ProofError) {
208198
bytes memory value_ = item.readBytes();
209199
if (value_.length == 0) return ("", ProofError.EMPTY_VALUE);
210-
if (i != trieProof.length - 1) return ("", ProofError.INVALID_EXTRA_PROOF_ELEMENT);
200+
if (i != trieProofLength - 1) return ("", ProofError.INVALID_EXTRA_PROOF_ELEMENT);
211201
return (value_, ProofError.NO_ERROR);
212202
}
213203

@@ -245,9 +235,9 @@ library TrieProof {
245235
* Used to determine how much of a path matches a key during trie traversal.
246236
*/
247237
function _sharedNibbleLength(bytes memory _a, bytes memory _b) private pure returns (uint256 shared_) {
248-
uint256 max = Math.max(_a.length, _b.length);
238+
uint256 min = Math.min(_a.length, _b.length);
249239
uint256 length;
250-
while (length < max && _a[length] == _b[length]) {
240+
while (length < min && _a[length] == _b[length]) {
251241
length++;
252242
}
253243
return length;

‎test/utils/cryptography/TrieProof.test.js‎

Lines changed: 49 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -76,25 +76,63 @@ describe('TrieProof', function () {
7676
}
7777
});
7878

79-
describe('verify', function () {
80-
it('returns true for a valid proof with leaf', async function () {
79+
describe('verify proof', function () {
80+
it('returns true with proof size 1 (even leaf [0x20])', async function () {
8181
const slot = ethers.ZeroHash;
82-
const tx = await call(this.storage, 'setUint256Slot', [slot, 42]);
83-
const { key, value, proof, storageHash } = await this.getProof(this.storage, slot, tx);
82+
await call(this.storage, 'setUint256Slot', [slot, 42]);
83+
const { key, value, proof, storageHash } = await this.getProof(this.storage, slot);
8484
const result = await this.mock.$verify(ethers.keccak256(key), value, proof, storageHash);
8585
expect(result).is.true;
8686
});
8787

88-
it('returns true for a valid proof with extension', async function () {
88+
it('returns true with proof size 2 (branch then odd leaf [0x3])', async function () {
8989
const slot0 = ethers.ZeroHash;
9090
const slot1 = '0x0000000000000000000000000000000000000000000000000000000000000001';
9191
await call(this.storage, 'setUint256Slot', [slot0, 42]);
92-
const tx = await call(this.storage, 'setUint256Slot', [slot1, 43]);
93-
const { key, value, proof, storageHash } = await this.getProof(this.storage, slot1, tx);
92+
await call(this.storage, 'setUint256Slot', [slot1, 43]);
93+
const { key, value, proof, storageHash } = await this.getProof(this.storage, slot1);
9494
const result = await this.mock.$verify(ethers.keccak256(key), value, proof, storageHash);
9595
expect(result).is.true;
9696
});
9797

98+
it('returns true with proof size 3 (even extension [0x00], branch then leaf)', async function () {
99+
const slots = [
100+
'0x0000000000000000000000000000000000000000000000000000000000001889', // 0xabc4243e220df4927f4d7b432d2d718dadbba652f6cee6a45bb90c077fa4e158
101+
'0x0000000000000000000000000000000000000000000000000000000000008b23', // 0xabd5ef9a39144905d28bd8554745ebae050359cf7e89079f49b66a6c06bd2bf9
102+
'0x0000000000000000000000000000000000000000000000000000000000002383', // 0xabe87cb73c1e15a89cfb0daa7fd0cc3eb1a762345fe15d668f5061a4900b22fa
103+
];
104+
await call(this.storage, 'setUint256Slot', [slots[0], 42]);
105+
await call(this.storage, 'setUint256Slot', [slots[1], 43]);
106+
await call(this.storage, 'setUint256Slot', [slots[2], 44]);
107+
for (let slot of slots) {
108+
const { key, value, proof, storageHash } = await this.getProof(this.storage, slot);
109+
const result = await this.mock.$verify(ethers.keccak256(key), value, proof, storageHash);
110+
expect(result).is.true;
111+
}
112+
});
113+
114+
it('returns true with proof size 3 (odd extension [0x1], branch then leaf)', async function () {
115+
const slots = [
116+
'0x0000000000000000000000000000000000000000000000000000000000004616', // 0xabcd2ce29d227a0aaaa2ea425df9d5c96a569b416fd0bb7e018b8c9ce9b9d15d
117+
'0x0000000000000000000000000000000000000000000000000000000000012dd3', // 0xabce7718834e2932319fc4642268a27405261f7d3826b19811d044bf2b56ebb1
118+
'0x000000000000000000000000000000000000000000000000000000000000ce8f', // 0xabcf8b375ce20d03da20a3f5efeb8f3666810beca66f729f995953f51559a4ff
119+
];
120+
await call(this.storage, 'setUint256Slot', [slots[0], 42]);
121+
await call(this.storage, 'setUint256Slot', [slots[1], 43]);
122+
await call(this.storage, 'setUint256Slot', [slots[2], 44]);
123+
for (let slot of slots) {
124+
const { key, value, proof, storageHash } = await this.getProof(this.storage, slot);
125+
const result = await this.mock.$verify(ethers.keccak256(key), value, proof, storageHash);
126+
expect(result).is.true;
127+
}
128+
});
129+
130+
it('returns false for invalid proof', async function () {
131+
await expect(this.mock.$verify('0x', '0x', [], ethers.ZeroHash)).to.eventually.be.false;
132+
});
133+
});
134+
135+
describe('process invalid proof', function () {
98136
it('fails to process proof with empty key', async function () {
99137
const [value, error] = await this.mock.$processProof('0x', [], ethers.ZeroHash);
100138
expect(value).to.equal('0x');
@@ -116,8 +154,8 @@ describe('TrieProof', function () {
116154
const slot0 = ethers.ZeroHash;
117155
const slot1 = '0x0000000000000000000000000000000000000000000000000000000000000001';
118156
await call(this.storage, 'setUint256Slot', [slot0, 42]);
119-
const tx = await call(this.storage, 'setUint256Slot', [slot1, 43]);
120-
const { key, proof, storageHash } = await this.getProof(this.storage, slot1, tx);
157+
await call(this.storage, 'setUint256Slot', [slot1, 43]);
158+
const { key, proof, storageHash } = await this.getProof(this.storage, slot1);
121159
proof[1] = ethers.toBeHex(BigInt(proof[1]) + 1n); // Corrupt internal large node hash
122160
const [processedValue, error] = await this.mock.$processProof(key, proof, storageHash);
123161
expect(processedValue).to.equal('0x');
@@ -127,7 +165,7 @@ describe('TrieProof', function () {
127165
it.skip('fails to process proof with invalid internal short node', async function () {}); // TODO: INVALID_INTERNAL_NODE_HASH
128166

129167
it('fails to process proof with empty value', async function () {
130-
const proof = [ethers.encodeRlp(['0x2000', '0x'])]; // Corrupt proof to yield empty value
168+
const proof = [ethers.encodeRlp(['0x2000', '0x'])];
131169
const [processedValue, error] = await this.mock.$processProof('0x00', proof, ethers.keccak256(proof[0]));
132170
expect(processedValue).to.equal('0x');
133171
expect(error).to.equal(ProofError.EMPTY_VALUE);
@@ -159,7 +197,7 @@ describe('TrieProof', function () {
159197
});
160198

161199
it('fails to process proof with invalid path remainder', async function () {
162-
const proof = [ethers.encodeRlp(['0x0011', '0x'])]; // Corrupt proof to yield invalid path remainder
200+
const proof = [ethers.encodeRlp(['0x0011', '0x'])];
163201
const [processedValue, error] = await this.mock.$processProof(ethers.ZeroHash, proof, ethers.keccak256(proof[0]));
164202
expect(processedValue).to.equal('0x');
165203
expect(error).to.equal(ProofError.INVALID_PATH_REMAINDER);

0 commit comments

Comments
 (0)