Skip to content

Commit 5a253fa

Browse files
nkrishangKrishang Nadgauda
nkrishang
and
Krishang Nadgauda
authored
Multichain Registry (#252)
* Multichain TWRegistry: accept chainId arg in add and remove * TWFactory: call registry with block chainid * Update TWRegsitry tests * run prettier * update tests for TWRegistry Co-authored-by: Krishang Nadgauda <[email protected]>
1 parent 0fdde7f commit 5a253fa

File tree

4 files changed

+182
-57
lines changed

4 files changed

+182
-57
lines changed

‎contracts/TWFactory.sol

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ contract TWFactory is Multicall, ERC2771Context, AccessControlEnumerable, IContr
7575

7676
emit ProxyDeployed(_implementation, deployedProxy, _msgSender());
7777

78-
registry.add(_msgSender(), deployedProxy);
78+
registry.add(_msgSender(), deployedProxy, block.chainid);
7979

8080
if (_data.length > 0) {
8181
// slither-disable-next-line unused-return

‎contracts/TWRegistry.sol

Lines changed: 54 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -6,47 +6,87 @@ import "@openzeppelin/contracts/access/AccessControlEnumerable.sol";
66
import "@openzeppelin/contracts/utils/Multicall.sol";
77
import "@openzeppelin/contracts/metatx/ERC2771Context.sol";
88

9-
contract TWRegistry is Multicall, ERC2771Context, AccessControlEnumerable {
9+
import "./interfaces/ITWRegistry.sol";
10+
11+
contract TWRegistry is ITWRegistry, Multicall, ERC2771Context, AccessControlEnumerable {
1012
bytes32 public constant OPERATOR_ROLE = keccak256("OPERATOR_ROLE");
1113

1214
using EnumerableSet for EnumerableSet.AddressSet;
15+
using EnumerableSet for EnumerableSet.UintSet;
1316

1417
/// @dev wallet address => [contract addresses]
15-
mapping(address => EnumerableSet.AddressSet) private deployments;
18+
mapping(address => mapping(uint256 => EnumerableSet.AddressSet)) private deployments;
1619

17-
event Added(address indexed deployer, address indexed deployment);
18-
event Deleted(address indexed deployer, address indexed deployment);
20+
EnumerableSet.UintSet private chainIds;
1921

2022
constructor(address _trustedForwarder) ERC2771Context(_trustedForwarder) {
2123
_setupRole(DEFAULT_ADMIN_ROLE, _msgSender());
2224
}
2325

2426
// slither-disable-next-line similar-names
25-
function add(address _deployer, address _deployment) external {
27+
function add(
28+
address _deployer,
29+
address _deployment,
30+
uint256 _chainId
31+
) external {
2632
require(hasRole(OPERATOR_ROLE, _msgSender()) || _deployer == _msgSender(), "not operator or deployer.");
2733

28-
bool added = deployments[_deployer].add(_deployment);
34+
bool added = deployments[_deployer][_chainId].add(_deployment);
2935
require(added, "failed to add");
3036

31-
emit Added(_deployer, _deployment);
37+
chainIds.add(_chainId);
38+
39+
emit Added(_deployer, _deployment, _chainId);
3240
}
3341

3442
// slither-disable-next-line similar-names
35-
function remove(address _deployer, address _deployment) external {
43+
function remove(
44+
address _deployer,
45+
address _deployment,
46+
uint256 _chainId
47+
) external {
3648
require(hasRole(OPERATOR_ROLE, _msgSender()) || _deployer == _msgSender(), "not operator or deployer.");
3749

38-
bool removed = deployments[_deployer].remove(_deployment);
50+
bool removed = deployments[_deployer][_chainId].remove(_deployment);
3951
require(removed, "failed to remove");
4052

41-
emit Deleted(_deployer, _deployment);
53+
emit Deleted(_deployer, _deployment, _chainId);
4254
}
4355

44-
function getAll(address _deployer) external view returns (address[] memory) {
45-
return deployments[_deployer].values();
56+
function getAll(address _deployer) external view returns (Deployment[] memory allDeployments) {
57+
uint256 totalDeployments;
58+
uint256 chainIdsLen = chainIds.length();
59+
60+
for (uint256 i = 0; i < chainIdsLen; i += 1) {
61+
uint256 chainId = chainIds.at(i);
62+
63+
totalDeployments += deployments[_deployer][chainId].length();
64+
}
65+
66+
allDeployments = new Deployment[](totalDeployments);
67+
uint256 idx;
68+
69+
for (uint256 j = 0; j < chainIdsLen; j += 1) {
70+
uint256 chainId = chainIds.at(j);
71+
72+
uint256 len = deployments[_deployer][chainId].length();
73+
address[] memory deploymentAddrs = deployments[_deployer][chainId].values();
74+
75+
for (uint256 k = 0; k < len; k += 1) {
76+
allDeployments[idx] = Deployment({ deploymentAddress: deploymentAddrs[k], chainId: chainId });
77+
idx += 1;
78+
}
79+
}
4680
}
4781

48-
function count(address _deployer) external view returns (uint256) {
49-
return deployments[_deployer].length();
82+
function count(address _deployer) external view returns (uint256 deploymentCount) {
83+
uint256 chainIdsLen = chainIds.length();
84+
85+
for (uint256 i = 0; i < chainIdsLen; i += 1) {
86+
uint256 chainId = chainIds.at(i);
87+
88+
deploymentCount += deployments[_deployer][chainId].length();
89+
}
5090
}
5191

5292
function _msgSender() internal view virtual override(Context, ERC2771Context) returns (address sender) {

‎contracts/interfaces/ITWRegistry.sol

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
pragma solidity ^0.8.11;
3+
4+
interface ITWRegistry {
5+
struct Deployment {
6+
address deploymentAddress;
7+
uint256 chainId;
8+
}
9+
10+
event Added(address indexed deployer, address indexed deployment, uint256 indexed chainId);
11+
event Deleted(address indexed deployer, address indexed deployment, uint256 indexed chainId);
12+
13+
/// @notice Add a deployment for a deployer.
14+
function add(
15+
address _deployer,
16+
address _deployment,
17+
uint256 _chainId
18+
) external;
19+
20+
/// @notice Remove a deployment for a deployer.
21+
function remove(
22+
address _deployer,
23+
address _deployment,
24+
uint256 _chainId
25+
) external;
26+
27+
/// @notice Get all deployments for a deployer.
28+
function getAll(address _deployer) external view returns (Deployment[] memory allDeployments);
29+
30+
/// @notice Get the total number of deployments for a deployer.
31+
function count(address _deployer) external view returns (uint256 deploymentCount);
32+
}

‎src/test/TWRegistry.t.sol

Lines changed: 95 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -3,26 +3,40 @@ pragma solidity ^0.8.11;
33

44
// Test imports
55
import "./utils/BaseTest.sol";
6+
import "contracts/interfaces/ITWRegistry.sol";
67
import "contracts/TWRegistry.sol";
8+
import "./mocks/MockThirdwebContract.sol";
79

810
interface ITWRegistryData {
9-
event Added(address indexed deployer, address indexed moduleAddress);
10-
event Deleted(address indexed deployer, address indexed moduleAddress);
11+
event Added(address indexed deployer, address indexed moduleAddress, uint256 indexed chainid);
12+
event Deleted(address indexed deployer, address indexed moduleAddress, uint256 indexed chainid);
1113
}
1214

1315
contract TWRegistryTest is ITWRegistryData, BaseTest {
1416
// Target contract
1517
TWRegistry internal _registry;
1618

1719
// Test params
18-
address internal mockModuleAddress = address(0x42);
19-
address internal actor;
20+
uint256[] internal chainIds;
21+
address[] internal deploymentAddresses;
22+
address internal deployer_;
23+
24+
uint256 total = 1000;
2025

2126
// ===== Set up =====
2227

2328
function setUp() public override {
2429
super.setUp();
25-
actor = getActor(0);
30+
31+
deployer_ = getActor(100);
32+
33+
for (uint256 i = 0; i < total; i += 1) {
34+
chainIds.push(i);
35+
vm.prank(deployer_);
36+
address depl = address(new MockThirdwebContract());
37+
deploymentAddresses.push(depl);
38+
}
39+
2640
_registry = TWRegistry(registry);
2741
}
2842

@@ -31,96 +45,135 @@ contract TWRegistryTest is ITWRegistryData, BaseTest {
3145
/// @dev Test `add`
3246

3347
function test_addFromFactory() public {
34-
vm.prank(factory);
35-
_registry.add(actor, mockModuleAddress);
48+
vm.startPrank(factory);
49+
for (uint256 i = 0; i < total; i += 1) {
50+
_registry.add(deployer_, deploymentAddresses[i], chainIds[i]);
51+
}
52+
vm.stopPrank();
53+
54+
ITWRegistry.Deployment[] memory modules = _registry.getAll(deployer_);
55+
56+
assertEq(modules.length, total);
57+
assertEq(_registry.count(deployer_), total);
3658

37-
address[] memory modules = _registry.getAll(actor);
38-
assertEq(modules.length, 1);
39-
assertEq(modules[0], mockModuleAddress);
40-
assertEq(_registry.count(actor), 1);
59+
for (uint256 i = 0; i < total; i += 1) {
60+
assertEq(modules[i].deploymentAddress, deploymentAddresses[i]);
61+
assertEq(modules[i].chainId, chainIds[i]);
62+
}
4163

4264
vm.prank(factory);
43-
_registry.add(actor, address(0x43));
65+
_registry.add(deployer_, address(0x43), 111);
4466

45-
modules = _registry.getAll(actor);
46-
assertEq(modules.length, 2);
47-
assertEq(_registry.count(actor), 2);
67+
modules = _registry.getAll(deployer_);
68+
assertEq(modules.length, total + 1);
69+
assertEq(_registry.count(deployer_), total + 1);
4870
}
4971

5072
function test_addFromSelf() public {
51-
vm.prank(actor);
52-
_registry.add(actor, mockModuleAddress);
73+
vm.startPrank(deployer_);
74+
for (uint256 i = 0; i < total; i += 1) {
75+
_registry.add(deployer_, deploymentAddresses[i], chainIds[i]);
76+
}
77+
vm.stopPrank();
78+
79+
ITWRegistry.Deployment[] memory modules = _registry.getAll(deployer_);
80+
81+
assertEq(modules.length, total);
82+
assertEq(_registry.count(deployer_), total);
5383

54-
address[] memory modules = _registry.getAll(actor);
84+
for (uint256 i = 0; i < total; i += 1) {
85+
assertEq(modules[i].deploymentAddress, deploymentAddresses[i]);
86+
assertEq(modules[i].chainId, chainIds[i]);
87+
}
5588

56-
assertEq(modules.length, 1);
57-
assertEq(modules[0], mockModuleAddress);
58-
assertEq(_registry.count(actor), 1);
89+
vm.prank(factory);
90+
_registry.add(deployer_, address(0x43), 111);
91+
92+
modules = _registry.getAll(deployer_);
93+
assertEq(modules.length, total + 1);
94+
assertEq(_registry.count(deployer_), total + 1);
5995
}
6096

6197
function test_add_emit_Added() public {
62-
vm.expectEmit(true, true, false, true);
63-
emit Added(actor, mockModuleAddress);
98+
vm.expectEmit(true, true, true, true);
99+
emit Added(deployer_, deploymentAddresses[0], chainIds[0]);
64100

65101
vm.prank(factory);
66-
_registry.add(actor, mockModuleAddress);
102+
_registry.add(deployer_, deploymentAddresses[0], chainIds[0]);
67103
}
68104

69105
// Test `remove`
70106

71107
function setUp_remove() public {
72-
vm.prank(factory);
73-
_registry.add(actor, mockModuleAddress);
108+
vm.startPrank(factory);
109+
for (uint256 i = 0; i < total; i += 1) {
110+
_registry.add(deployer_, deploymentAddresses[i], chainIds[i]);
111+
}
112+
vm.stopPrank();
74113
}
75114

76115
// ===== Functionality tests =====
77116
function test_removeFromFactory() public {
78117
setUp_remove();
79118
vm.prank(factory);
80-
_registry.remove(actor, mockModuleAddress);
119+
_registry.remove(deployer_, deploymentAddresses[0], chainIds[0]);
81120

82-
address[] memory modules = _registry.getAll(actor);
83-
assertEq(modules.length, 0);
121+
ITWRegistry.Deployment[] memory modules = _registry.getAll(deployer_);
122+
assertEq(modules.length, total - 1);
123+
124+
for (uint256 i = 0; i < total - 1; i += 1) {
125+
assertEq(modules[i].deploymentAddress, deploymentAddresses[i + 1]);
126+
assertEq(modules[i].chainId, chainIds[i + 1]);
127+
}
84128
}
85129

86130
function test_removeFromSelf() public {
87131
setUp_remove();
88-
vm.prank(actor);
89-
_registry.remove(actor, mockModuleAddress);
132+
vm.prank(factory);
133+
_registry.remove(deployer_, deploymentAddresses[0], chainIds[0]);
90134

91-
address[] memory modules = _registry.getAll(actor);
92-
assertEq(modules.length, 0);
135+
ITWRegistry.Deployment[] memory modules = _registry.getAll(deployer_);
136+
assertEq(modules.length, total - 1);
93137
}
94138

95139
function test_remove_revert_invalidCaller() public {
96140
setUp_remove();
97141
address invalidCaller = address(0x123);
98-
assertTrue(invalidCaller != factory || invalidCaller != actor);
142+
assertTrue(invalidCaller != factory || invalidCaller != deployer_);
99143

100144
vm.expectRevert("not operator or deployer.");
101145

102146
vm.prank(invalidCaller);
103-
_registry.remove(actor, mockModuleAddress);
147+
_registry.remove(deployer_, deploymentAddresses[0], chainIds[0]);
104148
}
105149

106150
function test_remove_revert_noModulesToRemove() public {
107151
setUp_remove();
108-
actor = getActor(1);
109-
address[] memory modules = _registry.getAll(actor);
152+
address actor = getActor(1);
153+
ITWRegistry.Deployment[] memory modules = _registry.getAll(actor);
110154
assertEq(modules.length, 0);
111155

112156
vm.expectRevert("failed to remove");
113157

114158
vm.prank(actor);
115-
_registry.remove(actor, mockModuleAddress);
159+
_registry.remove(actor, deploymentAddresses[0], chainIds[0]);
160+
}
161+
162+
function test_remove_revert_incorrectChainId() public {
163+
setUp_remove();
164+
165+
vm.expectRevert("failed to remove");
166+
167+
vm.prank(deployer_);
168+
_registry.remove(deployer_, deploymentAddresses[0], 12345);
116169
}
117170

118171
function test_remove_emit_Deleted() public {
119172
setUp_remove();
120-
vm.expectEmit(true, true, false, true);
121-
emit Deleted(actor, mockModuleAddress);
173+
vm.expectEmit(true, true, true, true);
174+
emit Deleted(deployer_, deploymentAddresses[0], chainIds[0]);
122175

123-
vm.prank(actor);
124-
_registry.remove(actor, mockModuleAddress);
176+
vm.prank(deployer_);
177+
_registry.remove(deployer_, deploymentAddresses[0], chainIds[0]);
125178
}
126179
}

0 commit comments

Comments
 (0)