Fix Shift Amount in Proxy (#868)

This PR fixes the shift amounts in the `SafeProxy` implementation. It
also slightly optimizes how we do the masking by 2 instructions (removes
the need to `PUSH 0` and `SHR`).

I added some additional tests to make sure the masking works as expected
and a comment explaining why we only mask for handling the
`masterCopy()` call.

---------

Co-authored-by: Mikhail <[email protected]>

contracts/proxies/SafeProxy.sol

@@ -40,7 +40,13 @@ contract SafeProxy {
             let _singleton := sload(0)
             // 0xa619486e == keccak("masterCopy()"). The value is right padded to 32-bytes with 0s
             if eq(calldataload(0), 0xa619486e00000000000000000000000000000000000000000000000000000000) {
-                mstore(0, shr(12, shl(12, _singleton)))
+                // We mask the singleton address when handling the `masterCopy()` call to ensure that it is correctly
+                // ABI-encoded. We do this by shifting the address left by 96 bits (or 12 bytes) and then storing it in
+                // memory with a 12 byte offset from where the return data starts. Note that we **intentionally** only
+                // do this for the `masterCopy()` call, since the EVM `DELEGATECALL` opcode ignores the most-significant
+                // 12 bytes from the address, so we do not need to make sure the top bytes are cleared when proxying
+                // calls to the `singleton`. This saves us a tiny amount of gas per proxied call.
+                mstore(0x0c, shl(96, _singleton))
                 return(0, 0x20)
             }
             calldatacopy(0, 0, calldatasize())

test/factory/Proxy.spec.ts

@@ -1,6 +1,7 @@
 import { expect } from "chai";
 import hre from "hardhat";
 import { AddressZero } from "@ethersproject/constants";
+import { deployContractFromSource } from "../utils/setup";
 
 describe("Proxy", () => {
     describe("constructor", () => {
@@ -9,4 +10,54 @@ describe("Proxy", () => {
             await expect(Proxy.deploy(AddressZero)).to.be.revertedWith("Invalid singleton address provided");
         });
     });
+
+    describe("masterCopy", () => {
+        const SINGLETON_SOURCE = `
+        contract Test {
+            uint256 _singletonSlot;
+            function masterCopy() public pure returns (address) {
+                return address(0);
+            }
+            function overwriteSingletonSlot(uint256 value) public {
+                _singletonSlot = value;
+            }
+            function theAnswerToLifeTheUniverseAndEverything() public pure returns (uint256) {
+                return 42;
+            }
+        }`;
+
+        const setupTests = hre.deployments.createFixture(async () => {
+            const [deployer] = await hre.ethers.getSigners();
+            const singleton = await deployContractFromSource(deployer, SINGLETON_SOURCE);
+            const Proxy = await hre.ethers.getContractFactory("SafeProxy");
+            const proxyDeployment = await Proxy.deploy(singleton.target);
+            const proxy = singleton.attach(proxyDeployment) as typeof singleton;
+            return {
+                singleton,
+                proxy,
+            };
+        });
+
+        it("should return the master copy address regardless of implementation", async () => {
+            const { singleton, proxy } = await setupTests();
+            expect(await singleton.masterCopy()).to.equal(hre.ethers.ZeroAddress);
+            expect(await proxy.masterCopy()).to.equal(await singleton.getAddress());
+        });
+
+        it("should correctly mask the address value", async () => {
+            const { proxy } = await setupTests();
+            await proxy.overwriteSingletonSlot(hre.ethers.MaxUint256);
+            expect(await proxy.masterCopy()).to.equal(hre.ethers.getAddress(`0x${"ff".repeat(20)}`));
+        });
+
+        it("should ignore most significant bits when calling singleton", async () => {
+            const { singleton, proxy } = await setupTests();
+            const singletonAddressAsUint = BigInt(await singleton.getAddress());
+            const mask = 0xffffffffffffffffffffffffn << 160n;
+
+            expect(await proxy.theAnswerToLifeTheUniverseAndEverything()).to.equal(42);
+            await proxy.overwriteSingletonSlot(singletonAddressAsUint | mask);
+            expect(await proxy.theAnswerToLifeTheUniverseAndEverything()).to.equal(42);
+        });
+    });
 });