all: Adds prelude with vmed modifier

verklegarden · Oct 21, 2023 · add51da · add51da
1 parent e46e7b5
commit add51da
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 14 deletions.
diff --git a/src/Random.sol b/src/Random.sol
@@ -21,13 +21,22 @@ import {Vm} from "forge-std/Vm.sol";
  * @dev Randomness is sourced from cast's `new wallet` command.
  */
 library Random {
-    Vm private constant vm =
-        Vm(address(uint160(uint(keccak256("hevm cheat code")))));
+    // ~~~~~~~ Prelude ~~~~~~~
+    // forgefmt: disable-start
+    Vm private constant vm = Vm(address(uint160(uint(keccak256("hevm cheat code")))));
+    modifier vmed() {
+        if (block.chainid != 31337) {
+            revert("requireVm");
+        }
+        _;
+    }
+    // forgefmt: disable-end
+    // ~~~~~~~~~~~~~~~~~~~~~~~
 
     /// @dev Returns 256 bit of cryptographically sound randomness.
     ///
     /// @custom:vm ffi `cast wallet new`
-    function readUint() internal returns (uint) {
+    function readUint() internal vmed returns (uint) {
         string[] memory inputs = new string[](3);
         inputs[0] = "cast";
         inputs[1] = "wallet";

diff --git a/src/curves/Secp256k1.sol b/src/curves/Secp256k1.sol
@@ -66,8 +66,20 @@ library Secp256k1 {
     using Secp256k1 for AffinePoint;
     using Secp256k1Arithmetic for AffinePoint;
 
-    Vm private constant vm =
-        Vm(address(uint160(uint(keccak256("hevm cheat code")))));
+    // ~~~~~~~ Prelude ~~~~~~~
+    // forgefmt: disable-start
+    Vm private constant vm = Vm(address(uint160(uint(keccak256("hevm cheat code")))));
+    modifier vmed() {
+        if (block.chainid != 31337) {
+            revert("requireVm");
+        }
+        _;
+    }
+    // forgefmt: disable-end
+    // ~~~~~~~~~~~~~~~~~~~~~~~
+
+    //--------------------------------------------------------------------------
+    // Private Constants
 
     uint private constant _ADDRESS_MASK =
         0x000000000000000000000000FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF;
@@ -93,7 +105,7 @@ library Secp256k1 {
     /// @dev Returns a new cryptographically secure private key.
     ///
     /// @custom:vm Random::readUint()(uint)
-    function newPrivateKey() internal returns (PrivateKey) {
+    function newPrivateKey() internal vmed returns (PrivateKey) {
         // Let scalar ∊ [1, Q) sourced cryptographically secure.
         uint scalar = (Random.readUint() % (Secp256k1Arithmetic.Q - 1)) + 1;
         return PrivateKey.wrap(scalar);
@@ -112,7 +124,11 @@ library Secp256k1 {
     ///      - Private key invalid
     ///
     /// @custom:vm vm.createWallet(uint)
-    function toPublicKey(PrivateKey self) internal returns (PublicKey memory) {
+    function toPublicKey(PrivateKey self)
+        internal
+        vmed
+        returns (PublicKey memory)
+    {
         if (!self.isValid()) {
             revert("PrivateKeyInvalid()");
         }

diff --git a/src/signatures/ECDSA.sol b/src/signatures/ECDSA.sol
@@ -67,8 +67,17 @@ library ECDSA {
     using Secp256k1 for PrivateKey;
     using Secp256k1 for PublicKey;
 
-    Vm private constant vm =
-        Vm(address(uint160(uint(keccak256("hevm cheat code")))));
+    // ~~~~~~~ Prelude ~~~~~~~
+    // forgefmt: disable-start
+    Vm private constant vm = Vm(address(uint160(uint(keccak256("hevm cheat code")))));
+    modifier vmed() {
+        if (block.chainid != 31337) {
+            revert("requireVm");
+        }
+        _;
+    }
+    // forgefmt: disable-end
+    // ~~~~~~~~~~~~~~~~~~~~~~~
 
     /// @dev Mask to receive an ECDSA's s value from an EIP-2098 compact
     ///      signature representation.
@@ -178,7 +187,7 @@ library ECDSA {
     ///      - Private key invalid
     function sign(PrivateKey privKey, bytes memory message)
         internal
-        pure
+        view
         returns (Signature memory)
     {
         bytes32 digest = keccak256(message);
@@ -195,7 +204,8 @@ library ECDSA {
     /// @custom:vm vm.sign(uint,bytes32)
     function sign(PrivateKey privKey, bytes32 digest)
         internal
-        pure
+        view
+        vmed
         returns (Signature memory)
     {
         if (!privKey.isValid()) {
@@ -215,7 +225,7 @@ library ECDSA {
 
     function signEthereumSignedMessage(PrivateKey privKey, bytes memory message)
         internal
-        pure
+        view
         returns (Signature memory)
     {
         bytes32 digest = Message.deriveEthereumSignedMessage(message);
@@ -225,7 +235,7 @@ library ECDSA {
 
     function signEthereumSignedMessageHash(PrivateKey privKey, bytes32 digest)
         internal
-        pure
+        view
         returns (Signature memory)
     {
         bytes32 digest2 = Message.deriveEthereumSignedMessageHash(digest);
@@ -249,7 +259,7 @@ library ECDSA {
     /// @custom:vm vm.toString(uint)
     function toString(Signature memory self)
         internal
-        pure
+        vmed
         returns (string memory)
     {
         string memory str = "ECDSA::Signature { \n";