1use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
4use alloy_primitives::{Address, B256, U256, keccak256};
5use alloy_signer::{Signer, SignerSync};
6use alloy_signer_local::{
7 LocalSigner, MnemonicBuilder, PrivateKeySigner,
8 coins_bip39::{
9 ChineseSimplified, ChineseTraditional, Czech, English, French, Italian, Japanese, Korean,
10 Portuguese, Spanish, Wordlist,
11 },
12};
13use alloy_sol_types::SolValue;
14use k256::{
15 FieldBytes, Scalar,
16 ecdsa::{SigningKey, hazmat},
17 elliptic_curve::{bigint::ArrayEncoding, sec1::ToEncodedPoint},
18};
19
20use p256::ecdsa::{
21 Signature as P256Signature, SigningKey as P256SigningKey, signature::hazmat::PrehashSigner,
22};
23
24const DEFAULT_DERIVATION_PATH_PREFIX: &str = "m/44'/60'/0'/0/";
26
27impl Cheatcode for createWallet_0Call {
28 fn apply(&self, state: &mut Cheatcodes) -> Result {
29 let Self { walletLabel } = self;
30 create_wallet(&U256::from_be_bytes(keccak256(walletLabel).0), Some(walletLabel), state)
31 }
32}
33
34impl Cheatcode for createWallet_1Call {
35 fn apply(&self, state: &mut Cheatcodes) -> Result {
36 let Self { privateKey } = self;
37 create_wallet(privateKey, None, state)
38 }
39}
40
41impl Cheatcode for createWallet_2Call {
42 fn apply(&self, state: &mut Cheatcodes) -> Result {
43 let Self { privateKey, walletLabel } = self;
44 create_wallet(privateKey, Some(walletLabel), state)
45 }
46}
47
48impl Cheatcode for sign_0Call {
49 fn apply(&self, _state: &mut Cheatcodes) -> Result {
50 let Self { wallet, digest } = self;
51 let sig = sign(&wallet.privateKey, digest)?;
52 Ok(encode_full_sig(sig))
53 }
54}
55
56impl Cheatcode for signWithNonceUnsafeCall {
57 fn apply(&self, _state: &mut Cheatcodes) -> Result {
58 let pk: U256 = self.privateKey;
59 let digest: B256 = self.digest;
60 let nonce: U256 = self.nonce;
61 let sig: alloy_primitives::Signature = sign_with_nonce(&pk, &digest, &nonce)?;
62 Ok(encode_full_sig(sig))
63 }
64}
65
66impl Cheatcode for signCompact_0Call {
67 fn apply(&self, _state: &mut Cheatcodes) -> Result {
68 let Self { wallet, digest } = self;
69 let sig = sign(&wallet.privateKey, digest)?;
70 Ok(encode_compact_sig(sig))
71 }
72}
73
74impl Cheatcode for deriveKey_0Call {
75 fn apply(&self, _state: &mut Cheatcodes) -> Result {
76 let Self { mnemonic, index } = self;
77 derive_key::<English>(mnemonic, DEFAULT_DERIVATION_PATH_PREFIX, *index)
78 }
79}
80
81impl Cheatcode for deriveKey_1Call {
82 fn apply(&self, _state: &mut Cheatcodes) -> Result {
83 let Self { mnemonic, derivationPath, index } = self;
84 derive_key::<English>(mnemonic, derivationPath, *index)
85 }
86}
87
88impl Cheatcode for deriveKey_2Call {
89 fn apply(&self, _state: &mut Cheatcodes) -> Result {
90 let Self { mnemonic, index, language } = self;
91 derive_key_str(mnemonic, DEFAULT_DERIVATION_PATH_PREFIX, *index, language)
92 }
93}
94
95impl Cheatcode for deriveKey_3Call {
96 fn apply(&self, _state: &mut Cheatcodes) -> Result {
97 let Self { mnemonic, derivationPath, index, language } = self;
98 derive_key_str(mnemonic, derivationPath, *index, language)
99 }
100}
101
102impl Cheatcode for rememberKeyCall {
103 fn apply(&self, state: &mut Cheatcodes) -> Result {
104 let Self { privateKey } = self;
105 let wallet = parse_wallet(privateKey)?;
106 let address = inject_wallet(state, wallet);
107 Ok(address.abi_encode())
108 }
109}
110
111impl Cheatcode for rememberKeys_0Call {
112 fn apply(&self, state: &mut Cheatcodes) -> Result {
113 let Self { mnemonic, derivationPath, count } = self;
114 let wallets = derive_wallets::<English>(mnemonic, derivationPath, *count)?;
115 let mut addresses = Vec::<Address>::with_capacity(wallets.len());
116 for wallet in wallets {
117 let addr = inject_wallet(state, wallet);
118 addresses.push(addr);
119 }
120
121 Ok(addresses.abi_encode())
122 }
123}
124
125impl Cheatcode for rememberKeys_1Call {
126 fn apply(&self, state: &mut Cheatcodes) -> Result {
127 let Self { mnemonic, derivationPath, language, count } = self;
128 let wallets = derive_wallets_str(mnemonic, derivationPath, language, *count)?;
129 let mut addresses = Vec::<Address>::with_capacity(wallets.len());
130 for wallet in wallets {
131 let addr = inject_wallet(state, wallet);
132 addresses.push(addr);
133 }
134
135 Ok(addresses.abi_encode())
136 }
137}
138
139fn inject_wallet(state: &mut Cheatcodes, wallet: LocalSigner<SigningKey>) -> Address {
140 let address = wallet.address();
141 state.wallets().add_local_signer(wallet);
142 address
143}
144
145impl Cheatcode for sign_1Call {
146 fn apply(&self, _state: &mut Cheatcodes) -> Result {
147 let Self { privateKey, digest } = self;
148 let sig = sign(privateKey, digest)?;
149 Ok(encode_full_sig(sig))
150 }
151}
152
153impl Cheatcode for signCompact_1Call {
154 fn apply(&self, _state: &mut Cheatcodes) -> Result {
155 let Self { privateKey, digest } = self;
156 let sig = sign(privateKey, digest)?;
157 Ok(encode_compact_sig(sig))
158 }
159}
160
161impl Cheatcode for sign_2Call {
162 fn apply(&self, state: &mut Cheatcodes) -> Result {
163 let Self { digest } = self;
164 let sig = sign_with_wallet(state, None, digest)?;
165 Ok(encode_full_sig(sig))
166 }
167}
168
169impl Cheatcode for signCompact_2Call {
170 fn apply(&self, state: &mut Cheatcodes) -> Result {
171 let Self { digest } = self;
172 let sig = sign_with_wallet(state, None, digest)?;
173 Ok(encode_compact_sig(sig))
174 }
175}
176
177impl Cheatcode for sign_3Call {
178 fn apply(&self, state: &mut Cheatcodes) -> Result {
179 let Self { signer, digest } = self;
180 let sig = sign_with_wallet(state, Some(*signer), digest)?;
181 Ok(encode_full_sig(sig))
182 }
183}
184
185impl Cheatcode for signCompact_3Call {
186 fn apply(&self, state: &mut Cheatcodes) -> Result {
187 let Self { signer, digest } = self;
188 let sig = sign_with_wallet(state, Some(*signer), digest)?;
189 Ok(encode_compact_sig(sig))
190 }
191}
192
193impl Cheatcode for signP256Call {
194 fn apply(&self, _state: &mut Cheatcodes) -> Result {
195 let Self { privateKey, digest } = self;
196 sign_p256(privateKey, digest)
197 }
198}
199
200impl Cheatcode for publicKeyP256Call {
201 fn apply(&self, _state: &mut Cheatcodes) -> Result {
202 let Self { privateKey } = self;
203 let pub_key =
204 parse_private_key_p256(privateKey)?.verifying_key().as_affine().to_encoded_point(false);
205 let pub_key_x = U256::from_be_bytes((*pub_key.x().unwrap()).into());
206 let pub_key_y = U256::from_be_bytes((*pub_key.y().unwrap()).into());
207
208 Ok((pub_key_x, pub_key_y).abi_encode())
209 }
210}
211
212fn create_wallet(private_key: &U256, label: Option<&str>, state: &mut Cheatcodes) -> Result {
217 let key = parse_private_key(private_key)?;
218 let addr = alloy_signer::utils::secret_key_to_address(&key);
219
220 let pub_key = key.verifying_key().as_affine().to_encoded_point(false);
221 let pub_key_x = U256::from_be_bytes((*pub_key.x().unwrap()).into());
222 let pub_key_y = U256::from_be_bytes((*pub_key.y().unwrap()).into());
223
224 if let Some(label) = label {
225 state.labels.insert(addr, label.into());
226 }
227
228 Ok(Wallet { addr, publicKeyX: pub_key_x, publicKeyY: pub_key_y, privateKey: *private_key }
229 .abi_encode())
230}
231
232fn encode_full_sig(sig: alloy_primitives::Signature) -> Vec<u8> {
233 let v = U256::from(sig.v() as u64 + 27);
235 let r = B256::from(sig.r());
236 let s = B256::from(sig.s());
237 (v, r, s).abi_encode()
238}
239
240fn encode_compact_sig(sig: alloy_primitives::Signature) -> Vec<u8> {
241 let r = B256::from(sig.r());
243 let mut vs = sig.s();
244 vs.set_bit(255, sig.v());
245 (r, vs).abi_encode()
246}
247
248fn sign(private_key: &U256, digest: &B256) -> Result<alloy_primitives::Signature> {
249 let wallet = parse_wallet(private_key)?;
251 let sig = wallet.sign_hash_sync(digest)?;
252 debug_assert_eq!(sig.recover_address_from_prehash(digest)?, wallet.address());
253 Ok(sig)
254}
255
256fn sign_with_nonce(
267 private_key: &U256,
268 digest: &B256,
269 nonce: &U256,
270) -> Result<alloy_primitives::Signature> {
271 let d_scalar: Scalar =
272 <Scalar as k256::elliptic_curve::PrimeField>::from_repr(private_key.to_be_bytes().into())
273 .into_option()
274 .ok_or_else(|| fmt_err!("invalid private key scalar"))?;
275 if bool::from(d_scalar.is_zero()) {
276 return Err(fmt_err!("private key cannot be 0"));
277 }
278
279 let k_scalar: Scalar =
280 <Scalar as k256::elliptic_curve::PrimeField>::from_repr(nonce.to_be_bytes().into())
281 .into_option()
282 .ok_or_else(|| fmt_err!("invalid nonce scalar"))?;
283 if bool::from(k_scalar.is_zero()) {
284 return Err(fmt_err!("nonce cannot be 0"));
285 }
286
287 let mut z = [0u8; 32];
288 z.copy_from_slice(digest.as_slice());
289 let z_fb: FieldBytes = FieldBytes::from(z);
290
291 let (sig_raw, recid_opt) =
294 <Scalar as hazmat::SignPrimitive<k256::Secp256k1>>::try_sign_prehashed(
295 &d_scalar, k_scalar, &z_fb,
296 )
297 .map_err(|e| fmt_err!("sign_prehashed failed: {e}"))?;
298
299 let (sig_low, flipped) =
301 if let Some(norm) = sig_raw.normalize_s() { (norm, true) } else { (sig_raw, false) };
302
303 let r_u256 = U256::from_be_bytes(sig_low.r().to_bytes().into());
304 let s_u256 = U256::from_be_bytes(sig_low.s().to_bytes().into());
305
306 let v_parity = if let Some(id) = recid_opt {
308 let mut v = id.to_byte() & 1;
309 if flipped {
310 v ^= 1;
311 }
312 v
313 } else {
314 let expected_addr = {
316 let sk: SigningKey = parse_private_key(private_key)?;
317 alloy_signer::utils::secret_key_to_address(&sk)
318 };
319 let cand0 = alloy_primitives::Signature::new(r_u256, s_u256, false);
321 if cand0.recover_address_from_prehash(digest).ok() == Some(expected_addr) {
322 return Ok(cand0);
323 }
324 let cand1 = alloy_primitives::Signature::new(r_u256, s_u256, true);
326 if cand1.recover_address_from_prehash(digest).ok() == Some(expected_addr) {
327 return Ok(cand1);
328 }
329 return Err(fmt_err!("failed to determine recovery id for signature"));
330 };
331
332 let y_parity = v_parity != 0;
333 Ok(alloy_primitives::Signature::new(r_u256, s_u256, y_parity))
334}
335
336fn sign_with_wallet(
337 state: &mut Cheatcodes,
338 signer: Option<Address>,
339 digest: &B256,
340) -> Result<alloy_primitives::Signature> {
341 if state.wallets().is_empty() {
342 bail!("no wallets available");
343 }
344
345 let mut wallets = state.wallets().inner.lock();
346 let maybe_provided_sender = wallets.provided_sender;
347 let signers = wallets.multi_wallet.signers()?;
348
349 let signer = if let Some(signer) = signer {
350 signer
351 } else if let Some(provided_sender) = maybe_provided_sender {
352 provided_sender
353 } else if signers.len() == 1 {
354 *signers.keys().next().unwrap()
355 } else {
356 bail!(
357 "could not determine signer, there are multiple signers available use vm.sign(signer, digest) to specify one"
358 );
359 };
360
361 let wallet = signers
362 .get(&signer)
363 .ok_or_else(|| fmt_err!("signer with address {signer} is not available"))?;
364
365 let sig = foundry_common::block_on(wallet.sign_hash(digest))?;
366 debug_assert_eq!(sig.recover_address_from_prehash(digest)?, signer);
367 Ok(sig)
368}
369
370fn sign_p256(private_key: &U256, digest: &B256) -> Result {
371 let signing_key = parse_private_key_p256(private_key)?;
372 let signature: P256Signature = signing_key.sign_prehash(digest.as_slice())?;
373 let signature = signature.normalize_s().unwrap_or(signature);
374 let r_bytes: [u8; 32] = signature.r().to_bytes().into();
375 let s_bytes: [u8; 32] = signature.s().to_bytes().into();
376
377 Ok((r_bytes, s_bytes).abi_encode())
378}
379
380fn validate_private_key<C: ecdsa::PrimeCurve>(private_key: &U256) -> Result<()> {
381 ensure!(*private_key != U256::ZERO, "private key cannot be 0");
382 let order = U256::from_be_slice(&C::ORDER.to_be_byte_array());
383 ensure!(
384 *private_key < order,
385 "private key must be less than the {curve:?} curve order ({order})",
386 curve = C::default(),
387 );
388
389 Ok(())
390}
391
392fn parse_private_key(private_key: &U256) -> Result<SigningKey> {
393 validate_private_key::<k256::Secp256k1>(private_key)?;
394 Ok(SigningKey::from_bytes((&private_key.to_be_bytes()).into())?)
395}
396
397fn parse_private_key_p256(private_key: &U256) -> Result<P256SigningKey> {
398 validate_private_key::<p256::NistP256>(private_key)?;
399 Ok(P256SigningKey::from_bytes((&private_key.to_be_bytes()).into())?)
400}
401
402pub(super) fn parse_wallet(private_key: &U256) -> Result<PrivateKeySigner> {
403 parse_private_key(private_key).map(PrivateKeySigner::from)
404}
405
406fn derive_key_str(mnemonic: &str, path: &str, index: u32, language: &str) -> Result {
407 match language {
408 "chinese_simplified" => derive_key::<ChineseSimplified>(mnemonic, path, index),
409 "chinese_traditional" => derive_key::<ChineseTraditional>(mnemonic, path, index),
410 "czech" => derive_key::<Czech>(mnemonic, path, index),
411 "english" => derive_key::<English>(mnemonic, path, index),
412 "french" => derive_key::<French>(mnemonic, path, index),
413 "italian" => derive_key::<Italian>(mnemonic, path, index),
414 "japanese" => derive_key::<Japanese>(mnemonic, path, index),
415 "korean" => derive_key::<Korean>(mnemonic, path, index),
416 "portuguese" => derive_key::<Portuguese>(mnemonic, path, index),
417 "spanish" => derive_key::<Spanish>(mnemonic, path, index),
418 _ => Err(fmt_err!("unsupported mnemonic language: {language:?}")),
419 }
420}
421
422fn derive_key<W: Wordlist>(mnemonic: &str, path: &str, index: u32) -> Result {
423 fn derive_key_path(path: &str, index: u32) -> String {
424 let mut out = path.to_string();
425 if !out.ends_with('/') {
426 out.push('/');
427 }
428 out.push_str(&index.to_string());
429 out
430 }
431
432 let wallet = MnemonicBuilder::<W>::default()
433 .phrase(mnemonic)
434 .derivation_path(derive_key_path(path, index))?
435 .build()?;
436 let private_key = U256::from_be_bytes(wallet.credential().to_bytes().into());
437 Ok(private_key.abi_encode())
438}
439
440fn derive_wallets_str(
441 mnemonic: &str,
442 path: &str,
443 language: &str,
444 count: u32,
445) -> Result<Vec<LocalSigner<SigningKey>>> {
446 match language {
447 "chinese_simplified" => derive_wallets::<ChineseSimplified>(mnemonic, path, count),
448 "chinese_traditional" => derive_wallets::<ChineseTraditional>(mnemonic, path, count),
449 "czech" => derive_wallets::<Czech>(mnemonic, path, count),
450 "english" => derive_wallets::<English>(mnemonic, path, count),
451 "french" => derive_wallets::<French>(mnemonic, path, count),
452 "italian" => derive_wallets::<Italian>(mnemonic, path, count),
453 "japanese" => derive_wallets::<Japanese>(mnemonic, path, count),
454 "korean" => derive_wallets::<Korean>(mnemonic, path, count),
455 "portuguese" => derive_wallets::<Portuguese>(mnemonic, path, count),
456 "spanish" => derive_wallets::<Spanish>(mnemonic, path, count),
457 _ => Err(fmt_err!("unsupported mnemonic language: {language:?}")),
458 }
459}
460
461fn derive_wallets<W: Wordlist>(
462 mnemonic: &str,
463 path: &str,
464 count: u32,
465) -> Result<Vec<LocalSigner<SigningKey>>> {
466 let mut out = path.to_string();
467
468 if !out.ends_with('/') {
469 out.push('/');
470 }
471
472 let mut wallets = Vec::with_capacity(count as usize);
473 for idx in 0..count {
474 let wallet = MnemonicBuilder::<W>::default()
475 .phrase(mnemonic)
476 .derivation_path(format!("{out}{idx}"))?
477 .build()?;
478 wallets.push(wallet);
479 }
480
481 Ok(wallets)
482}
483
484#[cfg(test)]
485mod tests {
486 use super::*;
487 use alloy_primitives::{FixedBytes, hex::FromHex};
488 use k256::elliptic_curve::Curve;
489 use p256::ecdsa::signature::hazmat::PrehashVerifier;
490
491 #[test]
492 fn test_sign_p256() {
493 use p256::ecdsa::VerifyingKey;
494
495 let pk_u256: U256 = "1".parse().unwrap();
496 let signing_key = P256SigningKey::from_bytes(&pk_u256.to_be_bytes().into()).unwrap();
497 let digest = FixedBytes::from_hex(
498 "0x44acf6b7e36c1342c2c5897204fe09504e1e2efb1a900377dbc4e7a6a133ec56",
499 )
500 .unwrap();
501
502 let result = sign_p256(&pk_u256, &digest).unwrap();
503 let result_bytes: [u8; 64] = result.try_into().unwrap();
504 let signature = P256Signature::from_bytes(&result_bytes.into()).unwrap();
505 let verifying_key = VerifyingKey::from(&signing_key);
506 assert!(verifying_key.verify_prehash(digest.as_slice(), &signature).is_ok());
507 }
508
509 #[test]
510 fn test_sign_p256_pk_too_large() {
511 let pk =
513 "0xffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551".parse().unwrap();
514 let digest = FixedBytes::from_hex(
515 "0x54705ba3baafdbdfba8c5f9a70f7a89bee98d906b53e31074da7baecdc0da9ad",
516 )
517 .unwrap();
518 let result = sign_p256(&pk, &digest);
519 assert_eq!(
520 result.err().unwrap().to_string(),
521 "private key must be less than the NistP256 curve order (115792089210356248762697446949407573529996955224135760342422259061068512044369)"
522 );
523 }
524
525 #[test]
526 fn test_sign_p256_pk_0() {
527 let digest = FixedBytes::from_hex(
528 "0x54705ba3baafdbdfba8c5f9a70f7a89bee98d906b53e31074da7baecdc0da9ad",
529 )
530 .unwrap();
531 let result = sign_p256(&U256::ZERO, &digest);
532 assert_eq!(result.err().unwrap().to_string(), "private key cannot be 0");
533 }
534
535 #[test]
536 fn test_sign_with_nonce_varies_and_recovers() {
537 let pk_u256: U256 = U256::from(1u64);
539 let digest = FixedBytes::from_hex(
540 "0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
541 )
542 .unwrap();
543
544 let n1: U256 = U256::from(123u64);
546 let n2: U256 = U256::from(456u64);
547
548 let sig1 = sign_with_nonce(&pk_u256, &digest, &n1).expect("sig1");
550 let sig2 = sign_with_nonce(&pk_u256, &digest, &n2).expect("sig2");
551
552 assert!(
554 sig1.r() != sig2.r() || sig1.s() != sig2.s(),
555 "signatures should differ with different nonces"
556 );
557
558 let sk = parse_private_key(&pk_u256).unwrap();
560 let expected = alloy_signer::utils::secret_key_to_address(&sk);
561
562 assert_eq!(sig1.recover_address_from_prehash(&digest).unwrap(), expected);
563 assert_eq!(sig2.recover_address_from_prehash(&digest).unwrap(), expected);
564 }
565
566 #[test]
567 fn test_sign_with_nonce_zero_nonce_errors() {
568 let pk_u256: U256 = U256::from(1u64);
570 let digest = FixedBytes::from_hex(
571 "0xbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
572 )
573 .unwrap();
574 let n0: U256 = U256::ZERO;
575
576 let err = sign_with_nonce(&pk_u256, &digest, &n0).unwrap_err();
577 let msg = err.to_string();
578 assert!(msg.contains("nonce cannot be 0"), "unexpected error: {msg}");
579 }
580
581 #[test]
582 fn test_sign_with_nonce_nonce_ge_order_errors() {
583 use k256::Secp256k1;
585 let n_u256 = U256::from_be_slice(&Secp256k1::ORDER.to_be_byte_array());
587
588 let pk_u256: U256 = U256::from(1u64);
589 let digest = FixedBytes::from_hex(
590 "0xcccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc",
591 )
592 .unwrap();
593
594 let err = sign_with_nonce(&pk_u256, &digest, &n_u256).unwrap_err();
596 let msg = err.to_string();
597 assert!(msg.contains("invalid nonce scalar"), "unexpected error: {msg}");
598 }
599}