Skip to main content

foundry_cheatcodes/
crypto.rs

1//! Implementations of [`Crypto`](spec::Group::Crypto) Cheatcodes.
2
3use 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
36/// The BIP32 default derivation path prefix.
37const 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
319/// Using a given private key, return its public ETH address, its public key affine x and y
320/// coordinates, and its private key (see the 'Wallet' struct)
321///
322/// If 'label' is set to 'Some()', assign that label to the associated ETH address in state
323fn 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    // Retrieve v, r and s from signature.
342    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    // Implement EIP-2098 compact signature.
350    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    // The `ecrecover` precompile does not use EIP-155. No chain ID is needed.
358    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
391/// Signs `digest` on secp256k1 using a user-supplied ephemeral nonce `k` (no RFC6979).
392/// - `private_key` and `nonce` must be in (0, n)
393/// - `digest` is a 32-byte prehash.
394///
395/// # Warning
396///
397/// Use [`sign_with_nonce`] with extreme caution!
398/// Reusing the same nonce (`k`) with the same private key in ECDSA will leak the private key.
399/// Always generate `nonce` with a cryptographically secure RNG, and never reuse it across
400/// signatures.
401fn 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    // Hazmat signing using the scalar `d` (SignPrimitive is implemented for `Scalar`)
427    // Note: returns (Signature, Option<RecoveryId>)
428    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    // Enforce low-s; if mirrored, parity flips (we’ll account for it below if we use recid)
435    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    // Determine v parity in {0,1}
442    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        // Fallback: choose parity by recovery to expected address
450        let expected_addr = {
451            let sk: SigningKey = parse_private_key(private_key)?;
452            alloy_signer::utils::secret_key_to_address(&sk)
453        };
454        // Try v = 0
455        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        // Try v = 1
460        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        // max n from https://neuromancer.sk/std/secg/secp256r1
759        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        // Given a fixed private key and digest
785        let pk_u256: U256 = U256::ONE;
786        let digest = FixedBytes::from_hex(
787            "0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
788        )
789        .unwrap();
790
791        // Two distinct nonces
792        let n1: U256 = U256::from(123u64);
793        let n2: U256 = U256::from(456u64);
794
795        // Sign with both nonces
796        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        // (r,s) must differ when nonce differs
800        assert!(
801            sig1.r() != sig2.r() || sig1.s() != sig2.s(),
802            "signatures should differ with different nonces"
803        );
804
805        // ecrecover must yield the address for both signatures
806        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        // nonce = 0 should be rejected
816        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        // nonce >= n should be rejected
831        use k256::Secp256k1;
832        // Curve order n as U256
833        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        // Try exactly n (>= n invalid)
842        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}