Add Replayables with Token

This code leads out the core changes for replayables with token. This was meant to be a WIP that just showed off the core features (it still is), but to get it to compile, I ended up converting most of the current tests to the new framework (e.g. mostly converting them to pick a semi-random nonce versus `nextNonce`, and checking `getNonceToken` as opposed to `getNextNonce`). Overall, the changes to the non-test code are very straight-forward. Most notable should be the changes to `QuarkStateManager` (adding new nonce and replay token code), and adding a new function to QuarkWallet `verifySigAndExecuteReplayableQuarkOperation`.

Note: I didn't spend much time working on the outer interface for this, and I need to consider how this works best with multi-quark operations, but I wanted to get the outline first so we could discuss. There are also failing test cases and absent test cases, but again, more of a discussion point than a final product here.

Patches:
  * Fix current tests

src/quark-core/src/QuarkStateManager.sol

@@ -9,76 +9,47 @@ import {IQuarkWallet} from "quark-core/src/interfaces/IQuarkWallet.sol";
  * @author Compound Labs, Inc.
  */
 contract QuarkStateManager {
-    error NonceAlreadySet();
-    error NoUnusedNonces();
+    error NonReplayableNonce(address wallet, bytes32 nonce, bytes32 replayToken);
+    error InvalidReplayToken(address wallet, bytes32 nonce, bytes32 replayToken);
 
-    /// @notice Bit-packed nonce values
-    mapping(address wallet => mapping(uint256 bucket => uint256 bitset)) public nonces;
+    event NonceSubmitted(address wallet, bytes32 nonce, bytes32 replayToken);
 
-    /**
-     * @notice Return whether a nonce has been exhausted
-     * @param wallet Address of the wallet owning the nonce
-     * @param nonce Nonce to check
-     * @return Whether the nonce has been exhausted
-     */
-    function isNonceSet(address wallet, uint96 nonce) public view returns (bool) {
-        (uint256 bucket, uint256 mask) = getBucket(nonce);
-        return isNonceSetInternal(wallet, bucket, mask);
-    }
+    /// @notice Represents the unclaimed bytes32 value.
+    bytes32 public constant CLAIMABLE_TOKEN = bytes32(uint256(0));
 
-    /// @dev Returns if a given nonce is set for a wallet, using the nonce's bucket and mask
-    function isNonceSetInternal(address wallet, uint256 bucket, uint256 mask) internal view returns (bool) {
-        return (nonces[wallet][bucket] & mask) != 0;
-    }
+    /// @notice A token that implies a Quark Operation is no longer replayable.
+    bytes32 public constant NO_REPLAY_TOKEN = bytes32(type(uint).max);
+
+    /// @notice Mapping from nonces to last used replay token.
+    mapping(address wallet => mapping(bytes32 nonce => bytes32 lastToken)) public nonceTokens;
 
     /**
-     * @notice Returns the next valid unset nonce for a given wallet
-     * @dev Any unset nonce is valid to use, but using this method
-     * increases the likelihood that the nonce you use will be in a bucket that
-     * has already been written to, which costs less gas
-     * @param wallet Address of the wallet to find the next nonce for
-     * @return The next unused nonce
+     * @notice Returns the nonce token (last replay token) for a given nonce. For finalized scripts, this will be `uint256(-1)`. For unclaimed nonces, this will be `uint256(0)`. Otherwise, it will be the next value in the replay chain.
+     * @param wallet The wallet for which to get the nonce token.
+     * @param nonce The nonce for the given request.
+     * @return replayToken The last used replay token, or 0 if unused or -1 if finalized.
      */
-    function nextNonce(address wallet) external view returns (uint96) {
-        // Any bucket larger than `type(uint88).max` will result in unsafe undercast when converting to nonce
-        for (uint256 bucket = 0; bucket <= type(uint88).max; ++bucket) {
-            uint96 bucketValue = uint96(bucket << 8);
-            uint256 bucketNonces = nonces[wallet][bucket];
-            // Move on to the next bucket if all bits in this bucket are already set
-            if (bucketNonces == type(uint256).max) continue;
-            for (uint256 maskOffset = 0; maskOffset < 256; ++maskOffset) {
-                uint256 mask = 1 << maskOffset;
-                if ((bucketNonces & mask) == 0) {
-                    uint96 nonce = uint96(bucketValue + maskOffset);
-                    return nonce;
-                }
-            }
-        }
-
-        revert NoUnusedNonces();
-    }
-
-    /// @dev Locate a nonce at a (bucket, mask) bitset position in the nonces mapping
-    function getBucket(uint96 nonce) internal pure returns (uint256, /* bucket */ uint256 /* mask */ ) {
-        uint256 bucket = nonce >> 8;
-        uint256 setMask = 1 << (nonce & 0xff);
-        return (bucket, setMask);
+    function getNonceToken(address wallet, bytes32 nonce) external view returns (bytes32 replayToken) {
+        return nonceTokens[wallet][nonce];
     }
 
     /**
-     * @notice Claim a given nonce for the calling wallet, reverting if the nonce is already set
-     * @param nonce Nonce to claim for the calling wallet
+     * @notice Attempts a first or subsequent submission of a given nonce from a wallet.
+     * @param nonce The nonce of the chain to submit.
+     * @param replayToken The replay token of the submission. For single-use operations, set `replayToken` to `uint256(-1)`. For first-use replayable operations, set `replayToken` = `nonce`.
      */
-    function claimNonce(uint96 nonce) external {
-        (uint256 bucket, uint256 setMask) = getBucket(nonce);
-        if (isNonceSetInternal(msg.sender, bucket, setMask)) {
-            revert NonceAlreadySet();
+    function submitNonceToken(bytes32 nonce, bytes32 replayToken) external {
+        bytes32 lastToken = nonceTokens[msg.sender][nonce];
+        if (lastToken == NO_REPLAY_TOKEN) {
+            revert NonReplayableNonce(msg.sender, nonce, replayToken);
+        }
+
+        bool validReplay = lastToken == CLAIMABLE_TOKEN || keccak256(abi.encodePacked(replayToken)) == lastToken;
+        if (!validReplay) {
+            revert InvalidReplayToken(msg.sender, nonce, replayToken);
         }
-        setNonceInternal(bucket, setMask);
-    }
 
-    /// @dev Set a nonce for the msg.sender, using the nonce's bucket and mask
-    function setNonceInternal(uint256 bucket, uint256 setMask) internal {
-        nonces[msg.sender][bucket] |= setMask;
+        nonceTokens[msg.sender][nonce] = replayToken;
+        emit NonceSubmitted(msg.sender, nonce, replayToken);
     }
 }

src/quark-core/src/QuarkWallet.sol

@@ -23,7 +23,7 @@ library QuarkWalletMetadata {
 
     /// @notice The EIP-712 typehash for authorizing an operation for this version of QuarkWallet
     bytes32 internal constant QUARK_OPERATION_TYPEHASH = keccak256(
-        "QuarkOperation(uint96 nonce,address scriptAddress,bytes[] scriptSources,bytes scriptCalldata,uint256 expiry)"
+        "QuarkOperation(bytes32 nonce,address scriptAddress,bytes[] scriptSources,bytes scriptCalldata,uint256 expiry)"
     );
 
     /// @notice The EIP-712 typehash for authorizing a MultiQuarkOperation for this version of QuarkWallet
@@ -65,7 +65,7 @@ contract QuarkWallet is IERC1271 {
 
     /// @notice Event emitted when a Quark script is executed by this Quark wallet
     event ExecuteQuarkScript(
-        address indexed executor, address indexed scriptAddress, uint96 indexed nonce, ExecutionType executionType
+        address indexed executor, address indexed scriptAddress, bytes32 indexed nonce, ExecutionType executionType
     );
 
     /// @notice Address of CodeJar contract used to deploy transaction script source code
@@ -113,13 +113,16 @@ contract QuarkWallet is IERC1271 {
     /// @notice Well-known storage slot for the currently executing script's address (if any)
     bytes32 public constant ACTIVE_SCRIPT_SLOT = bytes32(uint256(keccak256("quark.v1.active.script")) - 1);
 
+    /// @notice A token that implies a Quark Operation is no longer replayable.
+    bytes32 public constant NO_REPLAY_TOKEN = bytes32(type(uint).max);
+
     /// @notice The magic value to return for valid ERC1271 signature
     bytes4 internal constant EIP_1271_MAGIC_VALUE = 0x1626ba7e;
 
     /// @notice The structure of a signed operation to execute in the context of this wallet
     struct QuarkOperation {
         /// @notice Nonce identifier for the operation
-        uint96 nonce;
+        bytes32 nonce;
         /// @notice The address of the transaction script to run
         address scriptAddress;
         /// @notice Creation codes Quark must ensure are deployed before executing this operation
@@ -193,7 +196,7 @@ contract QuarkWallet is IERC1271 {
     }
 
     /**
-     * @notice Verify a signature and execute a QuarkOperation
+     * @notice Verify a signature and execute a single-use QuarkOperation
      * @param op A QuarkOperation struct
      * @param digest A EIP-712 digest for either a QuarkOperation or MultiQuarkOperation to verify the signature against
      * @param v EIP-712 signature v value
@@ -220,7 +223,44 @@ contract QuarkWallet is IERC1271 {
             codeJar.saveCode(op.scriptSources[i]);
         }
 
-        stateManager.claimNonce(op.nonce);
+        stateManager.submitNonceToken(op.nonce, NO_REPLAY_TOKEN);
+
+        emit ExecuteQuarkScript(msg.sender, op.scriptAddress, op.nonce, ExecutionType.Signature);
+
+        return executeScriptInternal(op.scriptAddress, op.scriptCalldata);
+    }
+
+    /**
+     * @notice Verify a signature and execute a replayable QuarkOperation
+     * @param op A QuarkOperation struct
+     * @param replayToken The replay token for the replayable quark operation. For the first submission, this is generally the `rootHash`.
+     * @param digest A EIP-712 digest for either a QuarkOperation or MultiQuarkOperation to verify the signature against
+     * @param v EIP-712 signature v value
+     * @param r EIP-712 signature r value
+     * @param s EIP-712 signature s value
+     * @return Return value from the executed operation
+     */
+    function verifySigAndExecuteReplayableQuarkOperation(
+        QuarkOperation calldata op,
+        bytes32 replayToken,
+        bytes32 digest,
+        uint8 v,
+        bytes32 r,
+        bytes32 s
+    ) internal returns (bytes memory) {
+        if (block.timestamp >= op.expiry) {
+            revert SignatureExpired();
+        }
+
+        // if the signature check does not revert, the signature is valid
+        checkValidSignatureInternal(IHasSignerExecutor(address(this)).signer(), digest, v, r, s);
+
+        // guarantee every script in scriptSources is deployed
+        for (uint256 i = 0; i < op.scriptSources.length; ++i) {
+            codeJar.saveCode(op.scriptSources[i]);
+        }
+
+        stateManager.submitNonceToken(op.nonce, replayToken);
 
         emit ExecuteQuarkScript(msg.sender, op.scriptAddress, op.nonce, ExecutionType.Signature);
 
@@ -237,7 +277,7 @@ contract QuarkWallet is IERC1271 {
      * @return Return value from the executed operation
      */
     function executeScript(
-        uint96 nonce,
+        bytes32 nonce,
         address scriptAddress,
         bytes calldata scriptCalldata,
         bytes[] calldata scriptSources
@@ -252,7 +292,7 @@ contract QuarkWallet is IERC1271 {
             codeJar.saveCode(scriptSources[i]);
         }
 
-        stateManager.claimNonce(nonce);
+        stateManager.submitNonceToken(nonce, NO_REPLAY_TOKEN);
 
         emit ExecuteQuarkScript(msg.sender, scriptAddress, nonce, ExecutionType.Direct);
 

src/quark-core/src/interfaces/IQuarkWallet.sol

@@ -10,7 +10,7 @@ interface IQuarkWallet {
     /// @notice The structure of a signed operation to execute in the context of this wallet
     struct QuarkOperation {
         /// @notice Nonce identifier for the operation
-        uint96 nonce;
+        bytes32 nonce;
         /// @notice The address of the transaction script to run
         address scriptAddress;
         /// @notice Creation codes Quark must ensure are deployed before executing this operation
@@ -32,7 +32,7 @@ interface IQuarkWallet {
         bytes32 s
     ) external returns (bytes memory);
     function executeScript(
-        uint96 nonce,
+        bytes32 nonce,
         address scriptAddress,
         bytes calldata scriptCalldata,
         bytes[] calldata scriptSources

test/lib/CancelOtherScript.sol

@@ -4,7 +4,7 @@ pragma solidity 0.8.23;
 import "quark-core/src/QuarkWallet.sol";
 
 contract CancelOtherScript {
-    function run(uint96 nonce) public {
-        return QuarkWallet(payable(address(this))).stateManager().claimNonce(nonce);
+    function run(bytes32 nonce) public {
+        return QuarkWallet(payable(address(this))).stateManager().submitNonceToken(nonce, bytes32(type(uint).max));
     }
 }

test/lib/ExecuteOnBehalf.sol

@@ -4,7 +4,7 @@ pragma solidity 0.8.23;
 import "quark-core/src/QuarkWallet.sol";
 
 contract ExecuteOnBehalf {
-    function run(QuarkWallet targetWallet, uint96 nonce, address scriptAddress, bytes calldata scriptCalldata)
+    function run(QuarkWallet targetWallet, bytes32 nonce, address scriptAddress, bytes calldata scriptCalldata)
         public
         returns (bytes memory)
     {

test/lib/ExecuteWithRequirements.sol

@@ -5,16 +5,16 @@ import "quark-core/src/QuarkWallet.sol";
 import "quark-core/src/QuarkStateManager.sol";
 
 contract ExecuteWithRequirements {
-    error RequirementNotMet(uint96 nonce);
+    error RequirementNotMet(bytes32 nonce);
 
-    function runWithRequirements(uint96[] memory requirements, address scriptAddress, bytes calldata scriptCalldata)
+    function runWithRequirements(bytes32[] memory requirements, address scriptAddress, bytes calldata scriptCalldata)
         public
         returns (bytes memory)
     {
         QuarkWallet wallet = QuarkWallet(payable(address(this)));
         QuarkStateManager stateManager = wallet.stateManager();
-        for (uint96 i = 0; i < requirements.length; i++) {
-            if (!stateManager.isNonceSet(address(wallet), requirements[i])) {
+        for (uint256 i = 0; i < requirements.length; i++) {
+            if (stateManager.getNonceToken(address(wallet), requirements[i]) == bytes32(uint256(0))) {
                 revert RequirementNotMet(requirements[i]);
             }
         }

test/lib/QuarkOperationHelper.sol

@@ -12,6 +12,9 @@ enum ScriptType {
 // TODO: QuarkOperationHelper ScriptType doesn't really make sense anymore, since scriptSource
 // has been replaced with scriptSources and scriptAddress is now always required.
 contract QuarkOperationHelper is Test {
+    error SemiRandomNonceRequiresQuarkStateManagerOrInitializedQuarkWallet(address quarkWallet);
+    error Impossible();
+
     function newBasicOp(QuarkWallet wallet, bytes memory scriptSource, ScriptType scriptType)
         external
         returns (QuarkWallet.QuarkOperation memory)
@@ -41,17 +44,43 @@ contract QuarkOperationHelper is Test {
                 scriptAddress: scriptAddress,
                 scriptSources: ensureScripts,
                 scriptCalldata: scriptCalldata,
-                nonce: wallet.stateManager().nextNonce(address(wallet)),
+                nonce: semiRandomNonce(wallet),
                 expiry: block.timestamp + 1000
             });
         } else {
             return QuarkWallet.QuarkOperation({
                 scriptAddress: scriptAddress,
                 scriptSources: ensureScripts,
                 scriptCalldata: scriptCalldata,
-                nonce: wallet.stateManager().nextNonce(address(wallet)),
+                nonce: semiRandomNonce(wallet),
                 expiry: block.timestamp + 1000
             });
         }
     }
+
+    /// @dev Note: not sufficiently random for non-test case usage.
+    function semiRandomNonce(QuarkWallet wallet) public view returns (bytes32) {
+        if (address(wallet).code.length == 0) {
+            revert SemiRandomNonceRequiresQuarkStateManagerOrInitializedQuarkWallet(address(wallet));
+        }
+
+        return semiRandomNonce(wallet.stateManager(), wallet);
+    }
+
+    /// @dev Note: not sufficiently random for non-test case usage.
+    function semiRandomNonce(QuarkStateManager quarkStateManager, QuarkWallet wallet) public view returns (bytes32) {
+        bytes32 nonce = bytes32(uint256(keccak256(abi.encodePacked(block.timestamp))) - 1);
+        while (true) {
+            if (quarkStateManager.getNonceToken(address(wallet), nonce) == bytes32(uint256(0))) {
+                return nonce;
+            }
+
+            nonce = bytes32(uint256(keccak256(abi.encodePacked(nonce))) - 1);
+        }
+        revert Impossible();
+    }
+
+    function incrementNonce(bytes32 nonce) public pure returns (bytes32) {
+        return bytes32(uint256(nonce) + 1);
+    }
 }

test/quark-core-scripts/Multicall.t.sol

@@ -391,8 +391,10 @@ contract MulticallTest is Test {
         // 1. transfer 0.5 WETH from wallet A to wallet B
         wallets[0] = address(walletA);
         walletCalls[0] = abi.encodeWithSignature(
-            "executeScript(uint96,address,bytes,bytes[])",
-            QuarkWallet(payable(factory.walletImplementation())).stateManager().nextNonce(address(walletA)),
+            "executeScript(bytes32,address,bytes,bytes[])",
+            new QuarkOperationHelper().semiRandomNonce(
+                QuarkWallet(payable(factory.walletImplementation())).stateManager(), walletA
+            ),
             ethcallAddress,
             abi.encodeWithSelector(
                 Ethcall.run.selector,
@@ -404,11 +406,12 @@ contract MulticallTest is Test {
         );
 
         // 2. approve Comet cUSDCv3 to receive 0.5 WETH from wallet B
-        uint96 walletBNextNonce =
-            QuarkWallet(payable(factory.walletImplementation())).stateManager().nextNonce(address(walletB));
+        bytes32 walletBNextNonce = new QuarkOperationHelper().semiRandomNonce(
+            QuarkWallet(payable(factory.walletImplementation())).stateManager(), walletB
+        );
         wallets[1] = address(walletB);
         walletCalls[1] = abi.encodeWithSignature(
-            "executeScript(uint96,address,bytes,bytes[])",
+            "executeScript(bytes32,address,bytes,bytes[])",
             walletBNextNonce,
             ethcallAddress,
             abi.encodeWithSelector(
@@ -423,8 +426,8 @@ contract MulticallTest is Test {
         // 3. supply 0.5 WETH from wallet B to Comet cUSDCv3
         wallets[2] = address(walletB);
         walletCalls[2] = abi.encodeWithSignature(
-            "executeScript(uint96,address,bytes,bytes[])",
-            walletBNextNonce + 1,
+            "executeScript(bytes32,address,bytes,bytes[])",
+            bytes32(uint256(walletBNextNonce) + 1),
             ethcallAddress,
             abi.encodeWithSelector(
                 Ethcall.run.selector,
@@ -493,7 +496,10 @@ contract MulticallTest is Test {
         deal(WETH, address(wallet), 100 ether);
 
         address subWallet1 = factory.walletAddressForSalt(alice, address(wallet), bytes32("1"));
-        uint96 nonce = QuarkWallet(payable(factory.walletImplementation())).stateManager().nextNonce(subWallet1);
+        bytes32 nonce = new QuarkOperationHelper().semiRandomNonce(
+            QuarkWallet(payable(factory.walletImplementation())).stateManager(), QuarkWallet(payable(subWallet1))
+        );
+
         // Steps: Wallet#1: Supply WETH to Comet -> Borrow USDC from Comet(USDC) to subwallet -> Create subwallet
         // -> Swap USDC to WETH on Uniswap -> Supply WETH to Comet(WETH)
         address[] memory callContracts = new address[](5);
@@ -534,7 +540,7 @@ contract MulticallTest is Test {
                                 path: abi.encodePacked(USDC, uint24(500), WETH) // Path: USDC - 0.05% -> WETH
                             })
                         )
-                        ),
+                    ),
                     new bytes[](0)
                 )
             ),
@@ -548,7 +554,7 @@ contract MulticallTest is Test {
             abi.encodeCall(
                 QuarkWallet.executeScript,
                 (
-                    nonce + 1,
+                    new QuarkOperationHelper().incrementNonce(nonce),
                     legendCometSupplyScriptAddress,
                     abi.encodeCall(CometSupplyActions.supply, (cWETHv3, WETH, 2 ether)),
                     new bytes[](0)