fix(Permit2Lib): add gas cap to DOMAIN_SEPERATOR staticcall (#166)

* test: add failing test for weth9 * fix: add gas cap to DOMAIN_SEPERATOR staticcall Prevents the call from eating a ton of gas w/ fallback tokens and the like. * perf(Permit2Lib): Skip checking DS if token address matches WETH (#167)
Uniswap · Jan 3, 2023 · 2141eef · 2141eef
1 parent a7cd186
commit 2141eef
Show file tree

Hide file tree

Showing 9 changed files with 307 additions and 47 deletions.
diff --git a/.forge-snapshots/permit2 + transferFrom2 with WETH9's mainnet address.snap b/.forge-snapshots/permit2 + transferFrom2 with WETH9's mainnet address.snap
@@ -0,0 +1 @@
+60281
diff --git a/.forge-snapshots/permit2 + transferFrom2 with a non EIP-2612 native token with fallback.snap b/.forge-snapshots/permit2 + transferFrom2 with a non EIP-2612 native token with fallback.snap
@@ -0,0 +1 @@
+65465
diff --git a/.forge-snapshots/permit2 + transferFrom2 with a non EIP-2612 native token.snap b/.forge-snapshots/permit2 + transferFrom2 with a non EIP-2612 native token.snap
@@ -1 +1 @@
-60711
+60743
diff --git a/.forge-snapshots/permit2 + transferFrom2 with an EIP-2612 native token.snap b/.forge-snapshots/permit2 + transferFrom2 with an EIP-2612 native token.snap
@@ -1 +1 @@
-46265
+46260
diff --git a/.gas-snapshot b/.gas-snapshot
@@ -45,28 +45,32 @@ NonceBitmapTest:testLowNonces() (gas: 41004)
 NonceBitmapTest:testNonceWordBoundary() (gas: 42203)
 NonceBitmapTest:testUseTwoRandomNonces(uint256,uint256) (runs: 256, μ: 49205, ~: 51640)
 NonceBitmapTest:testUsingNonceTwiceFails(uint256) (runs: 256, μ: 21866, ~: 21889)
-Permit2LibTest:testOZSafePermit() (gas: 24443)
-Permit2LibTest:testOZSafePermitPlusOZSafeTransferFrom() (gas: 129153)
-Permit2LibTest:testOZSafeTransferFrom() (gas: 38875)
-Permit2LibTest:testPermit2() (gas: 22737)
-Permit2LibTest:testPermit2DSLessToken() (gas: 6818)
-Permit2LibTest:testPermit2DSMore32Token() (gas: 6905)
-Permit2LibTest:testPermit2DSMoreToken() (gas: 6808)
-Permit2LibTest:testPermit2Full() (gas: 42125)
-Permit2LibTest:testPermit2InvalidAmount() (gas: 20513)
-Permit2LibTest:testPermit2LargerDS() (gas: 51128)
-Permit2LibTest:testPermit2LargerDSRevert() (gas: 32546)
-Permit2LibTest:testPermit2NonPermitToken() (gas: 31934)
-Permit2LibTest:testPermit2PlusTransferFrom2() (gas: 126854)
-Permit2LibTest:testPermit2PlusTransferFrom2WithNonPermit() (gas: 147888)
-Permit2LibTest:testPermit2SmallerDS() (gas: 77593)
-Permit2LibTest:testPermit2SmallerDSNoRevert() (gas: 59230)
-Permit2LibTest:testStandardPermit() (gas: 22252)
-Permit2LibTest:testStandardTransferFrom() (gas: 38077)
-Permit2LibTest:testTransferFrom2() (gas: 38514)
-Permit2LibTest:testTransferFrom2Full() (gas: 53214)
-Permit2LibTest:testTransferFrom2InvalidAmount() (gas: 12666)
-Permit2LibTest:testTransferFrom2NonPermitToken() (gas: 53060)
+Permit2LibTest:testOZSafePermit() (gas: 24509)
+Permit2LibTest:testOZSafePermitPlusOZSafeTransferFrom() (gas: 129197)
+Permit2LibTest:testOZSafeTransferFrom() (gas: 38919)
+Permit2LibTest:testPermit2() (gas: 22776)
+Permit2LibTest:testPermit2DSLessToken() (gas: 6989)
+Permit2LibTest:testPermit2DSMore32Token() (gas: 7076)
+Permit2LibTest:testPermit2DSMoreToken() (gas: 6957)
+Permit2LibTest:testPermit2Full() (gas: 42196)
+Permit2LibTest:testPermit2InvalidAmount() (gas: 20619)
+Permit2LibTest:testPermit2LargerDS() (gas: 51226)
+Permit2LibTest:testPermit2LargerDSRevert() (gas: 32650)
+Permit2LibTest:testPermit2NonPermitFallback() (gas: 37048)
+Permit2LibTest:testPermit2NonPermitToken() (gas: 32011)
+Permit2LibTest:testPermit2PlusTransferFrom2() (gas: 126893)
+Permit2LibTest:testPermit2PlusTransferFrom2WithNonPermit() (gas: 147999)
+Permit2LibTest:testPermit2PlusTransferFrom2WithNonPermitFallback() (gas: 174659)
+Permit2LibTest:testPermit2PlusTransferFrom2WithWETH9Mainnet() (gas: 147693)
+Permit2LibTest:testPermit2SmallerDS() (gas: 77619)
+Permit2LibTest:testPermit2SmallerDSNoRevert() (gas: 59269)
+Permit2LibTest:testPermit2WETH9Mainnet() (gas: 28712)
+Permit2LibTest:testStandardPermit() (gas: 22340)
+Permit2LibTest:testStandardTransferFrom() (gas: 38121)
+Permit2LibTest:testTransferFrom2() (gas: 38580)
+Permit2LibTest:testTransferFrom2Full() (gas: 53258)
+Permit2LibTest:testTransferFrom2InvalidAmount() (gas: 12710)
+Permit2LibTest:testTransferFrom2NonPermitToken() (gas: 53104)
 SignatureTransferTest:testCorrectWitnessTypehashes() (gas: 3075)
 SignatureTransferTest:testGasMultiplePermitBatchTransferFrom() (gas: 270919)
 SignatureTransferTest:testGasSinglePermitBatchTransferFrom() (gas: 186316)
@@ -102,4 +106,4 @@ TypehashGeneration:testPermitTransferFrom() (gas: 36520)
 TypehashGeneration:testPermitTransferFromWithWitness() (gas: 43369)
 TypehashGeneration:testPermitTransferFromWithWitnessIncorrectPermitData() (gas: 43430)
 TypehashGeneration:testPermitTransferFromWithWitnessIncorrectTypehashStub() (gas: 43833)
-MockPermit2Lib:testPermit2Code(address):(bool) (runs: 256, μ: 35465847065545473, ~: 2911)
+MockPermit2Lib:testPermit2Code(address):(bool) (runs: 256, μ: 3025, ~: 3016)
diff --git a/src/libraries/Permit2Lib.sol b/src/libraries/Permit2Lib.sol
@@ -20,6 +20,9 @@ library Permit2Lib {
     /// @dev The unique EIP-712 domain domain separator for the DAI token contract.
     bytes32 internal constant DAI_DOMAIN_SEPARATOR = 0xdbb8cf42e1ecb028be3f3dbc922e1d878b963f411dc388ced501601c60f7c6f7;
 
+    /// @dev The address for the WETH9 contract on Ethereum mainnet, encoded as a bytes32.
+    bytes32 internal constant WETH9_ADDRESS = 0x000000000000000000000000c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2;
+
     /// @dev The address of the Permit2 contract the library will use.
     Permit2 internal constant PERMIT2 = Permit2(address(0x000000000022D473030F116dDEE9F6B43aC78BA3));
 
@@ -79,18 +82,25 @@ library Permit2Lib {
 
         bool success; // Call the token contract as normal, capturing whether it succeeded.
         bytes32 domainSeparator; // If the call succeeded, we'll capture the return value here.
-        assembly {
-            success :=
-                and(
-                    // Should resolve false if its not 32 bytes or its first word is 0.
-                    and(iszero(iszero(mload(0))), eq(returndatasize(), 32)),
-                    // We use 0 and 32 to copy up to 32 bytes of return data into the scratch space.
-                    // Counterintuitively, this call must be positioned second to the and() call in the
-                    // surrounding and() call or else returndatasize() will be zero during the computation.
-                    staticcall(gas(), token, add(inputData, 32), mload(inputData), 0, 32)
-                )
 
-            domainSeparator := mload(0) // Copy the return value into the domainSeparator variable.
+        assembly {
+            // If the token is WETH9, we know it doesn't have a DOMAIN_SEPARATOR, and we'll skip this step.
+            // We make sure to mask the token address as its higher order bits aren't guaranteed to be clean.
+            if iszero(eq(and(token, 0xffffffffffffffffffffffffffffffffffffffff), WETH9_ADDRESS)) {
+                success :=
+                    and(
+                        // Should resolve false if its not 32 bytes or its first word is 0.
+                        and(iszero(iszero(mload(0))), eq(returndatasize(), 32)),
+                        // We use 0 and 32 to copy up to 32 bytes of return data into the scratch space.
+                        // Counterintuitively, this call must be positioned second to the and() call in the
+                        // surrounding and() call or else returndatasize() will be zero during the computation.
+                        // We send a maximum of 5000 gas to prevent tokens with fallbacks from using a ton of gas.
+                        // which should be plenty to allow tokens to fetch their DOMAIN_SEPARATOR from storage, etc.
+                        staticcall(5000, token, add(inputData, 32), mload(inputData), 0, 32)
+                    )
+
+                domainSeparator := mload(0) // Copy the return value into the domainSeparator variable.
+            }
         }
 
         // If the call to DOMAIN_SEPARATOR succeeded, try using permit on the token.

diff --git a/test/Permit2Lib.t.sol b/test/Permit2Lib.t.sol
@@ -5,7 +5,7 @@ import "forge-std/Test.sol";
 
 import {SafeERC20, IERC20, IERC20Permit} from "openzeppelin-contracts/contracts/token/ERC20/utils/SafeERC20.sol";
 import {DSTestPlus} from "solmate/src/test/utils/DSTestPlus.sol";
-import {MockERC20} from "solmate/src/test/utils/mocks/MockERC20.sol";
+import {MockERC20, ERC20} from "solmate/src/test/utils/mocks/MockERC20.sol";
 import {Permit2} from "../src/Permit2.sol";
 import {Permit2Lib} from "../src/libraries/Permit2Lib.sol";
 import {MockNonPermitERC20} from "./mocks/MockNonPermitERC20.sol";
@@ -17,6 +17,7 @@ import {SafeCast160} from "../src/libraries/SafeCast160.sol";
 import {MockPermitWithSmallDS, MockPermitWithLargerDS} from "./mocks/MockPermitWithDS.sol";
 import {MockNonPermitNonERC20WithDS} from "./mocks/MockNonPermitNonERC20WithDS.sol";
 import {SignatureVerification} from "../src/libraries/SignatureVerification.sol";
+import {MockFallbackERC20} from "./mocks/MockFallbackERC20.sol";
 
 contract Permit2LibTest is Test, PermitSignature, GasSnapshot {
     bytes32 constant PERMIT_TYPEHASH =
@@ -32,12 +33,15 @@ contract Permit2LibTest is Test, PermitSignature, GasSnapshot {
 
     Permit2 immutable permit2 = Permit2(0x000000000022D473030F116dDEE9F6B43aC78BA3);
 
+    ERC20 immutable weth9Mainnet = ERC20(payable(address(0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2)));
+
     // Use to test errors in Permit2Lib calls.
     MockPermit2Lib immutable permit2Lib = new MockPermit2Lib();
 
     MockERC20 immutable token = new MockERC20("Mock Token", "MOCK", 18);
 
     MockNonPermitERC20 immutable nonPermitToken = new MockNonPermitERC20("Mock NonPermit Token", "MOCK", 18);
+    MockFallbackERC20 immutable fallbackToken = new MockFallbackERC20("Mock Fallback Token", "MOCK", 18);
     MockPermitWithSmallDS immutable lessDSToken =
         new MockPermitWithSmallDS("Mock Permit Token Small Domain Sep", "MOCK", 18);
     MockPermitWithLargerDS immutable largerDSToken =
@@ -49,6 +53,7 @@ contract Permit2LibTest is Test, PermitSignature, GasSnapshot {
         PK_OWNER = vm.addr(PK);
         Permit2 tempPermit2 = new Permit2();
         vm.etch(address(permit2), address(tempPermit2).code);
+        vm.etch(address(weth9Mainnet), address(nonPermitToken).code);
 
         TOKEN_DOMAIN_SEPARATOR = token.DOMAIN_SEPARATOR();
         PERMIT2_DOMAIN_SEPARATOR = permit2.DOMAIN_SEPARATOR();
@@ -78,11 +83,29 @@ contract Permit2LibTest is Test, PermitSignature, GasSnapshot {
         nonPermitToken.mint(PK_OWNER, type(uint128).max);
         vm.prank(PK_OWNER);
         nonPermitToken.approve(address(permit2), type(uint128).max);
+
+        MockNonPermitERC20(address(weth9Mainnet)).mint(address(this), type(uint128).max);
+        weth9Mainnet.approve(address(this), type(uint128).max);
+        weth9Mainnet.approve(address(permit2), type(uint128).max);
+
+        MockNonPermitERC20(address(weth9Mainnet)).mint(PK_OWNER, type(uint128).max);
+        vm.prank(PK_OWNER);
+        weth9Mainnet.approve(address(permit2), type(uint128).max);
+
+        fallbackToken.mint(address(this), type(uint128).max);
+        fallbackToken.approve(address(this), type(uint128).max);
+        fallbackToken.approve(address(permit2), type(uint128).max);
+
+        fallbackToken.mint(PK_OWNER, type(uint128).max);
+        vm.prank(PK_OWNER);
+        fallbackToken.approve(address(permit2), type(uint128).max);
     }
 
     function setUp() public {
         testPermit2Full();
+        testPermit2NonPermitFallback();
         testPermit2NonPermitToken();
+        testPermit2WETH9Mainnet();
         testStandardPermit();
     }
 
@@ -224,6 +247,48 @@ contract Permit2LibTest is Test, PermitSignature, GasSnapshot {
         Permit2Lib.permit2(nonPermitToken, PK_OWNER, address(0xCAFE), 1e18, block.timestamp, v, r, s);
     }
 
+    function testPermit2WETH9Mainnet() public {
+        (,, uint48 nonce) = permit2.allowance(PK_OWNER, address(weth9Mainnet), address(0xCAFE));
+
+        IAllowanceTransfer.PermitSingle memory permit = IAllowanceTransfer.PermitSingle({
+            details: IAllowanceTransfer.PermitDetails({
+                token: address(weth9Mainnet),
+                amount: 1e18,
+                expiration: type(uint48).max,
+                nonce: nonce
+            }),
+            spender: address(0xCAFE),
+            sigDeadline: block.timestamp
+        });
+
+        (uint8 v, bytes32 r, bytes32 s) = getPermitSignatureRaw(permit, PK, PERMIT2_DOMAIN_SEPARATOR);
+
+        Permit2Lib.permit2(weth9Mainnet, PK_OWNER, address(0xCAFE), 1e18, block.timestamp, v, r, s);
+    }
+
+    function testPermit2NonPermitFallback() public {
+        (,, uint48 nonce) = permit2.allowance(PK_OWNER, address(fallbackToken), address(0xCAFE));
+
+        IAllowanceTransfer.PermitSingle memory permit = IAllowanceTransfer.PermitSingle({
+            details: IAllowanceTransfer.PermitDetails({
+                token: address(fallbackToken),
+                amount: 1e18,
+                expiration: type(uint48).max,
+                nonce: nonce
+            }),
+            spender: address(0xCAFE),
+            sigDeadline: block.timestamp
+        });
+
+        (uint8 v, bytes32 r, bytes32 s) = getPermitSignatureRaw(permit, PK, PERMIT2_DOMAIN_SEPARATOR);
+
+        uint256 gas1 = gasleft();
+
+        Permit2Lib.permit2(ERC20(address(fallbackToken)), PK_OWNER, address(0xCAFE), 1e18, block.timestamp, v, r, s);
+
+        assertLt(gas1 - gasleft(), 50000); // If unbounded the staticcall will consume a wild amount of gas.
+    }
+
     function testPermit2SmallerDS() public {
         (,, uint48 nonce) = permit2.allowance(PK_OWNER, address(lessDSToken), address(0xCAFE));
 
@@ -414,6 +479,58 @@ contract Permit2LibTest is Test, PermitSignature, GasSnapshot {
         snapEnd();
     }
 
+    function testPermit2PlusTransferFrom2WithNonPermitFallback() public {
+        (,, uint48 nonce) = permit2.allowance(PK_OWNER, address(fallbackToken), address(0xCAFE));
+
+        IAllowanceTransfer.PermitSingle memory permit = IAllowanceTransfer.PermitSingle({
+            details: IAllowanceTransfer.PermitDetails({
+                token: address(fallbackToken),
+                amount: 1e18,
+                expiration: type(uint48).max,
+                nonce: nonce
+            }),
+            spender: address(0xCAFE),
+            sigDeadline: block.timestamp
+        });
+
+        (uint8 v, bytes32 r, bytes32 s) = getPermitSignatureRaw(permit, PK, PERMIT2_DOMAIN_SEPARATOR);
+
+        vm.startPrank(address(0xCAFE));
+
+        snapStart("permit2 + transferFrom2 with a non EIP-2612 native token with fallback");
+
+        Permit2Lib.permit2(ERC20(address(fallbackToken)), PK_OWNER, address(0xCAFE), 1e18, block.timestamp, v, r, s);
+        Permit2Lib.transferFrom2(ERC20(address(fallbackToken)), PK_OWNER, address(0xB00B), 1e18);
+
+        snapEnd();
+    }
+
+    function testPermit2PlusTransferFrom2WithWETH9Mainnet() public {
+        (,, uint48 nonce) = permit2.allowance(PK_OWNER, address(weth9Mainnet), address(0xCAFE));
+
+        IAllowanceTransfer.PermitSingle memory permit = IAllowanceTransfer.PermitSingle({
+            details: IAllowanceTransfer.PermitDetails({
+                token: address(weth9Mainnet),
+                amount: 1e18,
+                expiration: type(uint48).max,
+                nonce: nonce
+            }),
+            spender: address(0xCAFE),
+            sigDeadline: block.timestamp
+        });
+
+        (uint8 v, bytes32 r, bytes32 s) = getPermitSignatureRaw(permit, PK, PERMIT2_DOMAIN_SEPARATOR);
+
+        vm.startPrank(address(0xCAFE));
+
+        snapStart("permit2 + transferFrom2 with WETH9's mainnet address");
+
+        Permit2Lib.permit2(weth9Mainnet, PK_OWNER, address(0xCAFE), 1e18, block.timestamp, v, r, s);
+        Permit2Lib.transferFrom2(weth9Mainnet, PK_OWNER, address(0xB00B), 1e18);
+
+        snapEnd();
+    }
+
     // mock tests
     function testPermit2DSLessToken() public {
         bool success = permit2Lib.testPermit2Code(MockERC20(address(lessDSToken)));