Restore transient slot values after calling script

.github/workflows/lint.yml

@@ -30,7 +30,7 @@ jobs:
       - name: Run Forge build
         run: |
           forge --version
-          forge build --sizes
+          forge build --via-ir --sizes
         id: build
 
       - name: Run Forge Format

src/quark-core/src/QuarkScript.sol

@@ -13,7 +13,6 @@ abstract contract QuarkScript {
     error ReentrantCall();
     error InvalidActiveNonce();
     error InvalidActiveSubmissionToken();
-    error NoActiveNonce();
 
     /// @notice Storage location for the re-entrancy guard
     bytes32 internal constant REENTRANCY_FLAG_SLOT =
@@ -125,9 +124,6 @@ abstract contract QuarkScript {
     // This provide cooperative isolation of storage between scripts.
     function getNonceIsolatedKey(bytes32 key) internal view returns (bytes32) {
         bytes32 nonce = getActiveNonce();
-        if (nonce == bytes32(0)) {
-            revert NoActiveNonce();
-        }
         return keccak256(abi.encodePacked(nonce, key));
     }
 

src/quark-core/src/QuarkWallet.sol

@@ -482,8 +482,17 @@ contract QuarkWallet is IERC1271 {
         bytes32 activeScriptSlot = ACTIVE_SCRIPT_SLOT;
         bytes32 activeNonceSlot = ACTIVE_NONCE_SLOT;
         bytes32 activeSubmissionTokenSlot = ACTIVE_SUBMISSION_TOKEN_SLOT;
+        bytes32 callbackSlot = CALLBACK_SLOT;
+        address oldActiveScript;
+        bytes32 oldActiveNonce;
+        bytes32 oldActiveSubmissionToken;
+        address oldCallback;
         assembly {
-            // TODO: TSTORE the callback slot to 0
+            // Cache the previous values in each of the transient slots so they can be restored after the callcode
+            oldActiveScript := tload(activeScriptSlot)
+            oldActiveNonce := tload(activeNonceSlot)
+            oldActiveSubmissionToken := tload(activeSubmissionTokenSlot)
+            oldCallback := tload(callbackSlot)
 
             // Transiently store the active script
             tstore(activeScriptSlot, scriptAddress)
@@ -499,14 +508,17 @@ contract QuarkWallet is IERC1271 {
                 callcode(gas(), scriptAddress, /* value */ 0, add(scriptCalldata, 0x20), scriptCalldataLen, 0x0, 0)
             returnSize := returndatasize()
 
-            // Transiently clear the active script
-            tstore(activeScriptSlot, 0)
+            // Transiently restore the active script
+            tstore(activeScriptSlot, oldActiveScript)
+
+            // Transiently restore the active nonce
+            tstore(activeNonceSlot, oldActiveNonce)
 
-            // Transiently clear the active nonce
-            tstore(activeNonceSlot, 0)
+            // Transiently restore the active submission token
+            tstore(activeSubmissionTokenSlot, oldActiveSubmissionToken)
 
-            // Transiently clear the active submission token
-            tstore(activeSubmissionTokenSlot, 0)
+            // Transiently restore the callback slot
+            tstore(callbackSlot, oldCallback)
         }
 
         bytes memory returnData = new bytes(returnSize);

test/lib/BatchCallback.sol

@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: BSD-3-Clause
+pragma solidity 0.8.27;
+
+import "quark-core/src/QuarkScript.sol";
+import "quark-core/src/QuarkWallet.sol";
+
+contract BatchSend {
+    function submitTwo(
+        QuarkWallet wallet1,
+        QuarkWallet.QuarkOperation memory op1,
+        uint8 v1,
+        bytes32 r1,
+        bytes32 s1,
+        QuarkWallet wallet2,
+        QuarkWallet.QuarkOperation memory op2,
+        uint8 v2,
+        bytes32 r2,
+        bytes32 s2
+    ) public returns (uint256) {
+        wallet1.executeQuarkOperation(op1, v1, r1, s1);
+        wallet2.executeQuarkOperation(op2, v2, r2, s2);
+        return IncrementByCallback(address(wallet1)).number();
+    }
+}
+
+contract IncrementByCallback is QuarkScript {
+    uint256 public number;
+
+    function run() public {
+        allowCallback();
+        IncrementByCallback(address(this)).increment();
+        IncrementByCallback(address(this)).increment();
+    }
+
+    function increment() external {
+        number++;
+    }
+}
+
+contract CallIncrement is QuarkScript {
+    function run(address wallet) public {
+        IncrementByCallback(wallet).increment();
+    }
+}

test/quark-core/BatchCallback.t.sol

@@ -0,0 +1,96 @@
+// SPDX-License-Identifier: BSD-3-Clause
+pragma solidity 0.8.27;
+
+import "forge-std/Test.sol";
+import "forge-std/console.sol";
+
+import {YulHelper} from "test/lib/YulHelper.sol";
+import {SignatureHelper} from "test/lib/SignatureHelper.sol";
+import {QuarkOperationHelper, ScriptType} from "test/lib/QuarkOperationHelper.sol";
+
+import {CodeJar} from "codejar/src/CodeJar.sol";
+
+import {QuarkScript} from "quark-core/src/QuarkScript.sol";
+import {QuarkNonceManager} from "quark-core/src/QuarkNonceManager.sol";
+import {QuarkWallet, QuarkWalletMetadata} from "quark-core/src/QuarkWallet.sol";
+import {QuarkWalletStandalone} from "quark-core/src/QuarkWalletStandalone.sol";
+import {IHasSignerExecutor} from "quark-core/src/interfaces/IHasSignerExecutor.sol";
+
+import {QuarkMinimalProxy} from "quark-proxy/src/QuarkMinimalProxy.sol";
+
+import {BatchSend} from "test/lib/BatchCallback.sol";
+
+contract BatchCallbackTest is Test {
+    enum ExecutionType {
+        Signature,
+        Direct
+    }
+
+    CodeJar public codeJar;
+    QuarkNonceManager public nonceManager;
+    QuarkWallet public walletImplementation;
+
+    uint256 alicePrivateKey = 0x8675309;
+    address aliceAccount = vm.addr(alicePrivateKey);
+    QuarkWallet aliceWallet; // see constructor()
+
+    uint256 bobPrivateKey = 0x8675309;
+    address bobAccount = vm.addr(bobPrivateKey);
+    QuarkWallet bobWallet; // see constructor()
+
+    bytes32 constant EXHAUSTED_TOKEN = bytes32(type(uint256).max);
+
+    // wallet proxy instantiation helper
+    function newWallet(address signer, address executor) internal returns (QuarkWallet) {
+        return QuarkWallet(payable(new QuarkMinimalProxy(address(walletImplementation), signer, executor)));
+    }
+
+    constructor() {
+        codeJar = new CodeJar();
+        console.log("CodeJar deployed to: %s", address(codeJar));
+
+        nonceManager = new QuarkNonceManager();
+        console.log("QuarkNonceManager deployed to: %s", address(nonceManager));
+
+        walletImplementation = new QuarkWallet(codeJar, nonceManager);
+        console.log("QuarkWallet implementation: %s", address(walletImplementation));
+
+        aliceWallet = newWallet(aliceAccount, address(0));
+        console.log("Alice signer: %s", aliceAccount);
+        console.log("Alice wallet at: %s", address(aliceWallet));
+
+        bobWallet = newWallet(bobAccount, address(0));
+        console.log("Bob signer: %s", bobAccount);
+        console.log("Bob wallet at: %s", address(bobWallet));
+    }
+
+    /**
+     * get active nonce, submission token, replay count ***************************
+     *
+     * single
+     */
+    function testBatchCallWithCallback() public {
+        // gas: do not meter set-up
+        vm.pauseGasMetering();
+        BatchSend batchSend = new BatchSend();
+        bytes memory incrementByCallbackScript = new YulHelper().getCode("BatchCallback.sol/IncrementByCallback.json");
+        QuarkWallet.QuarkOperation memory op1 = new QuarkOperationHelper().newBasicOpWithCalldata(
+            aliceWallet, incrementByCallbackScript, abi.encodeWithSignature("run()"), ScriptType.ScriptSource
+        );
+        (uint8 v1, bytes32 r1, bytes32 s1) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op1);
+
+        bytes memory callIncrementScript = new YulHelper().getCode("BatchCallback.sol/CallIncrement.json");
+        QuarkWallet.QuarkOperation memory op2 = new QuarkOperationHelper().newBasicOpWithCalldata(
+            bobWallet,
+            callIncrementScript,
+            abi.encodeWithSignature("run(address)", address(aliceWallet)),
+            ScriptType.ScriptSource
+        );
+        (uint8 v2, bytes32 r2, bytes32 s2) = new SignatureHelper().signOp(bobPrivateKey, bobWallet, op2);
+
+        // gas: meter execute
+        vm.resumeGasMetering();
+        vm.expectRevert(abi.encodeWithSelector(QuarkWallet.NoActiveCallback.selector));
+        batchSend.submitTwo(aliceWallet, op1, v1, r1, s1, bobWallet, op2, v2, r2, s2);
+    }
+}

test/quark-core/Noncer.t.sol

@@ -18,6 +18,8 @@ import {IHasSignerExecutor} from "quark-core/src/interfaces/IHasSignerExecutor.s
 
 import {QuarkMinimalProxy} from "quark-proxy/src/QuarkMinimalProxy.sol";
 
+import {Counter} from "test/lib/Counter.sol";
+import {MaxCounterScript} from "test/lib/MaxCounterScript.sol";
 import {Stow} from "test/lib/Noncer.sol";
 
 contract NoncerTest is Test {
@@ -27,6 +29,7 @@ contract NoncerTest is Test {
     }
 
     CodeJar public codeJar;
+    Counter public counter;
     QuarkNonceManager public nonceManager;
     QuarkWallet public walletImplementation;
 
@@ -45,6 +48,10 @@ contract NoncerTest is Test {
         codeJar = new CodeJar();
         console.log("CodeJar deployed to: %s", address(codeJar));
 
+        counter = new Counter();
+        counter.setNumber(0);
+        console.log("Counter deployed to: %s", address(counter));
+
         nonceManager = new QuarkNonceManager();
         console.log("QuarkNonceManager deployed to: %s", address(nonceManager));
 
@@ -113,7 +120,7 @@ contract NoncerTest is Test {
         assertEq(replayCount, 0);
     }
 
-    /* 
+    /*
      * nested
      */
 
@@ -148,7 +155,7 @@ contract NoncerTest is Test {
 
         (bytes32 pre, bytes32 post, bytes memory innerResult) = abi.decode(result, (bytes32, bytes32, bytes));
         assertEq(pre, op.nonce);
-        assertEq(post, bytes32(0));
+        assertEq(post, op.nonce);
         bytes32 innerNonce = abi.decode(innerResult, (bytes32));
         assertEq(innerNonce, nestedOp.nonce);
     }
@@ -184,7 +191,7 @@ contract NoncerTest is Test {
 
         (bytes32 pre, bytes32 post, bytes memory innerResult) = abi.decode(result, (bytes32, bytes32, bytes));
         assertEq(pre, op.nonce);
-        assertEq(post, bytes32(0));
+        assertEq(post, op.nonce);
         bytes32 innerNonce = abi.decode(innerResult, (bytes32));
         assertEq(innerNonce, nestedOp.nonce);
     }
@@ -248,12 +255,13 @@ contract NoncerTest is Test {
         assertEq(innerNonce, 0);
     }
 
-    function testPostNestReadFailure() public {
+    function testPostNestReadsCorrectValue() public {
         // gas: do not meter set-up
         vm.pauseGasMetering();
         bytes memory noncerScript = new YulHelper().getCode("Noncer.sol/Noncer.json");
+        bytes memory maxCounter = new YulHelper().getCode("MaxCounterScript.sol/MaxCounterScript.json");
         QuarkWallet.QuarkOperation memory nestedOp = new QuarkOperationHelper().newBasicOpWithCalldata(
-            aliceWallet, noncerScript, abi.encodeWithSignature("checkNonce()"), ScriptType.ScriptSource
+            aliceWallet, maxCounter, abi.encodeWithSignature("run(address)", address(counter)), ScriptType.ScriptSource
         );
         nestedOp.nonce = bytes32(uint256(keccak256(abi.encodePacked(block.timestamp))) - 2); // Don't overlap on nonces
         (uint8 nestedV, bytes32 nestedR, bytes32 nestedS) =
@@ -275,11 +283,22 @@ contract NoncerTest is Test {
 
         // gas: meter execute
         vm.resumeGasMetering();
-        vm.expectRevert(abi.encodeWithSelector(QuarkScript.NoActiveNonce.selector));
-        aliceWallet.executeQuarkOperation(op, v, r, s);
+        bytes memory result = aliceWallet.executeQuarkOperation(op, v, r, s);
+
+        uint256 value = abi.decode(result, (uint256));
+        assertEq(value, 0);
+        // Counter should be incremented in storage for the inner op, not the outer op
+        assertEq(
+            vm.load(address(aliceWallet), keccak256(abi.encodePacked(op.nonce, keccak256("count")))),
+            bytes32(uint256(0))
+        );
+        assertEq(
+            vm.load(address(aliceWallet), keccak256(abi.encodePacked(nestedOp.nonce, keccak256("count")))),
+            bytes32(uint256(1))
+        );
     }
 
-    /* 
+    /*
      * replayable
      */