Adding implementation() function to proxies

openfort-xyz · Jul 13, 2023 · 0177f98 · 0177f98
1 parent 8853d1e
commit 0177f98
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 12 deletions.
diff --git a/contracts/core/managed/OpenfortBeaconProxy.sol b/contracts/core/managed/OpenfortBeaconProxy.sol
@@ -12,4 +12,8 @@ import {BeaconProxy} from "@openzeppelin/contracts/proxy/beacon/BeaconProxy.sol"
  */
 contract OpenfortBeaconProxy is BeaconProxy {
     constructor(address beacon, bytes memory data) BeaconProxy(beacon, data) {}
+
+    function implementation() external view returns (address) {
+        return _implementation();
+    }
 }
diff --git a/contracts/core/upgradeable/OpenfortUpgradeableProxy.sol b/contracts/core/upgradeable/OpenfortUpgradeableProxy.sol
@@ -12,4 +12,8 @@ import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.s
  */
 contract OpenfortUpgradeableProxy is ERC1967Proxy {
     constructor(address _logic, bytes memory _data) ERC1967Proxy(_logic, _data) {}
+
+    function implementation() external view returns (address) {
+        return _getImplementation();
+    }
 }
diff --git a/script/deployAllChains.sh b/script/deployAllChains.sh
@@ -19,21 +19,21 @@ LOG_FILE=script/deployments/$(date +%Y-%m-%d_%H:%M)"-deploymentAllChains.log"
 # forge script StaticOpenfortDeploy --rpc-url $ARBITRUM_GOERLI_RPC -vvvv --verify --broadcast --slow --etherscan-api-key $ARBISCAN_API_KEY >> $LOG_FILE
 # sleep 3
 
-echo "------ UpgradeableOpenfortDeploy ------ (Goerli)"
-forge script UpgradeableOpenfortDeploy --rpc-url $GOERLI_RPC -vvvv --verify --broadcast --slow --etherscan-api-key $GOERLI_API_KEY >> $LOG_FILE
-sleep 3
-echo "------ UpgradeableOpenfortDeploy ------ (Mumbai)"
-forge script UpgradeableOpenfortDeploy --rpc-url $POLYGON_MUMBAI_RPC -vvvv --verify --broadcast --slow --legacy --etherscan-api-key $POLYGON_MUMBAI_KEY >> $LOG_FILE
-sleep 3
+# echo "------ UpgradeableOpenfortDeploy ------ (Goerli)"
+# forge script UpgradeableOpenfortDeploy --rpc-url $GOERLI_RPC -vvvv --verify --broadcast --slow -g 150 --etherscan-api-key $GOERLI_API_KEY >> $LOG_FILE
+# sleep 3
+# echo "------ UpgradeableOpenfortDeploy ------ (Mumbai)"
+# forge script UpgradeableOpenfortDeploy --rpc-url $POLYGON_MUMBAI_RPC -vvvv --verify --broadcast --slow --legacy --etherscan-api-key $POLYGON_MUMBAI_KEY >> $LOG_FILE
+# sleep 3
 echo "------ UpgradeableOpenfortDeploy ------ (Fuji)"
 forge script UpgradeableOpenfortDeploy --rpc-url $AVALANCHE_FUJI_RPC -vvvv --verify --broadcast --slow --etherscan-api-key $FUJI_API_KEY >> $LOG_FILE
 sleep 3
-echo "------ UpgradeableOpenfortDeploy ------ (BSC testnet)"
-forge script UpgradeableOpenfortDeploy --rpc-url $BSC_TESTNET_RPC -vvvv --verify --broadcast --slow --etherscan-api-key $BSCSCAN_TESTNET_API_KEY >> $LOG_FILE
-sleep 3
-echo "------ UpgradeableOpenfortDeploy ------ (Arbitrum Goerli testnet)"
-forge script UpgradeableOpenfortDeploy --rpc-url $ARBITRUM_GOERLI_RPC -vvvv --verify --broadcast --slow -g 200 --etherscan-api-key $ARBISCAN_API_KEY >> $LOG_FILE
-sleep 3
+# echo "------ UpgradeableOpenfortDeploy ------ (BSC testnet)"
+# forge script UpgradeableOpenfortDeploy --rpc-url $BSC_TESTNET_RPC -vvvv --verify --broadcast --slow --etherscan-api-key $BSCSCAN_TESTNET_API_KEY >> $LOG_FILE
+# sleep 3
+# echo "------ UpgradeableOpenfortDeploy ------ (Arbitrum Goerli testnet)"
+# forge script UpgradeableOpenfortDeploy --rpc-url $ARBITRUM_GOERLI_RPC -vvvv --verify --broadcast --slow -g 200 --etherscan-api-key $ARBISCAN_API_KEY >> $LOG_FILE
+# sleep 3
 
 # echo "------ ManagedOpenfortDeploy ------ (Goerli)"
 # forge script ManagedOpenfortDeploy --rpc-url $GOERLI_RPC -vvvv --verify --broadcast --slow --etherscan-api-key $GOERLI_API_KEY >> $LOG_FILE

