Skip to content

Commit

Permalink
feat: add signature validation hooks
Browse files Browse the repository at this point in the history
  • Loading branch information
adamegyed committed Aug 15, 2024
1 parent 2168a96 commit 0e510c6
Show file tree
Hide file tree
Showing 11 changed files with 154 additions and 26 deletions.
29 changes: 24 additions & 5 deletions src/account/UpgradeableModularAccount.sol
Original file line number Diff line number Diff line change
Expand Up @@ -234,8 +234,6 @@ contract UpgradeableModularAccount is
}

/// @notice Initializes the account with a validation function added to the global pool.
/// TODO: remove and merge with regular initialization, after we figure out a better install/uninstall workflow
/// with user install configs.
/// @dev This function is only callable once.
function initializeWithValidation(
ValidationConfig validationConfig,
Expand Down Expand Up @@ -298,14 +296,35 @@ contract UpgradeableModularAccount is
AccountStorage storage _storage = getAccountStorage();

ModuleEntity sigValidation = ModuleEntity.wrap(bytes24(signature));
signature = signature[24:];

(address module, uint32 entityId) = sigValidation.unpack();
if (!_storage.validationData[sigValidation].isSignatureValidation) {
revert SignatureValidationInvalid(module, entityId);
(address _module, uint32 _entityId) = sigValidation.unpack();
revert SignatureValidationInvalid(_module, _entityId);
}

ModuleEntity[] memory preSignatureValidationHooks =
getAccountStorage().validationData[sigValidation].preValidationHooks;

for (uint256 i = 0; i < preSignatureValidationHooks.length; ++i) {
(address hookModule, uint32 hookEntityId) = preSignatureValidationHooks[i].unpack();

bytes memory currentSignatureSegment;

(currentSignatureSegment, signature) = signature.advanceSegmentIfAtIndex(uint8(i));

// If this reverts, bubble up revert reason.
IValidationHookModule(hookModule).preSignatureValidationHook(
hookEntityId, msg.sender, hash, currentSignatureSegment
);
}

signature = signature.getFinalSegment();

(address module, uint32 entityId) = sigValidation.unpack();

if (
IValidationModule(module).validateSignature(address(this), entityId, msg.sender, hash, signature[24:])
IValidationModule(module).validateSignature(address(this), entityId, msg.sender, hash, signature)
== _1271_MAGIC_VALUE
) {
return _1271_MAGIC_VALUE;
Expand Down
2 changes: 0 additions & 2 deletions src/interfaces/IModuleManager.sol
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ interface IModuleManager {
/// path.
/// Installs a validation function across a set of execution selectors, and optionally mark it as a global
/// validation.
/// TODO: remove or update.
/// @dev This does not validate anything against the manifest - the caller must ensure validity.
/// @param validationConfig The validation function to install, along with configuration flags.
/// @param selectors The selectors to install the validation function for.
Expand All @@ -46,7 +45,6 @@ interface IModuleManager {
) external;

/// @notice Uninstall a validation function from a set of execution selectors.
/// TODO: remove or update.
/// @param validationFunction The validation function to uninstall.
/// @param uninstallData Optional data to be decoded and used by the module to clear module data for the
/// account.
Expand Down
10 changes: 3 additions & 7 deletions src/interfaces/IValidationHookModule.sol
Original file line number Diff line number Diff line change
Expand Up @@ -32,18 +32,14 @@ interface IValidationHookModule is IModule {
bytes calldata authorization
) external;

// TODO: support this hook type within the account & in the manifest

/// @notice Run the pre signature validation hook specified by the `entityId`.
/// @dev To indicate the call should revert, the function MUST revert.
/// @param entityId An identifier that routes the call to different internal implementations, should there
/// be more than one.
/// @param sender The caller address.
/// @param hash The hash of the message being signed.
/// @param signature The signature of the message.
// function preSignatureValidationHook(uint32 entityId, address sender, bytes32 hash, bytes calldata
// signature)
// external
// view
// returns (bytes4);
function preSignatureValidationHook(uint32 entityId, address sender, bytes32 hash, bytes calldata signature)
external
view;
}
4 changes: 4 additions & 0 deletions src/modules/NativeTokenLimitModule.sol
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,10 @@ contract NativeTokenLimitModule is BaseModule, IExecutionHookModule, IValidation
override
{} // solhint-disable-line no-empty-blocks

function preSignatureValidationHook(uint32, address, bytes32, bytes calldata) external pure override {
return;
}

/// @inheritdoc IModule
function moduleMetadata() external pure virtual override returns (ModuleMetadata memory) {
ModuleMetadata memory metadata;
Expand Down
4 changes: 4 additions & 0 deletions src/modules/permissionhooks/AllowlistModule.sol
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,10 @@ contract AllowlistModule is IValidationHookModule, BaseModule {
return;
}

function preSignatureValidationHook(uint32, address, bytes32, bytes calldata) external pure override {
return;
}

function moduleMetadata() external pure override returns (ModuleMetadata memory) {
ModuleMetadata memory metadata;
metadata.name = "Allowlist Module";
Expand Down
47 changes: 47 additions & 0 deletions test/account/PerHookData.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -352,6 +352,53 @@ contract PerHookDataTest is CustomValidationTestBase {
);
}

function test_pass1271AccessControl() public {
string memory message = "Hello, world!";

bytes32 messageHash = keccak256(abi.encodePacked(message));

(uint8 v, bytes32 r, bytes32 s) = vm.sign(owner1Key, messageHash);

PreValidationHookData[] memory preValidationHookData = new PreValidationHookData[](1);
preValidationHookData[0] = PreValidationHookData({index: 0, validationData: abi.encodePacked(message)});

bytes4 result = account1.isValidSignature(
messageHash, _encode1271Signature(_signerValidation, preValidationHookData, abi.encodePacked(r, s, v))
);

assertEq(result, bytes4(0x1626ba7e));
}

function test_fail1271AccessControl_badSigData() public {
string memory message = "Hello, world!";

bytes32 messageHash = keccak256(abi.encodePacked(message));

(uint8 v, bytes32 r, bytes32 s) = vm.sign(owner1Key, messageHash);

PreValidationHookData[] memory preValidationHookData = new PreValidationHookData[](1);
preValidationHookData[0] = PreValidationHookData({
index: 0,
validationData: abi.encodePacked(address(0x1234123412341234123412341234123412341234))
});

vm.expectRevert("Preimage not provided");
account1.isValidSignature(
messageHash, _encode1271Signature(_signerValidation, preValidationHookData, abi.encodePacked(r, s, v))
);
}

function test_fail1271AccessControl_noSigData() public {
string memory message = "Hello, world!";

bytes32 messageHash = keccak256(abi.encodePacked(message));

(uint8 v, bytes32 r, bytes32 s) = vm.sign(owner1Key, messageHash);

vm.expectRevert("Preimage not provided");
account1.isValidSignature(messageHash, _encode1271Signature(_signerValidation, abi.encodePacked(r, s, v)));
}

function _getCounterUserOP() internal view returns (PackedUserOperation memory, bytes32) {
PackedUserOperation memory userOp = PackedUserOperation({
sender: address(account1),
Expand Down
8 changes: 6 additions & 2 deletions test/account/UpgradeableModularAccount.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ import {MessageHashUtils} from "@openzeppelin/contracts/utils/cryptography/Messa
import {ModuleManagerInternals} from "../../src/account/ModuleManagerInternals.sol";
import {UpgradeableModularAccount} from "../../src/account/UpgradeableModularAccount.sol";

import {ModuleEntityLib} from "../../src/helpers/ModuleEntityLib.sol";

import {ExecutionDataView, IAccountLoupe} from "../../src/interfaces/IAccountLoupe.sol";
import {ExecutionManifest} from "../../src/interfaces/IExecutionModule.sol";
import {IModuleManager} from "../../src/interfaces/IModuleManager.sol";
Expand Down Expand Up @@ -381,8 +383,10 @@ contract UpgradeableModularAccountTest is AccountTestBase {

// singleSignerValidationModule.ownerOf(address(account1));

bytes memory signature =
abi.encodePacked(address(singleSignerValidationModule), TEST_DEFAULT_VALIDATION_ENTITY_ID, r, s, v);
bytes memory signature = _encode1271Signature(
ModuleEntityLib.pack(address(singleSignerValidationModule), TEST_DEFAULT_VALIDATION_ENTITY_ID),
abi.encodePacked(r, s, v)
);

bytes4 validationResult = IERC1271(address(account1)).isValidSignature(message, signature);

Expand Down
4 changes: 4 additions & 0 deletions test/mocks/modules/ComprehensiveModule.sol
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,10 @@ contract ComprehensiveModule is
revert NotImplemented();
}

function preSignatureValidationHook(uint32, address, bytes32, bytes calldata) external pure override {
return;
}

function validateSignature(address, uint32 entityId, address, bytes32, bytes calldata)
external
pure
Expand Down
12 changes: 12 additions & 0 deletions test/mocks/modules/MockAccessControlHookModule.sol
Original file line number Diff line number Diff line change
Expand Up @@ -72,5 +72,17 @@ contract MockAccessControlHookModule is IValidationHookModule, BaseModule {
revert NotImplemented();
}

function preSignatureValidationHook(uint32, address, bytes32 hash, bytes calldata signature)
external
pure
override
{
// Simulates some signature checking by requiring a preimage of the hash.

require(keccak256(signature) == hash, "Preimage not provided");

return;
}

function moduleMetadata() external pure override returns (ModuleMetadata memory) {}
}
4 changes: 4 additions & 0 deletions test/mocks/modules/ValidationModuleMocks.sol
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,10 @@ abstract contract MockBaseUserOpValidationModule is
revert NotImplemented();
}

function preSignatureValidationHook(uint32, address, bytes32, bytes calldata) external pure override {
return;
}

function validateSignature(address, uint32, address, bytes32, bytes calldata)
external
pure
Expand Down
56 changes: 46 additions & 10 deletions test/utils/AccountTestBase.sol
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,21 @@ abstract contract AccountTestBase is OptimizedTest {
return bytes32(uint256((g1 << 128) + uint128(g2)));
}

// helper function to encode a 1271 signature, according to the per-hook and per-validation data format.
function _encode1271Signature(
ModuleEntity validationFunction,
PreValidationHookData[] memory preValidationHookData,
bytes memory validationData
) internal pure returns (bytes memory) {
bytes memory sig = abi.encodePacked(validationFunction);

sig = abi.encodePacked(sig, _packPreHookDatas(preValidationHookData));

sig = abi.encodePacked(sig, _packValidationResWithIndex(255, validationData));

return sig;
}

// helper function to encode a signature, according to the per-hook and per-validation data format.
function _encodeSignature(
ModuleEntity validationFunction,
Expand All @@ -215,17 +230,8 @@ abstract contract AccountTestBase is OptimizedTest {
) internal pure returns (bytes memory) {
bytes memory sig = abi.encodePacked(validationFunction, globalOrNot);

for (uint256 i = 0; i < preValidationHookData.length; ++i) {
sig = abi.encodePacked(
sig,
_packValidationResWithIndex(
preValidationHookData[i].index, preValidationHookData[i].validationData
)
);
}
sig = abi.encodePacked(sig, _packPreHookDatas(preValidationHookData));

// Index of the actual validation data is the length of the preValidationHooksRetrieved - aka
// one-past-the-end
sig = abi.encodePacked(sig, _packValidationResWithIndex(255, validationData));

return sig;
Expand All @@ -241,6 +247,36 @@ abstract contract AccountTestBase is OptimizedTest {
return _encodeSignature(validationFunction, globalOrNot, emptyPreValidationHookData, validationData);
}

// overload for the case where there are no pre-validation hooks
function _encode1271Signature(ModuleEntity validationFunction, bytes memory validationData)
internal
pure
returns (bytes memory)
{
PreValidationHookData[] memory emptyPreValidationHookData = new PreValidationHookData[](0);
return _encode1271Signature(validationFunction, emptyPreValidationHookData, validationData);
}

// helper function to pack pre-validation hook datas, according to the sparse calldata segment spec.
function _packPreHookDatas(PreValidationHookData[] memory preValidationHookData)
internal
pure
returns (bytes memory)
{
bytes memory res = "";

for (uint256 i = 0; i < preValidationHookData.length; ++i) {
res = abi.encodePacked(
res,
_packValidationResWithIndex(
preValidationHookData[i].index, preValidationHookData[i].validationData
)
);
}

return res;
}

// helper function to pack validation data with an index, according to the sparse calldata segment spec.
function _packValidationResWithIndex(uint8 index, bytes memory validationData)
internal
Expand Down

0 comments on commit 0e510c6

Please sign in to comment.