Skip to content

Commit 6b24bc5

Browse files
EndymionJkbEndymionJkb
authored
Merge pull request #6 from balancer/angstrom-updates-v2
Angstrom review / refactor
2 parents 42f0474 + d5b569c commit 6b24bc5

File tree

6 files changed

+357
-323
lines changed

6 files changed

+357
-323
lines changed

‎contracts/AngstromBalancer.sol‎

Lines changed: 213 additions & 145 deletions
Large diffs are not rendered by default.

‎contracts/test/AngstromBalancerMock.sol‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,16 +19,8 @@ contract AngstromBalancerMock is AngstromBalancer {
1919
// solhint-disable-previous-line no-empty-blocks
2020
}
2121

22-
function manualUnlockAngstromWithRouter() external {
23-
_unlockAngstromWithRouter();
24-
}
25-
26-
function getLastUnlockBlockNumber() external view returns (uint256) {
27-
return _lastUnlockBlockNumber;
28-
}
29-
30-
function isNode(address account) external view returns (bool) {
31-
return _isNode(account);
22+
function manualUnlockAngstrom() external {
23+
_unlockAngstrom();
3224
}
3325

3426
function getDigest() external view returns (bytes32) {

‎test/foundry/AngstromBalancerUnit.t.sol‎

Lines changed: 16 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -7,78 +7,51 @@ import "forge-std/Test.sol";
77
import { IAuthentication } from "@balancer-labs/v3-interfaces/contracts/solidity-utils/helpers/IAuthentication.sol";
88

99
import { ArrayHelpers } from "@balancer-labs/v3-solidity-utils/contracts/test/ArrayHelpers.sol";
10-
import { BaseVaultTest } from "@balancer-labs/v3-vault/test/foundry/utils/BaseVaultTest.sol";
1110

1211
import { AngstromBalancerMock } from "../../contracts/test/AngstromBalancerMock.sol";
1312
import { AngstromBalancer } from "../../contracts/AngstromBalancer.sol";
13+
import { BaseAngstromTest } from "./utils/BaseAngstromTest.sol";
1414

15-
contract AngstromBalancerUnitTest is BaseVaultTest {
15+
contract AngstromBalancerUnitTest is BaseAngstromTest {
1616
using ArrayHelpers for *;
1717

18-
AngstromBalancerMock private _angstromBalancer;
19-
20-
function setUp() public virtual override {
21-
super.setUp();
22-
23-
_angstromBalancer = new AngstromBalancerMock(vault, weth, permit2, "AngstromBalancer Mock v1");
24-
25-
authorizer.grantRole(_angstromBalancer.getActionId(AngstromBalancer.toggleNodes.selector), admin);
26-
}
27-
2818
function testToggleNodesIsAuthenticated() public {
2919
vm.expectRevert(IAuthentication.SenderNotAllowed.selector);
30-
_angstromBalancer.toggleNodes([bob].toMemoryArray());
20+
angstromBalancer.toggleNodes([bob].toMemoryArray());
3121
}
3222

3323
function testToggleNodes() public {
34-
vm.prank(admin);
35-
_angstromBalancer.toggleNodes([bob].toMemoryArray());
36-
assertTrue(_angstromBalancer.isNode(bob), "Bob is not a node");
24+
makeAngstromNode(bob);
25+
assertTrue(angstromBalancer.isRegisteredNode(bob), "Bob is not a node");
3726
}
3827

3928
function testAddAndRemoveNodes() public {
40-
vm.startPrank(admin);
41-
_angstromBalancer.toggleNodes([bob, alice].toMemoryArray());
42-
assertTrue(_angstromBalancer.isNode(bob), "Bob is not a node");
43-
assertTrue(_angstromBalancer.isNode(alice), "Alice is not a node");
44-
_angstromBalancer.toggleNodes([bob].toMemoryArray());
45-
vm.stopPrank();
46-
assertFalse(_angstromBalancer.isNode(bob), "Bob is still a node");
47-
assertTrue(_angstromBalancer.isNode(alice), "Alice is not a node after bob was removed");
48-
}
29+
makeAngstromNodes([bob, alice].toMemoryArray());
30+
assertTrue(angstromBalancer.isRegisteredNode(bob), "Bob is not a node");
31+
assertTrue(angstromBalancer.isRegisteredNode(alice), "Alice is not a node");
4932

50-
function testUnlockAngstromNotNode() public {
51-
vm.expectRevert(AngstromBalancer.NotNode.selector);
52-
_angstromBalancer.manualUnlockAngstromWithRouter();
53-
}
54-
55-
function testUnlockAngstromTwice() public {
5633
vm.prank(admin);
57-
_angstromBalancer.toggleNodes([bob].toMemoryArray());
34+
angstromBalancer.toggleNodes([bob].toMemoryArray());
5835

59-
vm.startPrank(bob);
60-
_angstromBalancer.manualUnlockAngstromWithRouter();
61-
vm.expectRevert(AngstromBalancer.OnlyOncePerBlock.selector);
62-
_angstromBalancer.manualUnlockAngstromWithRouter();
63-
vm.stopPrank();
36+
assertFalse(angstromBalancer.isRegisteredNode(bob), "Bob is still a node");
37+
assertTrue(angstromBalancer.isRegisteredNode(alice), "Alice is not a node after bob was removed");
6438
}
6539

6640
function testUnlockAngstromSetsLastUnlockBlockNumber() public {
67-
vm.prank(admin);
68-
_angstromBalancer.toggleNodes([bob].toMemoryArray());
41+
makeAngstromNode(bob);
6942

70-
assertEq(_angstromBalancer.getLastUnlockBlockNumber(), 0, "Last unlock block number is not 0");
43+
assertEq(angstromBalancer.getLastUnlockBlockNumber(), 0, "Last unlock block number is not 0");
7144

7245
vm.prank(bob);
73-
_angstromBalancer.manualUnlockAngstromWithRouter();
46+
angstromBalancer.manualUnlockAngstrom();
7447
assertEq(
75-
_angstromBalancer.getLastUnlockBlockNumber(),
48+
angstromBalancer.getLastUnlockBlockNumber(),
7649
block.number,
7750
"Last unlock block number is not the current block number"
7851
);
7952
}
8053

8154
function testGetVault() public view {
82-
assertEq(address(_angstromBalancer.getVault()), address(vault), "Wrong vault address");
55+
assertEq(address(angstromBalancer.getVault()), address(vault), "Wrong vault address");
8356
}
8457
}

‎test/foundry/AngstromHook.t.sol‎

Lines changed: 29 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -13,38 +13,23 @@ import {
1313
} from "@balancer-labs/v3-interfaces/contracts/vault/BatchRouterTypes.sol";
1414

1515
import { ArrayHelpers } from "@balancer-labs/v3-solidity-utils/contracts/test/ArrayHelpers.sol";
16-
import { BaseVaultTest } from "@balancer-labs/v3-vault/test/foundry/utils/BaseVaultTest.sol";
1716
import { FixedPoint } from "@balancer-labs/v3-solidity-utils/contracts/math/FixedPoint.sol";
1817

1918
import { AngstromBalancerMock } from "../../contracts/test/AngstromBalancerMock.sol";
2019
import { AngstromBalancer } from "../../contracts/AngstromBalancer.sol";
20+
import { BaseAngstromTest } from "./utils/BaseAngstromTest.sol";
2121

22-
contract AngstromHookTest is BaseVaultTest {
22+
contract AngstromHookTest is BaseAngstromTest {
2323
using ArrayHelpers for *;
2424

25-
AngstromBalancerMock private _angstromBalancer;
26-
27-
function setUp() public virtual override {
28-
super.setUp();
29-
authorizer.grantRole(_angstromBalancer.getActionId(AngstromBalancer.toggleNodes.selector), admin);
30-
}
31-
32-
function createHook() internal override returns (address) {
33-
// Creating the router and hook in this function ensures that "pool" has the hook set correctly.
34-
_angstromBalancer = new AngstromBalancerMock(vault, weth, permit2, "AngstromBalancer Mock v1");
35-
return address(_angstromBalancer);
36-
}
37-
3825
/***************************************************************************
3926
Swaps
4027
***************************************************************************/
4128

4229
function testOnBeforeSwapNotNode() public {
43-
(, bytes memory userData) = _generateSignatureAndUserData(bob, bobKey);
44-
4530
vm.expectRevert(AngstromBalancer.NotNode.selector);
4631
vm.prank(bob);
47-
router.swapSingleTokenExactIn(pool, dai, usdc, 1e18, 0, MAX_UINT256, false, userData);
32+
router.swapSingleTokenExactIn(pool, dai, usdc, 1e18, 0, MAX_UINT256, false, bobUserData);
4833
}
4934

5035
function testOnBeforeSwapCannotSwapWhileLocked() public {
@@ -64,55 +49,37 @@ contract AngstromHookTest is BaseVaultTest {
6449

6550
function testOnBeforeSwapInvalidSignature() public {
6651
// If the signature is invalid, the hook reverts (in this case, signer and key do not match).
67-
(, bytes memory userData) = _generateSignatureAndUserData(bob, aliceKey);
52+
(, bytes memory userData) = generateSignatureAndUserData(bob, aliceKey);
6853

6954
vm.prank(admin);
70-
_angstromBalancer.toggleNodes([bob].toMemoryArray());
55+
angstromBalancer.toggleNodes([bob].toMemoryArray());
7156

7257
vm.expectRevert(AngstromBalancer.InvalidSignature.selector);
7358
vm.prank(bob);
7459
router.swapSingleTokenExactIn(pool, dai, usdc, 1e18, 0, MAX_UINT256, false, userData);
7560
}
7661

7762
function testOnBeforeSwapSucceedsAndSetBlockNumber() public {
78-
(, bytes memory userData) = _generateSignatureAndUserData(bob, bobKey);
79-
80-
vm.prank(admin);
81-
_angstromBalancer.toggleNodes([bob].toMemoryArray());
63+
makeAngstromNode(bob);
8264

8365
vm.prank(bob);
84-
router.swapSingleTokenExactIn(pool, dai, usdc, 1e18, 0, MAX_UINT256, false, userData);
66+
router.swapSingleTokenExactIn(pool, dai, usdc, 1e18, 0, MAX_UINT256, false, bobUserData);
8567
assertEq(
86-
_angstromBalancer.getLastUnlockBlockNumber(),
68+
angstromBalancer.getLastUnlockBlockNumber(),
8769
block.number,
8870
"Last unlock block number is not the current block number"
8971
);
9072
}
9173

9274
function testOnlyOncePerBlock() public {
93-
vm.prank(admin);
94-
_angstromBalancer.toggleNodes([bob].toMemoryArray());
95-
96-
(bytes memory signature, ) = _generateSignatureAndUserData(bob, bobKey);
97-
vm.prank(bob);
98-
_angstromBalancer.unlockWithEmptyAttestation(bob, signature);
99-
100-
vm.expectRevert(AngstromBalancer.OnlyOncePerBlock.selector);
101-
vm.prank(bob);
102-
_angstromBalancer.unlockWithEmptyAttestation(bob, signature);
103-
}
104-
105-
function testOnlyOncePerBlockCalldata() public {
106-
vm.prank(admin);
107-
_angstromBalancer.toggleNodes([bob].toMemoryArray());
75+
makeAngstromNode(bob);
10876

109-
(bytes memory signature, ) = _generateSignatureAndUserData(bob, bobKey);
11077
vm.prank(bob);
111-
_angstromBalancer.unlockWithEmptyAttestationCalldata(bob, signature);
78+
angstromBalancer.unlockWithEmptyAttestation(bob, bobSignature);
11279

11380
vm.expectRevert(AngstromBalancer.OnlyOncePerBlock.selector);
11481
vm.prank(bob);
115-
_angstromBalancer.unlockWithEmptyAttestationCalldata(bob, signature);
82+
angstromBalancer.unlockWithEmptyAttestation(bob, bobSignature);
11683
}
11784

11885
/***************************************************************************
@@ -138,40 +105,40 @@ contract AngstromHookTest is BaseVaultTest {
138105
}
139106

140107
function testOnBeforeAddLiquidityUnbalancedNotNode() public {
141-
(, bytes memory userData) = _generateSignatureAndUserData(alice, aliceKey);
142-
143108
vm.expectRevert(AngstromBalancer.NotNode.selector);
144109
vm.prank(alice);
145-
router.addLiquidityUnbalanced(pool, [FixedPoint.ONE, FixedPoint.ONE].toMemoryArray(), 1e18, false, userData);
110+
router.addLiquidityUnbalanced(
111+
pool,
112+
[FixedPoint.ONE, FixedPoint.ONE].toMemoryArray(),
113+
1e18,
114+
false,
115+
aliceUserData
116+
);
146117
}
147118

148119
function testOnBeforeAddLiquidityUnbalancedInvalidSignature() public {
149-
(, bytes memory userData) = _generateSignatureAndUserData(alice, bobKey);
120+
(, bytes memory userData) = generateSignatureAndUserData(alice, bobKey);
150121

151-
vm.prank(admin);
152-
_angstromBalancer.toggleNodes([alice].toMemoryArray());
122+
makeAngstromNode(alice);
153123

154124
vm.expectRevert(AngstromBalancer.InvalidSignature.selector);
155125
vm.prank(alice);
156126
router.addLiquidityUnbalanced(pool, [FixedPoint.ONE, FixedPoint.ONE].toMemoryArray(), 1e18, false, userData);
157127
}
158128

159129
function testOnBeforeAddLiquidityUnbalancedSucceedsAndSetBlockNumber() public {
160-
(, bytes memory userData) = _generateSignatureAndUserData(alice, aliceKey);
161-
162-
vm.prank(admin);
163-
_angstromBalancer.toggleNodes([alice].toMemoryArray());
130+
makeAngstromNode(alice);
164131

165132
vm.prank(alice);
166133
uint256 expectedBptAmountOut = router.addLiquidityUnbalanced(
167134
pool,
168135
[FixedPoint.ONE, FixedPoint.ONE].toMemoryArray(),
169136
1e18,
170137
false,
171-
userData
138+
aliceUserData
172139
);
173140
assertEq(
174-
_angstromBalancer.getLastUnlockBlockNumber(),
141+
angstromBalancer.getLastUnlockBlockNumber(),
175142
block.number,
176143
"Last unlock block number is not the current block number"
177144
);
@@ -206,52 +173,36 @@ contract AngstromHookTest is BaseVaultTest {
206173
}
207174

208175
function testOnBeforeRemoveLiquidityUnbalancedNotNode() public {
209-
(, bytes memory userData) = _generateSignatureAndUserData(lp, lpKey);
210-
211176
vm.expectRevert(AngstromBalancer.NotNode.selector);
212177
vm.prank(lp);
213-
router.removeLiquiditySingleTokenExactIn(pool, 1e18, dai, 0.1e18, false, userData);
178+
router.removeLiquiditySingleTokenExactIn(pool, 1e18, dai, 0.1e18, false, lpUserData);
214179
}
215180

216181
function testOnBeforeRemoveLiquidityUnbalancedInvalidSignature() public {
217-
(, bytes memory userData) = _generateSignatureAndUserData(lp, bobKey);
182+
(, bytes memory userData) = generateSignatureAndUserData(lp, bobKey);
218183

219-
vm.prank(admin);
220-
_angstromBalancer.toggleNodes([lp].toMemoryArray());
184+
makeAngstromNode(lp);
221185

222186
vm.expectRevert(AngstromBalancer.InvalidSignature.selector);
223187
vm.prank(lp);
224188
router.removeLiquiditySingleTokenExactIn(pool, 1e18, dai, 0.1e18, false, userData);
225189
}
226190

227191
function testOnBeforeRemoveLiquidityUnbalancedSucceedsAndSetBlockNumber() public {
228-
(, bytes memory userData) = _generateSignatureAndUserData(lp, lpKey);
229-
230-
vm.prank(admin);
231-
_angstromBalancer.toggleNodes([lp].toMemoryArray());
192+
makeAngstromNode(lp);
232193

233194
Balances memory balancesBefore = getBalances(lp);
234195

235196
vm.prank(lp);
236-
router.removeLiquiditySingleTokenExactIn(pool, 1e18, dai, 0.1e18, false, userData);
197+
router.removeLiquiditySingleTokenExactIn(pool, 1e18, dai, 0.1e18, false, lpUserData);
237198

238199
Balances memory balancesAfter = getBalances(lp);
239200

240201
assertEq(
241-
_angstromBalancer.getLastUnlockBlockNumber(),
202+
angstromBalancer.getLastUnlockBlockNumber(),
242203
block.number,
243204
"Last unlock block number is not the current block number"
244205
);
245206
assertEq(balancesAfter.lpBpt, balancesBefore.lpBpt - 1e18, "LP did not burn BPTs");
246207
}
247-
248-
function _generateSignatureAndUserData(
249-
address signer,
250-
uint256 privateKey
251-
) private view returns (bytes memory signature, bytes memory userData) {
252-
bytes32 hash = _angstromBalancer.getDigest();
253-
(uint8 v, bytes32 r, bytes32 s) = vm.sign(privateKey, hash);
254-
signature = abi.encodePacked(r, s, v);
255-
userData = abi.encodePacked(signer, signature);
256-
}
257208
}

0 commit comments

Comments
 (0)