Skip to content

Commit

Permalink
First true tests for replay
Browse files Browse the repository at this point in the history
  • Loading branch information
@hayesgm
hayesgm committed Sep 9, 2024
1 parent 9950311 commit ba3546e
Show file tree
Hide file tree
Showing 6 changed files with 201 additions and 123 deletions.
7 changes: 2 additions & 5 deletions src/quark-core/src/QuarkNonceManager.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,10 +44,6 @@ contract QuarkNonceManager {
if (lastTokenSubmission == EXHAUSTED) {
revert NonReplayableNonce(msg.sender, nonce, submissionToken, true);
}
// Defense-in-depth check for non-replayable operations
if (!isReplayable && lastTokenSubmission != FREE) {
revert NonReplayableNonce(msg.sender, nonce, submissionToken, false);
}
// Defense-in-deptch check for `submissionToken != FREE`
if (submissionToken == FREE) {
revert InvalidSubmissionToken(msg.sender, nonce, submissionToken);
Expand All @@ -62,7 +58,8 @@ contract QuarkNonceManager {
revert InvalidSubmissionToken(msg.sender, nonce, submissionToken);
}

nonceSubmissions[msg.sender][nonce] = submissionToken;
// Note: even with a valid submission token, we always set non-replayables to exhausted (e.g. for cancelations)
nonceSubmissions[msg.sender][nonce] = isReplayable ? submissionToken : EXHAUSTED;
emit NonceSubmitted(msg.sender, nonce, submissionToken);
}
}
7 changes: 4 additions & 3 deletions test/lib/CancelOtherScript.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,9 @@ pragma solidity 0.8.23;
import "quark-core/src/QuarkWallet.sol";

