fix: Fix old ProxyAdmin cannot upgrade openzeppelin-v5's TransparentU…

…pgradeableProxy with upgradeTo & null check ProxyAdmin & expose get sender pk for signing
axieinfinity · Dec 14, 2023 · ed44275 · ed44275
1 parent 6b75a10
commit ed44275
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 23 deletions.
diff --git a/script/BaseMigration.s.sol b/script/BaseMigration.s.sol
@@ -166,11 +166,7 @@ abstract contract BaseMigration is ScriptExtended {
     _mockUpgradeRaw(proxy.getProxyAdmin(), proxy, logic, args);
   }
 
-  function _upgradeProxy(TContract contractType)
-    internal
-    virtual
-    returns (address payable proxy)
-  {
+  function _upgradeProxy(TContract contractType) internal virtual returns (address payable proxy) {
     proxy = _upgradeProxy(contractType, arguments());
   }
 
@@ -189,41 +185,73 @@ abstract contract BaseMigration is ScriptExtended {
     internal
     virtual
   {
+    ITransparentUpgradeableProxy iProxy = ITransparentUpgradeableProxy(proxy);
+    ProxyAdmin wProxyAdmin = ProxyAdmin(proxyAdmin);
+    // if proxyAdmin is External Owned Wallet
     if (proxyAdmin.code.length == 0) {
       vm.prank(proxyAdmin);
-      if (args.length == 0) ITransparentUpgradeableProxy(proxy).upgradeTo(logic);
-      else ITransparentUpgradeableProxy(proxy).upgradeToAndCall(logic, args);
+      if (args.length == 0) iProxy.upgradeTo(logic);
+      else iProxy.upgradeToAndCall(logic, args);
     } else {
-      try ProxyAdmin(proxyAdmin).owner() returns (address owner) {
-        vm.prank(owner);
+      try wProxyAdmin.owner() returns (address owner) {
         if (args.length == 0) {
-          ProxyAdmin(proxyAdmin).upgrade(ITransparentUpgradeableProxy(proxy), logic);
+          // try `upgrade` function
+          vm.prank(owner);
+          (bool success,) = proxyAdmin.call(abi.encodeCall(ProxyAdmin.upgrade, (iProxy, logic)));
+          if (success) {
+            vm.prank(owner);
+            wProxyAdmin.upgrade(iProxy, logic);
+          } else {
+            console.log(
+              StdStyle.yellow(
+                "`ProxyAdmin:upgrade` failed!. Retrying with `ProxyAdmin:upgradeAndCall` with emty args..."
+              )
+            );
+            vm.prank(owner);
+            wProxyAdmin.upgradeAndCall(iProxy, logic, args);
+          }
         } else {
-          ProxyAdmin(proxyAdmin).upgradeAndCall(ITransparentUpgradeableProxy(proxy), logic, args);
+          vm.prank(owner);
+          wProxyAdmin.upgradeAndCall(iProxy, logic, args);
         }
       } catch {
-        vm.prank(proxyAdmin);
-        if (args.length == 0) ITransparentUpgradeableProxy(proxy).upgradeTo(logic);
-        else ITransparentUpgradeableProxy(proxy).upgradeToAndCall(logic, args);
+        revert("BaseMigration: Unknown ProxyAdmin contract!");
       }
     }
   }
 
   function _upgradeRaw(address proxyAdmin, address payable proxy, address logic, bytes memory args) internal virtual {
+    ITransparentUpgradeableProxy iProxy = ITransparentUpgradeableProxy(proxy);
+    ProxyAdmin wProxyAdmin = ProxyAdmin(proxyAdmin);
+    // if proxyAdmin is External Owned Wallet
     if (proxyAdmin.code.length == 0) {
       vm.broadcast(proxyAdmin);
-      if (args.length == 0) ITransparentUpgradeableProxy(proxy).upgradeTo(logic);
-      else ITransparentUpgradeableProxy(proxy).upgradeToAndCall(logic, args);
+      if (args.length == 0) iProxy.upgradeTo(logic);
+      else iProxy.upgradeToAndCall(logic, args);
     } else {
-      try ProxyAdmin(proxyAdmin).owner() returns (address owner) {
-        vm.broadcast(owner);
+      try wProxyAdmin.owner() returns (address owner) {
         if (args.length == 0) {
-          ProxyAdmin(proxyAdmin).upgrade(ITransparentUpgradeableProxy(proxy), logic);
+          // try `upgrade` function
+          vm.prank(owner);
+          (bool success,) = proxyAdmin.call(abi.encodeCall(ProxyAdmin.upgrade, (iProxy, logic)));
+          if (success) {
+            vm.broadcast(owner);
+            wProxyAdmin.upgrade(iProxy, logic);
+          } else {
+            console.log(
+              StdStyle.yellow(
+                "`ProxyAdmin:upgrade` failed!. Retrying with `ProxyAdmin:upgradeAndCall` with emty args..."
+              )
+            );
+            vm.broadcast(owner);
+            wProxyAdmin.upgradeAndCall(iProxy, logic, args);
+          }
         } else {
-          ProxyAdmin(proxyAdmin).upgradeAndCall(ITransparentUpgradeableProxy(proxy), logic, args);
+          vm.broadcast(owner);
+          wProxyAdmin.upgradeAndCall(iProxy, logic, args);
         }
       } catch {
-        revert("BaseMigration: Unhandled case for upgrading proxy!");
+        revert("BaseMigration: Unknown ProxyAdmin contract!");
       }
     }
   }

diff --git a/script/configs/WalletConfig.sol b/script/configs/WalletConfig.sol
@@ -10,6 +10,10 @@ abstract contract WalletConfig is IWalletConfig {
 
   function getSender() public view virtual returns (address payable sender);
 
+  function getSenderPk() public view virtual returns (uint256) {
+    return _envPk;
+  }
+
   function trezorPrefix() public view virtual returns (string memory) {
     return "trezor://";
   }

diff --git a/script/interfaces/configs/IWalletConfig.sol b/script/interfaces/configs/IWalletConfig.sol
@@ -2,6 +2,8 @@
 pragma solidity ^0.8.19;
 
 interface IWalletConfig {
+  function getSenderPk() external view returns (uint256);
+
   function getSender() external view returns (address payable sender);
 
   function trezorPrefix() external view returns (string memory);

diff --git a/script/libraries/LibProxy.sol b/script/libraries/LibProxy.sol
@@ -9,19 +9,29 @@ library LibProxy {
   bytes32 internal constant ADMIN_SLOT = 0xb53127684a568b3173ae13b9f8a6016e243e63b6e8ee1178d6a717850b5d6103;
   bytes32 internal constant IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
 
-  function getProxyAdmin(address payable proxy) internal returns (address payable proxyAdmin) {
+  function getProxyAdmin(address payable proxy, bool nullCheck) internal returns (address payable proxyAdmin) {
     proxyAdmin = payable(address(uint160(uint256(vm.load(address(proxy), ADMIN_SLOT)))));
+    if (!nullCheck) return proxyAdmin;
     require(
       proxyAdmin != address(0x0),
       string.concat("LibProxy: Null ProxyAdmin, Provided address: ", vm.getLabel(proxy), " is not EIP1967 Proxy")
     );
   }
 
-  function getProxyImplementation(address payable proxy) internal returns (address payable impl) {
+  function getProxyAdmin(address payable proxy) internal returns (address payable proxyAdmin) {
+    proxyAdmin = getProxyAdmin({ proxy: proxy, nullCheck: true });
+  }
+
+  function getProxyImplementation(address payable proxy, bool nullCheck) internal returns (address payable impl) {
     impl = payable(address(uint160(uint256(vm.load(address(proxy), IMPLEMENTATION_SLOT)))));
+    if (!nullCheck) return impl;
     require(
       impl != address(0x0),
       string.concat("LibProxy: Null Implementation, Provided address: ", vm.getLabel(proxy), " is not EIP1967 Proxy")
     );
   }
+
+  function getProxyImplementation(address payable proxy) internal returns (address payable impl) {
+    impl = getProxyImplementation({ proxy: proxy, nullCheck: true });
+  }
 }