From 1f6690ab7e353d0e9e569235f81e2611c5b7097d Mon Sep 17 00:00:00 2001
From: Eric Lau <ericglau@outlook.com>
Date: Wed, 14 Aug 2024 16:56:01 -0400
Subject: [PATCH] Fix revert when getting upgrade interface version with
 OpenZeppelin Contracts v4 (#65)

---
 CHANGELOG.md                                |  4 +++
 src/internal/Core.sol                       | 17 ++++++---
 test/contracts/UpgradeInterfaceVersions.sol | 22 ++++++++++++
 test/internal/Core.t.sol                    | 38 +++++++++++++++++++++
 4 files changed, 76 insertions(+), 5 deletions(-)
 create mode 100644 test/contracts/UpgradeInterfaceVersions.sol
 create mode 100644 test/internal/Core.t.sol

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 14d14ef..e60a653 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## 0.3.2 (2024-08-14)
+
+- Fix simulation failure due to revert when upgrading deployments using OpenZeppelin Contracts v4. ([#65](https://github.com/OpenZeppelin/openzeppelin-foundry-upgrades/pull/65))
+
 ## 0.3.1 (2024-05-21)
 
 - Fix upgrade interface version detection in `upgradeProxy` function. ([#53](https://github.com/OpenZeppelin/openzeppelin-foundry-upgrades/pull/53))
diff --git a/src/internal/Core.sol b/src/internal/Core.sol
index de292e8..6d2b966 100644
--- a/src/internal/Core.sol
+++ b/src/internal/Core.sol
@@ -73,7 +73,7 @@ library Core {
 
         bytes32 adminSlot = vm.load(proxy, ADMIN_SLOT);
         if (adminSlot == bytes32(0)) {
-            string memory upgradeInterfaceVersion = _getUpgradeInterfaceVersion(proxy);
+            string memory upgradeInterfaceVersion = getUpgradeInterfaceVersion(proxy);
             if (upgradeInterfaceVersion.toSlice().equals("5.0.0".toSlice()) || data.length > 0) {
                 IUpgradeableProxy(proxy).upgradeToAndCall(newImpl, data);
             } else {
@@ -81,7 +81,7 @@ library Core {
             }
         } else {
             address admin = address(uint160(uint256(adminSlot)));
-            string memory upgradeInterfaceVersion = _getUpgradeInterfaceVersion(admin);
+            string memory upgradeInterfaceVersion = getUpgradeInterfaceVersion(admin);
             if (upgradeInterfaceVersion.toSlice().equals("5.0.0".toSlice()) || data.length > 0) {
                 IProxyAdmin(admin).upgradeAndCall(proxy, newImpl, data);
             } else {
@@ -302,9 +302,16 @@ library Core {
 
     using strings for *;
 
-    function _getUpgradeInterfaceVersion(address addr) private returns (string memory) {
-        (bool success, bytes memory returndata) = addr.call(abi.encodeWithSignature("UPGRADE_INTERFACE_VERSION()"));
-        if (success) {
+    /**
+     * @dev Gets the upgrade interface version string from a proxy or admin contract using the `UPGRADE_INTERFACE_VERSION()` getter.
+     * If the contract does not have the getter or the return data does not look like a string, this function returns an empty string.
+     */
+    function getUpgradeInterfaceVersion(address addr) internal view returns (string memory) {
+        // Use staticcall to prevent forge from broadcasting it
+        (bool success, bytes memory returndata) = addr.staticcall(
+            abi.encodeWithSignature("UPGRADE_INTERFACE_VERSION()")
+        );
+        if (success && returndata.length > 32) {
             return abi.decode(returndata, (string));
         } else {
             return "";
diff --git a/test/contracts/UpgradeInterfaceVersions.sol b/test/contracts/UpgradeInterfaceVersions.sol
new file mode 100644
index 0000000..aedd50c
--- /dev/null
+++ b/test/contracts/UpgradeInterfaceVersions.sol
@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+// These contracts are for testing only.
+
+contract UpgradeInterfaceVersionString {
+    string public constant UPGRADE_INTERFACE_VERSION = "5.0.0";
+}
+
+contract UpgradeInterfaceVersionNoGetter {}
+
+contract UpgradeInterfaceVersionEmpty {
+    string public constant UPGRADE_INTERFACE_VERSION = "";
+}
+
+contract UpgradeInterfaceVersionInteger {
+    uint256 public constant UPGRADE_INTERFACE_VERSION = 5;
+}
+
+contract UpgradeInterfaceVersionVoid {
+    function UPGRADE_INTERFACE_VERSION() external pure {}
+}
diff --git a/test/internal/Core.t.sol b/test/internal/Core.t.sol
new file mode 100644
index 0000000..cf50e33
--- /dev/null
+++ b/test/internal/Core.t.sol
@@ -0,0 +1,38 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import {Test} from "forge-std/Test.sol";
+
+import {Core} from "openzeppelin-foundry-upgrades/internal/Core.sol";
+
+import {UpgradeInterfaceVersionString, UpgradeInterfaceVersionNoGetter, UpgradeInterfaceVersionEmpty, UpgradeInterfaceVersionInteger, UpgradeInterfaceVersionVoid} from "../contracts/UpgradeInterfaceVersions.sol";
+
+/**
+ * @dev Tests the Core internal library.
+ */
+contract CoreTest is Test {
+    function testGetUpgradeInterfaceVersion_string() public {
+        UpgradeInterfaceVersionString u = new UpgradeInterfaceVersionString();
+        assertEq(Core.getUpgradeInterfaceVersion(address(u)), "5.0.0");
+    }
+
+    function testGetUpgradeInterfaceVersion_noGetter() public {
+        UpgradeInterfaceVersionNoGetter u = new UpgradeInterfaceVersionNoGetter();
+        assertEq(Core.getUpgradeInterfaceVersion(address(u)), "");
+    }
+
+    function testGetUpgradeInterfaceVersion_empty() public {
+        UpgradeInterfaceVersionEmpty u = new UpgradeInterfaceVersionEmpty();
+        assertEq(Core.getUpgradeInterfaceVersion(address(u)), "");
+    }
+
+    function testGetUpgradeInterfaceVersion_integer() public {
+        UpgradeInterfaceVersionInteger u = new UpgradeInterfaceVersionInteger();
+        assertEq(Core.getUpgradeInterfaceVersion(address(u)), "");
+    }
+
+    function testGetUpgradeInterfaceVersion_void() public {
+        UpgradeInterfaceVersionVoid u = new UpgradeInterfaceVersionVoid();
+        assertEq(Core.getUpgradeInterfaceVersion(address(u)), "");
+    }
+}