feat(cheatcodes): implement EIP-2098 in vm.sign (#8538)

* feat: implement EIP-2098 in vm.sign

* fix: clippy

* test: fix

* fix: forge-fmt

* chore: implement new cheatcode vm.signEIP2098

* fix: typos

* fix: another typo

* fix

* test: update sign tests

* fix: typo

* chore: rename `vm.signEIP2098()` into `vm.signCompact()`

* chore: update `encode_compact_signature` impl

* chore: update `signCompact` tests

* chore: factorize

* chore: rename fns

* chore: nit

* fix: issue

* chore: nits

crates/cheatcodes/spec/src/vm.rs

@@ -286,6 +286,14 @@ interface Vm {
     #[cheatcode(group = Evm, safety = Safe)]
     function sign(uint256 privateKey, bytes32 digest) external pure returns (uint8 v, bytes32 r, bytes32 s);
 
+    /// Signs `digest` with `privateKey` using the secp256k1 curve.
+    ///
+    /// Returns a compact signature (`r`, `vs`) as per EIP-2098, where `vs` encodes both the
+    /// signature's `s` value, and the recovery id `v` in a single bytes32.
+    /// This format reduces the signature size from 65 to 64 bytes.
+    #[cheatcode(group = Evm, safety = Safe)]
+    function signCompact(uint256 privateKey, bytes32 digest) external pure returns (bytes32 r, bytes32 vs);
+
     /// Signs `digest` with signer provided to script using the secp256k1 curve.
     ///
     /// If `--sender` is provided, the signer with provided address is used, otherwise,
@@ -296,12 +304,36 @@ interface Vm {
     #[cheatcode(group = Evm, safety = Safe)]
     function sign(bytes32 digest) external pure returns (uint8 v, bytes32 r, bytes32 s);
 
+    /// Signs `digest` with signer provided to script using the secp256k1 curve.
+    ///
+    /// Returns a compact signature (`r`, `vs`) as per EIP-2098, where `vs` encodes both the
+    /// signature's `s` value, and the recovery id `v` in a single bytes32.
+    /// This format reduces the signature size from 65 to 64 bytes.
+    ///
+    /// If `--sender` is provided, the signer with provided address is used, otherwise,
+    /// if exactly one signer is provided to the script, that signer is used.
+    ///
+    /// Raises error if signer passed through `--sender` does not match any unlocked signers or
+    /// if `--sender` is not provided and not exactly one signer is passed to the script.
+    #[cheatcode(group = Evm, safety = Safe)]
+    function signCompact(bytes32 digest) external pure returns (bytes32 r, bytes32 vs);
+
     /// Signs `digest` with signer provided to script using the secp256k1 curve.
     ///
     /// Raises error if none of the signers passed into the script have provided address.
     #[cheatcode(group = Evm, safety = Safe)]
     function sign(address signer, bytes32 digest) external pure returns (uint8 v, bytes32 r, bytes32 s);
 
+    /// Signs `digest` with signer provided to script using the secp256k1 curve.
+    ///
+    /// Returns a compact signature (`r`, `vs`) as per EIP-2098, where `vs` encodes both the
+    /// signature's `s` value, and the recovery id `v` in a single bytes32.
+    /// This format reduces the signature size from 65 to 64 bytes.
+    ///
+    /// Raises error if none of the signers passed into the script have provided address.
+    #[cheatcode(group = Evm, safety = Safe)]
+    function signCompact(address signer, bytes32 digest) external pure returns (bytes32 r, bytes32 vs);
+
     /// Signs `digest` with `privateKey` using the secp256r1 curve.
     #[cheatcode(group = Evm, safety = Safe)]
     function signP256(uint256 privateKey, bytes32 digest) external pure returns (bytes32 r, bytes32 s);
@@ -2153,6 +2185,14 @@ interface Vm {
     #[cheatcode(group = Utilities)]
     function sign(Wallet calldata wallet, bytes32 digest) external returns (uint8 v, bytes32 r, bytes32 s);
 
+    /// Signs data with a `Wallet`.
+    ///
+    /// Returns a compact signature (`r`, `vs`) as per EIP-2098, where `vs` encodes both the
+    /// signature's `s` value, and the recovery id `v` in a single bytes32.
+    /// This format reduces the signature size from 65 to 64 bytes.
+    #[cheatcode(group = Utilities)]
+    function signCompact(Wallet calldata wallet, bytes32 digest) external returns (bytes32 r, bytes32 vs);
+
     /// Derive a private key from a provided mnenomic string (or mnenomic file path)
     /// at the derivation path `m/44'/60'/0'/0/{index}`.
     #[cheatcode(group = Utilities)]

crates/cheatcodes/src/evm.rs

@@ -162,21 +162,48 @@ impl Cheatcode for dumpStateCall {
 impl Cheatcode for sign_0Call {
     fn apply_stateful<DB: DatabaseExt>(&self, _: &mut CheatsCtxt<DB>) -> Result {
         let Self { privateKey, digest } = self;
-        super::utils::sign(privateKey, digest)
+        let sig = super::utils::sign(privateKey, digest)?;
+        Ok(super::utils::encode_full_sig(sig))
+    }
+}
+
+impl Cheatcode for signCompact_0Call {
+    fn apply_stateful<DB: DatabaseExt>(&self, _: &mut CheatsCtxt<DB>) -> Result {
+        let Self { privateKey, digest } = self;
+        let sig = super::utils::sign(privateKey, digest)?;
+        Ok(super::utils::encode_compact_sig(sig))
     }
 }
 
 impl Cheatcode for sign_1Call {
     fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { digest } = self;
-        super::utils::sign_with_wallet(ccx, None, digest)
+        let sig = super::utils::sign_with_wallet(ccx, None, digest)?;
+        Ok(super::utils::encode_full_sig(sig))
+    }
+}
+
+impl Cheatcode for signCompact_1Call {
+    fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { digest } = self;
+        let sig = super::utils::sign_with_wallet(ccx, None, digest)?;
+        Ok(super::utils::encode_compact_sig(sig))
     }
 }
 
 impl Cheatcode for sign_2Call {
     fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { signer, digest } = self;
-        super::utils::sign_with_wallet(ccx, Some(*signer), digest)
+        let sig = super::utils::sign_with_wallet(ccx, Some(*signer), digest)?;
+        Ok(super::utils::encode_full_sig(sig))
+    }
+}
+
+impl Cheatcode for signCompact_2Call {
+    fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { signer, digest } = self;
+        let sig = super::utils::sign_with_wallet(ccx, Some(*signer), digest)?;
+        Ok(super::utils::encode_compact_sig(sig))
     }
 }
 

crates/cheatcodes/src/utils.rs

@@ -55,7 +55,16 @@ impl Cheatcode for getNonce_1Call {
 impl Cheatcode for sign_3Call {
     fn apply_stateful<DB: DatabaseExt>(&self, _: &mut CheatsCtxt<DB>) -> Result {
         let Self { wallet, digest } = self;
-        sign(&wallet.privateKey, digest)
+        let sig = sign(&wallet.privateKey, digest)?;
+        Ok(encode_full_sig(sig))
+    }
+}
+
+impl Cheatcode for signCompact_3Call {
+    fn apply_stateful<DB: DatabaseExt>(&self, _: &mut CheatsCtxt<DB>) -> Result {
+        let Self { wallet, digest } = self;
+        let sig = sign(&wallet.privateKey, digest)?;
+        Ok(encode_compact_sig(sig))
     }
 }
 
@@ -201,25 +210,35 @@ fn create_wallet(private_key: &U256, label: Option<&str>, state: &mut Cheatcodes
         .abi_encode())
 }
 
-fn encode_vrs(sig: alloy_primitives::Signature) -> Vec<u8> {
-    let v = sig.v().y_parity_byte_non_eip155().unwrap_or(sig.v().y_parity_byte());
+pub(super) fn encode_full_sig(sig: alloy_primitives::Signature) -> Vec<u8> {
+    // Retrieve v, r and s from signature.
+    let v = U256::from(sig.v().y_parity_byte_non_eip155().unwrap_or(sig.v().y_parity_byte()));
+    let r = B256::from(sig.r());
+    let s = B256::from(sig.s());
+    (v, r, s).abi_encode()
+}
 
-    (U256::from(v), B256::from(sig.r()), B256::from(sig.s())).abi_encode()
+pub(super) fn encode_compact_sig(sig: alloy_primitives::Signature) -> Vec<u8> {
+    // Implement EIP-2098 compact signature.
+    let r = B256::from(sig.r());
+    let mut vs = sig.s();
+    vs.set_bit(255, sig.v().y_parity());
+    (r, vs).abi_encode()
 }
 
-pub(super) fn sign(private_key: &U256, digest: &B256) -> Result {
+pub(super) fn sign(private_key: &U256, digest: &B256) -> Result<alloy_primitives::Signature> {
     // The `ecrecover` precompile does not use EIP-155. No chain ID is needed.
     let wallet = parse_wallet(private_key)?;
     let sig = wallet.sign_hash_sync(digest)?;
     debug_assert_eq!(sig.recover_address_from_prehash(digest)?, wallet.address());
-    Ok(encode_vrs(sig))
+    Ok(sig)
 }
 
 pub(super) fn sign_with_wallet<DB: DatabaseExt>(
     ccx: &mut CheatsCtxt<DB>,
     signer: Option<Address>,
     digest: &B256,
-) -> Result {
+) -> Result<alloy_primitives::Signature> {
     let Some(script_wallets) = ccx.state.script_wallets() else {
         bail!("no wallets are available");
     };
@@ -244,7 +263,7 @@ pub(super) fn sign_with_wallet<DB: DatabaseExt>(
 
     let sig = foundry_common::block_on(wallet.sign_hash(digest))?;
     debug_assert_eq!(sig.recover_address_from_prehash(digest)?, signer);
-    Ok(encode_vrs(sig))
+    Ok(sig)
 }
 
 pub(super) fn sign_p256(private_key: &U256, digest: &B256, _state: &mut Cheatcodes) -> Result {

testdata/default/cheats/Sign.t.sol

@@ -13,11 +13,35 @@ contract SignTest is DSTest {
         (uint8 v, bytes32 r, bytes32 s) = vm.sign(pk, digest);
         address expected = vm.addr(pk);
         address actual = ecrecover(digest, v, r, s);
+        assertEq(actual, expected, "digest signer did not match");
+    }
+
+    function testSignCompactDigest(uint248 pk, bytes32 digest) public {
+        vm.assume(pk != 0);
+
+        (bytes32 r, bytes32 vs) = vm.signCompact(pk, digest);
+
+        // Extract `s` from `vs`.
+        // Shift left by 1 bit to clear the leftmost bit, then shift right by 1 bit to restore the original position.
+        // This effectively clears the leftmost bit of `vs`, giving us `s`.
+        bytes32 s = bytes32((uint256(vs) << 1) >> 1);
 
+        // Extract `v` from `vs`.
+        // We shift `vs` right by 255 bits to isolate the leftmost bit.
+        // Converting this to uint8 gives us the parity bit (0 or 1).
+        // Adding 27 converts this parity bit to the correct `v` value (27 or 28).
+        uint8 v = uint8(uint256(vs) >> 255) + 27;
+
+        address expected = vm.addr(pk);
+        address actual = ecrecover(digest, v, r, s);
         assertEq(actual, expected, "digest signer did not match");
     }
 
     function testSignMessage(uint248 pk, bytes memory message) public {
         testSignDigest(pk, keccak256(message));
     }
+
+    function testSignCompactMessage(uint248 pk, bytes memory message) public {
+        testSignCompactDigest(pk, keccak256(message));
+    }
 }

testdata/default/cheats/Wallet.t.sol

@@ -104,10 +104,36 @@ contract WalletTest is DSTest {
         assertEq(recovered, wallet.addr);
     }
 
+    function testSignCompactWithWalletDigest(uint256 pkSeed, bytes32 digest) public {
+        uint256 pk = bound(pkSeed, 1, Q - 1);
+
+        Vm.Wallet memory wallet = vm.createWallet(pk);
+
+        (bytes32 r, bytes32 vs) = vm.signCompact(wallet, digest);
+
+        // Extract `s` from `vs`.
+        // Shift left by 1 bit to clear the leftmost bit, then shift right by 1 bit to restore the original position.
+        // This effectively clears the leftmost bit of `vs`, giving us `s`.
+        bytes32 s = bytes32((uint256(vs) << 1) >> 1);
+
+        // Extract `v` from `vs`.
+        // We shift `vs` right by 255 bits to isolate the leftmost bit.
+        // Converting this to uint8 gives us the parity bit (0 or 1).
+        // Adding 27 converts this parity bit to the correct `v` value (27 or 28).
+        uint8 v = uint8(uint256(vs) >> 255) + 27;
+
+        address recovered = ecrecover(digest, v, r, s);
+        assertEq(recovered, wallet.addr);
+    }
+
     function testSignWithWalletMessage(uint256 pkSeed, bytes memory message) public {
         testSignWithWalletDigest(pkSeed, keccak256(message));
     }
 
+    function testSignCompactWithWalletMessage(uint256 pkSeed, bytes memory message) public {
+        testSignCompactWithWalletDigest(pkSeed, keccak256(message));
+    }
+
     function testGetNonceWallet(uint256 pkSeed) public {
         uint256 pk = bound(pkSeed, 1, Q - 1);