Skip to content

Commit

Permalink
Remove isolated storage
Browse files Browse the repository at this point in the history
  • Loading branch information
@kevincheng96
kevincheng96 committed Aug 24, 2024
1 parent e36c8cc commit 3661014
Show file tree
Hide file tree
Showing 16 changed files with 1,010 additions and 1,191 deletions.
263 changes: 122 additions & 141 deletions .gas-snapshot
View file

Large diffs are not rendered by default.

45 changes: 34 additions & 11 deletions src/quark-core/src/QuarkScript.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,24 @@ abstract contract QuarkScript {

/// @notice A safer, but gassier reentrancy guard that writes the flag to the QuarkStateManager
modifier nonReentrant() {
if (read(REENTRANCY_FLAG_SLOT) == bytes32(uint256(1))) {
bytes32 slot = REENTRANCY_FLAG_SLOT;
bytes32 flag;
// TODO: Move to TSTORE after updating Solidity version to >=0.8.24
assembly {
flag := sload(slot)
}
if (flag == bytes32(uint256(1))) {
revert ReentrantCall();
}
write(REENTRANCY_FLAG_SLOT, bytes32(uint256(1)));
assembly {
sstore(slot, 1)
}

_;

write(REENTRANCY_FLAG_SLOT, bytes32(uint256(0)));
assembly {
sstore(slot, 0)
}
}

/**
Expand Down Expand Up @@ -58,16 +68,23 @@ abstract contract QuarkScript {

function allowCallback() internal {
QuarkWallet self = QuarkWallet(payable(address(this)));
self.stateManager().write(self.CALLBACK_KEY(), bytes32(uint256(uint160(self.stateManager().getActiveScript()))));
// TODO: Can save gas by just having the constant in QuarkScript
bytes32 callbackSlot = self.CALLBACK_SLOT();
bytes32 activeScriptSlot = self.ACTIVE_SCRIPT_SLOT();
assembly {
// TODO: Move to TLOAD/TSTORE after updating Solidity version to >=0.8.24
let activeScript := sload(activeScriptSlot)
sstore(callbackSlot, activeScript)
}
}

function clearCallback() internal {
QuarkWallet self = QuarkWallet(payable(address(this)));
self.stateManager().write(self.CALLBACK_KEY(), bytes32(0));
}

function allowReplay() internal {
return QuarkWallet(payable(address(this))).stateManager().clearNonce();
bytes32 callbackSlot = self.CALLBACK_SLOT();
assembly {
// TODO: Move to TSTORE after updating Solidity version to >=0.8.24
sstore(callbackSlot, 0)
}
}

function readU256(string memory key) internal view returns (uint256) {
Expand All @@ -79,7 +96,11 @@ abstract contract QuarkScript {
}

function read(bytes32 key) internal view returns (bytes32) {
return QuarkWallet(payable(address(this))).stateManager().read(key);
bytes32 value;
assembly {
value := sload(key)
}
return value;
}

function writeU256(string memory key, uint256 value) internal {
Expand All @@ -91,6 +112,8 @@ abstract contract QuarkScript {
}

function write(bytes32 key, bytes32 value) internal {
return QuarkWallet(payable(address(this))).stateManager().write(key, value);
assembly {
sstore(key, value)
}
}
}
126 changes: 8 additions & 118 deletions src/quark-core/src/QuarkStateManager.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,38 +5,18 @@ import {IQuarkWallet} from "quark-core/src/interfaces/IQuarkWallet.sol";

/**
* @title Quark State Manager
* @notice Contract for managing nonces and storage for Quark wallets, guaranteeing storage isolation across wallets
* and Quark operations
* @notice Contract for managing nonces for Quark wallets
* @author Compound Labs, Inc.
*/
contract QuarkStateManager {
event ClearNonce(address indexed wallet, uint96 nonce);

error NoActiveNonce();
error NoUnusedNonces();
error NonceAlreadySet();
error NonceScriptMismatch();

/// @notice Bit-packed structure of a nonce-script pair
struct NonceScript {
uint96 nonce;
address scriptAddress;
}
error NoUnusedNonces();

/// @notice Bit-packed nonce values
mapping(address wallet => mapping(uint256 bucket => uint256 bitset)) public nonces;

/// @notice Per-wallet-nonce address for preventing replays with changed script address
mapping(address wallet => mapping(uint96 nonce => address scriptAddress)) public nonceScriptAddress;

/// @notice Per-wallet-nonce storage space that can be utilized while a nonce is active
mapping(address wallet => mapping(uint96 nonce => mapping(bytes32 key => bytes32 value))) public walletStorage;

/// @notice Currently active nonce-script pair for a wallet, if any, for which storage is accessible
mapping(address wallet => NonceScript) internal activeNonceScript;

/**
* @notice Return whether a nonce has been exhausted; note that if a nonce is not set, that does not mean it has not been used before
* @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
Expand Down Expand Up @@ -70,125 +50,35 @@ contract QuarkStateManager {
uint256 mask = 1 << maskOffset;
if ((bucketNonces & mask) == 0) {
uint96 nonce = uint96(bucketValue + maskOffset);
// The next available nonce should not be reserved for a replayable transaction
if (nonceScriptAddress[wallet][nonce] == address(0)) {
return nonce;
}
return nonce;
}
}
}

revert NoUnusedNonces();
}

/**
* @notice Return the script address associated with the currently active nonce; revert if none
* @return Currently active script address
*/
function getActiveScript() external view returns (address) {
address scriptAddress = activeNonceScript[msg.sender].scriptAddress;
if (scriptAddress == address(0)) {
revert NoActiveNonce();
}
return scriptAddress;
}

/// @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);
}

/// @notice Clears (un-sets) the active nonce to allow its reuse; allows a script to be replayed
function clearNonce() external {
if (activeNonceScript[msg.sender].scriptAddress == address(0)) {
revert NoActiveNonce();
}

uint96 nonce = activeNonceScript[msg.sender].nonce;
(uint256 bucket, uint256 setMask) = getBucket(nonce);
nonces[msg.sender][bucket] &= ~setMask;
emit ClearNonce(msg.sender, nonce);
}

/**
* @notice Set a given nonce for the calling wallet; effectively cancels any replayable script using that nonce
* @notice Set a given nonce for the calling wallet, reverting if the nonce is already set
* @param nonce Nonce to set for the calling wallet
*/
function setNonce(uint96 nonce) external {
// TODO: should we check whether there exists a nonceScriptAddress?
(uint256 bucket, uint256 setMask) = getBucket(nonce);
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;
}

/**
* @notice Set a wallet nonce as the active nonce and yield control back to the wallet by calling into callback
* @param nonce Nonce to activate for the transaction
* @param scriptAddress Address of script to invoke with nonce lock
* @param scriptCalldata Calldata for script call to invoke with nonce lock
* @return Return value from the executed operation
* @dev The script is expected to clearNonce() if it wishes to be replayable
*/
function setActiveNonceAndCallback(uint96 nonce, address scriptAddress, bytes calldata scriptCalldata)
external
returns (bytes memory)
{
// retrieve the (bucket, mask) pair that addresses the nonce in memory
(uint256 bucket, uint256 setMask) = getBucket(nonce);

// ensure nonce is not already set
if (isNonceSetInternal(msg.sender, bucket, setMask)) {
revert NonceAlreadySet();
}

address cachedScriptAddress = nonceScriptAddress[msg.sender][nonce];
// if the nonce has been used before, check if the script address matches, and revert if not
if ((cachedScriptAddress != address(0)) && (cachedScriptAddress != scriptAddress)) {
revert NonceScriptMismatch();
}

// spend the nonce; only if the callee chooses to clear it will it get un-set and become replayable
setNonceInternal(bucket, setMask);

// set the nonce-script pair active and yield to the wallet callback
NonceScript memory previousNonceScript = activeNonceScript[msg.sender];
activeNonceScript[msg.sender] = NonceScript({nonce: nonce, scriptAddress: scriptAddress});

bytes memory result = IQuarkWallet(msg.sender).executeScriptWithNonceLock(scriptAddress, scriptCalldata);

// if a nonce was cleared, set the nonceScriptAddress to lock nonce re-use to the same script address
if (cachedScriptAddress == address(0) && !isNonceSetInternal(msg.sender, bucket, setMask)) {
nonceScriptAddress[msg.sender][nonce] = scriptAddress;
}

// release the nonce when the wallet finishes executing callback
activeNonceScript[msg.sender] = previousNonceScript;

return result;
}

/// @notice Write arbitrary bytes to storage namespaced by the currently active nonce; reverts if no nonce is currently active
function write(bytes32 key, bytes32 value) external {
if (activeNonceScript[msg.sender].scriptAddress == address(0)) {
revert NoActiveNonce();
}
walletStorage[msg.sender][activeNonceScript[msg.sender].nonce][key] = value;
}

/**
* @notice Read from storage namespaced by the currently active nonce; reverts if no nonce is currently active
* @return Value at the nonce storage location, as bytes
*/
function read(bytes32 key) external view returns (bytes32) {
if (activeNonceScript[msg.sender].scriptAddress == address(0)) {
revert NoActiveNonce();
}
return walletStorage[msg.sender][activeNonceScript[msg.sender].nonce][key];
/// @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;
}
}
38 changes: 28 additions & 10 deletions src/quark-core/src/QuarkWallet.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,8 +107,11 @@ contract QuarkWallet is IERC1271 {
)
);

