1use crate::{Cheatcode, CheatcodesCore, Result, Vm::*};
4use alloy_primitives::{Address, B256, U256, keccak256};
5use alloy_signer::{Signer, SignerSync};
6use alloy_signer_local::{
7 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 foundry_common::wallet::{derive_private_key, derive_private_key_with_language};
15use k256::{
16 AffinePoint, EncodedPoint, FieldBytes, FieldElement, ProjectivePoint, Scalar,
17 ecdsa::{SigningKey, hazmat},
18 elliptic_curve::{
19 bigint::{ArrayEncoding, U256 as K256U256},
20 group::Group,
21 ops::Reduce,
22 sec1::{FromEncodedPoint, ToEncodedPoint},
23 },
24};
25
26use p256::ecdsa::{
27 Signature as P256Signature, SigningKey as P256SigningKey, signature::hazmat::PrehashSigner,
28};
29
30use ed25519_consensus::{
31 Signature as Ed25519Signature, SigningKey as Ed25519SigningKey,
32 VerificationKey as Ed25519VerificationKey,
33};
34use tempo_primitives::transaction::{KeychainSignature, PrimitiveSignature, TempoSignature};
35
36const DEFAULT_DERIVATION_PATH_PREFIX: &str = "m/44'/60'/0'/0/";
38const PRIVATE_KEY_SIGNER_CACHE_LIMIT: usize = 64;
39
40impl Cheatcode for createWallet_0Call {
41 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
42 let Self { walletLabel } = self;
43 create_wallet(&keccak256(walletLabel).into(), Some(walletLabel), state)
44 }
45}
46
47impl Cheatcode for createWallet_1Call {
48 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
49 let Self { privateKey } = self;
50 create_wallet(privateKey, None, state)
51 }
52}
53
54impl Cheatcode for createWallet_2Call {
55 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
56 let Self { privateKey, walletLabel } = self;
57 create_wallet(privateKey, Some(walletLabel), state)
58 }
59}
60
61impl Cheatcode for sign_0Call {
62 fn apply_context_free(&self) -> Result {
63 let Self { wallet, digest } = self;
64 let sig = sign(&wallet.privateKey, digest)?;
65 Ok(encode_full_sig(sig))
66 }
67}
68
69impl Cheatcode for signWithNonceUnsafeCall {
70 fn apply_context_free(&self) -> Result {
71 let pk: U256 = self.privateKey;
72 let digest: B256 = self.digest;
73 let nonce: U256 = self.nonce;
74 let sig: alloy_primitives::Signature = sign_with_nonce(&pk, &digest, &nonce)?;
75 Ok(encode_full_sig(sig))
76 }
77}
78
79impl Cheatcode for signKeychainCall {
80 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
81 let Self { privateKey, account, digest } = self;
82 sign_keychain(state, privateKey, account, digest)
83 }
84}
85
86impl Cheatcode for signKeychainAdminCall {
87 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
88 let Self { privateKey, account, digest } = self;
89 sign_keychain(state, privateKey, account, digest)
90 }
91}
92
93impl Cheatcode for signCompact_0Call {
94 fn apply_context_free(&self) -> Result {
95 let Self { wallet, digest } = self;
96 let sig = sign(&wallet.privateKey, digest)?;
97 Ok(encode_compact_sig(sig))
98 }
99}
100
101impl Cheatcode for deriveKey_0Call {
102 fn apply_context_free(&self) -> Result {
103 let Self { mnemonic, index } = self;
104 derive_key::<English>(mnemonic, DEFAULT_DERIVATION_PATH_PREFIX, *index)
105 }
106}
107
108impl Cheatcode for deriveKey_1Call {
109 fn apply_context_free(&self) -> Result {
110 let Self { mnemonic, derivationPath, index } = self;
111 derive_key::<English>(mnemonic, derivationPath, *index)
112 }
113}
114
115impl Cheatcode for deriveKey_2Call {
116 fn apply_context_free(&self) -> Result {
117 let Self { mnemonic, index, language } = self;
118 derive_key_str(mnemonic, DEFAULT_DERIVATION_PATH_PREFIX, *index, language)
119 }
120}
121
122impl Cheatcode for deriveKey_3Call {
123 fn apply_context_free(&self) -> Result {
124 let Self { mnemonic, derivationPath, index, language } = self;
125 derive_key_str(mnemonic, derivationPath, *index, language)
126 }
127}
128
129impl Cheatcode for rememberKeyCall {
130 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
131 let Self { privateKey } = self;
132 let wallet = with_private_key_signer(state, privateKey, |wallet| Ok(wallet.clone()))?;
133 let address = inject_wallet(state, wallet);
134 Ok(address.abi_encode())
135 }
136}
137
138impl Cheatcode for rememberKeys_0Call {
139 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
140 let Self { mnemonic, derivationPath, count } = self;
141 let wallets = derive_wallets::<English>(mnemonic, derivationPath, *count)?;
142 let mut addresses = Vec::<Address>::with_capacity(wallets.len());
143 for wallet in wallets {
144 let addr = inject_wallet(state, wallet);
145 addresses.push(addr);
146 }
147
148 Ok(addresses.abi_encode())
149 }
150}
151
152impl Cheatcode for rememberKeys_1Call {
153 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
154 let Self { mnemonic, derivationPath, language, count } = self;
155 let wallets = derive_wallets_str(mnemonic, derivationPath, language, *count)?;
156 let mut addresses = Vec::<Address>::with_capacity(wallets.len());
157 for wallet in wallets {
158 let addr = inject_wallet(state, wallet);
159 addresses.push(addr);
160 }
161
162 Ok(addresses.abi_encode())
163 }
164}
165
166fn inject_wallet(state: &mut CheatcodesCore, wallet: PrivateKeySigner) -> Address {
167 let address = wallet.address();
168 state.wallets().add_local_signer(wallet);
169 address
170}
171
172impl Cheatcode for sign_1Call {
173 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
174 let Self { privateKey, digest } = self;
175 let sig = sign_cached(state, privateKey, digest)?;
176 Ok(encode_full_sig(sig))
177 }
178}
179
180impl Cheatcode for signCompact_1Call {
181 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
182 let Self { privateKey, digest } = self;
183 let sig = sign_cached(state, privateKey, digest)?;
184 Ok(encode_compact_sig(sig))
185 }
186}
187
188impl Cheatcode for sign_2Call {
189 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
190 let Self { digest } = self;
191 let sig = sign_with_wallet(state, None, digest)?;
192 Ok(encode_full_sig(sig))
193 }
194}
195
196impl Cheatcode for signCompact_2Call {
197 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
198 let Self { digest } = self;
199 let sig = sign_with_wallet(state, None, digest)?;
200 Ok(encode_compact_sig(sig))
201 }
202}
203
204impl Cheatcode for sign_3Call {
205 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
206 let Self { signer, digest } = self;
207 let sig = sign_with_wallet(state, Some(*signer), digest)?;
208 Ok(encode_full_sig(sig))
209 }
210}
211
212impl Cheatcode for signCompact_3Call {
213 fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
214 let Self { signer, digest } = self;
215 let sig = sign_with_wallet(state, Some(*signer), digest)?;
216 Ok(encode_compact_sig(sig))
217 }
218}
219
220impl Cheatcode for signP256Call {
221 fn apply_context_free(&self) -> Result {
222 let Self { privateKey, digest } = self;
223 sign_p256(privateKey, digest)
224 }
225}
226
227impl Cheatcode for publicKeyP256Call {
228 fn apply_context_free(&self) -> Result {
229 let Self { privateKey } = self;
230 let pub_key =
231 parse_private_key_p256(privateKey)?.verifying_key().as_affine().to_encoded_point(false);
232 let pub_key_x = U256::from_be_bytes((*pub_key.x().unwrap()).into());
233 let pub_key_y = U256::from_be_bytes((*pub_key.y().unwrap()).into());
234
235 Ok((pub_key_x, pub_key_y).abi_encode())
236 }
237}
238
239impl Cheatcode for ecAffineToProjectiveCall {
240 fn apply_context_free(&self) -> Result {
241 let Self { pointX, pointY } = self;
242 let point = parse_affine_point(pointX, pointY, "point")?;
243 encode_projective_point(ProjectivePoint::from(point))
244 }
245}
246
247impl Cheatcode for ecProjectiveToAffineCall {
248 fn apply_context_free(&self) -> Result {
249 let Self { pointX, pointY, pointZ } = self;
250 let point = parse_projective_point(pointX, pointY, pointZ, "point")?;
251 encode_affine_point(point)
252 }
253}
254
255impl Cheatcode for ecAddAffineCall {
256 fn apply_context_free(&self) -> Result {
257 let Self { pointX1, pointY1, pointX2, pointY2 } = self;
258 let lhs = parse_affine_point(pointX1, pointY1, "first point")?;
259 let rhs = parse_affine_point(pointX2, pointY2, "second point")?;
260 encode_affine_point(ProjectivePoint::from(lhs) + rhs)
261 }
262}
263
264impl Cheatcode for ecAddProjectiveCall {
265 fn apply_context_free(&self) -> Result {
266 let Self { pointX1, pointY1, pointZ1, pointX2, pointY2, pointZ2 } = self;
267 let lhs = parse_projective_point(pointX1, pointY1, pointZ1, "first point")?;
268 let rhs = parse_projective_point(pointX2, pointY2, pointZ2, "second point")?;
269 encode_projective_point(lhs + rhs)
270 }
271}
272
273impl Cheatcode for ecMulAffineCall {
274 fn apply_context_free(&self) -> Result {
275 let Self { pointX, pointY, scalar } = self;
276 let point = parse_affine_point(pointX, pointY, "point")?;
277 let scalar = reduce_ec_scalar(scalar);
278 encode_affine_point(ProjectivePoint::from(point) * scalar)
279 }
280}
281
282impl Cheatcode for ecMulProjectiveCall {
283 fn apply_context_free(&self) -> Result {
284 let Self { pointX, pointY, pointZ, scalar } = self;
285 let point = parse_projective_point(pointX, pointY, pointZ, "point")?;
286 let scalar = reduce_ec_scalar(scalar);
287 encode_projective_point(point * scalar)
288 }
289}
290
291impl Cheatcode for createEd25519KeyCall {
292 fn apply_context_free(&self) -> Result {
293 let Self { salt } = self;
294 create_ed25519_key(salt)
295 }
296}
297
298impl Cheatcode for publicKeyEd25519Call {
299 fn apply_context_free(&self) -> Result {
300 let Self { privateKey } = self;
301 public_key_ed25519(privateKey)
302 }
303}
304
305impl Cheatcode for signEd25519Call {
306 fn apply_context_free(&self) -> Result {
307 let Self { namespace, message, privateKey } = self;
308 sign_ed25519(namespace, message, privateKey)
309 }
310}
311
312impl Cheatcode for verifyEd25519Call {
313 fn apply_context_free(&self) -> Result {
314 let Self { signature, namespace, message, publicKey } = self;
315 verify_ed25519(signature, namespace, message, publicKey)
316 }
317}
318
319fn create_wallet(private_key: &U256, label: Option<&str>, state: &mut CheatcodesCore) -> Result {
324 let (addr, pub_key_x, pub_key_y) = with_private_key_signer(state, private_key, |wallet| {
325 let addr = wallet.address();
326 let pub_key = wallet.credential().verifying_key().as_affine().to_encoded_point(false);
327 let pub_key_x = U256::from_be_bytes((*pub_key.x().unwrap()).into());
328 let pub_key_y = U256::from_be_bytes((*pub_key.y().unwrap()).into());
329 Ok((addr, pub_key_x, pub_key_y))
330 })?;
331
332 if let Some(label) = label {
333 state.labels.insert(addr, label.into());
334 }
335
336 Ok(Wallet { addr, publicKeyX: pub_key_x, publicKeyY: pub_key_y, privateKey: *private_key }
337 .abi_encode())
338}
339
340fn encode_full_sig(sig: alloy_primitives::Signature) -> Vec<u8> {
341 let v = U256::from(sig.v_byte());
343 let r = B256::from(sig.r());
344 let s = B256::from(sig.s());
345 (v, r, s).abi_encode()
346}
347
348fn encode_compact_sig(sig: alloy_primitives::Signature) -> Vec<u8> {
349 let r = B256::from(sig.r());
351 let mut vs = sig.s();
352 vs.set_bit(255, sig.v());
353 (r, vs).abi_encode()
354}
355
356fn sign(private_key: &U256, digest: &B256) -> Result<alloy_primitives::Signature> {
357 let wallet = parse_wallet(private_key)?;
359 let sig = wallet.sign_hash_sync(digest)?;
360 debug_assert_eq!(sig.recover_address_from_prehash(digest)?, wallet.address());
361 Ok(sig)
362}
363
364fn sign_cached(
365 state: &mut CheatcodesCore,
366 private_key: &U256,
367 digest: &B256,
368) -> Result<alloy_primitives::Signature> {
369 with_private_key_signer(state, private_key, |wallet| {
370 let sig = wallet.sign_hash_sync(digest)?;
371 debug_assert_eq!(sig.recover_address_from_prehash(digest)?, wallet.address());
372 Ok(sig)
373 })
374}
375
376fn sign_keychain(
377 state: &mut CheatcodesCore,
378 private_key: &U256,
379 account: &Address,
380 digest: &B256,
381) -> Result {
382 let signing_hash = KeychainSignature::signing_hash(*digest, *account);
383 let inner = sign_cached(state, private_key, &signing_hash)?;
384 let signature = TempoSignature::Keychain(KeychainSignature::new(
385 *account,
386 PrimitiveSignature::Secp256k1(inner),
387 ));
388 Ok(signature.to_bytes().abi_encode())
389}
390
391fn sign_with_nonce(
402 private_key: &U256,
403 digest: &B256,
404 nonce: &U256,
405) -> Result<alloy_primitives::Signature> {
406 let d_scalar: Scalar =
407 <Scalar as k256::elliptic_curve::PrimeField>::from_repr(private_key.to_be_bytes().into())
408 .into_option()
409 .ok_or_else(|| fmt_err!("invalid private key scalar"))?;
410 if bool::from(d_scalar.is_zero()) {
411 return Err(fmt_err!("private key cannot be 0"));
412 }
413
414 let k_scalar: Scalar =
415 <Scalar as k256::elliptic_curve::PrimeField>::from_repr(nonce.to_be_bytes().into())
416 .into_option()
417 .ok_or_else(|| fmt_err!("invalid nonce scalar"))?;
418 if bool::from(k_scalar.is_zero()) {
419 return Err(fmt_err!("nonce cannot be 0"));
420 }
421
422 let mut z = [0u8; 32];
423 z.copy_from_slice(digest.as_slice());
424 let z_fb: FieldBytes = FieldBytes::from(z);
425
426 let (sig_raw, recid_opt) =
429 <Scalar as hazmat::SignPrimitive<k256::Secp256k1>>::try_sign_prehashed(
430 &d_scalar, k_scalar, &z_fb,
431 )
432 .map_err(|e| fmt_err!("sign_prehashed failed: {e}"))?;
433
434 let (sig_low, flipped) =
436 if let Some(norm) = sig_raw.normalize_s() { (norm, true) } else { (sig_raw, false) };
437
438 let r_u256 = U256::from_be_bytes(sig_low.r().to_bytes().into());
439 let s_u256 = U256::from_be_bytes(sig_low.s().to_bytes().into());
440
441 let v_parity = if let Some(id) = recid_opt {
443 let mut v = id.to_byte() & 1;
444 if flipped {
445 v ^= 1;
446 }
447 v
448 } else {
449 let expected_addr = {
451 let sk: SigningKey = parse_private_key(private_key)?;
452 alloy_signer::utils::secret_key_to_address(&sk)
453 };
454 let cand0 = alloy_primitives::Signature::new(r_u256, s_u256, false);
456 if cand0.recover_address_from_prehash(digest).ok() == Some(expected_addr) {
457 return Ok(cand0);
458 }
459 let cand1 = alloy_primitives::Signature::new(r_u256, s_u256, true);
461 if cand1.recover_address_from_prehash(digest).ok() == Some(expected_addr) {
462 return Ok(cand1);
463 }
464 return Err(fmt_err!("failed to determine recovery id for signature"));
465 };
466
467 let y_parity = v_parity != 0;
468 Ok(alloy_primitives::Signature::new(r_u256, s_u256, y_parity))
469}
470
471fn sign_with_wallet(
472 state: &mut CheatcodesCore,
473 signer: Option<Address>,
474 digest: &B256,
475) -> Result<alloy_primitives::Signature> {
476 let mut wallets = state.wallets().inner.lock();
477 let maybe_provided_sender = wallets.provided_sender;
478 let signers = wallets.multi_wallet.signers()?;
479 if signers.is_empty() {
480 bail!("no wallets available");
481 }
482
483 let signer = if let Some(signer) = signer {
484 signer
485 } else if let Some(provided_sender) = maybe_provided_sender {
486 provided_sender
487 } else if signers.len() == 1 {
488 *signers.keys().next().unwrap()
489 } else {
490 bail!(
491 "could not determine signer, there are multiple signers available use vm.sign(signer, digest) to specify one"
492 );
493 };
494
495 let wallet = signers
496 .get(&signer)
497 .ok_or_else(|| fmt_err!("signer with address {signer} is not available"))?;
498
499 let sig = foundry_common::block_on(wallet.sign_hash(digest))?;
500 debug_assert_eq!(sig.recover_address_from_prehash(digest)?, signer);
501 Ok(sig)
502}
503
504fn sign_p256(private_key: &U256, digest: &B256) -> Result {
505 let signing_key = parse_private_key_p256(private_key)?;
506 let signature: P256Signature = signing_key.sign_prehash(digest.as_slice())?;
507 let signature = signature.normalize_s().unwrap_or(signature);
508 let r_bytes: [u8; 32] = signature.r().to_bytes().into();
509 let s_bytes: [u8; 32] = signature.s().to_bytes().into();
510
511 Ok((r_bytes, s_bytes).abi_encode())
512}
513
514fn parse_affine_point(x: &U256, y: &U256, name: &str) -> Result<AffinePoint> {
515 if x.is_zero() && y.is_zero() {
516 return Ok(AffinePoint::IDENTITY);
517 }
518
519 let encoded = EncodedPoint::from_affine_coordinates(
520 &FieldBytes::from(x.to_be_bytes()),
521 &FieldBytes::from(y.to_be_bytes()),
522 false,
523 );
524 AffinePoint::from_encoded_point(&encoded)
525 .into_option()
526 .ok_or_else(|| fmt_err!("invalid secp256k1 {name}"))
527}
528
529fn parse_projective_point(x: &U256, y: &U256, z: &U256, name: &str) -> Result<ProjectivePoint> {
530 let x_field = parse_field_element(x, name)?;
531 let y_field = parse_field_element(y, name)?;
532 let z_field = parse_field_element(z, name)?;
533
534 if bool::from(z_field.is_zero()) {
535 ensure!(
536 bool::from(x_field.is_zero()) && !bool::from(y_field.is_zero()),
537 "invalid secp256k1 {name}"
538 );
539 return Ok(ProjectivePoint::IDENTITY);
540 }
541
542 let z_inv = z_field.invert().expect("non-zero field element is invertible");
543 let affine_x = U256::from_be_slice(&(x_field * z_inv).to_bytes());
544 let affine_y = U256::from_be_slice(&(y_field * z_inv).to_bytes());
545
546 Ok(ProjectivePoint::from(parse_affine_point(&affine_x, &affine_y, name)?))
547}
548
549fn parse_field_element(value: &U256, name: &str) -> Result<FieldElement> {
550 FieldElement::from_bytes(&FieldBytes::from(value.to_be_bytes()))
551 .into_option()
552 .ok_or_else(|| fmt_err!("invalid secp256k1 {name}"))
553}
554
555fn reduce_ec_scalar(scalar: &U256) -> Scalar {
556 <Scalar as Reduce<K256U256>>::reduce_bytes(&scalar.to_be_bytes().into())
557}
558
559fn encode_affine_point(point: ProjectivePoint) -> Result {
560 if bool::from(point.is_identity()) {
561 return Ok((U256::ZERO, U256::ZERO).abi_encode());
562 }
563
564 let encoded = point.to_affine().to_encoded_point(false);
565 let x = U256::from_be_slice(encoded.x().expect("non-identity point has x coordinate"));
566 let y = U256::from_be_slice(encoded.y().expect("non-identity point has y coordinate"));
567
568 Ok((x, y).abi_encode())
569}
570
571fn encode_projective_point(point: ProjectivePoint) -> Result {
572 if bool::from(point.is_identity()) {
573 return Ok((U256::ZERO, U256::ONE, U256::ZERO).abi_encode());
574 }
575
576 let encoded = point.to_affine().to_encoded_point(false);
577 let x = U256::from_be_slice(encoded.x().expect("non-identity point has x coordinate"));
578 let y = U256::from_be_slice(encoded.y().expect("non-identity point has y coordinate"));
579
580 Ok((x, y, U256::ONE).abi_encode())
581}
582
583fn validate_private_key<C: ecdsa::PrimeCurve>(private_key: &U256) -> Result<()> {
584 ensure!(!private_key.is_zero(), "private key cannot be 0");
585 let order = U256::from_be_slice(&C::ORDER.to_be_byte_array());
586 ensure!(
587 *private_key < order,
588 "private key must be less than the {curve:?} curve order ({order})",
589 curve = C::default(),
590 );
591
592 Ok(())
593}
594
595fn parse_private_key(private_key: &U256) -> Result<SigningKey> {
596 validate_private_key::<k256::Secp256k1>(private_key)?;
597 Ok(SigningKey::from_bytes((&private_key.to_be_bytes()).into())?)
598}
599
600fn parse_private_key_p256(private_key: &U256) -> Result<P256SigningKey> {
601 validate_private_key::<p256::NistP256>(private_key)?;
602 Ok(P256SigningKey::from_bytes((&private_key.to_be_bytes()).into())?)
603}
604
605fn parse_signing_key_ed25519(private_key: &B256) -> Result<Ed25519SigningKey> {
606 Ed25519SigningKey::try_from(private_key.as_slice())
607 .map_err(|e| fmt_err!("invalid Ed25519 private key: {e}"))
608}
609
610fn create_ed25519_key(salt: &B256) -> Result {
611 let signing_key = parse_signing_key_ed25519(salt)?;
612 let public_key = B256::from_slice(signing_key.verification_key().as_ref());
613 Ok((public_key, *salt).abi_encode())
614}
615
616fn public_key_ed25519(private_key: &B256) -> Result {
617 let signing_key = parse_signing_key_ed25519(private_key)?;
618 Ok(B256::from_slice(signing_key.verification_key().as_ref()).abi_encode())
619}
620
621fn sign_ed25519(namespace: &[u8], message: &[u8], private_key: &B256) -> Result {
622 let signing_key = parse_signing_key_ed25519(private_key)?;
623 let combined = [namespace, message].concat();
624 let signature: [u8; 64] = signing_key.sign(&combined).into();
625 Ok(signature.to_vec().abi_encode())
626}
627
628fn verify_ed25519(signature: &[u8], namespace: &[u8], message: &[u8], public_key: &B256) -> Result {
629 if signature.len() != 64 {
630 return Ok(false.abi_encode());
631 }
632
633 let Ok(verification_key) = Ed25519VerificationKey::try_from(public_key.as_slice()) else {
634 return Ok(false.abi_encode());
635 };
636
637 let Ok(sig_bytes): Result<[u8; 64], _> = signature.try_into() else {
638 return Ok(false.abi_encode());
639 };
640
641 let combined = [namespace, message].concat();
642 let valid = verification_key.verify(&Ed25519Signature::from(sig_bytes), &combined).is_ok();
643 Ok(valid.abi_encode())
644}
645
646pub(super) fn parse_wallet(private_key: &U256) -> Result<PrivateKeySigner> {
647 parse_private_key(private_key).map(PrivateKeySigner::from)
648}
649
650pub(super) fn with_private_key_signer<R>(
651 state: &mut CheatcodesCore,
652 private_key: &U256,
653 f: impl FnOnce(&PrivateKeySigner) -> Result<R>,
654) -> Result<R> {
655 if !state.private_key_signers.contains_key(private_key)
656 && state.private_key_signers.len() < PRIVATE_KEY_SIGNER_CACHE_LIMIT
657 {
658 let wallet = parse_wallet(private_key)?;
659 state.private_key_signers.insert(*private_key, wallet);
660 }
661
662 if let Some(wallet) = state.private_key_signers.get(private_key) {
663 f(wallet)
664 } else {
665 let wallet = parse_wallet(private_key)?;
666 f(&wallet)
667 }
668}
669
670fn derive_key_str(mnemonic: &str, path: &str, index: u32, language: &str) -> Result {
671 let private_key = derive_private_key_with_language(mnemonic, path, index, language)
672 .map_err(|e| fmt_err!("{e}"))?;
673 Ok(private_key.abi_encode())
674}
675
676fn derive_key<W: Wordlist>(mnemonic: &str, path: &str, index: u32) -> Result {
677 let private_key =
678 derive_private_key::<W>(mnemonic, path, index).map_err(|e| fmt_err!("{e}"))?;
679 Ok(private_key.abi_encode())
680}
681
682fn derive_wallets_str(
683 mnemonic: &str,
684 path: &str,
685 language: &str,
686 count: u32,
687) -> Result<Vec<PrivateKeySigner>> {
688 match language {
689 "chinese_simplified" => derive_wallets::<ChineseSimplified>(mnemonic, path, count),
690 "chinese_traditional" => derive_wallets::<ChineseTraditional>(mnemonic, path, count),
691 "czech" => derive_wallets::<Czech>(mnemonic, path, count),
692 "english" => derive_wallets::<English>(mnemonic, path, count),
693 "french" => derive_wallets::<French>(mnemonic, path, count),
694 "italian" => derive_wallets::<Italian>(mnemonic, path, count),
695 "japanese" => derive_wallets::<Japanese>(mnemonic, path, count),
696 "korean" => derive_wallets::<Korean>(mnemonic, path, count),
697 "portuguese" => derive_wallets::<Portuguese>(mnemonic, path, count),
698 "spanish" => derive_wallets::<Spanish>(mnemonic, path, count),
699 _ => Err(fmt_err!("unsupported mnemonic language: {language:?}")),
700 }
701}
702
703fn derive_wallets<W: Wordlist>(
704 mnemonic: &str,
705 path: &str,
706 count: u32,
707) -> Result<Vec<PrivateKeySigner>> {
708 foundry_common::wallet::validate_bip32_path(path).map_err(|e| fmt_err!("{e}"))?;
709
710 let mut wallets = Vec::with_capacity(count as usize);
711 for idx in 0..count {
712 let full_path = foundry_common::wallet::derive_key_path_checked(path, idx)
713 .map_err(|e| fmt_err!("{e}"))?;
714 let wallet =
715 MnemonicBuilder::<W>::default().phrase(mnemonic).derivation_path(full_path)?.build()?;
716 wallets.push(wallet);
717 }
718
719 Ok(wallets)
720}
721
722#[cfg(test)]
723mod tests {
724 use super::*;
725 use alloy_primitives::{FixedBytes, hex::FromHex};
726 use alloy_sol_types::SolCall;
727 use k256::elliptic_curve::Curve;
728 use p256::ecdsa::signature::hazmat::PrehashVerifier;
729 use tempo_contracts::precompiles::{IAccountKeychain, ISignatureVerifier};
730 use tempo_hardfork::TempoHardfork;
731 use tempo_precompiles::{
732 Precompile,
733 account_keychain::{AccountKeychain, KeyRestrictions, SignatureType},
734 signature_verifier::SignatureVerifier,
735 storage::{StorageCtx, hashmap::HashMapStorageProvider},
736 };
737
738 #[test]
739 fn test_sign_p256() {
740 use p256::ecdsa::VerifyingKey;
741
742 let pk_u256: U256 = "1".parse().unwrap();
743 let signing_key = P256SigningKey::from_bytes(&pk_u256.to_be_bytes().into()).unwrap();
744 let digest = FixedBytes::from_hex(
745 "0x44acf6b7e36c1342c2c5897204fe09504e1e2efb1a900377dbc4e7a6a133ec56",
746 )
747 .unwrap();
748
749 let result = sign_p256(&pk_u256, &digest).unwrap();
750 let result_bytes: [u8; 64] = result.try_into().unwrap();
751 let signature = P256Signature::from_bytes(&result_bytes.into()).unwrap();
752 let verifying_key = VerifyingKey::from(&signing_key);
753 assert!(verifying_key.verify_prehash(digest.as_slice(), &signature).is_ok());
754 }
755
756 #[test]
757 fn test_sign_p256_pk_too_large() {
758 let pk =
760 "0xffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551".parse().unwrap();
761 let digest = FixedBytes::from_hex(
762 "0x54705ba3baafdbdfba8c5f9a70f7a89bee98d906b53e31074da7baecdc0da9ad",
763 )
764 .unwrap();
765 let result = sign_p256(&pk, &digest);
766 assert_eq!(
767 result.err().unwrap().to_string(),
768 "private key must be less than the NistP256 curve order (115792089210356248762697446949407573529996955224135760342422259061068512044369)"
769 );
770 }
771
772 #[test]
773 fn test_sign_p256_pk_0() {
774 let digest = FixedBytes::from_hex(
775 "0x54705ba3baafdbdfba8c5f9a70f7a89bee98d906b53e31074da7baecdc0da9ad",
776 )
777 .unwrap();
778 let result = sign_p256(&U256::ZERO, &digest);
779 assert_eq!(result.err().unwrap().to_string(), "private key cannot be 0");
780 }
781
782 #[test]
783 fn test_sign_with_nonce_varies_and_recovers() {
784 let pk_u256: U256 = U256::ONE;
786 let digest = FixedBytes::from_hex(
787 "0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
788 )
789 .unwrap();
790
791 let n1: U256 = U256::from(123u64);
793 let n2: U256 = U256::from(456u64);
794
795 let sig1 = sign_with_nonce(&pk_u256, &digest, &n1).expect("sig1");
797 let sig2 = sign_with_nonce(&pk_u256, &digest, &n2).expect("sig2");
798
799 assert!(
801 sig1.r() != sig2.r() || sig1.s() != sig2.s(),
802 "signatures should differ with different nonces"
803 );
804
805 let sk = parse_private_key(&pk_u256).unwrap();
807 let expected = alloy_signer::utils::secret_key_to_address(&sk);
808
809 assert_eq!(sig1.recover_address_from_prehash(&digest).unwrap(), expected);
810 assert_eq!(sig2.recover_address_from_prehash(&digest).unwrap(), expected);
811 }
812
813 #[test]
814 fn test_sign_with_nonce_zero_nonce_errors() {
815 let pk_u256: U256 = U256::ONE;
817 let digest = FixedBytes::from_hex(
818 "0xbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
819 )
820 .unwrap();
821 let n0: U256 = U256::ZERO;
822
823 let err = sign_with_nonce(&pk_u256, &digest, &n0).unwrap_err();
824 let msg = err.to_string();
825 assert!(msg.contains("nonce cannot be 0"), "unexpected error: {msg}");
826 }
827
828 #[test]
829 fn test_sign_with_nonce_nonce_ge_order_errors() {
830 use k256::Secp256k1;
832 let n_u256 = U256::from_be_slice(&Secp256k1::ORDER.to_be_byte_array());
834
835 let pk_u256: U256 = U256::ONE;
836 let digest = FixedBytes::from_hex(
837 "0xcccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc",
838 )
839 .unwrap();
840
841 let err = sign_with_nonce(&pk_u256, &digest, &n_u256).unwrap_err();
843 let msg = err.to_string();
844 assert!(msg.contains("invalid nonce scalar"), "unexpected error: {msg}");
845 }
846
847 #[test]
848 fn test_sign_keychain_encodes_v2_signature_for_account() {
849 let private_key = U256::from(0xB0Bu64);
850 let account = Address::repeat_byte(0x11);
851 let digest = B256::repeat_byte(0x22);
852 let mut state = CheatcodesCore::default();
853
854 let result = sign_keychain(&mut state, &private_key, &account, &digest).unwrap();
855 let signature = Vec::<u8>::abi_decode(&result).unwrap();
856
857 assert_eq!(signature.len(), 86);
858 assert_eq!(signature[0], 0x04);
859 assert_eq!(Address::from_slice(&signature[1..21]), account);
860
861 let parsed = TempoSignature::from_bytes(&signature).unwrap();
862 assert!(parsed.is_v2_keychain());
863
864 let keychain = parsed.as_keychain().unwrap();
865 let expected_key = parse_wallet(&private_key).unwrap().address();
866 assert_eq!(keychain.user_address, account);
867 assert_eq!(keychain.key_id(&digest).unwrap(), expected_key);
868 }
869
870 #[test]
871 fn private_key_signers_are_cached_for_repeated_lookup() {
872 let private_key = U256::from(0xB0Bu64);
873 let mut state = CheatcodesCore::default();
874
875 let first =
876 with_private_key_signer(&mut state, &private_key, |wallet| Ok(wallet.address()))
877 .unwrap();
878 assert_eq!(state.private_key_signers.len(), 1);
879
880 let second =
881 with_private_key_signer(&mut state, &private_key, |wallet| Ok(wallet.address()))
882 .unwrap();
883 assert_eq!(state.private_key_signers.len(), 1);
884 assert_eq!(first, second);
885 }
886
887 #[test]
888 fn test_sign_keychain_matches_t6_signature_verifier_state() {
889 let root_pk = U256::from(0xA11CEu64);
890 let access_pk = U256::from(0xB0Bu64);
891 let admin_pk = U256::from(0xC0FFEEu64);
892 let revoked_pk = U256::from(0xBADu64);
893 let expired_pk = U256::from(0xE441u64);
894 let unknown_pk = U256::from(0xFACEu64);
895
896 let root = parse_wallet(&root_pk).unwrap().address();
897 let access_key = parse_wallet(&access_pk).unwrap().address();
898 let admin_key = parse_wallet(&admin_pk).unwrap().address();
899 let revoked_key = parse_wallet(&revoked_pk).unwrap().address();
900 let expired_key = parse_wallet(&expired_pk).unwrap().address();
901
902 let hash = B256::repeat_byte(0x44);
903 let admin_hash = B256::repeat_byte(0x66);
904
905 let mut storage = HashMapStorageProvider::new_with_spec(1, TempoHardfork::T6);
906 storage.set_timestamp(U256::from(1_000u64));
907 StorageCtx::enter(&mut storage, || {
908 let mut keychain = AccountKeychain::new();
909 keychain.initialize()?;
910 keychain.set_tx_origin(root)?;
911
912 authorize_t6_access_key(&mut keychain, root, access_key, u64::MAX)?;
913 authorize_t6_access_key(&mut keychain, root, revoked_key, u64::MAX)?;
914 authorize_t6_access_key(&mut keychain, root, expired_key, 1_005)?;
915 keychain.authorize_admin_key(root, admin_key, SignatureType::Secp256k1, None)?;
916 keychain.revoke_key(root, IAccountKeychain::revokeKeyCall { keyId: revoked_key })?;
917
918 assert!(verify_keychain(root, hash, keychain_signature(&access_pk, root, hash)));
919 assert!(!verify_keychain(root, hash, keychain_signature(&revoked_pk, root, hash)));
920 assert!(!verify_keychain(root, hash, keychain_signature(&unknown_pk, root, hash)));
921 assert!(!verify_keychain(
922 Address::repeat_byte(0x99),
923 hash,
924 keychain_signature(&access_pk, root, hash)
925 ));
926 assert!(verify_keychain_admin(
927 root,
928 admin_hash,
929 keychain_signature(&admin_pk, root, admin_hash)
930 ));
931 assert!(verify_keychain_admin(
932 root,
933 admin_hash,
934 keychain_signature(&root_pk, root, admin_hash)
935 ));
936 assert!(!verify_keychain_admin(
937 root,
938 admin_hash,
939 keychain_signature(&access_pk, root, admin_hash)
940 ));
941 assert!(!verify_keychain_admin(
942 Address::repeat_byte(0x88),
943 admin_hash,
944 keychain_signature(&admin_pk, root, admin_hash)
945 ));
946 assert_keychain_signature_reverts(root, hash, vec![0x04]);
947
948 Ok::<_, eyre::Report>(())
949 })
950 .unwrap();
951
952 storage.set_timestamp(U256::from(1_006u64));
953 StorageCtx::enter(&mut storage, || {
954 assert!(!verify_keychain(root, hash, keychain_signature(&expired_pk, root, hash)));
955 Ok::<_, eyre::Report>(())
956 })
957 .unwrap();
958 }
959
960 fn authorize_t6_access_key(
961 keychain: &mut AccountKeychain,
962 account: Address,
963 key_id: Address,
964 expiry: u64,
965 ) -> eyre::Result<()> {
966 keychain.authorize_key(
967 account,
968 key_id,
969 SignatureType::Secp256k1,
970 KeyRestrictions {
971 expiry,
972 enforceLimits: false,
973 limits: vec![],
974 allowAnyCalls: true,
975 allowedCalls: vec![],
976 },
977 None,
978 )?;
979 Ok(())
980 }
981
982 fn keychain_signature(private_key: &U256, account: Address, hash: B256) -> Vec<u8> {
983 let mut state = CheatcodesCore::default();
984 Vec::<u8>::abi_decode(&sign_keychain(&mut state, private_key, &account, &hash).unwrap())
985 .unwrap()
986 }
987
988 fn verify_keychain(account: Address, hash: B256, signature: Vec<u8>) -> bool {
989 let calldata =
990 ISignatureVerifier::verifyKeychainCall { account, hash, signature: signature.into() }
991 .abi_encode();
992
993 let output = SignatureVerifier::new().call(&calldata, Address::ZERO).unwrap();
994 assert!(!output.is_revert(), "verifyKeychain reverted: {:?}", output.bytes);
995 ISignatureVerifier::verifyKeychainCall::abi_decode_returns(&output.bytes).unwrap()
996 }
997
998 fn verify_keychain_admin(account: Address, hash: B256, signature: Vec<u8>) -> bool {
999 let calldata = ISignatureVerifier::verifyKeychainAdminCall {
1000 account,
1001 hash,
1002 signature: signature.into(),
1003 }
1004 .abi_encode();
1005
1006 let output = SignatureVerifier::new().call(&calldata, Address::ZERO).unwrap();
1007 assert!(!output.is_revert(), "verifyKeychainAdmin reverted: {:?}", output.bytes);
1008 ISignatureVerifier::verifyKeychainAdminCall::abi_decode_returns(&output.bytes).unwrap()
1009 }
1010
1011 fn assert_keychain_signature_reverts(account: Address, hash: B256, signature: Vec<u8>) {
1012 let calldata =
1013 ISignatureVerifier::verifyKeychainCall { account, hash, signature: signature.into() }
1014 .abi_encode();
1015
1016 let output = SignatureVerifier::new().call(&calldata, Address::ZERO).unwrap();
1017 assert!(output.is_revert(), "malformed keychain signature should revert");
1018 }
1019
1020 #[test]
1021 fn test_create_ed25519_key_determinism() {
1022 let salt = B256::repeat_byte(1u8);
1023 let result1 = create_ed25519_key(&salt).unwrap();
1024 let result2 = create_ed25519_key(&salt).unwrap();
1025 assert_eq!(result1, result2, "same salt should produce same keys");
1026 }
1027
1028 #[test]
1029 fn test_create_ed25519_key_different_salts() {
1030 let salt1 = B256::repeat_byte(1u8);
1031 let salt2 = B256::repeat_byte(2u8);
1032 let result1 = create_ed25519_key(&salt1).unwrap();
1033 let result2 = create_ed25519_key(&salt2).unwrap();
1034 assert_ne!(result1, result2, "different salts should produce different keys");
1035 }
1036
1037 #[test]
1038 fn test_public_key_ed25519_consistency() {
1039 let salt = B256::repeat_byte(42u8);
1040 let create_result = create_ed25519_key(&salt).unwrap();
1041 let (expected_public, private): (B256, B256) =
1042 <(B256, B256)>::abi_decode(&create_result).unwrap();
1043
1044 let derived_public_result = public_key_ed25519(&private).unwrap();
1045 let derived_public = B256::abi_decode(&derived_public_result).unwrap();
1046
1047 assert_eq!(expected_public, derived_public, "derived public key should match");
1048 }
1049
1050 #[test]
1051 fn test_sign_and_verify_ed25519_valid() {
1052 let salt = B256::repeat_byte(123u8);
1053 let create_result = create_ed25519_key(&salt).unwrap();
1054 let (public_key, private_key): (B256, B256) =
1055 <(B256, B256)>::abi_decode(&create_result).unwrap();
1056
1057 let namespace = b"test.namespace";
1058 let message = b"hello world";
1059 let sig_result = sign_ed25519(namespace, message, &private_key).unwrap();
1060 let sig_bytes: Vec<u8> = Vec::abi_decode(&sig_result).unwrap();
1061
1062 let verify_result = verify_ed25519(&sig_bytes, namespace, message, &public_key).unwrap();
1063 let valid = bool::abi_decode(&verify_result).unwrap();
1064
1065 assert!(valid, "signature should be valid");
1066 }
1067
1068 #[test]
1069 fn test_verify_ed25519_invalid_signature() {
1070 let salt = B256::repeat_byte(123u8);
1071 let create_result = create_ed25519_key(&salt).unwrap();
1072 let (public_key, _): (B256, B256) = <(B256, B256)>::abi_decode(&create_result).unwrap();
1073
1074 let invalid_sig = [0u8; 64];
1075 let namespace = b"test.namespace";
1076 let message = b"hello world";
1077
1078 let verify_result = verify_ed25519(&invalid_sig, namespace, message, &public_key).unwrap();
1079 let valid = bool::abi_decode(&verify_result).unwrap();
1080
1081 assert!(!valid, "invalid signature should not verify");
1082 }
1083
1084 #[test]
1085 fn test_verify_ed25519_namespace_separation() {
1086 let salt = B256::repeat_byte(123u8);
1087 let create_result = create_ed25519_key(&salt).unwrap();
1088 let (public_key, private_key): (B256, B256) =
1089 <(B256, B256)>::abi_decode(&create_result).unwrap();
1090
1091 let namespace_a = b"namespace.a";
1092 let message = b"message";
1093 let sig_result = sign_ed25519(namespace_a, message, &private_key).unwrap();
1094 let sig_bytes: Vec<u8> = Vec::abi_decode(&sig_result).unwrap();
1095
1096 let namespace_b = b"namespace.b";
1097 let verify_result = verify_ed25519(&sig_bytes, namespace_b, message, &public_key).unwrap();
1098 let valid = bool::abi_decode(&verify_result).unwrap();
1099 assert!(!valid, "signature with namespace A should not verify with namespace B");
1100
1101 let verify_result = verify_ed25519(&sig_bytes, namespace_a, message, &public_key).unwrap();
1102 let valid = bool::abi_decode(&verify_result).unwrap();
1103 assert!(valid, "signature should verify with correct namespace");
1104 }
1105
1106 #[test]
1107 fn test_verify_ed25519_invalid_signature_length() {
1108 let salt = B256::repeat_byte(123u8);
1109 let create_result = create_ed25519_key(&salt).unwrap();
1110 let (public_key, _): (B256, B256) = <(B256, B256)>::abi_decode(&create_result).unwrap();
1111
1112 let invalid_sig = [0u8; 32];
1113 let namespace = b"test";
1114 let message = b"message";
1115
1116 let verify_result = verify_ed25519(&invalid_sig, namespace, message, &public_key).unwrap();
1117 let valid = bool::abi_decode(&verify_result).unwrap();
1118 assert!(!valid, "signature with wrong length should not verify");
1119 }
1120
1121 const MNEMONIC: &str = "test test test test test test test test test test test junk";
1122
1123 #[test]
1124 fn derive_key_rejects_harden_bit_overflow() {
1125 let err = derive_key::<English>(MNEMONIC, "m/44'/60'/0'/0/2147483648'", 0)
1126 .unwrap_err()
1127 .to_string();
1128 assert!(err.contains("harden bit"), "{err}");
1129
1130 let err = derive_key::<English>(MNEMONIC, "m/44'/60'/0'/0/2147483648h", 0)
1131 .unwrap_err()
1132 .to_string();
1133 assert!(err.contains("harden bit"), "{err}");
1134
1135 let err =
1136 derive_key::<English>(MNEMONIC, "m/44'/60'/0'/0", foundry_common::wallet::BIP32_HARDEN)
1137 .unwrap_err()
1138 .to_string();
1139 assert!(err.contains("harden bit"), "{err}");
1140
1141 assert!(derive_key::<English>(MNEMONIC, "m/44'/60'/0'/0", 0).is_ok());
1142 assert!(
1143 derive_key::<English>(
1144 MNEMONIC,
1145 &format!("m/44'/60'/0'/0/{}", foundry_common::wallet::BIP32_HARDEN - 1),
1146 0
1147 )
1148 .is_ok()
1149 );
1150 }
1151
1152 #[test]
1153 fn remember_keys_rejects_harden_bit_overflow() {
1154 let err = derive_wallets::<English>(MNEMONIC, "m/44'/60'/0'/0/2147483648'", 1)
1155 .unwrap_err()
1156 .to_string();
1157 assert!(err.contains("harden bit"), "{err}");
1158 assert!(derive_wallets::<English>(MNEMONIC, "m/44'/60'/0'/0", 1).is_ok());
1159 }
1160}