add audit fixes

wormhole-foundation · Jan 14, 2025 · 2be79b9 · 2be79b9
1 parent 6d8b638
commit 2be79b9
Show file tree

Hide file tree

Showing 8 changed files with 551 additions and 54 deletions.
diff --git a/src/components/dispatcher/AccessControl.sol b/src/components/dispatcher/AccessControl.sol
@@ -14,7 +14,8 @@ import {
   ADD_ADMIN_ID,
   PROPOSE_OWNERSHIP_TRANSFER_ID,
   ACQUIRE_OWNERSHIP_ID,
-  RELINQUISH_OWNERSHIP_ID
+  RELINQUISH_OWNERSHIP_ID,
+  CANCEL_OWNERSHIP_TRANSFER_ID
 } from "wormhole-sdk/components/dispatcher/Ids.sol";
 
 //rationale for different roles (owner, admin):
@@ -83,33 +84,42 @@ abstract contract AccessControl {
     address owner,
     address[] memory admins
   ) internal {
-    accessControlState().owner = owner;
+    AccessControlState storage state = accessControlState();
+    state.owner = owner;
     for (uint i = 0; i < admins.length; ++i)
-      _updateAdmins(admins[i], true);
+      _updateAdmins(state, admins[i], true);
   }
 
   // ---- external -----
 
+  //selector: f2fde38b
   function transferOwnership(address newOwner) external {
     AccessControlState storage state = accessControlState();
-    failAuthIf(msg.sender != state.owner);
+    if (msg.sender != state.owner)
+      revert NotAuthorized();
 
-    state.pendingOwner = newOwner;
+    _proposeOwnershipTransfer(state, newOwner);
   }
 
+  //selector: 23452b9c 
   function cancelOwnershipTransfer() external {
     AccessControlState storage state = accessControlState();
-    failAuthIf(msg.sender != state.owner);
+    if (msg.sender != state.owner)
+      revert NotAuthorized();
 
-    state.pendingOwner = address(0);
+    _cancelOwnershipTransfer(state);
   }
 
+  //selector: 1c74a301
   function receiveOwnership() external {
     _acquireOwnership();
   }
 
   // ---- internals ----
 
