refactor(common): use underlying VK types instead of Vec<u8>

Cargo.toml

@@ -81,7 +81,11 @@ sha2 = "0.10.8"
 auto_impl = "1.2.0"
 bincode = "1.3.3"
 ed25519-consensus = "2.1.0"
-secp256k1 = { version = "0.29.0", features = ["global-context", "rand-std"] }
+secp256k1 = { version = "0.29.0", features = [
+    "global-context",
+    "rand-std",
+    "serde",
+] }
 sp1-zkvm = { version = "3.0.0-rc1" }
 sp1-sdk = { version = "3.0.0-rc1" }
 prism-common = { path = "crates/common" }

crates/bin/src/main.rs

@@ -3,7 +3,6 @@ mod node_types;
 
 use cfg::{initialize_da_layer, load_config, CommandLineArgs, Commands};
 use clap::Parser;
-use ed25519_consensus::VerificationKey as Ed25519VerifyingKey;
 use keystore_rs::{KeyChain, KeyStore, KeyStoreType};
 use prism_common::keys::VerifyingKey;
 
@@ -41,7 +40,10 @@ async fn main() -> std::io::Result<()> {
             let prover_vk = config
                 .verifying_key
                 .and_then(|s| s.try_into().ok())
-                .and_then(|vk: VerifyingKey| Ed25519VerifyingKey::try_from(vk.as_bytes()).ok());
+                .and_then(|vk: VerifyingKey| match vk {
+                    VerifyingKey::Ed25519(key) => Some(key),
+                    _ => None,
+                });
 
             Arc::new(LightClient::new(da, celestia_config, prover_vk))
         }

crates/common/src/keys.rs

