From e49313325491a3cc1cf4eaf3c247e6cde7f9e2e9 Mon Sep 17 00:00:00 2001
From: catalyst17 <37663786+catalyst17@users.noreply.github.com>
Date: Thu, 5 Jun 2025 13:58:25 +0200
Subject: [PATCH] fix: handle short storage values in beacon proxy pattern

---
 .changeset/cold-icons-shop.md                 |  5 ++
 .../bytecode/handleBeaconProxyPattern.test.ts | 52 +++++++++++++++++++
 2 files changed, 57 insertions(+)
 create mode 100644 .changeset/cold-icons-shop.md
 create mode 100644 packages/thirdweb/src/utils/bytecode/handleBeaconProxyPattern.test.ts

diff --git a/.changeset/cold-icons-shop.md b/.changeset/cold-icons-shop.md
new file mode 100644
index 00000000000..9616f24a69d
--- /dev/null
+++ b/.changeset/cold-icons-shop.md
@@ -0,0 +1,5 @@
+---
+"thirdweb": patch
+---
+
+fix implementation resolution for Beacon
diff --git a/packages/thirdweb/src/utils/bytecode/handleBeaconProxyPattern.test.ts b/packages/thirdweb/src/utils/bytecode/handleBeaconProxyPattern.test.ts
new file mode 100644
index 00000000000..00fd883a915
--- /dev/null
+++ b/packages/thirdweb/src/utils/bytecode/handleBeaconProxyPattern.test.ts
@@ -0,0 +1,52 @@
+import type { EIP1193RequestFn, EIP1474Methods } from "viem";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { ANVIL_CHAIN } from "../../../test/src/chains.js";
+import { TEST_CLIENT } from "../../../test/src/test-clients.js";
+import { getBytecode } from "../../contract/actions/get-bytecode.js";
+import { eth_getStorageAt } from "../../rpc/actions/eth_getStorageAt.js";
+import { getRpcClient } from "../../rpc/rpc.js";
+import { resolveImplementation } from "./resolveImplementation.js";
+
+// Mock dependencies
+vi.mock("../../contract/actions/get-bytecode.js");
+vi.mock("../../rpc/rpc.js");
+vi.mock("../../rpc/actions/eth_getStorageAt.js");
+
+describe.runIf(process.env.TW_SECRET_KEY)(
+  "Handle beacon proxy pattern",
+  async () => {
+    beforeEach(() => {
+      vi.resetAllMocks();
+    });
+    it("should handle beacon proxy pattern with short storage values", async () => {
+      // This test verifies that the implementation can handle short storage values
+      // which was the original issue we were trying to fix
+      const mockContract = {
+        address: "0x0000000000000000000000000000000000000001",
+        client: TEST_CLIENT,
+        chain: ANVIL_CHAIN,
+      } as const;
+
+      // Mock getBytecode to return a minimal proxy bytecode that will trigger the beacon check
+      const mockGetBytecode = vi
+        .fn()
+        .mockResolvedValueOnce("") // First call for proxy (empty bytecode triggers beacon check)
+        .mockResolvedValueOnce("0x1234"); // Second call for implementation
+
+      vi.mocked(getBytecode).mockImplementation(mockGetBytecode);
+
+      // Mock eth_getStorageAt to return a short value
+      const mockRpcRequest = { request: vi.fn() };
+      vi.mocked(getRpcClient).mockReturnValue(
+        mockRpcRequest as unknown as EIP1193RequestFn<EIP1474Methods>,
+      );
+      vi.mocked(eth_getStorageAt).mockResolvedValue("0x1234"); // Short storage value
+
+      const result = await resolveImplementation(mockContract);
+
+      // Should not throw and should return some result (even if it's the original address)
+      expect(result).toBeDefined();
+      expect(result.address).toBeDefined();
+    });
+  },
+);