From 6e4a375e225c09820b7705047101c624e53d1dbe Mon Sep 17 00:00:00 2001
From: Kyrylo Riabov <kyryl.ryabov@gmail.com>
Date: Wed, 15 May 2024 16:54:09 +0300
Subject: [PATCH] Add `upgradeToWithSigAndCall` function to the
 UUPSSignableUpgradeable (#7)

* Added `upgradeToWithSigAndCall` function to the UUPSSignableUpgradeable

* Changed `forceCall` to false
---
 contracts/utils/UUPSSignableUpgradeable.sol | 17 +++++++++++++++++
 test/bridge/Upgradeable.test.ts             |  8 ++++++++
 2 files changed, 25 insertions(+)

diff --git a/contracts/utils/UUPSSignableUpgradeable.sol b/contracts/utils/UUPSSignableUpgradeable.sol
index 71bf894..dd2e4ae 100644
--- a/contracts/utils/UUPSSignableUpgradeable.sol
+++ b/contracts/utils/UUPSSignableUpgradeable.sol
@@ -30,4 +30,21 @@ abstract contract UUPSSignableUpgradeable is UUPSUpgradeable {
         _authorizeUpgrade(newImplementation_, signatures_);
         _upgradeToAndCallUUPS(newImplementation_, new bytes(0), false);
     }
+
+    /**
+     * @notice Upgrade the implementation of the proxy to `newImplementation`, and subsequently execute the function call
+     * encoded in `data`.
+     *
+     * Calls {_authorizeUpgrade} with the provided `newImplementation` and `signatures`.
+     *
+     * Emits an {Upgraded} event.
+     */
+    function upgradeToWithSigAndCall(
+        address newImplementation_,
+        bytes[] calldata signatures_,
+        bytes calldata data_
+    ) external virtual onlyProxy {
+        _authorizeUpgrade(newImplementation_, signatures_);
+        _upgradeToAndCallUUPS(newImplementation_, data_, false);
+    }
 }
diff --git a/test/bridge/Upgradeable.test.ts b/test/bridge/Upgradeable.test.ts
index d20228a..e066ab5 100644
--- a/test/bridge/Upgradeable.test.ts
+++ b/test/bridge/Upgradeable.test.ts
@@ -50,6 +50,10 @@ describe("Upgradeable", () => {
     await expect(newBridge.upgradeToWithSig(await newBridge.getAddress(), [])).to.be.rejectedWith(
       "Function must be called through delegatecall",
     );
+
+    await expect(newBridge.upgradeToWithSigAndCall(await newBridge.getAddress(), [], "0x")).to.be.rejectedWith(
+      "Function must be called through delegatecall",
+    );
   });
 
   it("should upgrade implementation", async () => {
@@ -85,6 +89,10 @@ describe("Upgradeable", () => {
     await expect(proxyBridge.connect(SECOND).upgradeToWithSig(await newBridge.getAddress(), [])).to.be.rejectedWith(
       "Ownable: caller is not the owner",
     );
+
+    await expect(
+      proxyBridge.connect(SECOND).upgradeToWithSigAndCall(await newBridge.getAddress(), [], "0x"),
+    ).to.be.rejectedWith("Ownable: caller is not the owner");
   });
 
   it("should receive ether through proxy", async () => {