From 8b3be091d0420fbd2119397ddc66da4cf70c6d84 Mon Sep 17 00:00:00 2001
From: umadayal <uma.dayal@yahoo.com>
Date: Wed, 4 Dec 2024 12:46:34 -0800
Subject: [PATCH] hook r1 ecrecover

---
 Cargo.lock                                |  1 +
 crates/core/executor/src/hook.rs          | 44 ++++++++++++++++++-----
 crates/curves/Cargo.toml                  |  1 +
 crates/curves/src/lib.rs                  |  4 +++
 crates/test-artifacts/programs/Cargo.lock |  1 +
 5 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index b09a6db73..189b57b43 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5700,6 +5700,7 @@ version = "4.0.0-rc.1"
 dependencies = [
  "cfg-if",
  "dashu",
+ "ecdsa",
  "elliptic-curve",
  "generic-array 1.1.1",
  "itertools 0.13.0",
diff --git a/crates/core/executor/src/hook.rs b/crates/core/executor/src/hook.rs
index 875479b55..186599b54 100644
--- a/crates/core/executor/src/hook.rs
+++ b/crates/core/executor/src/hook.rs
@@ -4,8 +4,9 @@ use std::sync::{Arc, RwLock, RwLockWriteGuard};
 
 use hashbrown::HashMap;
 use sp1_curves::{
+    ecdsa::RecoveryId as ecdsaRecoveryId,
     k256::{Invert, RecoveryId, Signature, VerifyingKey},
-    p256::Signature as p256Signature,
+    p256::{Signature as p256Signature, VerifyingKey as p256VerifyingKey},
 };
 
 use crate::Executor;
@@ -187,20 +188,47 @@ pub fn hook_ed_decompress(_: HookEnv, buf: &[u8]) -> Vec<Vec<u8>> {
 /// # Arguments
 ///
 /// * `env` - The environment in which the hook is invoked.
-/// * `buf` - The buffer containing the signature.
-///     - The signature is 64 bytes.
+/// * `buf` - The buffer containing the signature and message hash.
+///     - The signature is 65 bytes, the first 64 bytes are the signature and the last byte is the
+///       recovery ID.
+///     - The message hash is 32 bytes.
 ///
-/// The result is a single 32 byte vector containing s inverse.
+/// The result is returned as a status and a pair of bytes, where the first 32 bytes are the X coordinate
+/// and the second 32 bytes are the Y coordinate of the decompressed point.
+///
+/// A status of 0 indicates that the public key could not be recovered.
+///
+/// WARNING: This function is used to recover the public key outside of the zkVM context. These
+/// values must be constrained by the zkVM for correctness.
 #[must_use]
 pub fn hook_r1_ecrecover(_: HookEnv, buf: &[u8]) -> Vec<Vec<u8>> {
-    assert_eq!(buf.len(), 64, "ecrecover input should have length 64");
-    let sig: &[u8; 64] = buf.try_into().unwrap();
-    let sig = p256Signature::from_slice(sig).unwrap();
+    assert_eq!(buf.len(), 65 + 32, "ecrecover input should have length 65 + 32, this is a bug.");
+    let (sig, msg_hash) = buf.split_at(65);
+    let sig: &[u8; 65] = sig.try_into().unwrap();
+    let msg_hash: &[u8; 32] = msg_hash.try_into().unwrap();
+
+    let mut recovery_id = sig[64];
+    let mut sig = p256Signature::from_slice(&sig[..64]).unwrap();
+
+    if let Some(sig_normalized) = sig.normalize_s() {
+        sig = sig_normalized;
+        recovery_id ^= 1;
+    };
+    let recid = ecdsaRecoveryId::from_byte(recovery_id)
+        .expect("Computed recovery ID is invalid, this is a bug.");
+
+    // Attempting to recover the public key has failed, write a 0 to indicate to the caller.
+    let Ok(recovered_key) = p256VerifyingKey::recover_from_prehash(&msg_hash[..], &sig, recid)
+    else {
+        return vec![vec![0]];
+    };
+
+    let bytes = recovered_key.to_sec1_bytes();
 
     let (_, s) = sig.split_scalars();
     let s_inverse = s.invert();
 
-    vec![s_inverse.to_bytes().to_vec()]
+    vec![vec![1], bytes.to_vec(), s_inverse.to_bytes().to_vec()]
 }
 
 #[cfg(test)]
diff --git a/crates/curves/Cargo.toml b/crates/curves/Cargo.toml
index 8242decfa..268260c23 100644
--- a/crates/curves/Cargo.toml
+++ b/crates/curves/Cargo.toml
@@ -19,6 +19,7 @@ generic-array = { version = "1.1.0", features = ["alloc", "serde"] }
 amcl = { package = "snowbridge-amcl", version = "1.0.2", default-features = false, features = [
   "bls381",
 ] }
+ecdsa = "0.16.9"
 elliptic-curve = "0.13.8"
 dashu = "0.4.2"
 
diff --git a/crates/curves/src/lib.rs b/crates/curves/src/lib.rs
index 84a1aad0b..73682a74d 100644
--- a/crates/curves/src/lib.rs
+++ b/crates/curves/src/lib.rs
@@ -56,6 +56,10 @@ pub mod p256 {
     };
 }
 
+pub mod ecdsa {
+    pub use ecdsa::RecoveryId;
+}
+
 use params::{FieldParameters, NumWords};
 use sp1_primitives::consts::WORD_SIZE;
 use std::{
diff --git a/crates/test-artifacts/programs/Cargo.lock b/crates/test-artifacts/programs/Cargo.lock
index 9e549c335..cd3a76201 100644
--- a/crates/test-artifacts/programs/Cargo.lock
+++ b/crates/test-artifacts/programs/Cargo.lock
@@ -1821,6 +1821,7 @@ version = "4.0.0-rc.1"
 dependencies = [
  "cfg-if",
  "dashu",
+ "ecdsa",
  "elliptic-curve 0.13.8",
  "generic-array 1.1.1",
  "itertools 0.13.0",