From fa686b5246f9a27fdf627e59ca0f751256ac4e35 Mon Sep 17 00:00:00 2001
From: nonergodic <nonergodic@gmail.com>
Date: Thu, 19 Dec 2024 11:22:04 -0800
Subject: [PATCH] bugfix: prevent proxy upgrade to addresses with no code

---
 src/proxy/ProxyBase.sol |  4 ++++
 test/Proxy.t.sol        | 17 +++++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/src/proxy/ProxyBase.sol b/src/proxy/ProxyBase.sol
index 43f4e92..a7e234a 100644
--- a/src/proxy/ProxyBase.sol
+++ b/src/proxy/ProxyBase.sol
@@ -8,6 +8,7 @@ error InvalidSender();
 error IdempotentUpgrade();
 error InvalidMsgValue();
 error InvalidData();
+error InvalidImplementation();
 error UpgradeFailed(bytes revertData);
 
 event Upgraded(address indexed implementation);
@@ -40,6 +41,9 @@ abstract contract ProxyBase {
     if (newImplementation == implementationState().implementation)
       revert IdempotentUpgrade();
 
+    if (newImplementation.code.length == 0)
+      revert InvalidImplementation();
+
     implementationState().implementation = newImplementation;
 
     (bool success, bytes memory revertData) =
diff --git a/test/Proxy.t.sol b/test/Proxy.t.sol
index b49a804..805ce77 100644
--- a/test/Proxy.t.sol
+++ b/test/Proxy.t.sol
@@ -5,7 +5,12 @@ pragma solidity ^0.8.24;
 import "forge-std/Test.sol";
 
 import { adminState } from "wormhole-sdk/proxy/Eip1967Admin.sol";
-import { ProxyBase, UpgradeFailed, InvalidData } from "wormhole-sdk/proxy/ProxyBase.sol";
+import {
+  ProxyBase,
+  UpgradeFailed,
+  InvalidData,
+  InvalidImplementation
+} from "wormhole-sdk/proxy/ProxyBase.sol";
 import { Proxy, ProxyConstructionFailed } from "wormhole-sdk/proxy/Proxy.sol";
 
 error NotAuthorized();
@@ -104,4 +109,12 @@ contract TestProxy is Test {
     assertEq(contrct.immutableNum(), 1);
     assertEq(contrct.message(), "v2");
   }
-}
\ No newline at end of file
+
+  function testProxyInvalidUpgradeFails() public {
+    address logic1 = address(new LogicContractV1(1));
+    LogicContractV1 contrct = LogicContractV1(address(new Proxy(logic1, abi.encode("v1"))));
+
+    vm.expectRevert(InvalidImplementation.selector);
+    contrct.customUpgradeFun(makeAddr("wrongAddress"), abi.encode("oops"));
+  }
+}