https://github.com/Cyfrin/2024-07-zaros/blob/d687fe96bb7ace8652778797052a38763fbcbb1b/src/tree-proxy/leaves/RootUpgrade.sol#L177-L204
Summary
If the owner proceeds to upgrade the UpgradeBranch::upgrade function via the RootUpgrade::removeBranch function, it would render the addition of a new UpgradeBranch::upgrade function to the Protocol through the RootUpgrade::addBranch function impossible. Without the presence of the UpgradeBranch::upgrade function the Protocol loses upgradability.
Vulnerability Details
The UpgradeBranch uses for upgrade the Protocol. To do that, owner can call UpgradeBranch::upgrade function with next params Link to code:
FILE: src/tree-proxy/branches/UpgradeBranch.sol
function upgrade(
RootProxy.BranchUpgrade[] memory branchUpgrades,
address[] memory initializables,
bytes[] memory initializePayloads
) { ... }
FILE: src/tree-proxy/RootProxy.sol
enum BranchUpgradeAction {
Add,
Replace,
Remove
}
...
struct BranchUpgrade {
address branch;
BranchUpgradeAction action;
bytes4[] selectors;
}
In the current realization owner has two ways to upgrade any function:
First: call upgrade function with BranchUpgradeAction.Replace param
Second: call upgrade function with BranchUpgradeAction.Remove param and than call upgrade function with BranchUpgradeAction.Add param
RootUpgrade::removeBranch function link to code is internal function that provides the logic for removing branch selectors when BranchUpgradeAction.Remove param was choosen. There is one check for branch in this function: if (branch == address(this)) { revert Errors.ImmutableBranch(); }.
Since this function calls trought the Proxy, address(this) is address of the Proxy, i.e. address of PerpsEngine contract.
There is two checks for branch selector in RootUpgrade::removeBranch function: selector != bytes4(0) and self.selectorToBranch[selector] == branch.
So, owner can remove functions from any branch incude UpgradeBranch::upgrade.
FILE: src/tree-proxy/leaves/RootUpgrade.sol
function removeBranch(Data storage self, address branch, bytes4[] memory selectors) internal {
if (branch == address(this)) {
revert Errors.ImmutableBranch();
}
uint256 cachedSelectorsLength = selectors.length;
for (uint256 i; i < cachedSelectorsLength; i++) {
bytes4 selector = selectors[i];
if (selector == bytes4(0)) {
revert Errors.SelectorIsZero();
}
if (self.selectorToBranch[selector] != branch) {
revert Errors.CannotRemoveFromOtherBranch(branch, selector);
}
delete self.selectorToBranch[selector];
self.branchSelectors[branch].remove(selector);
}
if (self.branchSelectors[branch].length() == 0) {
self.branches.remove(branch);
}
}
If owner chooses second way to upgrade UpgradeBranch::upgrade function itself, he will find themselves unable to call upgrade function with BranchUpgradeAction.Add param following its removal. Since the upgrade function will be removed from the list of functions available for calling. Owner will no longer be able to call the upgrade function and update the protocol.
Proof of Concept
Let's add a new test to the upgrade.t.sol test file. We remove an upgrade function from the Protocol. Then we try to upgrade the protocol in several ways.
Also, add import { UpgradeBranch } from "@zaros/tree-proxy/branches/UpgradeBranch.sol"; to this file for corretly testing.
function test_WhenRemovingAUpgradeBranch() external givenTheSenderIsTheOwner {
changePrank({ msgSender: users.owner.account });
bytes4 selector = UpgradeBranch.upgrade.selector;
address upgradeBranch = perpsEngine.branchAddress(selector);
address[] memory branches = new address[](1);
branches[0] = upgradeBranch;
bytes4[] memory upgradeBranchSelectors = new bytes4[](1);
upgradeBranchSelectors[0] = selector;
bytes4[][] memory branchesSelectors = new bytes4[][](1);
branchesSelectors[0] = upgradeBranchSelectors;
RootProxy.BranchUpgrade[] memory branchUpgrades =
getBranchUpgrades(branches, branchesSelectors, RootProxy.BranchUpgradeAction.Remove);
perpsEngine.upgrade(branchUpgrades, new address[](0), new bytes[](0));
branchUpgrades =
getBranchUpgrades(branches, branchesSelectors, RootProxy.BranchUpgradeAction.Add);
vm.expectRevert({
revertData: abi.encodeWithSelector(
Errors.UnsupportedFunction.selector, UpgradeBranch.upgrade.selector
)
});
perpsEngine.upgrade(branchUpgrades, new address[](0), new bytes[](0));
Branch.Data[] memory perpsEngineBranches = perpsEngine.branches();
address orderBranchAddress;
for (uint256 i; i < perpsEngineBranches.length; i++) {
if (i == 4) {
orderBranchAddress = perpsEngineBranches[i].branch;
}
}
branches[0] = orderBranchAddress;
bytes4[] memory orderBranchSelectors = new bytes4[](1);
orderBranchSelectors[0] = OrderBranch.getActiveMarketOrder.selector;
branchesSelectors[0] = orderBranchSelectors;
branchUpgrades =
getBranchUpgrades(branches, branchesSelectors, RootProxy.BranchUpgradeAction.Remove);
vm.expectRevert({
revertData: abi.encodeWithSelector(
Errors.UnsupportedFunction.selector, UpgradeBranch.upgrade.selector
)
});
perpsEngine.upgrade(branchUpgrades, new address[](0), new bytes[](0));
address newOrderBranch = address(new NewOrderBranch());
branches[0] = newOrderBranch;
bytes4[] memory newOrderBranchSelectors = new bytes4[](1);
newOrderBranchSelectors[0] = NewOrderBranch.getName.selector;
branchesSelectors[0] = newOrderBranchSelectors;
branchUpgrades =
getBranchUpgrades(branches, branchesSelectors, RootProxy.BranchUpgradeAction.Replace);
vm.expectRevert({
revertData: abi.encodeWithSelector(
Errors.UnsupportedFunction.selector, UpgradeBranch.upgrade.selector
)
});
perpsEngine.upgrade(branchUpgrades, new address[](0), new bytes[](0));
}
Let's run the test: forge test --mt test_WhenRemovingAUpgradeBranch
2024-07-zaros$ forge test --mt test_WhenRemovingAUpgradeBranch
[⠒] Compiling...
[⠆] Compiling 1 files with Solc 0.8.25
[⠔] Solc 0.8.25 finished in 7.13s
Compiler run successful!
Ran 1 test for test/integration/tree-proxy/upgrade-branch/upgrade/upgrade.t.sol:Upgrade_Integration_Test
[PASS] test_WhenRemovingAUpgradeBranch() (gas: 2775890)
Suite result: ok. 1 passed; 0 failed; 0 skipped; finished in 16.25ms (1.22ms CPU time)
Ran 1 test suite in 16.95ms (16.25ms CPU time): 1 tests passed, 0 failed, 0 skipped (1 total tests)
We can see the test passing.
Impact
Vulnerability requires owner actions. But the impact of these actions is High. The protocol will completely lose the upgradability. To regain this functionality, it becomes necessary to redeploy the Protocol.
Tools Used
Manual Review
Recommendations
Add the require in the RootUpgrade::removeBranch function to check that branch selector doesn't equal UpgradeBranch::upgrade function selector.
FILE: src/tree-proxy/leaves/RootUpgrade.sol
+ import { UpgradeBranch } from "@zaros/tree-proxy/branches/UpgradeBranch.sol";
...
function removeBranch(Data storage self, address branch, bytes4[] memory selectors) internal {
if (branch == address(this)) {
revert Errors.ImmutableBranch();
}
uint256 cachedSelectorsLength = selectors.length;
for (uint256 i; i < cachedSelectorsLength; i++) {
bytes4 selector = selectors[i];
// also reverts if left side returns zero address
if (selector == bytes4(0)) {
revert Errors.SelectorIsZero();
}
if (self.selectorToBranch[selector] != branch) {
revert Errors.CannotRemoveFromOtherBranch(branch, selector);
}
+ if (selector == UpgradeBranch.upgrade.selector) {
+ revert("Don't remove upgrade selector");
+ }
delete self.selectorToBranch[selector];
// slither-disable-next-line unused-return
self.branchSelectors[branch].remove(selector);
}
// if no more selectors in branch, remove branch address
if (self.branchSelectors[branch].length() == 0) {
// slither-disable-next-line unused-return
self.branches.remove(branch);
}
}