contract CancelOtherScript {
function run(bytes32 nonce) public {
return
QuarkWallet(payable(address(this))).nonceManager().submitNonceToken(nonce, true, bytes32(type(uint256).max));
event CancelNonce();

function run() public {
emit CancelNonce();
}
}
11 changes: 6 additions & 5 deletions test/lib/Incrementer.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,12 @@ contract Incrementer {
Counter(counter).increment();
}

// TODO: Uncomment when replay tokens are supported
// function incrementCounterReplayable(Counter counter) public {
// incrementCounter(counter);
// QuarkWallet(payable(address(this))).nonceManager().clearNonce();
// }
function incrementCounter2(Counter counter) public {
Counter(counter).increment();
Counter(counter).increment();
Counter(counter).increment();
Counter(counter).increment();
}

fallback() external {
// Counter
Expand Down
49 changes: 49 additions & 0 deletions test/lib/QuarkOperationHelper.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ pragma solidity 0.8.23;

import "forge-std/Test.sol";
import "quark-core/src/QuarkWallet.sol";
import {YulHelper} from "test/lib/YulHelper.sol";

enum ScriptType {
ScriptAddress,
Expand Down Expand Up @@ -60,6 +61,54 @@ contract QuarkOperationHelper is Test {
}
}

function newReplayableOpWithCalldata(
QuarkWallet wallet,
bytes memory scriptSource,
bytes memory scriptCalldata,
ScriptType scriptType,
uint256 replays
) public returns (QuarkWallet.QuarkOperation memory, bytes32 nonceSecret) {
return newReplayableOpWithCalldata(wallet, scriptSource, scriptCalldata, new bytes[](0), scriptType, replays);
}

function newReplayableOpWithCalldata(
QuarkWallet wallet,
bytes memory scriptSource,
bytes memory scriptCalldata,
bytes[] memory ensureScripts,
ScriptType scriptType,
uint256 replays
) public returns (QuarkWallet.QuarkOperation memory, bytes32 nonceSecret) {
QuarkWallet.QuarkOperation memory operation =
newBasicOpWithCalldata(wallet, scriptSource, scriptCalldata, ensureScripts, scriptType);
nonceSecret = operation.nonce;
bytes32 nonce = operation.nonce;
for (uint256 i = 0; i < replays; i++) {
nonce = keccak256(abi.encodePacked(nonce));
}
operation.nonce = nonce;
operation.isReplayable = true;
return (operation, nonceSecret);
}

function cancelReplayable(QuarkWallet wallet, QuarkWallet.QuarkOperation memory quarkOperation)
public
returns (QuarkWallet.QuarkOperation memory)
{
bytes memory cancelOtherScript = new YulHelper().getCode("CancelOtherScript.sol/CancelOtherScript.json");
address scriptAddress = wallet.codeJar().saveCode(cancelOtherScript);
bytes[] memory scriptSources = new bytes[](1);
scriptSources[0] = cancelOtherScript;
return QuarkWallet.QuarkOperation({
scriptAddress: scriptAddress,
scriptSources: scriptSources,
scriptCalldata: abi.encodeWithSignature("run()"),
nonce: quarkOperation.nonce,
isReplayable: false,
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) {
Expand Down
13 changes: 7 additions & 6 deletions test/quark-core/QuarkNonceManager.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,22 +106,23 @@ contract QuarkNonceManagerTest is Test {
nonceManager.submitNonceToken(nonce, false, nonce);
}

function testRevertsDefenseInDepthChangingReplayableness() public {
function testDefenseInDepthChangingReplayableness() public {
bytes32 EXHAUSTED = nonceManager.EXHAUSTED();
bytes32 nonceSecret = bytes32(uint256(99));
bytes32 nonce = keccak256(abi.encodePacked(nonceSecret));

nonceManager.submitNonceToken(nonce, true, nonce);

// Reverts when isReplayable set to false
// Accepts as a cancel
nonceManager.submitNonceToken(nonce, false, nonceSecret);

assertEq(nonceManager.getNonceSubmission(address(this), nonce), EXHAUSTED_TOKEN);

vm.expectRevert(
abi.encodeWithSelector(
QuarkNonceManager.NonReplayableNonce.selector, address(this), nonce, nonceSecret, false
QuarkNonceManager.NonReplayableNonce.selector, address(this), nonce, nonceSecret, true
)
);
nonceManager.submitNonceToken(nonce, false, nonceSecret);

// Accepts when isReplayable set back to true
nonceManager.submitNonceToken(nonce, true, nonceSecret);
}

Expand Down
237 changes: 133 additions & 104 deletions test/quark-core/QuarkWallet.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -297,121 +297,150 @@ contract QuarkWalletTest is Test {

/* ===== replayability tests ===== */

// TODO: Uncomment when replay tokens are supported all of these tests
// function testCanReplaySameScriptWithDifferentCall() public {
// // gas: disable gas metering except while executing operations
// vm.pauseGasMetering();
// bytes memory incrementer = new YulHelper().getCode("Incrementer.sol/Incrementer.json");

// // 1. use nonce to increment a counter
// QuarkWallet.QuarkOperation memory op1 = new QuarkOperationHelper().newBasicOpWithCalldata(
// aliceWallet,
// incrementer,
// abi.encodeWithSignature("incrementCounterReplayable(address)", address(counter)),
// ScriptType.ScriptAddress
// );
// (uint8 v1, bytes32 r1, bytes32 s1) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op1);
function testCanReplaySameScriptWithDifferentCall() public {
// gas: disable gas metering except while executing operations
vm.pauseGasMetering();
bytes memory incrementer = new YulHelper().getCode("Incrementer.sol/Incrementer.json");

// address incrementerAddress = codeJar.saveCode(incrementer);
// 1. use nonce to increment a counter
(QuarkWallet.QuarkOperation memory op1, bytes32 nonceSecret) = new QuarkOperationHelper()
.newReplayableOpWithCalldata(
aliceWallet,
incrementer,
abi.encodeWithSignature("incrementCounter(address)", address(counter)),
ScriptType.ScriptAddress,
1
);
(uint8 v1, bytes32 r1, bytes32 s1) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op1);

// QuarkWallet.QuarkOperation memory op2 = QuarkWallet.QuarkOperation({
// nonce: op1.nonce,
// scriptAddress: incrementerAddress,
// scriptSources: new bytes[](0),
// scriptCalldata: abi.encodeWithSignature("incrementCounter(address)", address(counter)),
// expiry: block.timestamp + 1000
// });
// (uint8 v2, bytes32 r2, bytes32 s2) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op2);
address incrementerAddress = codeJar.saveCode(incrementer);

// // gas: meter execute
// vm.resumeGasMetering();
// vm.expectEmit(true, true, true, true);
// emit ClearNonce(address(aliceWallet), op1.nonce);
// aliceWallet.executeQuarkOperation(op1, v1, r1, s1);
// // incrementer increments the counter thrice
// assertEq(counter.number(), 3);
// // when reusing the nonce, you can change the call
// aliceWallet.executeQuarkOperation(op2, v2, r2, s2);
// // incrementer increments the counter thrice
// assertEq(counter.number(), 6);
// // but now that we did not use a replayable call, it is canceled
// vm.expectRevert(abi.encodeWithSelector(QuarkNonceManager.NonceAlreadySet.selector));
// aliceWallet.executeQuarkOperation(op1, v1, r1, s1);
// }
QuarkWallet.QuarkOperation memory op2 = QuarkWallet.QuarkOperation({
nonce: op1.nonce,
isReplayable: true,
scriptAddress: incrementerAddress,
scriptSources: new bytes[](0),
scriptCalldata: abi.encodeWithSignature("incrementCounter2(address)", address(counter)),
expiry: block.timestamp + 1000
});
(uint8 v2, bytes32 r2, bytes32 s2) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op2);

// function testRevertsForReusedNonceWithChangedScript() public {
// // gas: disable gas metering except while executing operations
// vm.pauseGasMetering();
// bytes memory incrementer = new YulHelper().getCode("Incrementer.sol/Incrementer.json");

// // 1. use nonce to increment a counter
// QuarkWallet.QuarkOperation memory op1 = new QuarkOperationHelper().newBasicOpWithCalldata(
// aliceWallet,
// incrementer,
// abi.encodeWithSignature("incrementCounterReplayable(address)", address(counter)),
// ScriptType.ScriptAddress
// );
// (uint8 v1, bytes32 r1, bytes32 s1) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op1);
// gas: meter execute
vm.resumeGasMetering();
aliceWallet.executeQuarkOperation(op1, v1, r1, s1);
// incrementer increments the counter thrice
assertEq(counter.number(), 3);
// when executing a replayable operation, you can change the call
aliceWallet.executeQuarkOperationWithSubmissionToken(op2, nonceSecret, v2, r2, s2);
// incrementer increments the counter frice
assertEq(counter.number(), 7);
// but now both operations are exhausted
vm.expectRevert(
abi.encodeWithSelector(QuarkNonceManager.InvalidSubmissionToken.selector, aliceWallet, op1.nonce, op1.nonce)
);
aliceWallet.executeQuarkOperation(op1, v1, r1, s1);
vm.expectRevert(
abi.encodeWithSelector(
QuarkNonceManager.InvalidSubmissionToken.selector, aliceWallet, op1.nonce, nonceSecret
)
);
aliceWallet.executeQuarkOperationWithSubmissionToken(op1, nonceSecret, v1, r1, s1);
vm.expectRevert(
abi.encodeWithSelector(QuarkNonceManager.InvalidSubmissionToken.selector, aliceWallet, op1.nonce, op2.nonce)
);
aliceWallet.executeQuarkOperation(op2, v2, r2, s2);
vm.expectRevert(
abi.encodeWithSelector(
QuarkNonceManager.InvalidSubmissionToken.selector, aliceWallet, op1.nonce, nonceSecret
)
);
aliceWallet.executeQuarkOperationWithSubmissionToken(op2, nonceSecret, v2, r2, s2);
}

// QuarkWallet.QuarkOperation memory op2 = QuarkWallet.QuarkOperation({
// nonce: op1.nonce,
// scriptAddress: address(counter),
// scriptSources: new bytes[](0),
// scriptCalldata: bytes(""),
// expiry: op1.expiry
// });
// (uint8 v2, bytes32 r2, bytes32 s2) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op2);
function testAllowsForReusedNonceWithChangedScript() public {
// gas: disable gas metering except while executing operations
vm.pauseGasMetering();
bytes memory incrementer = new YulHelper().getCode("Incrementer.sol/Incrementer.json");

// // gas: meter execute
// vm.resumeGasMetering();
// vm.expectEmit(true, true, true, true);
// emit ClearNonce(address(aliceWallet), op1.nonce);
// aliceWallet.executeQuarkOperation(op1, v1, r1, s1);
// // incrementer increments the counter thrice
// assertEq(counter.number(), 3);
// // when reusing the nonce but changing the script, revert
// vm.expectRevert(abi.encodeWithSelector(QuarkNonceManager.NonceScriptMismatch.selector));
// aliceWallet.executeQuarkOperation(op2, v2, r2, s2);
// }
// 1. use nonce to increment a counter
(QuarkWallet.QuarkOperation memory op1, bytes32 nonceSecret) = new QuarkOperationHelper()
.newReplayableOpWithCalldata(
aliceWallet,
incrementer,
abi.encodeWithSignature("incrementCounter(address)", address(counter)),
ScriptType.ScriptAddress,
1
);
(uint8 v1, bytes32 r1, bytes32 s1) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op1);

// function testRevertsForReplayOfCanceledScript() public {
// // gas: disable gas metering except while executing operations
// vm.pauseGasMetering();
// bytes memory incrementer = new YulHelper().getCode("Incrementer.sol/Incrementer.json");
// bytes memory cancelOtherScript = new YulHelper().getCode("CancelOtherScript.sol/CancelOtherScript.json");
QuarkWallet.QuarkOperation memory op2 = QuarkWallet.QuarkOperation({
nonce: op1.nonce,
isReplayable: true,
scriptAddress: address(counter),
scriptSources: new bytes[](0),
scriptCalldata: abi.encodeWithSignature("setNumber(uint256)", 999),
expiry: op1.expiry
});
(uint8 v2, bytes32 r2, bytes32 s2) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op2);

// QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata(
// aliceWallet,
// incrementer,
// abi.encodeWithSignature("incrementCounterReplayable(address)", address(counter)),
// ScriptType.ScriptAddress
// );
// (uint8 v, bytes32 r, bytes32 s) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op);
// gas: meter execute
vm.resumeGasMetering();
aliceWallet.executeQuarkOperation(op1, v1, r1, s1);
// incrementer increments the counter thrice
assertEq(counter.number(), 3);
// when reusing the nonce but changing the script, allow
aliceWallet.executeQuarkOperationWithSubmissionToken(op2, nonceSecret, v2, r2, s2);
// code didn't change counter since it was callcode
assertEq(counter.number(), 3);
}

// // gas: meter execute
// vm.resumeGasMetering();
// vm.expectEmit(true, true, true, true);
// emit ClearNonce(address(aliceWallet), op.nonce);
// aliceWallet.executeQuarkOperation(op, v, r, s);
// assertEq(counter.number(), 3);
// // can replay the same operation...
// aliceWallet.executeQuarkOperation(op, v, r, s);
// assertEq(counter.number(), 6);
function testScriptCanBeCanceled() public {
// gas: disable gas metering except while executing operations
vm.pauseGasMetering();
bytes memory incrementer = new YulHelper().getCode("Incrementer.sol/Incrementer.json");

// // can cancel the replayable nonce...
// vm.pauseGasMetering();
// QuarkWallet.QuarkOperation memory cancelOtherOp = new QuarkOperationHelper().newBasicOpWithCalldata(
// aliceWallet, cancelOtherScript, abi.encodeWithSignature("run(bytes32)", op.nonce), ScriptType.ScriptAddress
// );
// (uint8 cancel_v, bytes32 cancel_r, bytes32 cancel_s) =
// new SignatureHelper().signOp(alicePrivateKey, aliceWallet, cancelOtherOp);
// vm.resumeGasMetering();
// aliceWallet.executeQuarkOperation(cancelOtherOp, cancel_v, cancel_r, cancel_s);
(QuarkWallet.QuarkOperation memory op, bytes32 nonceSecret) = new QuarkOperationHelper()
.newReplayableOpWithCalldata(
aliceWallet,
incrementer,
abi.encodeWithSignature("incrementCounter(address)", address(counter)),
ScriptType.ScriptAddress,
1
);
(uint8 v, bytes32 r, bytes32 s) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op);

// // and now you can no longer replay
// vm.expectRevert(abi.encodeWithSelector(QuarkNonceManager.NonceAlreadySet.selector));
// aliceWallet.executeQuarkOperation(op, v, r, s);
// }
// gas: meter execute
vm.resumeGasMetering();
aliceWallet.executeQuarkOperation(op, v, r, s);
assertEq(counter.number(), 3);
// cannot replay the same operation directly...
vm.expectRevert(
abi.encodeWithSelector(QuarkNonceManager.InvalidSubmissionToken.selector, aliceWallet, op.nonce, op.nonce)
);
aliceWallet.executeQuarkOperation(op, v, r, s);
assertEq(counter.number(), 3);

// can cancel the replayable nonce...
vm.pauseGasMetering();
QuarkWallet.QuarkOperation memory cancelOtherOp = new QuarkOperationHelper().cancelReplayable(aliceWallet, op);
(uint8 cancel_v, bytes32 cancel_r, bytes32 cancel_s) =
new SignatureHelper().signOp(alicePrivateKey, aliceWallet, cancelOtherOp);
vm.resumeGasMetering();
vm.expectEmit(true, true, true, true);
emit CancelOtherScript.CancelNonce();
aliceWallet.executeQuarkOperationWithSubmissionToken(cancelOtherOp, nonceSecret, cancel_v, cancel_r, cancel_s);

// and now you can no longer replay
vm.expectRevert(
abi.encodeWithSelector(
QuarkNonceManager.NonReplayableNonce.selector, address(aliceWallet), op.nonce, nonceSecret, true
)
);
aliceWallet.executeQuarkOperationWithSubmissionToken(op, nonceSecret, v, r, s);

// Ensure exhausted
assertEq(nonceManager.getNonceSubmission(address(aliceWallet), op.nonce), bytes32(type(uint256).max));
}

/* ===== direct execution path tests ===== */

Expand Down

0 comments on commit ba3546e

Please sign in to comment.