+  /**
+   * Dispatch an execute function. Execute functions almost always modify contract state.
+   */
   function dispatchExecAccessControl(
     bytes calldata data,
     uint offset,
@@ -125,6 +135,9 @@ abstract contract AccessControl {
     return (true, offset);
   }
 
+  /**
+   * Dispatch a query function. Query functions never modify contract state.
+   */
   function dispatchQueryAccessControl(
     bytes calldata data,
     uint offset,
@@ -154,7 +167,7 @@ abstract contract AccessControl {
       if (command == REVOKE_ADMIN_ID) {
         address admin;
         (admin, offset) = commands.asAddressCdUnchecked(offset);
-        _updateAdmins(admin, false);
+        _updateAdmins(state, admin, false);
       }
       else {
         if (!isOwner)
@@ -163,20 +176,22 @@ abstract contract AccessControl {
         if (command == ADD_ADMIN_ID) {
           address newAdmin;
           (newAdmin, offset) = commands.asAddressCdUnchecked(offset);
-
-          _updateAdmins(newAdmin, true);
+          _updateAdmins(state, newAdmin, true);
         }
         else if (command == PROPOSE_OWNERSHIP_TRANSFER_ID) {
           address newOwner;
           (newOwner, offset) = commands.asAddressCdUnchecked(offset);
 
-          state.pendingOwner = newOwner;
+          _proposeOwnershipTransfer(state, newOwner);
+        }
+        else if (command == CANCEL_OWNERSHIP_TRANSFER_ID) {
+          _cancelOwnershipTransfer(state);
         }
         else if (command == RELINQUISH_OWNERSHIP_ID) {
-          _updateOwner(address(0));
+          _relinquishOwnership(state);
 
           //ownership relinquishment must be the last command in the batch
-          commands.checkLengthCd(offset);
+          BytesParsing.checkLength(offset, commands.length);
         }
         else
           revert InvalidAccessControlCommand(command);
@@ -224,25 +239,41 @@ abstract contract AccessControl {
     return (ret, offset);
   }
 
-  function _acquireOwnership() internal {
+  // ---- private ----
+
+  function _acquireOwnership() private {
     AccessControlState storage state = accessControlState();
-    if (msg.sender !=state.pendingOwner)
+    if (state.pendingOwner != msg.sender)
       revert NotAuthorized();
 
-    state.pendingOwner = address(0);
-    _updateOwner(msg.sender);
+    _updateOwner(state, msg.sender);
   }
 
-  // ---- private ----
+  function _relinquishOwnership(AccessControlState storage state) private {
+    _updateOwner(state, address(0));
+  }
+
+  function _updateOwner(AccessControlState storage state, address newOwner) private {
+    address oldAddress = state.owner;
+    state.owner = newOwner;
+    state.pendingOwner = address(0);
 
-  function _updateOwner(address newOwner) private {
-    address oldAddress;
-    accessControlState().owner = newOwner;
     emit OwnerUpdated(oldAddress, newOwner, block.timestamp);
   }
 
-  function _updateAdmins(address admin, bool authorization) private { unchecked {
-    AccessControlState storage state = accessControlState();
+  function _proposeOwnershipTransfer(AccessControlState storage state, address newOwner) private {
+    state.pendingOwner = newOwner;
+  }
+
+  function _cancelOwnershipTransfer(AccessControlState storage state) private {
+    state.pendingOwner = address(0);
+  }
+
+  function _updateAdmins(
+    AccessControlState storage state, 
+    address admin, 
+    bool authorization
+  ) private { unchecked {
     if ((state.isAdmin[admin] != 0) == authorization)
       return;
 
@@ -252,8 +283,11 @@ abstract contract AccessControl {
     }
     else {
       uint256 rawIndex = state.isAdmin[admin];
-      if (rawIndex != state.admins.length)
-        state.admins[rawIndex - 1] = state.admins[state.admins.length - 1];
+      if (rawIndex != state.admins.length) {
+        address tmpAdmin = state.admins[state.admins.length - 1];
+        state.isAdmin[tmpAdmin] = rawIndex;
+        state.admins[rawIndex - 1] = tmpAdmin;
+      }
 
       state.isAdmin[admin] = 0;
       state.admins.pop();

diff --git a/src/components/dispatcher/Ids.sol b/src/components/dispatcher/Ids.sol
@@ -27,6 +27,7 @@ uint8 constant REVOKE_ADMIN_ID = 0x00;
 uint8 constant PROPOSE_OWNERSHIP_TRANSFER_ID = 0x10;
 uint8 constant RELINQUISH_OWNERSHIP_ID = 0x11;
 uint8 constant ADD_ADMIN_ID = 0x12;
+uint8 constant CANCEL_OWNERSHIP_TRANSFER_ID = 0x13;
 
 // Query commands
 

diff --git a/src/components/dispatcher/SweepTokens.sol b/src/components/dispatcher/SweepTokens.sol
@@ -24,7 +24,7 @@ abstract contract SweepTokens {
     bytes calldata commands,
     uint offset
   ) internal returns (uint) {
-    senderAtLeastAdmin();
+    sweepTokenDoAuth();
 
     address token;
     uint256 amount;
@@ -34,4 +34,8 @@ abstract contract SweepTokens {
     tokenOrNativeTransfer(token, msg.sender, amount);
     return offset;
   }
+
+  function sweepTokenDoAuth() view internal virtual {
+    senderAtLeastAdmin();
+  }
 }
diff --git a/src/components/dispatcher/Upgrade.sol b/src/components/dispatcher/Upgrade.sol
@@ -13,6 +13,20 @@ error InvalidGovernanceQuery(uint8 query);
 abstract contract Upgrade is ProxyBase {
   using BytesParsing for bytes;
 
+  // ------ external ------
+
+  //selector: c987336c
+  function upgrade(address implementation, bytes calldata data) external {
+    failAuthIf(senderRole() != Role.Owner);
+
+    _upgradeTo(implementation, data);
+  }
+
+  // ------ internal ------
+
+  /**
+   * Dispatch an execute function. Execute functions almost always modify contract state.
+   */
   function dispatchExecUpgrade(
     bytes calldata data,
     uint offset,
@@ -23,6 +37,9 @@ abstract contract Upgrade is ProxyBase {
       : (false, offset);
   }
 
+  /**
+   * Dispatch a query function. Query functions never modify contract state.
+   */
   function dispatchQueryUpgrade(
     bytes calldata,
     uint offset,
@@ -33,12 +50,6 @@ abstract contract Upgrade is ProxyBase {
       : (false, new bytes(0), offset);
   }
 
-  function upgrade(address implementation, bytes calldata data) external {
-    failAuthIf(senderRole() != Role.Owner);
-
-    _upgradeTo(implementation, data);
-  }
-
   function _upgradeContract(
     bytes calldata commands,
     uint offset
@@ -47,10 +58,13 @@ abstract contract Upgrade is ProxyBase {
 
     address newImplementation;
     (newImplementation, offset) = commands.asAddressCdUnchecked(offset);
+    bytes calldata data;
+    (data, offset) = commands.sliceCdUnchecked(offset, commands.length - offset);
+
     //contract upgrades must be the last command in the batch
-    commands.checkLengthCd(offset);
+    BytesParsing.checkLength(offset, commands.length);
 
-    _upgradeTo(newImplementation, new bytes(0));
+    _upgradeTo(newImplementation, data);
 
     return offset;
   }