Skip to content

Commit 83b904f

Browse files
klkvrklkvr
authored
feat: make find_by_deployed_code_exact smarter (#11560)
* feat: make find_by_deployed_code_exact smarter * clippy * prefer exact match * clippy * forge fmt * typo
1 parent 1b3b9b4 commit 83b904f

File tree

2 files changed

+119
-72
lines changed

2 files changed

+119
-72
lines changed

‎crates/common/src/contracts.rs‎

Lines changed: 102 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
//! Commonly used contract types and functions.
22
3-
use crate::{compile::PathOrContractInfo, strip_bytecode_placeholders};
3+
use crate::{compile::PathOrContractInfo, find_metadata_start, strip_bytecode_placeholders};
44
use alloy_dyn_abi::JsonAbiExt;
55
use alloy_json_abi::{Event, Function, JsonAbi};
66
use alloy_primitives::{Address, B256, Bytes, Selector, hex};
@@ -241,94 +241,130 @@ impl ContractsByArtifact {
241241
return None;
242242
}
243243

244-
self.iter().find(|(_, contract)| {
245-
let Some(deployed_bytecode) = &contract.deployed_bytecode else {
246-
return false;
247-
};
248-
let Some(deployed_code) = &deployed_bytecode.object else {
249-
return false;
250-
};
251-
252-
let len = match deployed_code {
253-
BytecodeObject::Bytecode(bytes) => bytes.len(),
254-
BytecodeObject::Unlinked(bytes) => bytes.len() / 2,
255-
};
256-
257-
if len != code.len() {
258-
return false;
259-
}
244+
let mut partial_match = None;
245+
self.iter()
246+
.find(|(id, contract)| {
247+
let Some(deployed_bytecode) = &contract.deployed_bytecode else {
248+
return false;
249+
};
250+
let Some(deployed_code) = &deployed_bytecode.object else {
251+
return false;
252+
};
253+
254+
let len = match deployed_code {
255+
BytecodeObject::Bytecode(bytes) => bytes.len(),
256+
BytecodeObject::Unlinked(bytes) => bytes.len() / 2,
257+
};
260258

261-
// Collect ignored offsets by chaining link and immutable references.
262-
let mut ignored = deployed_bytecode
263-
.immutable_references
264-
.values()
265-
.chain(deployed_bytecode.link_references.values().flat_map(|v| v.values()))
266-
.flatten()
267-
.cloned()
268-
.collect::<Vec<_>>();
269-
270-
// For libraries solidity adds a call protection prefix to the bytecode. We need to
271-
// ignore it as it includes library address determined at runtime.
272-
// See https://docs.soliditylang.org/en/latest/contracts.html#call-protection-for-libraries and
273-
// https://github.com/NomicFoundation/hardhat/blob/af7807cf38842a4f56e7f4b966b806e39631568a/packages/hardhat-verify/src/internal/solc/bytecode.ts#L172
274-
let has_call_protection = match deployed_code {
275-
BytecodeObject::Bytecode(bytes) => {
276-
bytes.starts_with(&CALL_PROTECTION_BYTECODE_PREFIX)
259+
if len != code.len() {
260+
return false;
277261
}
278-
BytecodeObject::Unlinked(bytes) => {
279-
if let Ok(bytes) =
280-
Bytes::from_str(&bytes[..CALL_PROTECTION_BYTECODE_PREFIX.len() * 2])
281-
{
262+
263+
// Collect ignored offsets by chaining link and immutable references.
264+
let mut ignored = deployed_bytecode
265+
.immutable_references
266+
.values()
267+
.chain(deployed_bytecode.link_references.values().flat_map(|v| v.values()))
268+
.flatten()
269+
.cloned()
270+
.collect::<Vec<_>>();
271+
272+
// For libraries solidity adds a call protection prefix to the bytecode. We need to
273+
// ignore it as it includes library address determined at runtime.
274+
// See https://docs.soliditylang.org/en/latest/contracts.html#call-protection-for-libraries and
275+
// https://github.com/NomicFoundation/hardhat/blob/af7807cf38842a4f56e7f4b966b806e39631568a/packages/hardhat-verify/src/internal/solc/bytecode.ts#L172
276+
let has_call_protection = match deployed_code {
277+
BytecodeObject::Bytecode(bytes) => {
282278
bytes.starts_with(&CALL_PROTECTION_BYTECODE_PREFIX)
283-
} else {
284-
false
285279
}
280+
BytecodeObject::Unlinked(bytes) => {
281+
if let Ok(bytes) =
282+
Bytes::from_str(&bytes[..CALL_PROTECTION_BYTECODE_PREFIX.len() * 2])
283+
{
284+
bytes.starts_with(&CALL_PROTECTION_BYTECODE_PREFIX)
285+
} else {
286+
false
287+
}
288+
}
289+
};
290+
291+
if has_call_protection {
292+
ignored.push(Offsets { start: 1, length: 20 });
286293
}
287-
};
288294

289-
if has_call_protection {
290-
ignored.push(Offsets { start: 1, length: 20 });
291-
}
295+
let metadata_start = find_metadata_start(code);
292296

293-
ignored.sort_by_key(|o| o.start);
297+
if let Some(metadata) = metadata_start {
298+
ignored.push(Offsets {
299+
start: metadata as u32,
300+
length: (code.len() - metadata) as u32,
301+
});
302+
}
294303

295-
let mut left = 0;
296-
for offset in ignored {
297-
let right = offset.start as usize;
304+
ignored.sort_by_key(|o| o.start);
298305

299-
let matched = match deployed_code {
300-
BytecodeObject::Bytecode(bytes) => bytes[left..right] == code[left..right],
301-
BytecodeObject::Unlinked(bytes) => {
302-
if let Ok(bytes) = Bytes::from_str(&bytes[left * 2..right * 2]) {
303-
bytes == code[left..right]
304-
} else {
305-
false
306+
let mut left = 0;
307+
for offset in ignored {
308+
let right = offset.start as usize;
309+
310+
let matched = match deployed_code {
311+
BytecodeObject::Bytecode(bytes) => bytes[left..right] == code[left..right],
312+
BytecodeObject::Unlinked(bytes) => {
313+
if let Ok(bytes) = Bytes::from_str(&bytes[left * 2..right * 2]) {
314+
bytes == code[left..right]
315+
} else {
316+
false
317+
}
306318
}
319+
};
320+
321+
if !matched {
322+
return false;
307323
}
324+
325+
left = right + offset.length as usize;
326+
}
327+
328+
let is_partial = if left < code.len() {
329+
match deployed_code {
330+
BytecodeObject::Bytecode(bytes) => bytes[left..] == code[left..],
331+
BytecodeObject::Unlinked(bytes) => {
332+
if let Ok(bytes) = Bytes::from_str(&bytes[left * 2..]) {
333+
bytes == code[left..]
334+
} else {
335+
false
336+
}
337+
}
338+
}
339+
} else {
340+
true
308341
};
309342

310-
if !matched {
343+
if !is_partial {
311344
return false;
312345
}
313346

314-
left = right + offset.length as usize;
315-
}
347+
let Some(metadata) = metadata_start else { return true };
316348

317-
if left < code.len() {
318-
match deployed_code {
319-
BytecodeObject::Bytecode(bytes) => bytes[left..] == code[left..],
349+
let exact_match = match deployed_code {
350+
BytecodeObject::Bytecode(bytes) => bytes[metadata..] == code[metadata..],
320351
BytecodeObject::Unlinked(bytes) => {
321-
if let Ok(bytes) = Bytes::from_str(&bytes[left * 2..]) {
322-
bytes == code[left..]
352+
if let Ok(bytes) = Bytes::from_str(&bytes[metadata * 2..]) {
353+
bytes == code[metadata..]
323354
} else {
324355
false
325356
}
326357
}
358+
};
359+
360+
if exact_match {
361+
true
362+
} else {
363+
partial_match = Some((*id, *contract));
364+
false
327365
}
328-
} else {
329-
true
330-
}
331-
})
366+
})
367+
.or(partial_match)
332368
}
333369

334370
/// Finds a contract which has the same contract name or identifier as `id`. If more than one is

‎crates/common/src/utils.rs‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -45,17 +45,28 @@ pub fn erc7201(id: &str) -> B256 {
4545
keccak256(x.to_be_bytes::<32>()) & B256::from(!U256::from(0xff))
4646
}
4747

48-
/// Utility function to ignore metadata hash of the given bytecode.
48+
/// Utility function to find the start of the metadata in the bytecode.
4949
/// This assumes that the metadata is at the end of the bytecode.
50-
pub fn ignore_metadata_hash(bytecode: &[u8]) -> &[u8] {
50+
pub fn find_metadata_start(bytecode: &[u8]) -> Option<usize> {
5151
// Get the last two bytes of the bytecode to find the length of CBOR metadata.
52-
let Some((rest, metadata_len_bytes)) = bytecode.split_last_chunk() else { return bytecode };
52+
let (rest, metadata_len_bytes) = bytecode.split_last_chunk()?;
5353
let metadata_len = u16::from_be_bytes(*metadata_len_bytes) as usize;
5454
if metadata_len > rest.len() {
55-
return bytecode;
55+
return None;
56+
}
57+
ciborium::from_reader::<ciborium::Value, _>(&rest[rest.len() - metadata_len..])
58+
.is_ok()
59+
.then(|| rest.len() - metadata_len)
60+
}
61+
62+
/// Utility function to ignore metadata hash of the given bytecode.
63+
/// This assumes that the metadata is at the end of the bytecode.
64+
pub fn ignore_metadata_hash(bytecode: &[u8]) -> &[u8] {
65+
if let Some(metadata) = find_metadata_start(bytecode) {
66+
&bytecode[..metadata]
67+
} else {
68+
bytecode
5669
}
57-
let (rest, metadata) = rest.split_at(rest.len() - metadata_len);
58-
if ciborium::from_reader::<ciborium::Value, _>(metadata).is_ok() { rest } else { bytecode }
5970
}
6071

6172
/// Strips all __$xxx$__ placeholders from the bytecode if it's an unlinked bytecode.

0 commit comments

Comments
 (0)