Skip to content

Commit

Permalink
Merge pull request #140 from BeanstalkFarms/well-upgr-token-check
Browse files Browse the repository at this point in the history
Add check for new well tokens in `authorizeUpgrade` [#23]
  • Loading branch information
nickkatsios authored Aug 21, 2024
2 parents 163b53e + 54ac748 commit e824771
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 9 deletions.
3 changes: 1 addition & 2 deletions foundry.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ out = 'out'
libs = ['lib', 'node_modules']
fuzz = { runs = 256 }
optimizer = true
optimizer_runs = 400
optimizer_runs = 200
remappings = [
'@openzeppelin/=node_modules/@openzeppelin/',
]
Expand All @@ -27,7 +27,6 @@ ignore = ["src/libraries/LibClone.sol", "src/utils/Clone.sol", "src/libraries/AB
int_types = "long"
line_length = 120
multiline_func_header = "params_first"
number_underscore = "thousands"
override_spacing = false
quote_style = "double"
tab_width = 4
Expand Down
8 changes: 8 additions & 0 deletions src/WellUpgradeable.sol
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,14 @@ contract WellUpgradeable is Well, UUPSUpgradeable, OwnableUpgradeable {
"New implementation must be a well implmentation"
);

// verify the new well uses the same tokens in the same order.
IERC20[] memory _tokens = tokens();
IERC20[] memory newTokens = WellUpgradeable(newImplmentation).tokens();
require(_tokens.length == newTokens.length, "New well must use the same number of tokens");
for (uint256 i; i < _tokens.length; ++i) {
require(_tokens[i] == newTokens[i], "New well must use the same tokens in the same order");
}

// verify the new implmentation is a valid ERC-1967 implmentation.
require(
UUPSUpgradeable(newImplmentation).proxiableUUID() == _IMPLEMENTATION_SLOT,
Expand Down
57 changes: 50 additions & 7 deletions test/WellUpgradeable.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -28,12 +28,14 @@ contract WellUpgradeTest is Test, WellDeployer {
address token2Address;
address wellAddress;
address wellImplementation;
IERC20[] tokens = new IERC20[](2);

function setUp() public {
// Tokens
IERC20[] memory tokens = new IERC20[](2);
tokens[0] = new MockToken("BEAN", "BEAN", 6);
tokens[1] = new MockToken("WETH", "WETH", 18);
IERC20 token0 = new MockToken("BEAN", "BEAN", 6);
IERC20 token1 = new MockToken("WETH", "WETH", 18);
tokens[0] = token0;
tokens[1] = token1;

token1Address = address(tokens[0]);
vm.label(token1Address, "token1");
Expand Down Expand Up @@ -182,18 +184,15 @@ contract WellUpgradeTest is Test, WellDeployer {
////////////////////// Upgrade Tests //////////////////////

function testUpgradeToNewImplementation() public {
IERC20[] memory tokens = new IERC20[](2);
tokens[0] = new MockToken("BEAN", "BEAN", 6);
tokens[1] = new MockToken("WETH", "WETH", 18);
Call memory wellFunction = Call(wellFunctionAddress, abi.encode("2"));
Call[] memory pumps = new Call[](1);
pumps[0] = Call(mockPumpAddress, abi.encode("2"));
// create new mock Well Implementation:
address wellImpl = address(new MockWellUpgradeable());
// bore new well with the same 2 tokens
WellUpgradeable well2 =
encodeAndBoreWellUpgradeable(aquifer, wellImpl, tokens, wellFunction, pumps, bytes32(abi.encode("2")));
vm.label(address(well2), "upgradeableWell2");

vm.startPrank(initialOwner);
WellUpgradeable proxy = WellUpgradeable(payable(proxyAddress));
proxy.upgradeTo(address(well2));
Expand All @@ -204,4 +203,48 @@ contract WellUpgradeTest is Test, WellDeployer {
assertEq(100, MockWellUpgradeable(proxyAddress).getVersion(100));
vm.stopPrank();
}

function testUpgradeToNewImplementationDiffTokens() public {
// create 2 new tokens with new addresses
IERC20[] memory newTokens = new IERC20[](2);
newTokens[0] = new MockToken("WBTC", "WBTC", 6);
newTokens[1] = new MockToken("WETH2", "WETH2", 18);
Call memory wellFunction = Call(wellFunctionAddress, abi.encode("2"));
Call[] memory pumps = new Call[](1);
pumps[0] = Call(mockPumpAddress, abi.encode("2"));
// create new mock Well Implementation:
address wellImpl = address(new MockWellUpgradeable());
// bore new well with the different tokens
WellUpgradeable well2 =
encodeAndBoreWellUpgradeable(aquifer, wellImpl, newTokens, wellFunction, pumps, bytes32(abi.encode("2")));
vm.label(address(well2), "upgradeableWell2");
vm.startPrank(initialOwner);
WellUpgradeable proxy = WellUpgradeable(payable(proxyAddress));
// expect revert since new well uses different tokens
vm.expectRevert("New well must use the same tokens in the same order");
proxy.upgradeTo(address(well2));
vm.stopPrank();
}

function testUpgradeToNewImplementationDiffTokenOrder() public {
// create 2 new tokens with new addresses
IERC20[] memory newTokens = new IERC20[](2);
newTokens[0] = tokens[1];
newTokens[1] = tokens[0];
Call memory wellFunction = Call(wellFunctionAddress, abi.encode("2"));
Call[] memory pumps = new Call[](1);
pumps[0] = Call(mockPumpAddress, abi.encode("2"));
// create new mock Well Implementation:
address wellImpl = address(new MockWellUpgradeable());
// bore new well with the different tokens
WellUpgradeable well2 =
encodeAndBoreWellUpgradeable(aquifer, wellImpl, newTokens, wellFunction, pumps, bytes32(abi.encode("2")));
vm.label(address(well2), "upgradeableWell2");
vm.startPrank(initialOwner);
WellUpgradeable proxy = WellUpgradeable(payable(proxyAddress));
// expect revert since new well uses different tokens
vm.expectRevert("New well must use the same tokens in the same order");
proxy.upgradeTo(address(well2));
vm.stopPrank();
}
}

0 comments on commit e824771

Please sign in to comment.