From 1a99984aa13656ef0c70f826085842354c85bb4e Mon Sep 17 00:00:00 2001
From: TuDo1403 <tudo.dev@gmail.com>
Date: Thu, 28 Dec 2023 18:19:06 +0700
Subject: [PATCH] feat: allow logging multisig proposal on mainnet when upgrade
 proxy

---
 script/BaseMigration.s.sol | 83 +++++++++++++++++++++++++++++++++-----
 1 file changed, 72 insertions(+), 11 deletions(-)

diff --git a/script/BaseMigration.s.sol b/script/BaseMigration.s.sol
index 416d2f1..0df9558 100644
--- a/script/BaseMigration.s.sol
+++ b/script/BaseMigration.s.sol
@@ -16,7 +16,7 @@ import { DefaultContract } from "./utils/DefaultContract.sol";
 import { TContract } from "./types/Types.sol";
 
 abstract contract BaseMigration is ScriptExtended {
-  using StdStyle for string;
+  using StdStyle for *;
   using LibString for bytes32;
   using LibProxy for address payable;
 
@@ -232,6 +232,7 @@ abstract contract BaseMigration is ScriptExtended {
   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);
@@ -240,24 +241,34 @@ abstract contract BaseMigration is ScriptExtended {
     } else {
       try wProxyAdmin.owner() returns (address owner) {
         if (args.length == 0) {
-          // try `upgrade` function
+          // try `upgrade(address,address)` function
           vm.prank(owner);
           (bool success,) = proxyAdmin.call(abi.encodeCall(ProxyAdmin.upgrade, (iProxy, logic)));
           if (success) {
-            vm.broadcast(owner);
-            wProxyAdmin.upgrade(iProxy, logic);
+            if (owner.code.length != 0) {
+              _cheatUpgrade(owner, wProxyAdmin, iProxy, logic);
+            } else {
+              vm.broadcast(owner);
+              wProxyAdmin.upgrade(iProxy, logic);
+            }
           } else {
             console.log(
-              StdStyle.yellow(
-                "`ProxyAdmin:upgrade` failed!. Retrying with `ProxyAdmin:upgradeAndCall` with emty args..."
-              )
+              "`ProxyAdmin:upgrade` failed!. Retrying with `ProxyAdmin:upgradeAndCall` with emty args...".yellow()
             );
+            if (owner.code.length != 0) {
+              _cheatUpgradeAndCall(owner, wProxyAdmin, iProxy, logic, args);
+            } else {
+              vm.broadcast(owner);
+              wProxyAdmin.upgradeAndCall(iProxy, logic, args);
+            }
+          }
+        } else {
+          if (owner.code.length != 0) {
+            _cheatUpgradeAndCall(owner, wProxyAdmin, iProxy, logic, args);
+          } else {
             vm.broadcast(owner);
             wProxyAdmin.upgradeAndCall(iProxy, logic, args);
           }
-        } else {
-          vm.broadcast(owner);
-          wProxyAdmin.upgradeAndCall(iProxy, logic, args);
         }
       } catch {
         revert("BaseMigration: Unknown ProxyAdmin contract!");
@@ -265,6 +276,56 @@ abstract contract BaseMigration is ScriptExtended {
     }
   }
 
+  function _cheatUpgrade(address owner, ProxyAdmin wProxyAdmin, ITransparentUpgradeableProxy iProxy, address logic)
+    internal
+    virtual
+  {
+    bytes memory callData = abi.encodeCall(ProxyAdmin.upgrade, (iProxy, logic));
+    console.log(
+      "------------------------------------------------------------------------------- Multi-Sig Proposal -------------------------------------------------------------------------------"
+    );
+    console.log("To:".cyan(), vm.getLabel(address(wProxyAdmin)));
+    console.log(
+      "Method:\n".cyan(),
+      string.concat(" - upgrade(address,address)\n  - ", vm.getLabel(address(iProxy)), "\n  - ", vm.getLabel(logic))
+    );
+    console.log("Raw Calldata Data:\n".cyan(), string.concat(" - ", vm.toString(callData)));
+    console.log(
+      "----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------"
+    );
+
+    // cheat prank to update `implementation slot` for next call
+    vm.prank(owner);
+    wProxyAdmin.upgrade(iProxy, logic);
+  }
+
+  function _cheatUpgradeAndCall(
+    address owner,
+    ProxyAdmin wProxyAdmin,
+    ITransparentUpgradeableProxy iProxy,
+    address logic,
+    bytes memory args
+  ) internal virtual {
+    bytes memory callData = abi.encodeCall(ProxyAdmin.upgradeAndCall, (iProxy, logic, args));
+    console.log(
+      "------------------------------------------------------------------------------- Multi-Sig Proposal -------------------------------------------------------------------------------"
+    );
+    console.log("To:".cyan(), vm.getLabel(address(wProxyAdmin)));
+    console.log(
+      "Method:\n".cyan(),
+      " - upgradeAndCall(address,address,bytes)\n",
+      string.concat(" - ", vm.getLabel(address(iProxy)), "\n  - ", vm.getLabel(logic), "\n  - ", vm.toString(args))
+    );
+    console.log("Raw Call Data:\n".cyan(), " - ", vm.toString(callData));
+    console.log(
+      "----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------"
+    );
+
+    // cheat prank to update `implementation slot` for next call
+    vm.prank(owner);
+    wProxyAdmin.upgradeAndCall(iProxy, logic, args);
+  }
+
   function _setDependencyDeployScript(TContract contractType, IScriptExtended deployScript) internal virtual {
     _deployScript[contractType] = IMigrationScript(address(deployScript));
   }
@@ -272,4 +333,4 @@ abstract contract BaseMigration is ScriptExtended {
   function _setDependencyDeployScript(TContract contractType, address deployScript) internal virtual {
     _deployScript[contractType] = IMigrationScript(deployScript);
   }
-}
\ No newline at end of file
+}