From e671bd2726f1b6cc6c951b9e4817c4318a2fb8f2 Mon Sep 17 00:00:00 2001
From: Nalin Bhardwaj <github@nibnalin.me>
Date: Sat, 18 May 2024 22:00:43 -0700
Subject: [PATCH] add fallback test

---
 src/P256.sol    |  9 +++++++-
 test/P256.t.sol | 59 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/src/P256.sol b/src/P256.sol
index 17d7230..c7792ea 100644
--- a/src/P256.sol
+++ b/src/P256.sol
@@ -3,6 +3,8 @@ pragma solidity 0.8.21;
 
 /**
  * Helper library for external contracts to verify P256 signatures.
+ * Tries to use RIP-7212 precompile if available on the chain, and if not falls
+ * back to more expensive Solidity implementation.
  **/
 library P256 {
     address constant PRECOMPILE = address(0x100);
@@ -18,7 +20,12 @@ library P256 {
         bytes memory args = abi.encode(message_hash, r, s, x, y);
 
         (bool success, bytes memory ret) = PRECOMPILE.staticcall(args);
-        if (success && ret.length > 0) return abi.decode(ret, (uint256)) == 1;
+        if (success && ret.length > 0) {
+            // RIP-7212 precompile returns 1 if signature is valid
+            // and nothing if signature is invalid, so those fall back to
+            // more expensive Solidity implementation.
+            return abi.decode(ret, (uint256)) == 1;
+        }
 
         (bool fallbackSuccess, bytes memory fallbackRet) = VERIFIER.staticcall(
             args
diff --git a/test/P256.t.sol b/test/P256.t.sol
index f64820b..5308d1a 100644
--- a/test/P256.t.sol
+++ b/test/P256.t.sol
@@ -6,6 +6,20 @@ import {stdJson} from "forge-std/StdJson.sol";
 import {P256} from "../src/P256.sol";
 import {P256Verifier} from "../src/P256Verifier.sol";
 
+contract FakePrecompile {
+    fallback(bytes calldata input) external returns (bytes memory) {
+        (bool success, bytes memory ret) = P256.VERIFIER.staticcall(input);
+        assert(success);
+
+        uint256 retUint = abi.decode(ret, (uint256));
+        if (retUint == 1) {
+            return abi.encode(1);
+        } else {
+            return abi.encode();
+        }
+    }
+}
+
 contract P256Test is Test {
     uint256[2] public pubKey;
 
@@ -57,4 +71,49 @@ contract P256Test is Test {
         res = P256.verifySignature(hash, r, s, pubKey[0], pubKey[1]);
         assertEq(res, true);
     }
+
+    function testPrecompileUsage() public {
+        uint256 r = 0x01655c1753db6b61a9717e4ccc5d6c4bf7681623dd54c2d6babc55125756661c;
+        uint256 s = 7033802732221576339889804108463427183539365869906989872244893535944704590394;
+
+        bytes32 hash = 0x267f9ea080b54bbea2443dff8aa543604564329783b6a515c6663a691c555490;
+
+        uint gasBefore = gasleft();
+        bool res = P256.verifySignatureAllowMalleability(
+            hash,
+            r,
+            s,
+            pubKey[0],
+            pubKey[1]
+        );
+        assertEq(res, true);
+        uint gasUsedFallback = gasBefore - gasleft();
+        assert(gasUsedFallback > 100_000); // no precompile, used Solidity implementation
+
+        vm.etch(P256.PRECOMPILE, type(FakePrecompile).runtimeCode);
+
+        gasBefore = gasleft();
+        res = P256.verifySignatureAllowMalleability(
+            hash,
+            r,
+            s,
+            pubKey[0],
+            pubKey[1]
+        );
+        assertEq(res, true);
+        uint gasUsedPrecompile = gasBefore - gasleft();
+        assert(gasUsedPrecompile < gasUsedFallback); // precompile, used precompile
+
+        gasBefore = gasleft();
+        res = P256.verifySignatureAllowMalleability(
+            hash,
+            r + 1,
+            s,
+            pubKey[0],
+            pubKey[1]
+        );
+        assertEq(res, false);
+        uint gasUsedInvalidSignature = gasBefore - gasleft();
+        assert(gasUsedInvalidSignature > gasUsedFallback); // invalid signature, used precompile and failed, so fall back to Solidity implementation and failed with nearly double gas
+    }
 }