diff --git a/test/foundry/core/managed/ManagedOpenfortAccountTest.t.sol b/test/foundry/core/managed/ManagedOpenfortAccountTest.t.sol
@@ -9,6 +9,7 @@ import {TestToken} from "account-abstraction/test/TestToken.sol";
 import {OpenfortBeacon} from "contracts/core/managed/OpenfortBeacon.sol";
 import {ManagedOpenfortAccount} from "contracts/core/managed/ManagedOpenfortAccount.sol";
 import {ManagedOpenfortFactory} from "contracts/core/managed/ManagedOpenfortFactory.sol";
+import {OpenfortBeaconProxy} from "contracts/core/managed/OpenfortBeaconProxy.sol";
 import {MockedV2ManagedOpenfortAccount} from "contracts/mock/MockedV2ManagedOpenfortAccount.sol";
 
 contract ManagedOpenfortAccountTest is Test {
@@ -1036,6 +1037,11 @@ contract ManagedOpenfortAccountTest is Test {
         ManagedOpenfortAccount managedAccount = ManagedOpenfortAccount(accountOld);
         assertEq(managedAccount.version(), 1);
 
+        OpenfortBeaconProxy p = OpenfortBeaconProxy(payable(account));
+        // Printing account address and the implementation address
+        console.log(account);
+        console.log(p.implementation());
+
         // Deploy the new implementation
         MockedV2ManagedOpenfortAccount newImplementation = new MockedV2ManagedOpenfortAccount();
         address newImplementationAddress = address(newImplementation);
@@ -1056,6 +1062,10 @@ contract ManagedOpenfortAccountTest is Test {
         address payable account3 = payable(managedOpenfortFactory.createAccountWithNonce(accountAdmin, "3"));
         ManagedOpenfortAccount managedAccount3 = ManagedOpenfortAccount(account3);
         managedAccount3.version();
+
+        // Printing account address and the implementation address. Impl address should have changed
+        console.log(account);
+        console.log(p.implementation());
     }
 
     /*

diff --git a/test/foundry/core/upgradeable/UpgradeableOpenfortAccountTest.t.sol b/test/foundry/core/upgradeable/UpgradeableOpenfortAccountTest.t.sol
@@ -8,6 +8,7 @@ import {TestCounter} from "account-abstraction/test/TestCounter.sol";
 import {TestToken} from "account-abstraction/test/TestToken.sol";
 import {UpgradeableOpenfortAccount} from "contracts/core/upgradeable/UpgradeableOpenfortAccount.sol";
 import {UpgradeableOpenfortFactory} from "contracts/core/upgradeable/UpgradeableOpenfortFactory.sol";
+import {OpenfortUpgradeableProxy} from "contracts/core/upgradeable/OpenfortUpgradeableProxy.sol";
 import {MockedV2UpgradeableOpenfortAccount} from "contracts/mock/MockedV2UpgradeableOpenfortAccount.sol";
 
 contract UpgradeableOpenfortAccountTest is Test {
@@ -988,6 +989,10 @@ contract UpgradeableOpenfortAccountTest is Test {
     function testUpgradeAccount() public {
         assertEq(UpgradeableOpenfortAccount(payable(account)).version(), 1);
         MockedV2UpgradeableOpenfortAccount newAccountImplementation = new MockedV2UpgradeableOpenfortAccount();
+        OpenfortUpgradeableProxy p = OpenfortUpgradeableProxy(payable(account));
+        // Printing account address and the implementation address
+        console.log(account);
+        console.log(p.implementation());
 
         vm.expectRevert("Ownable: caller is not the owner");
         UpgradeableOpenfortAccount(payable(account)).upgradeTo(address(newAccountImplementation));
@@ -997,6 +1002,10 @@ contract UpgradeableOpenfortAccountTest is Test {
 
         // Notice that, even though we bind the address to the old implementation, version() now returns 2
         assertEq(UpgradeableOpenfortAccount(payable(account)).version(), 2);
+
+        // Printing account address and the implementation address. Impl address should have changed
+        console.log(account);
+        console.log(p.implementation());
     }
 
     /*