feat: add initializer to be run after forking

mento-protocol · Aug 21, 2024 · a797a8f · a797a8f
1 parent 14c3dd8
commit a797a8f
Showing 2 changed files with 71 additions and 13 deletions.
diff --git a/src/CeloPrecompiles.sol b/src/CeloPrecompiles.sol
@@ -7,15 +7,25 @@ import "./Constants.sol";
 abstract contract CeloPrecompiles {
     Vm private constant vm = Vm(VM_ADDRESS);
 
-    bytes4 constant TRANSFER_SIG = bytes4(keccak256("transfer(address,address,uint256)"));
+    bytes4 constant TRANSFER_SIG =
+        bytes4(keccak256("transfer(address,address,uint256)"));
     bytes4 constant CATCHALL_SIG = bytes4(keccak256("catchAll()"));
 
     constructor() {
+        __CeloPrecompiles_init();
+    }
+
+    /// @notice Initialize the Celo precompiles.
+    function __CeloPrecompiles_init() internal {
         vm.etch(TRANSFER_PRECOMPILE, proxyTo(TRANSFER_SIG));
         vm.label(TRANSFER_PRECOMPILE, "TRANSFER_PRECOMPILE");
     }
 
-    function transfer(address from, address to, uint256 amount) public returns (bool) {
+    function transfer(
+        address from,
+        address to,
+        uint256 amount
+    ) public returns (bool) {
         if (from.balance < amount) {
             revert("TransferPrecompile: insufficient balance");
         }
@@ -47,11 +57,20 @@ abstract contract CeloPrecompiles {
             let sigPrefix := shl(0x50, 0x63) // shift the signature prefix
             let shiftedSig := shl(0x30, shr(0xe0, selector)) // shift the signature
             let suffix := 0x600060043601
-            mstore(mc, or(addrPrefix, or(addr, or(sigPrefix, or(shiftedSig, suffix))))) // constructor the masked first 32 bytes
+            mstore(
+                mc,
+                or(addrPrefix, or(addr, or(sigPrefix, or(shiftedSig, suffix))))
+            ) // constructor the masked first 32 bytes
             mc := add(mc, 0x20) // move mc to the next 32 bytes
-            mstore(mc, 0x8260e01b82523660006004840137600080828434885af13d6000816000823e82) // store the rest of the bytecode
+            mstore(
+                mc,
+                0x8260e01b82523660006004840137600080828434885af13d6000816000823e82
+            ) // store the rest of the bytecode
             mc := add(mc, 0x20) // move mc to the next 32 bytes
-            mstore(mc, 0x60008114604a578282f35b8282fd000000000000000000000000000000000000) // store the rest of the bytecode
+            mstore(
+                mc,
+                0x60008114604a578282f35b8282fd000000000000000000000000000000000000
+            ) // store the rest of the bytecode
             mstore(0x40, add(proxyByteCode, 0x80)) // increment the free memory pointer
         }
 

diff --git a/test/CeloPrecompiles.t.sol b/test/CeloPrecompiles.t.sol
@@ -4,27 +4,66 @@ pragma solidity ^0.8;
 import {console} from "forge-std/console.sol";
 import {Test} from "src/Test.sol";
 
-import {TRANSFER_PRECOMPILE} from "src/Constants.sol";
+import {TRANSFER_PRECOMPILE, CELO_ID} from "src/Constants.sol";
 
 contract CeloPrecompilesTest is Test {
+    function setUp() public {
+        vm.setEnv("CELO_RPC_URL", "https://forno.celo.org");
+    }
+
     function test_transferPrecompile() public {
-        address from = address(0x1);
-        address to = address(0x2);
+        address from = makeAddr("from");
+        address to = makeAddr("to");
+
+        vm.deal(from, 100);
+        (bool ok, ) = TRANSFER_PRECOMPILE.call(abi.encode(from, to, 100));
+
+        assertEq(ok, true);
+        assertEq(from.balance, 0);
+        assertEq(to.balance, 100);
+    }
+
+    function test_transferPrecompileAfterFork() public {
+        address from = makeAddr("from");
+        address to = makeAddr("to");
+
+        vm.deal(from, 100);
+        (bool ok, ) = TRANSFER_PRECOMPILE.call(abi.encode(from, to, 100));
+
+        assertEq(ok, true);
+        assertEq(from.balance, 0);
+        assertEq(to.balance, 100);
+
+        fork(CELO_ID);
+        (ok, ) = TRANSFER_PRECOMPILE.call(abi.encode(from, to, 100));
+        assertTrue(ok);
+        assertEq(from.balance, 0);
+        assertEq(to.balance, 0);
+
+        __CeloPrecompiles_init();
 
         vm.deal(from, 100);
-        (bool ok,) = TRANSFER_PRECOMPILE.call(abi.encode(from, to, 100));
+        (ok, ) = TRANSFER_PRECOMPILE.call(abi.encode(from, to, 100));
 
         assertEq(ok, true);
         assertEq(from.balance, 0);
         assertEq(to.balance, 100);
     }
 
     function test_transferPrecompileFails() public {
-        address from = address(0x1);
-        address to = address(0x2);
+        address from = makeAddr("from");
+        address to = makeAddr("to");
 
-        (bool ok, bytes memory data) = TRANSFER_PRECOMPILE.call(abi.encode(from, to, 100));
+        (bool ok, bytes memory data) = TRANSFER_PRECOMPILE.call(
+            abi.encode(from, to, 100)
+        );
         assertEq(ok, false);
-        assertEq(data, abi.encodeWithSignature("Error(string)", "TransferPrecompile: insufficient balance"));
+        assertEq(
+            data,
+            abi.encodeWithSignature(
+                "Error(string)",
+                "TransferPrecompile: insufficient balance"
+            )
+        );
     }
 }