foundry_cheatcodes/
crypto.rs

Help
1//! Implementations of [`Crypto`](spec::Group::Crypto) Cheatcodes.
2
3use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
4use alloy_primitives::{keccak256, Address, B256, U256};
5use alloy_signer::{Signer, SignerSync};
6use alloy_signer_local::{
7    coins_bip39::{
8        ChineseSimplified, ChineseTraditional, Czech, English, French, Italian, Japanese, Korean,
9        Portuguese, Spanish, Wordlist,
10    },
11    LocalSigner, MnemonicBuilder, PrivateKeySigner,
12};
13use alloy_sol_types::SolValue;
14use k256::{
15    ecdsa::SigningKey,
16    elliptic_curve::{bigint::ArrayEncoding, sec1::ToEncodedPoint},
17};
18use p256::ecdsa::{
19    signature::hazmat::PrehashSigner, Signature as P256Signature, SigningKey as P256SigningKey,
20};
21
22/// The BIP32 default derivation path prefix.
23const DEFAULT_DERIVATION_PATH_PREFIX: &str = "m/44'/60'/0'/0/";
24
25impl Cheatcode for createWallet_0Call {
26    fn apply(&self, state: &mut Cheatcodes) -> Result {
27        let Self { walletLabel } = self;
28        create_wallet(&U256::from_be_bytes(keccak256(walletLabel).0), Some(walletLabel), state)
29    }
30}
31
32impl Cheatcode for createWallet_1Call {
33    fn apply(&self, state: &mut Cheatcodes) -> Result {
34        let Self { privateKey } = self;
35        create_wallet(privateKey, None, state)
36    }
37}
38
39impl Cheatcode for createWallet_2Call {
40    fn apply(&self, state: &mut Cheatcodes) -> Result {
41        let Self { privateKey, walletLabel } = self;
42        create_wallet(privateKey, Some(walletLabel), state)
43    }
44}
45
46impl Cheatcode for sign_0Call {
47    fn apply(&self, _state: &mut Cheatcodes) -> Result {
48        let Self { wallet, digest } = self;
49        let sig = sign(&wallet.privateKey, digest)?;
50        Ok(encode_full_sig(sig))
51    }
52}
53
54impl Cheatcode for signCompact_0Call {
55    fn apply(&self, _state: &mut Cheatcodes) -> Result {
56        let Self { wallet, digest } = self;
57        let sig = sign(&wallet.privateKey, digest)?;
58        Ok(encode_compact_sig(sig))
59    }
60}
61
62impl Cheatcode for deriveKey_0Call {
63    fn apply(&self, _state: &mut Cheatcodes) -> Result {
64        let Self { mnemonic, index } = self;
65        derive_key::<English>(mnemonic, DEFAULT_DERIVATION_PATH_PREFIX, *index)
66    }
67}
68
69impl Cheatcode for deriveKey_1Call {
70    fn apply(&self, _state: &mut Cheatcodes) -> Result {
71        let Self { mnemonic, derivationPath, index } = self;
72        derive_key::<English>(mnemonic, derivationPath, *index)
73    }
74}
75
76impl Cheatcode for deriveKey_2Call {
77    fn apply(&self, _state: &mut Cheatcodes) -> Result {
78        let Self { mnemonic, index, language } = self;
79        derive_key_str(mnemonic, DEFAULT_DERIVATION_PATH_PREFIX, *index, language)
80    }
81}
82
83impl Cheatcode for deriveKey_3Call {
84    fn apply(&self, _state: &mut Cheatcodes) -> Result {
85        let Self { mnemonic, derivationPath, index, language } = self;
86        derive_key_str(mnemonic, derivationPath, *index, language)
87    }
88}
89
90impl Cheatcode for rememberKeyCall {
91    fn apply(&self, state: &mut Cheatcodes) -> Result {
92        let Self { privateKey } = self;
93        let wallet = parse_wallet(privateKey)?;
94        let address = inject_wallet(state, wallet);
95        Ok(address.abi_encode())
96    }
97}
98
99impl Cheatcode for rememberKeys_0Call {
100    fn apply(&self, state: &mut Cheatcodes) -> Result {
101        let Self { mnemonic, derivationPath, count } = self;
102        let wallets = derive_wallets::<English>(mnemonic, derivationPath, *count)?;
103        let mut addresses = Vec::<Address>::with_capacity(wallets.len());
104        for wallet in wallets {
105            let addr = inject_wallet(state, wallet);
106            addresses.push(addr);
107        }
108
109        Ok(addresses.abi_encode())
110    }
111}
112
113impl Cheatcode for rememberKeys_1Call {
114    fn apply(&self, state: &mut Cheatcodes) -> Result {
115        let Self { mnemonic, derivationPath, language, count } = self;
116        let wallets = derive_wallets_str(mnemonic, derivationPath, language, *count)?;
117        let mut addresses = Vec::<Address>::with_capacity(wallets.len());
118        for wallet in wallets {
119            let addr = inject_wallet(state, wallet);
120            addresses.push(addr);
121        }
122
123        Ok(addresses.abi_encode())
124    }
125}
126
127fn inject_wallet(state: &mut Cheatcodes, wallet: LocalSigner<SigningKey>) -> Address {
128    let address = wallet.address();
129    state.wallets().add_local_signer(wallet);
130    address
131}
132
133impl Cheatcode for sign_1Call {
134    fn apply(&self, _state: &mut Cheatcodes) -> Result {
135        let Self { privateKey, digest } = self;
136        let sig = sign(privateKey, digest)?;
137        Ok(encode_full_sig(sig))
138    }
139}
140
141impl Cheatcode for signCompact_1Call {
142    fn apply(&self, _state: &mut Cheatcodes) -> Result {
143        let Self { privateKey, digest } = self;
144        let sig = sign(privateKey, digest)?;
145        Ok(encode_compact_sig(sig))
146    }
147}
148
149impl Cheatcode for sign_2Call {
150    fn apply(&self, state: &mut Cheatcodes) -> Result {
151        let Self { digest } = self;
152        let sig = sign_with_wallet(state, None, digest)?;
153        Ok(encode_full_sig(sig))
154    }
155}
156
157impl Cheatcode for signCompact_2Call {
158    fn apply(&self, state: &mut Cheatcodes) -> Result {
159        let Self { digest } = self;
160        let sig = sign_with_wallet(state, None, digest)?;
161        Ok(encode_compact_sig(sig))
162    }
163}
164
165impl Cheatcode for sign_3Call {
166    fn apply(&self, state: &mut Cheatcodes) -> Result {
167        let Self { signer, digest } = self;
168        let sig = sign_with_wallet(state, Some(*signer), digest)?;
169        Ok(encode_full_sig(sig))
170    }
171}
172
173impl Cheatcode for signCompact_3Call {
174    fn apply(&self, state: &mut Cheatcodes) -> Result {
175        let Self { signer, digest } = self;
176        let sig = sign_with_wallet(state, Some(*signer), digest)?;
177        Ok(encode_compact_sig(sig))
178    }
179}
180
181impl Cheatcode for signP256Call {
182    fn apply(&self, _state: &mut Cheatcodes) -> Result {
183        let Self { privateKey, digest } = self;
184        sign_p256(privateKey, digest)
185    }
186}
187
188impl Cheatcode for publicKeyP256Call {
189    fn apply(&self, _state: &mut Cheatcodes) -> Result {
190        let Self { privateKey } = self;
191        let pub_key =
192            parse_private_key_p256(privateKey)?.verifying_key().as_affine().to_encoded_point(false);
193        let pub_key_x = U256::from_be_bytes((*pub_key.x().unwrap()).into());
194        let pub_key_y = U256::from_be_bytes((*pub_key.y().unwrap()).into());
195
196        Ok((pub_key_x, pub_key_y).abi_encode())
197    }
198}
199
200/// Using a given private key, return its public ETH address, its public key affine x and y
201/// coordinates, and its private key (see the 'Wallet' struct)
202///
203/// If 'label' is set to 'Some()', assign that label to the associated ETH address in state
204fn create_wallet(private_key: &U256, label: Option<&str>, state: &mut Cheatcodes) -> Result {
205    let key = parse_private_key(private_key)?;
206    let addr = alloy_signer::utils::secret_key_to_address(&key);
207
208    let pub_key = key.verifying_key().as_affine().to_encoded_point(false);
209    let pub_key_x = U256::from_be_bytes((*pub_key.x().unwrap()).into());
210    let pub_key_y = U256::from_be_bytes((*pub_key.y().unwrap()).into());
211
212    if let Some(label) = label {
213        state.labels.insert(addr, label.into());
214    }
215
216    Ok(Wallet { addr, publicKeyX: pub_key_x, publicKeyY: pub_key_y, privateKey: *private_key }
217        .abi_encode())
218}
219
220fn encode_full_sig(sig: alloy_primitives::PrimitiveSignature) -> Vec<u8> {
221    // Retrieve v, r and s from signature.
222    let v = U256::from(sig.v() as u64 + 27);
223    let r = B256::from(sig.r());
224    let s = B256::from(sig.s());
225    (v, r, s).abi_encode()
226}
227
228fn encode_compact_sig(sig: alloy_primitives::PrimitiveSignature) -> Vec<u8> {
229    // Implement EIP-2098 compact signature.
230    let r = B256::from(sig.r());
231    let mut vs = sig.s();
232    vs.set_bit(255, sig.v());
233    (r, vs).abi_encode()
234}
235
236fn sign(private_key: &U256, digest: &B256) -> Result<alloy_primitives::PrimitiveSignature> {
237    // The `ecrecover` precompile does not use EIP-155. No chain ID is needed.
238    let wallet = parse_wallet(private_key)?;
239    let sig = wallet.sign_hash_sync(digest)?;
240    debug_assert_eq!(sig.recover_address_from_prehash(digest)?, wallet.address());
241    Ok(sig)
242}
243
244fn sign_with_wallet(
245    state: &mut Cheatcodes,
246    signer: Option<Address>,
247    digest: &B256,
248) -> Result<alloy_primitives::PrimitiveSignature> {
249    if state.wallets().is_empty() {
250        bail!("no wallets available");
251    }
252
253    let mut wallets = state.wallets().inner.lock();
254    let maybe_provided_sender = wallets.provided_sender;
255    let signers = wallets.multi_wallet.signers()?;
256
257    let signer = if let Some(signer) = signer {
258        signer
259    } else if let Some(provided_sender) = maybe_provided_sender {
260        provided_sender
261    } else if signers.len() == 1 {
262        *signers.keys().next().unwrap()
263    } else {
264        bail!("could not determine signer, there are multiple signers available use vm.sign(signer, digest) to specify one");
265    };
266
267    let wallet = signers
268        .get(&signer)
269        .ok_or_else(|| fmt_err!("signer with address {signer} is not available"))?;
270
271    let sig = foundry_common::block_on(wallet.sign_hash(digest))?;
272    debug_assert_eq!(sig.recover_address_from_prehash(digest)?, signer);
273    Ok(sig)
274}
275
276fn sign_p256(private_key: &U256, digest: &B256) -> Result {
277    let signing_key = parse_private_key_p256(private_key)?;
278    let signature: P256Signature = signing_key.sign_prehash(digest.as_slice())?;
279    let signature = signature.normalize_s().unwrap_or(signature);
280    let r_bytes: [u8; 32] = signature.r().to_bytes().into();
281    let s_bytes: [u8; 32] = signature.s().to_bytes().into();
282
283    Ok((r_bytes, s_bytes).abi_encode())
284}
285
286fn validate_private_key<C: ecdsa::PrimeCurve>(private_key: &U256) -> Result<()> {
287    ensure!(*private_key != U256::ZERO, "private key cannot be 0");
288    let order = U256::from_be_slice(&C::ORDER.to_be_byte_array());
289    ensure!(
290        *private_key < U256::from_be_slice(&C::ORDER.to_be_byte_array()),
291        "private key must be less than the {curve:?} curve order ({order})",
292        curve = C::default(),
293    );
294
295    Ok(())
296}
297
298fn parse_private_key(private_key: &U256) -> Result<SigningKey> {
299    validate_private_key::<k256::Secp256k1>(private_key)?;
300    Ok(SigningKey::from_bytes((&private_key.to_be_bytes()).into())?)
301}
302
303fn parse_private_key_p256(private_key: &U256) -> Result<P256SigningKey> {
304    validate_private_key::<p256::NistP256>(private_key)?;
305    Ok(P256SigningKey::from_bytes((&private_key.to_be_bytes()).into())?)
306}
307
308pub(super) fn parse_wallet(private_key: &U256) -> Result<PrivateKeySigner> {
309    parse_private_key(private_key).map(PrivateKeySigner::from)
310}
311
312fn derive_key_str(mnemonic: &str, path: &str, index: u32, language: &str) -> Result {
313    match language {
314        "chinese_simplified" => derive_key::<ChineseSimplified>(mnemonic, path, index),
315        "chinese_traditional" => derive_key::<ChineseTraditional>(mnemonic, path, index),
316        "czech" => derive_key::<Czech>(mnemonic, path, index),
317        "english" => derive_key::<English>(mnemonic, path, index),
318        "french" => derive_key::<French>(mnemonic, path, index),
319        "italian" => derive_key::<Italian>(mnemonic, path, index),
320        "japanese" => derive_key::<Japanese>(mnemonic, path, index),
321        "korean" => derive_key::<Korean>(mnemonic, path, index),
322        "portuguese" => derive_key::<Portuguese>(mnemonic, path, index),
323        "spanish" => derive_key::<Spanish>(mnemonic, path, index),
324        _ => Err(fmt_err!("unsupported mnemonic language: {language:?}")),
325    }
326}
327
328fn derive_key<W: Wordlist>(mnemonic: &str, path: &str, index: u32) -> Result {
329    fn derive_key_path(path: &str, index: u32) -> String {
330        let mut out = path.to_string();
331        if !out.ends_with('/') {
332            out.push('/');
333        }
334        out.push_str(&index.to_string());
335        out
336    }
337
338    let wallet = MnemonicBuilder::<W>::default()
339        .phrase(mnemonic)
340        .derivation_path(derive_key_path(path, index))?
341        .build()?;
342    let private_key = U256::from_be_bytes(wallet.credential().to_bytes().into());
343    Ok(private_key.abi_encode())
344}
345
346fn derive_wallets_str(
347    mnemonic: &str,
348    path: &str,
349    language: &str,
350    count: u32,
351) -> Result<Vec<LocalSigner<SigningKey>>> {
352    match language {
353        "chinese_simplified" => derive_wallets::<ChineseSimplified>(mnemonic, path, count),
354        "chinese_traditional" => derive_wallets::<ChineseTraditional>(mnemonic, path, count),
355        "czech" => derive_wallets::<Czech>(mnemonic, path, count),
356        "english" => derive_wallets::<English>(mnemonic, path, count),
357        "french" => derive_wallets::<French>(mnemonic, path, count),
358        "italian" => derive_wallets::<Italian>(mnemonic, path, count),
359        "japanese" => derive_wallets::<Japanese>(mnemonic, path, count),
360        "korean" => derive_wallets::<Korean>(mnemonic, path, count),
361        "portuguese" => derive_wallets::<Portuguese>(mnemonic, path, count),
362        "spanish" => derive_wallets::<Spanish>(mnemonic, path, count),
363        _ => Err(fmt_err!("unsupported mnemonic language: {language:?}")),
364    }
365}
366
367fn derive_wallets<W: Wordlist>(
368    mnemonic: &str,
369    path: &str,
370    count: u32,
371) -> Result<Vec<LocalSigner<SigningKey>>> {
372    let mut out = path.to_string();
373
374    if !out.ends_with('/') {
375        out.push('/');
376    }
377
378    let mut wallets = Vec::with_capacity(count as usize);
379    for idx in 0..count {
380        let wallet = MnemonicBuilder::<W>::default()
381            .phrase(mnemonic)
382            .derivation_path(format!("{out}{idx}"))?
383            .build()?;
384        wallets.push(wallet);
385    }
386
387    Ok(wallets)
388}
389
390#[cfg(test)]
391mod tests {
392    use super::*;
393    use alloy_primitives::{hex::FromHex, FixedBytes};
394    use p256::ecdsa::signature::hazmat::PrehashVerifier;
395
396    #[test]
397    fn test_sign_p256() {
398        use p256::ecdsa::VerifyingKey;
399
400        let pk_u256: U256 = "1".parse().unwrap();
401        let signing_key = P256SigningKey::from_bytes(&pk_u256.to_be_bytes().into()).unwrap();
402        let digest = FixedBytes::from_hex(
403            "0x44acf6b7e36c1342c2c5897204fe09504e1e2efb1a900377dbc4e7a6a133ec56",
404        )
405        .unwrap();
406
407        let result = sign_p256(&pk_u256, &digest).unwrap();
408        let result_bytes: [u8; 64] = result.try_into().unwrap();
409        let signature = P256Signature::from_bytes(&result_bytes.into()).unwrap();
410        let verifying_key = VerifyingKey::from(&signing_key);
411        assert!(verifying_key.verify_prehash(digest.as_slice(), &signature).is_ok());
412    }
413
414    #[test]
415    fn test_sign_p256_pk_too_large() {
416        // max n from https://neuromancer.sk/std/secg/secp256r1
417        let pk =
418            "0xffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551".parse().unwrap();
419        let digest = FixedBytes::from_hex(
420            "0x54705ba3baafdbdfba8c5f9a70f7a89bee98d906b53e31074da7baecdc0da9ad",
421        )
422        .unwrap();
423        let result = sign_p256(&pk, &digest);
424        assert_eq!(result.err().unwrap().to_string(), "private key must be less than the NistP256 curve order (115792089210356248762697446949407573529996955224135760342422259061068512044369)");
425    }
426
427    #[test]
428    fn test_sign_p256_pk_0() {
429        let digest = FixedBytes::from_hex(
430            "0x54705ba3baafdbdfba8c5f9a70f7a89bee98d906b53e31074da7baecdc0da9ad",
431        )
432        .unwrap();
433        let result = sign_p256(&U256::ZERO, &digest);
434        assert_eq!(result.err().unwrap().to_string(), "private key cannot be 0");
435    }
436}
foundry_cheatcodes/crypto.rs

foundry_cheatcodes/
crypto.rs
