From 13a7bd95a8aa2f374c8da035f4bf268978cf2cd0 Mon Sep 17 00:00:00 2001 From: zeroknots Date: Tue, 28 May 2024 07:43:19 +0700 Subject: [PATCH] fix: attester storage fixed threshold finding using associated storage to store attesters to be 4337 compliant --- remappings.txt | 4 ++++ src/DataTypes.sol | 2 +- src/core/TrustManager.sol | 18 ++++++++++---- test/TrustDelegation.t.sol | 49 +++++++++++++++++++++++++++++++++++++- 4 files changed, 67 insertions(+), 6 deletions(-) diff --git a/remappings.txt b/remappings.txt index e40312c4..2a842303 100644 --- a/remappings.txt +++ b/remappings.txt @@ -3,3 +3,7 @@ solmate/=node_modules/solmate/src/ solady/=node_modules/solady/src/ forge-std/=node_modules/forge-std/src/ ds-test/=node_modules/ds-test/src/ +erc4337-validation/=node_modules/@rhinestone/erc4337-validation/src/ +account-abstraction/=node_modules/account-abstraction/contracts/ +account-abstraction-v0.6/=node_modules/account-abstraction-v0.6/contracts/ +@openzeppelin/=node_modules/@openzeppelin/ diff --git a/src/DataTypes.sol b/src/DataTypes.sol index 4813f234..483be60f 100644 --- a/src/DataTypes.sol +++ b/src/DataTypes.sol @@ -41,7 +41,7 @@ struct TrustedAttesterRecord { uint8 attesterCount; // number of attesters in the linked list uint8 threshold; // minimum number of attesters required address attester; // first attester in linked list. (packed to save gas) - mapping(address attester => address linkedAttester) linkedAttesters; // linked list + mapping(address attester => mapping(address account => address linkedAttester)) linkedAttesters; } /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ diff --git a/src/core/TrustManager.sol b/src/core/TrustManager.sol index 982d28e7..705e3b3b 100644 --- a/src/core/TrustManager.sol +++ b/src/core/TrustManager.sol @@ -62,7 +62,7 @@ abstract contract TrustManager is IRegistry { address _attester = attesters[i]; // user could have set attester to address(0) if (_attester == ZERO_ADDRESS) revert InvalidTrustedAttesterInput(); - $trustedAttester.linkedAttesters[_attester] = attesters[i + 1]; + $trustedAttester.linkedAttesters[_attester][msg.sender] = attesters[i + 1]; } emit NewTrustedAttesters(); @@ -119,7 +119,17 @@ abstract contract TrustManager is IRegistry { // use this condition to save gas else if (threshold == 1) { AttestationRecord storage $attestation = $getAttestation({ module: module, attester: attester }); - $attestation.enforceValid(moduleType); + if ($attestation.checkValid(moduleType)) return; + + // if first attestation is not valid, iterate over the linked list of attesters + // and check if the attestation is valid + for (uint256 i; i < attesterCount; i++) { + attester = $trustedAttesters.linkedAttesters[attester][smartAccount]; + $attestation = $getAttestation({ module: module, attester: attester }); + if ($attestation.checkValid(moduleType)) return; + } + // if no valid attestations were found in the for loop. the module is not valid + revert InsufficientAttestations(); } // smart account has more than one trusted attester else { @@ -129,7 +139,7 @@ abstract contract TrustManager is IRegistry { for (uint256 i = 1; i < attesterCount; i++) { // get next attester from linked List - attester = $trustedAttesters.linkedAttesters[attester]; + attester = $trustedAttesters.linkedAttesters[attester][smartAccount]; $attestation = $getAttestation({ module: module, attester: attester }); if ($attestation.checkValid(moduleType)) threshold--; @@ -154,7 +164,7 @@ abstract contract TrustManager is IRegistry { for (uint256 i = 1; i < count; i++) { // get next attester from linked List - attesters[i] = $trustedAttesters.linkedAttesters[attesters[i - 1]]; + attesters[i] = $trustedAttesters.linkedAttesters[attesters[i - 1]][smartAccount]; } } diff --git a/test/TrustDelegation.t.sol b/test/TrustDelegation.t.sol index ceac7e5f..ac617520 100644 --- a/test/TrustDelegation.t.sol +++ b/test/TrustDelegation.t.sol @@ -4,9 +4,12 @@ pragma solidity ^0.8.0; import "./Attestation.t.sol"; import "src/DataTypes.sol"; import { LibSort } from "solady/utils/LibSort.sol"; +import { VmSafe } from "forge-std/Vm.sol"; +import { ERC4337SpecsParser } from "erc4337-validation/SpecsParser.sol"; contract TrustTest is AttestationTest { using LibSort for address[]; + using ERC4337SpecsParser for VmSafe.AccountAccess; function setUp() public override { super.setUp(); @@ -31,7 +34,7 @@ contract TrustTest is AttestationTest { uint8 threshold, address[] memory attesters ) - external + public whenSettingAttester prankWithAccount(smartAccount1) { @@ -123,4 +126,48 @@ contract TrustTest is AttestationTest { vm.expectRevert(); registry.check(address(module1), ModuleType.wrap(3)); } + + function test_WhenAttestersSetCheckOnlyOneThreshold() external whenQueryingRegistry { + test_WhenUsingValidECDSA(); + + vm.startPrank(smartAccount1.addr); + // It should not revert. + address[] memory attesters = new address[](2); + attesters[0] = address(makeAddr("foo")); + attesters[1] = address(attester1.addr); + registry.trustAttesters(1, attesters); + + registry.check(address(module1), ModuleType.wrap(1)); + registry.check(address(module1), ModuleType.wrap(2)); + vm.expectRevert(); + registry.check(address(module1), ModuleType.wrap(3)); + } + + function test_WhenSupplyingManyAttesters_ShouldBe4337Compliant(uint8 threshold, address[] memory attesters) public { + vm.startMappingRecording(); + vm.startStateDiffRecording(); + + test_WhenSupplyingManyAttesters(threshold, attesters); + + VmSafe.AccountAccess[] memory accesses = vm.stopAndReturnStateDiff(); + + ERC4337SpecsParser.Entities memory entities = ERC4337SpecsParser.Entities({ + account: smartAccount1.addr, + factory: address(0), + isFactoryStaked: false, + paymaster: address(0), + isPaymasterStaked: false, + aggregator: address(0), + isAggregatorStaked: false + }); + + for (uint256 i; i < accesses.length; i++) { + VmSafe.AccountAccess memory currentAccess = accesses[i]; + if (currentAccess.account != address(this) && currentAccess.accessor != address(this)) { + currentAccess.validateBannedStorageLocations(entities); + } + } + + vm.stopMappingRecording(); + } }