@@ -17,17 +17,17 @@ use crate::tree::Digest;
 /// Represents a public key supported by the system.
 pub enum VerifyingKey {
     /// Bitcoin, Ethereum
-    Secp256k1(Vec<u8>),
+    Secp256k1(Secp256k1VerifyingKey),
     /// Cosmos, OpenSSH, GnuPG
-    Ed25519(Vec<u8>),
+    Ed25519(Ed25519VerifyingKey),
 }
 
 impl VerifyingKey {
     /// Returns the byte representation of the public key.
-    pub fn as_bytes(&self) -> &[u8] {
+    pub fn as_bytes(&self) -> Vec<u8> {
         match self {
-            VerifyingKey::Ed25519(bytes) => bytes,
-            VerifyingKey::Secp256k1(bytes) => bytes,
+            VerifyingKey::Ed25519(vk) => vk.to_bytes().to_vec(),
+            VerifyingKey::Secp256k1(vk) => vk.serialize().to_vec(),
         }
     }
 
@@ -36,14 +36,12 @@ impl VerifyingKey {
             return Err(anyhow!("Invalid signature length"));
         }
         match self {
-            VerifyingKey::Ed25519(bytes) => {
-                let vk = Ed25519VerifyingKey::try_from(bytes.as_slice()).map_err(|e| anyhow!(e))?;
+            VerifyingKey::Ed25519(vk) => {
                 let signature = Ed25519Signature::try_from(signature).map_err(|e| anyhow!(e))?;
                 vk.verify(&signature, message).map_err(|e| anyhow!(e))
             }
-            VerifyingKey::Secp256k1(bytes) => {
+            VerifyingKey::Secp256k1(vk) => {
                 let hashed_message = Digest::hash(message).to_bytes();
-                let vk = Secp256k1VerifyingKey::from_slice(bytes.as_slice())?;
                 let message = Secp256k1Message::from_digest(hashed_message);
                 let signature = Secp256k1Signature::from_compact(signature)?;
 
@@ -56,13 +54,13 @@ impl VerifyingKey {
 
 impl From<Ed25519SigningKey> for VerifyingKey {
     fn from(sk: Ed25519SigningKey) -> Self {
-        VerifyingKey::Ed25519(sk.verification_key().to_bytes().to_vec())
+        VerifyingKey::Ed25519(sk.verification_key())
     }
 }
 
 impl From<Ed25519VerifyingKey> for VerifyingKey {
     fn from(vk: Ed25519VerifyingKey) -> Self {
-        VerifyingKey::Ed25519(vk.to_bytes().to_vec())
+        VerifyingKey::Ed25519(vk)
     }
 }
 
@@ -74,7 +72,7 @@ impl From<Secp256k1SigningKey> for VerifyingKey {
 
 impl From<Secp256k1VerifyingKey> for VerifyingKey {
     fn from(vk: Secp256k1VerifyingKey) -> Self {
-        VerifyingKey::Secp256k1(vk.serialize().to_vec())
+        VerifyingKey::Secp256k1(vk)
     }
 }
 
@@ -102,8 +100,16 @@ impl TryFrom<String> for VerifyingKey {
             .map_err(|e| anyhow!("Failed to decode base64 string: {}", e))?;
 
         match bytes.len() {
-            32 => Ok(VerifyingKey::Ed25519(bytes)),
-            33 | 65 => Ok(VerifyingKey::Secp256k1(bytes)),
+            32 => {
+                let vk = Ed25519VerifyingKey::try_from(bytes.as_slice())
+                    .map_err(|e| anyhow!("Invalid Ed25519 key: {}", e))?;
+                Ok(VerifyingKey::Ed25519(vk))
+            }
+            33 | 65 => {
+                let vk = Secp256k1VerifyingKey::from_slice(bytes.as_slice())
+                    .map_err(|e| anyhow!("Invalid Secp256k1 key: {}", e))?;
+                Ok(VerifyingKey::Secp256k1(vk))
+            }
             _ => Err(anyhow!("Invalid public key length")),
         }
     }
@@ -143,53 +149,28 @@ mod tests {
 
     #[test]
     fn test_verifying_key_from_string_ed25519() {
-        let ed25519_vk =
+        let original_key =
             SigningKey::Ed25519(Box::new(Ed25519SigningKey::new(OsRng))).verifying_key();
-        let encoded = engine.encode(ed25519_vk.as_bytes());
+        let encoded = engine.encode(original_key.as_bytes());
 
         let result = VerifyingKey::try_from(encoded);
         assert!(result.is_ok());
 
-        if let Ok(VerifyingKey::Ed25519(key_bytes)) = result {
-            assert_eq!(key_bytes.len(), 32);
-            assert_eq!(key_bytes, ed25519_vk.as_bytes());
-        } else {
-            panic!("Expected Ed25519 key");
-        }
+        let decoded_key = result.unwrap();
+        assert_eq!(decoded_key.as_bytes(), original_key.as_bytes());
     }
 
     #[test]
-    fn test_verifying_key_from_string_secp256k1_compressed() {
-        let secp256k1_vk =
+    fn test_verifying_key_from_string_secp256k1() {
+        let original_key =
             SigningKey::Secp256k1(Secp256k1SigningKey::new(&mut OsRng)).verifying_key();
-        let secp256k1_bytes = secp256k1_vk.as_bytes();
-        let encoded = engine.encode(secp256k1_bytes);
-
-        let result = VerifyingKey::try_from(encoded);
-        assert!(result.is_ok());
-
-        if let Ok(VerifyingKey::Secp256k1(key_bytes)) = result {
-            dbg!(key_bytes.len());
-            assert_eq!(key_bytes, secp256k1_bytes);
-        } else {
-            panic!("Expected Secp256k1 key");
-        }
-    }
-
-    #[test]
-    fn test_verifying_key_from_string_secp256k1_uncompressed() {
-        let secp256k1_bytes = [0; 65];
-        let encoded = engine.encode(secp256k1_bytes);
+        let encoded = engine.encode(original_key.as_bytes());
 
         let result = VerifyingKey::try_from(encoded);
         assert!(result.is_ok());
 
-        if let Ok(VerifyingKey::Secp256k1(key_bytes)) = result {
-            assert_eq!(key_bytes.len(), 65);
-            assert_eq!(key_bytes, secp256k1_bytes);
-        } else {
-            panic!("Expected Secp256k1 key");
-        }
+        let decoded_key = result.unwrap();
+        assert_eq!(decoded_key.as_bytes(), original_key.as_bytes());
     }
 
     #[test]

crates/common/src/test_utils.rs

@@ -91,7 +91,7 @@ impl TestTreeState {
             .process_operation(&account.hashchain.last().unwrap().operation)?;
         if let Proof::Insert(insert_proof) = proof {
             self.inserted_keys.insert(account.key_hash);
-            return Ok(insert_proof);
+            return Ok(*insert_proof);
         }
         Err(anyhow!("Insert proof not returned"))
     }
@@ -105,7 +105,7 @@ impl TestTreeState {
             .tree
             .process_operation(&account.hashchain.last().unwrap().operation)?;
         if let Proof::Update(update_proof) = proof {
-            return Ok(update_proof);
+            return Ok(*update_proof);
         }
         Err(anyhow!("Update proof not returned"))
     }

crates/common/src/tree.rs

@@ -185,8 +185,8 @@ pub struct Batch {
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum Proof {
-    Update(UpdateProof),
-    Insert(InsertProof),
+    Update(Box<UpdateProof>),
+    Insert(Box<InsertProof>),
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -399,7 +399,7 @@ where
                         debug!("updating hashchain for user id {}", id.clone());
                         let proof = self.update(key_hash, new_entry.clone())?;
 
-                        Ok(Proof::Update(proof))
+                        Ok(Proof::Update(Box::new(proof)))
                     }
                     NotFound(_) => {
                         bail!("Failed to get hashchain for ID {}", id)
@@ -470,9 +470,9 @@ where
 
                 debug!("creating new hashchain for user ID {}", id);
 
-                Ok(Proof::Insert(
+                Ok(Proof::Insert(Box::new(
                     self.insert(account_key_hash, new_account_chain)?,
-                ))
+                )))
             }
             Operation::RegisterService(RegisterServiceArgs {
                 id, creation_gate, ..
@@ -491,9 +491,9 @@ where
                 debug!("creating new hashchain for service id {}", id);
                 let chain = Hashchain::register_service(id.clone(), creation_gate.clone())?;
 
-                Ok(Proof::Insert(
+                Ok(Proof::Insert(Box::new(
                     self.insert(KeyHash::with::<Hasher>(hashed_id), chain)?,
-                ))
+                )))
             }
         }
     }