diff --git a/script/BaseMigration.s.sol b/script/BaseMigration.s.sol index 60b1772..5466b23 100644 --- 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 index d690afc..8aafef6 100644 --- 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 index 86879aa..1532dfa 100644 --- 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 index b4dee57..77198c1 100644 --- 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 }); + } }