From c801dd7f4f3c08e08835a48797ea3b6439fcbdd2 Mon Sep 17 00:00:00 2001 From: Patrick Beza Date: Mon, 13 Jan 2025 22:47:24 +0100 Subject: [PATCH] Address code review comments --- Cargo.lock | 3 +- core/bin/zksync_tee_prover/Cargo.toml | 3 +- core/bin/zksync_tee_prover/src/api_client.rs | 3 +- core/bin/zksync_tee_prover/src/config.rs | 22 ++- core/bin/zksync_tee_prover/src/tee_prover.rs | 141 +++++++++++-------- 5 files changed, 105 insertions(+), 67 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bce35e14ebf3..26db94496ae5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12957,7 +12957,6 @@ dependencies = [ "reqwest 0.12.9", "secp256k1", "serde", - "sha3 0.10.8", "thiserror 1.0.69", "tokio", "tracing", @@ -12965,11 +12964,11 @@ dependencies = [ "vise", "zksync_basic_types", "zksync_config", + "zksync_crypto_primitives", "zksync_env_config", "zksync_node_framework", "zksync_prover_interface", "zksync_tee_verifier", - "zksync_types", "zksync_vlog", ] diff --git a/core/bin/zksync_tee_prover/Cargo.toml b/core/bin/zksync_tee_prover/Cargo.toml index 2628673845c9..aed8cfb49f90 100644 --- a/core/bin/zksync_tee_prover/Cargo.toml +++ b/core/bin/zksync_tee_prover/Cargo.toml @@ -29,13 +29,12 @@ url.workspace = true vise.workspace = true zksync_basic_types.workspace = true zksync_config = { workspace = true, features = ["observability_ext"] } +zksync_crypto_primitives.workspace = true zksync_env_config.workspace = true zksync_node_framework.workspace = true zksync_prover_interface.workspace = true zksync_tee_verifier.workspace = true -zksync_types.workspace = true zksync_vlog.workspace = true [dev-dependencies] hex.workspace = true -sha3.workspace = true diff --git a/core/bin/zksync_tee_prover/src/api_client.rs b/core/bin/zksync_tee_prover/src/api_client.rs index 186acc9f952a..fdcb13fb2992 100644 --- a/core/bin/zksync_tee_prover/src/api_client.rs +++ b/core/bin/zksync_tee_prover/src/api_client.rs @@ -2,13 +2,12 @@ use reqwest::{Client, Response, StatusCode}; use secp256k1::PublicKey; use serde::Serialize; use url::Url; -use zksync_basic_types::H256; +use zksync_basic_types::{tee_types::TeeType, L1BatchNumber, H256}; use zksync_prover_interface::{ api::{RegisterTeeAttestationRequest, SubmitTeeProofRequest, TeeProofGenerationDataRequest}, inputs::TeeVerifierInput, outputs::L1BatchTeeProofForL1, }; -use zksync_types::{tee_types::TeeType, L1BatchNumber}; use crate::{error::TeeProverError, metrics::METRICS}; diff --git a/core/bin/zksync_tee_prover/src/config.rs b/core/bin/zksync_tee_prover/src/config.rs index 1c2eb229d616..5f702575d616 100644 --- a/core/bin/zksync_tee_prover/src/config.rs +++ b/core/bin/zksync_tee_prover/src/config.rs @@ -3,8 +3,8 @@ use std::{path::PathBuf, time::Duration}; use secp256k1::SecretKey; use serde::Deserialize; use url::Url; +use zksync_basic_types::tee_types::TeeType; use zksync_env_config::FromEnv; -use zksync_types::tee_types::TeeType; /// Configuration for the TEE prover. #[derive(Debug, Clone, Deserialize)] @@ -19,13 +19,17 @@ pub(crate) struct TeeProverConfig { pub api_url: Url, /// Number of retries for retriable errors before giving up on recovery (i.e., returning an error /// from [`Self::run()`]). + #[serde(default = "TeeProverConfig::default_max_retries")] pub max_retries: usize, /// Initial back-off interval when retrying recovery on a retriable error. Each subsequent retry interval /// will be multiplied by [`Self.retry_backoff_multiplier`]. + #[serde(default = "TeeProverConfig::default_initial_retry_backoff_sec")] pub initial_retry_backoff_sec: u64, /// Multiplier for the back-off interval when retrying recovery on a retriable error. + #[serde(default = "TeeProverConfig::default_retry_backoff_multiplier")] pub retry_backoff_multiplier: f32, /// Maximum back-off interval when retrying recovery on a retriable error. + #[serde(default = "TeeProverConfig::default_max_backoff_sec")] pub max_backoff_sec: u64, } @@ -37,6 +41,22 @@ impl TeeProverConfig { pub fn max_backoff(&self) -> Duration { Duration::from_secs(self.max_backoff_sec) } + + pub const fn default_max_retries() -> usize { + 10 + } + + pub const fn default_initial_retry_backoff_sec() -> u64 { + 1 + } + + pub const fn default_retry_backoff_multiplier() -> f32 { + 2.0 + } + + pub const fn default_max_backoff_sec() -> u64 { + 128 + } } impl FromEnv for TeeProverConfig { diff --git a/core/bin/zksync_tee_prover/src/tee_prover.rs b/core/bin/zksync_tee_prover/src/tee_prover.rs index 92547dd9f301..1651c2cce71c 100644 --- a/core/bin/zksync_tee_prover/src/tee_prover.rs +++ b/core/bin/zksync_tee_prover/src/tee_prover.rs @@ -1,7 +1,8 @@ use std::fmt; -use secp256k1::{Message, PublicKey, Secp256k1, SecretKey, SECP256K1}; -use zksync_basic_types::H256; +use secp256k1::{PublicKey, Secp256k1}; +use zksync_basic_types::{web3, L1BatchNumber, H256}; +use zksync_crypto_primitives::K256PrivateKey; use zksync_node_framework::{ service::StopReceiver, task::{Task, TaskId}, @@ -10,12 +11,13 @@ use zksync_node_framework::{ }; use zksync_prover_interface::inputs::TeeVerifierInput; use zksync_tee_verifier::Verify; -use zksync_types::L1BatchNumber; use crate::{ api_client::TeeApiClient, config::TeeProverConfig, error::TeeProverError, metrics::METRICS, }; +type Signature = [u8; 65]; + /// Wiring layer for `TeeProver` #[derive(Debug)] pub(crate) struct TeeProverLayer { @@ -65,35 +67,40 @@ impl fmt::Debug for TeeProver { .finish() } } +pub trait SignatureSerialization { + fn to_bytes(&self) -> Signature; +} + +impl SignatureSerialization for web3::Signature { + fn to_bytes(&self) -> Signature { + let mut bytes = [0u8; 65]; + bytes[..32].copy_from_slice(self.r.as_bytes()); + bytes[32..64].copy_from_slice(self.s.as_bytes()); + bytes[64] = self.v as u8; + bytes + } +} impl TeeProver { /// Signs the message in Ethereum-compatible format for on-chain verification. - pub fn sign_message(sec: &SecretKey, message: Message) -> Result<[u8; 65], TeeProverError> { - let s = SECP256K1.sign_ecdsa_recoverable(&message, sec); - let (rec_id, data) = s.serialize_compact(); - - let mut signature = [0u8; 65]; - signature[..64].copy_from_slice(&data); - // as defined in the Ethereum Yellow Paper (Appendix F) - // https://ethereum.github.io/yellowpaper/paper.pdf - signature[64] = 27 + rec_id.to_i32() as u8; - + pub fn sign_message(&self, message: &H256) -> Result { + let secret_bytes = self.config.signing_key.secret_bytes(); + let private_key = K256PrivateKey::from_bytes(secret_bytes.into()) + .map_err(|e| TeeProverError::Verification(e.into()))?; + let signature = private_key.sign_web3(message, None).to_bytes(); Ok(signature) } fn verify( &self, tvi: TeeVerifierInput, - ) -> Result<([u8; 65], L1BatchNumber, H256), TeeProverError> { + ) -> Result<(Signature, L1BatchNumber, H256), TeeProverError> { match tvi { TeeVerifierInput::V1(tvi) => { let observer = METRICS.proof_generation_time.start(); let verification_result = tvi.verify().map_err(TeeProverError::Verification)?; - let root_hash_bytes = verification_result.value_hash.as_bytes(); let batch_number = verification_result.batch_number; - let msg_to_sign = Message::from_slice(root_hash_bytes) - .map_err(|e| TeeProverError::Verification(e.into()))?; - let signature = TeeProver::sign_message(&self.config.signing_key, msg_to_sign)?; + let signature = self.sign_message(&verification_result.value_hash)?; let duration = observer.observe(); tracing::info!( proof_generation_time = duration.as_secs_f64(), @@ -199,62 +206,76 @@ impl Task for TeeProver { #[cfg(test)] mod tests { - use anyhow::Result; - use secp256k1::ecdsa::{RecoverableSignature, RecoveryId}; - use sha3::{Digest, Keccak256}; - use super::*; + use anyhow::Result; + use secp256k1::{ + ecdsa::{RecoverableSignature, RecoveryId}, + Message, SecretKey, SECP256K1, + }; + use std::path::PathBuf; + use url::Url; + use zksync_basic_types::{self, tee_types::TeeType, web3::keccak256, Address, H512}; - /// Converts a public key into an Ethereum address by hashing the encoded public key with Keccak256. - pub fn public_key_to_ethereum_address(public: &PublicKey) -> [u8; 20] { - let public_key_bytes = public.serialize_uncompressed(); - - // Skip the first byte (0x04) which indicates uncompressed key - let hash: [u8; 32] = Keccak256::digest(&public_key_bytes[1..]).into(); + type Public = H512; - // Take the last 20 bytes of the hash to get the Ethereum address - let mut address = [0u8; 20]; - address.copy_from_slice(&hash[12..]); - address + /// Recovers the public key from the signature for the message + fn recover(signature: &Signature, message: &Message) -> Result { + let rsig = RecoverableSignature::from_compact( + &signature[0..64], + RecoveryId::from_i32(signature[64] as i32 - 27)?, + )?; + let pubkey = &SECP256K1.recover_ecdsa(&Message::from_slice(&message[..])?, &rsig)?; + let serialized = pubkey.serialize_uncompressed(); + let mut public = Public::default(); + public.as_bytes_mut().copy_from_slice(&serialized[1..65]); + Ok(public) } - /// Equivalent to the ecrecover precompile, ensuring that the signatures we produce off-chain - /// can be recovered on-chain. - pub fn recover_signer(sig: &[u8; 65], msg: &Message) -> Result<[u8; 20]> { - let sig = RecoverableSignature::from_compact( - &sig[0..64], - RecoveryId::from_i32(sig[64] as i32 - 27)?, - )?; - let public = SECP256K1.recover_ecdsa(msg, &sig)?; - Ok(public_key_to_ethereum_address(&public)) + /// Convert public key into the address + fn public_to_address(public: &Public) -> Address { + let hash = keccak256(public.as_bytes()); + let mut result = Address::zero(); + result.as_bytes_mut().copy_from_slice(&hash[12..]); + result } #[test] - fn recover() { - // Decode the sample secret key, generate the public key, and derive the Ethereum address - // from the public key - let secp = Secp256k1::new(); - let secret_key_bytes = - hex::decode("c87509a1c067bbde78beb793e6fa76530b6382a4c0241e5e4a9ec0a0f44dc0d3") - .unwrap(); - let secret_key = SecretKey::from_slice(&secret_key_bytes).unwrap(); - let public_key = PublicKey::from_secret_key(&secp, &secret_key); - let expected_address = hex::decode("627306090abaB3A6e1400e9345bC60c78a8BEf57").unwrap(); - let address = public_key_to_ethereum_address(&public_key); - - assert_eq!(address, expected_address.as_slice()); + fn test_recover() { + let signing_key = SecretKey::from_slice( + &hex::decode("c87509a1c067bbde78beb793e6fa76530b6382a4c0241e5e4a9ec0a0f44dc0d3") + .unwrap(), + ) + .unwrap(); + let tee_prover_config = TeeProverConfig { + signing_key, + attestation_quote_file_path: PathBuf::from("/tmp/mock"), + tee_type: TeeType::Sgx, + api_url: Url::parse("http://mock").unwrap(), + max_retries: TeeProverConfig::default_max_retries(), + initial_retry_backoff_sec: TeeProverConfig::default_initial_retry_backoff_sec(), + retry_backoff_multiplier: TeeProverConfig::default_retry_backoff_multiplier(), + max_backoff_sec: TeeProverConfig::default_max_backoff_sec(), + }; + let tee_prover = TeeProver { + config: tee_prover_config, + api_client: TeeApiClient::new(Url::parse("http://mock").unwrap()), + }; + let private_key = K256PrivateKey::from_bytes(signing_key.secret_bytes().into()).unwrap(); + let expected_address = + Address::from_slice(&hex::decode("627306090abaB3A6e1400e9345bC60c78a8BEf57").unwrap()); + assert!(private_key.address() == expected_address); // Generate a random root hash, create a message from the hash, and sign the message using // the secret key - let root_hash = H256::random(); - let root_hash_bytes = root_hash.as_bytes(); - let msg_to_sign = Message::from_slice(root_hash_bytes).unwrap(); - let signature = TeeProver::sign_message(&secret_key, msg_to_sign).unwrap(); + let random_root_hash = H256::random(); + let signature = tee_prover.sign_message(&random_root_hash).unwrap(); // Recover the signer's Ethereum address from the signature and the message, and verify it // matches the expected address - let proof_addr = recover_signer(&signature, &msg_to_sign).unwrap(); + let message = Message::from_slice(random_root_hash.as_bytes()).unwrap(); + let recovered_pubkey = recover(&signature, &message).unwrap(); + let proof_address = public_to_address(&recovered_pubkey); - assert_eq!(proof_addr, expected_address.as_slice()); + assert_eq!(proof_address, expected_address); } }