/// @notice Well-known stateManager key for the currently executing script's callback address (if any)
bytes32 public constant CALLBACK_KEY = keccak256("callback.v1.quark");
/// @notice Well-known storage slot for the currently executing script's callback address (if any)
bytes32 public constant CALLBACK_SLOT = bytes32(uint256(keccak256("quark.v1.callback")) - 1);

/// @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 The magic value to return for valid ERC1271 signature
bytes4 internal constant EIP_1271_MAGIC_VALUE = 0x1626ba7e;
Expand Down Expand Up @@ -219,7 +222,9 @@ contract QuarkWallet is IERC1271 {

emit ExecuteQuarkScript(msg.sender, op.scriptAddress, op.nonce, ExecutionType.Signature);

return stateManager.setActiveNonceAndCallback(op.nonce, op.scriptAddress, op.scriptCalldata);
stateManager.setNonce(op.nonce);

return executeScriptInternal(op.scriptAddress, op.scriptCalldata);
}

/**
Expand Down Expand Up @@ -249,7 +254,9 @@ contract QuarkWallet is IERC1271 {

emit ExecuteQuarkScript(msg.sender, scriptAddress, nonce, ExecutionType.Direct);

return stateManager.setActiveNonceAndCallback(nonce, scriptAddress, scriptCalldata);
stateManager.setNonce(nonce);

return executeScriptInternal(scriptAddress, scriptCalldata);
}

/**
Expand Down Expand Up @@ -378,29 +385,35 @@ contract QuarkWallet is IERC1271 {
}

/**
* @notice Execute a QuarkOperation with a lock acquired on nonce-namespaced storage
* @dev Can only be called by stateManager during setActiveNonceAndCallback()
* @notice Execute a script using the given calldata
* @param scriptAddress Address of script to execute
* @param scriptCalldata Encoded calldata for the call to execute on the scriptAddress
* @return Result of executing the script, encoded as bytes
*/
function executeScriptWithNonceLock(address scriptAddress, bytes memory scriptCalldata)
external
function executeScriptInternal(address scriptAddress, bytes memory scriptCalldata)
internal
returns (bytes memory)
{
require(msg.sender == address(stateManager));
if (scriptAddress.code.length == 0) {
revert EmptyCode();
}

bool success;
uint256 returnSize;
uint256 scriptCalldataLen = scriptCalldata.length;
bytes32 activeScriptSlot = ACTIVE_SCRIPT_SLOT;
assembly {
// Store the active script
// TODO: Move to TSTORE after updating Solidity version to >=0.8.24
sstore(activeScriptSlot, scriptAddress)

// Note: CALLCODE is used to set the QuarkWallet as the `msg.sender`
success :=
callcode(gas(), scriptAddress, /* value */ 0, add(scriptCalldata, 0x20), scriptCalldataLen, 0x0, 0)
returnSize := returndatasize()

// TODO: Move to TSTORE after updating Solidity version to >=0.8.24
sstore(activeScriptSlot, 0)
}

bytes memory returnData = new bytes(returnSize);
Expand All @@ -422,7 +435,12 @@ contract QuarkWallet is IERC1271 {
* @dev Reverts if callback is not enabled by the script
*/
fallback(bytes calldata data) external payable returns (bytes memory) {
address callback = address(uint160(uint256(stateManager.read(CALLBACK_KEY))));
bytes32 callbackSlot = CALLBACK_SLOT;
address callback;
assembly {
// TODO: Move to TLOAD after updating Solidity version to >=0.8.24
callback := sload(callbackSlot)
}
if (callback != address(0)) {
(bool success, bytes memory result) = callback.delegatecall(data);
if (!success) {
Expand Down
Loading

0 comments on commit 3661014

Please sign in to comment.