From 2aa4bf9190f047c1d41224ddec5baed98600334d Mon Sep 17 00:00:00 2001
From: Vectorized <webby1111@hotmail.com>
Date: Mon, 17 Feb 2025 20:31:03 +0000
Subject: [PATCH] Tidy

---
 src/accounts/EIP7702Proxy.sol |  4 ++--
 src/accounts/LibEIP7702.sol   | 28 ++++++++++++++++------------
 test/EIP7702Proxy.t.sol       |  9 +++++----
 3 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/src/accounts/EIP7702Proxy.sol b/src/accounts/EIP7702Proxy.sol
index d11c11b7a..820882f90 100644
--- a/src/accounts/EIP7702Proxy.sol
+++ b/src/accounts/EIP7702Proxy.sol
@@ -77,7 +77,7 @@ contract EIP7702Proxy {
                 }
                 // Admin workflow.
                 if eq(caller(), admin) {
-                    let addr := shr(96, shl(96, calldataload(0x04)))
+                    let addr := shr(96, calldataload(0x10))
                     // `changeAdmin(address)`.
                     if eq(0x8f283970, fnSel) {
                         sstore(_ERC1967_ADMIN_SLOT, addr)
@@ -119,7 +119,7 @@ contract EIP7702Proxy {
                 // to signal that we should replace it with the actual implementation.
                 // This also gives the flexibility on whether to let the proxy auto-upgrade,
                 // or let the authority manually upgrade.
-                if iszero(xor(sload(_ERC1967_IMPLEMENTATION_SLOT), _ERC1967_IMPLEMENTATION_SLOT)) {
+                if eq(sload(_ERC1967_IMPLEMENTATION_SLOT), _ERC1967_IMPLEMENTATION_SLOT) {
                     // The `implementation` is still at `calldatasize()` in memory.
                     sstore(_ERC1967_IMPLEMENTATION_SLOT, mload(calldatasize()))
                 }
diff --git a/src/accounts/LibEIP7702.sol b/src/accounts/LibEIP7702.sol
index a9b214492..94664d7c3 100644
--- a/src/accounts/LibEIP7702.sol
+++ b/src/accounts/LibEIP7702.sol
@@ -8,8 +8,11 @@ library LibEIP7702 {
     /*                       CUSTOM ERRORS                        */
     /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
 
-    /// @dev Failed to perform the proxy query.
-    error ProxyQueryFailed();
+    /// @dev Failed to change the proxy admin.
+    error ChangeProxyAdminFailed();
+
+    /// @dev Failed to upgrade the proxy.
+    error UpgradeProxyFailed();
 
     /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
     /*                         CONSTANTS                          */
@@ -41,31 +44,31 @@ library LibEIP7702 {
     /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
 
     /// @dev Returns the implementation of the proxy.
+    /// If the `proxy` does not exist, returns `address(0)`.
+    /// Assumes that the proxy is a proper EIP7702Proxy, if it exists.
     function proxyImplementation(address proxy) internal view returns (address result) {
         /// @solidity memory-safe-assembly
         assembly {
-            if iszero(mul(returndatasize(), staticcall(gas(), proxy, 0x00, 0x00, 0x00, 0x20))) {
-                mstore(0x00, 0x26ec9b6a) // `ProxyQueryFailed()`.
-                revert(0x1c, 0x04)
-            }
-            result := mload(0x00)
+            mstore(0x00, 0x5c60da1b) // `implementation()`.
+            let t := staticcall(gas(), proxy, 0x1c, 0x04, 0x00, 0x20)
+            result := mul(mload(0x00), and(gt(returndatasize(), 0x1f), t))
         }
     }
 
     /// @dev Returns the admin of the proxy.
+    /// If the `proxy` does not exist, returns `address(0)`.
+    /// Assumes that the proxy is a proper EIP7702Proxy, if it exists.
     function proxyAdmin(address proxy) internal view returns (address result) {
         /// @solidity memory-safe-assembly
         assembly {
             mstore(0x00, 0xf851a440) // `admin()`.
-            if iszero(mul(returndatasize(), staticcall(gas(), proxy, 0x1c, 0x04, 0x00, 0x20))) {
-                mstore(0x00, 0x26ec9b6a) // `ProxyQueryFailed()`.
-                revert(0x1c, 0x04)
-            }
-            result := mload(0x00)
+            let t := staticcall(gas(), proxy, 0x1c, 0x04, 0x00, 0x20)
+            result := mul(mload(0x00), and(gt(returndatasize(), 0x1f), t))
         }
     }
 
     /// @dev Changes the admin on the proxy. The caller must be the admin.
+    /// Assumes that the proxy is a proper EIP7702Proxy, if it exists.
     function changeProxyAdmin(address proxy, address newAdmin) internal {
         /// @solidity memory-safe-assembly
         assembly {
@@ -79,6 +82,7 @@ library LibEIP7702 {
     }
 
     /// @dev Changes the implementation on the proxy. The caller must be the admin.
+    /// Assumes that the proxy is a proper EIP7702Proxy, if it exists.
     function upgradeProxy(address proxy, address newImplementation) internal {
         /// @solidity memory-safe-assembly
         assembly {
diff --git a/test/EIP7702Proxy.t.sol b/test/EIP7702Proxy.t.sol
index d379628d7..784595945 100644
--- a/test/EIP7702Proxy.t.sol
+++ b/test/EIP7702Proxy.t.sol
@@ -32,14 +32,15 @@ contract EIP7702ProxyTest is SoladyTest {
 
     function setValue(uint256 value_) public {
         value = value_;
+        LibEIP7702.requestProxyDelegationInitialization();
     }
 
     function revertWithError() public view {
         revert CustomError(value);
     }
 
-    function requestProxyDelegationInitialization() public {
-        LibEIP7702.requestProxyDelegationInitialization();
+    function version() external pure returns (uint256) {
+        return 1;
     }
 
     function upgradeToLatestProxyDelegation() public {
@@ -49,6 +50,8 @@ contract EIP7702ProxyTest is SoladyTest {
     function _checkBehavesLikeProxy(address instance) internal {
         assertTrue(instance != address(0));
 
+        assertEq(EIP7702ProxyTest(instance).version(), 1);
+
         uint256 v = _random();
         uint256 thisValue = this.value();
         if (thisValue == v) {
@@ -130,8 +133,6 @@ contract EIP7702ProxyTest is SoladyTest {
 
         _checkBehavesLikeProxy(authority);
 
-        EIP7702ProxyTest(authority).requestProxyDelegationInitialization();
-
         // Check that upgrading the proxy won't cause the authority's implementation to change.
         if (!f && _randomChance(2)) {
             vm.startPrank(admin);