Skip to content

Commit

Permalink
Restore transient slot values after calling script (#217)
Browse files Browse the repository at this point in the history
We now cache the old transient slot values (e.g. for `activeNonce`,
`callback`, etc) in order to restore them after calling the Quark
script.

This also adds a new test to test batched submission for a contract with
callbacks to check its behavior.

---------

Co-authored-by: kevincheng96 <[email protected]>
  • Loading branch information
hayesgm and kevincheng96 authored Sep 13, 2024
1 parent b00f3f6 commit 95da9ec
Show file tree
Hide file tree
Showing 6 changed files with 187 additions and 20 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/lint.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 0 additions & 4 deletions src/quark-core/src/QuarkScript.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down Expand Up @@ -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));
}

Expand Down
26 changes: 19 additions & 7 deletions src/quark-core/src/QuarkWallet.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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);
Expand Down
44 changes: 44 additions & 0 deletions test/lib/BatchCallback.sol
Original file line number Diff line number Diff line change
@@ -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();
}
}
96 changes: 96 additions & 0 deletions test/quark-core/BatchCallback.t.sol
Original file line number Diff line number Diff line change
@@ -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);
}
}
35 changes: 27 additions & 8 deletions test/quark-core/Noncer.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -27,6 +29,7 @@ contract NoncerTest is Test {
}

CodeJar public codeJar;
Counter public counter;
QuarkNonceManager public nonceManager;
QuarkWallet public walletImplementation;

Expand All @@ -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));

Expand Down Expand Up @@ -113,7 +120,7 @@ contract NoncerTest is Test {
assertEq(replayCount, 0);
}

/*
/*
* nested
*/

Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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) =
Expand All @@ -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
*/

Expand Down

0 comments on commit 95da9ec

Please sign in to comment.