Address code review comments

core/bin/zksync_tee_prover/src/config.rs

@@ -1,39 +1,71 @@
-use std::path::PathBuf;
+use std::{path::PathBuf, time::Duration};
 
-use secp256k1::SecretKey;
+use secp256k1::{PublicKey, Secp256k1, SecretKey};
 use url::Url;
 use zksync_env_config::FromEnv;
 use zksync_types::tee_types::TeeType;
 
 /// Configuration for the TEE prover.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) struct TeeProverConfig {
     /// The private key used to sign the proofs.
     pub signing_key: SecretKey,
+    /// The public key used to verify the proofs.
+    pub public_key: PublicKey,
     /// The path to the file containing the TEE quote.
     pub attestation_quote_file_path: PathBuf,
     /// Attestation quote file.
     pub tee_type: TeeType,
     /// TEE proof data handler API.
     pub api_url: Url,
+    /// Number of retries for retriable errors before giving up on recovery (i.e., returning an error
+    /// from [`Self::run()`]).
+    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`].
+    pub initial_retry_backoff: Duration,
+    /// Multiplier for the back-off interval when retrying recovery on a retriable error.
+    pub retry_backoff_multiplier: f32,
+    /// Maximum back-off interval when retrying recovery on a retriable error.
+    pub max_backoff: Duration,
 }
 
 impl FromEnv for TeeProverConfig {
     /// Constructs the TEE Prover configuration from environment variables.
     ///
     /// Example usage of environment variables for tests:
     /// ```
-    /// export TEE_SIGNING_KEY="b50b38c8d396c88728fc032ece558ebda96907a0b1a9340289715eef7bf29deb"
-    /// export TEE_QUOTE_FILE="/tmp/test"  # run `echo test > /tmp/test` beforehand
-    /// export TEE_TYPE="sgx"
-    /// export TEE_API_URL="http://127.0.0.1:3320"
+    /// export TEE_PROVER_SIGNING_KEY="b50b38c8d396c88728fc032ece558ebda96907a0b1a9340289715eef7bf29deb"
+    /// export TEE_PROVER_QUOTE_FILE="/tmp/test"  # run `echo test > /tmp/test` beforehand
+    /// export TEE_PROVER_TYPE="sgx"
+    /// export TEE_PROVER_API_URL="http://127.0.0.1:3320"
+    /// export TEE_PROVER_MAX_RETRIES=10
+    /// export TEE_PROVER_INITIAL_RETRY_BACKOFF_SECONDS=1
+    /// export TEE_PROVER_RETRY_BACKOFF_MULTIPLIER=2.0
+    /// export TEE_PROVER_MAX_BACKOFF_SECONDS=128
     /// ```
     fn from_env() -> anyhow::Result<Self> {
+        let signing_key = std::env::var("TEE_PROVER_SIGNING_KEY")?.parse()?;
         Ok(Self {
-            signing_key: std::env::var("TEE_SIGNING_KEY")?.parse()?,
-            attestation_quote_file_path: std::env::var("TEE_QUOTE_FILE")?.parse()?,
-            tee_type: std::env::var("TEE_TYPE")?.parse()?,
-            api_url: std::env::var("TEE_API_URL")?.parse()?,
+            signing_key,
+            public_key: signing_key.public_key(&Secp256k1::new()),
+            attestation_quote_file_path: std::env::var("TEE_PROVER_QUOTE_FILE")?.parse()?,
+            tee_type: std::env::var("TEE_PROVER_TYPE")?.parse()?,
+            api_url: std::env::var("TEE_PROVER_API_URL")?.parse()?,
+            max_retries: std::env::var("TEE_PROVER_MAX_RETRIES")?.parse()?,
+            initial_retry_backoff: Duration::from_secs(
+                std::env::var("TEE_PROVER_INITIAL_RETRY_BACKOFF_SECONDS")
+                    .unwrap_or_else(|_| "1".to_string())
+                    .parse()?,
+            ),
+            retry_backoff_multiplier: std::env::var("TEE_PROVER_RETRY_BACKOFF_MULTIPLIER")
+                .unwrap_or("2.0".to_string())
+                .parse()?,
+            max_backoff: Duration::from_secs(
+                std::env::var("TEE_PROVER_MAX_BACKOFF_SECONDS")
+                    .unwrap_or_else(|_| "128".to_string())
+                    .parse()?,
+            ),
         })
     }
 }

core/bin/zksync_tee_prover/src/main.rs

@@ -32,8 +32,6 @@ fn main() -> anyhow::Result<()> {
         ObservabilityConfig::from_env().context("ObservabilityConfig::from_env()")?;
 
     let tee_prover_config = TeeProverConfig::from_env()?;
-    let attestation_quote_bytes = std::fs::read(tee_prover_config.attestation_quote_file_path)?;
-
     let prometheus_config = PrometheusConfig::from_env()?;
 
     let mut builder = ZkStackServiceBuilder::new()?;
@@ -45,12 +43,7 @@ fn main() -> anyhow::Result<()> {
 
     builder
         .add_layer(SigintHandlerLayer)
-        .add_layer(TeeProverLayer::new(
-            tee_prover_config.api_url,
-            tee_prover_config.signing_key,
-            attestation_quote_bytes,
-            tee_prover_config.tee_type,
-        ));
+        .add_layer(TeeProverLayer::new(tee_prover_config));
 
     if let Some(gateway) = prometheus_config.gateway_endpoint() {
         let exporter_config =

core/bin/zksync_tee_prover/src/tee_prover.rs

@@ -1,7 +1,6 @@
-use std::{fmt, time::Duration};
+use std::fmt;
 
-use secp256k1::{ecdsa::Signature, Message, PublicKey, Secp256k1, SecretKey};
-use url::Url;
+use secp256k1::{ecdsa::Signature, Message};
 use zksync_basic_types::H256;
 use zksync_node_framework::{
     service::StopReceiver,
@@ -11,32 +10,21 @@ use zksync_node_framework::{
 };
 use zksync_prover_interface::inputs::TeeVerifierInput;
 use zksync_tee_verifier::Verify;
-use zksync_types::{tee_types::TeeType, L1BatchNumber};
+use zksync_types::L1BatchNumber;
 
-use crate::{api_client::TeeApiClient, error::TeeProverError, metrics::METRICS};
+use crate::{
+    api_client::TeeApiClient, config::TeeProverConfig, error::TeeProverError, metrics::METRICS,
+};
 
 /// Wiring layer for `TeeProver`
 #[derive(Debug)]
 pub(crate) struct TeeProverLayer {
-    api_url: Url,
-    signing_key: SecretKey,
-    attestation_quote_bytes: Vec<u8>,
-    tee_type: TeeType,
+    config: TeeProverConfig,
 }
 
 impl TeeProverLayer {
-    pub fn new(
-        api_url: Url,
-        signing_key: SecretKey,
-        attestation_quote_bytes: Vec<u8>,
-        tee_type: TeeType,
-    ) -> Self {
-        Self {
-            api_url,
-            signing_key,
-            attestation_quote_bytes,
-            tee_type,
-        }
+    pub fn new(config: TeeProverConfig) -> Self {
+        Self { config }
     }
 }
 
@@ -56,34 +44,25 @@ impl WiringLayer for TeeProverLayer {
     }
 
     async fn wire(self, _input: Self::Input) -> Result<Self::Output, WiringError> {
+        let api_url = self.config.api_url.clone();
         let tee_prover = TeeProver {
-            config: Default::default(),
-            signing_key: self.signing_key,
-            public_key: self.signing_key.public_key(&Secp256k1::new()),
-            attestation_quote_bytes: self.attestation_quote_bytes,
-            tee_type: self.tee_type,
-            api_client: TeeApiClient::new(self.api_url),
+            config: self.config,
+            api_client: TeeApiClient::new(api_url),
         };
         Ok(LayerOutput { tee_prover })
     }
 }
 
 pub(crate) struct TeeProver {
     config: TeeProverConfig,
-    signing_key: SecretKey,
-    public_key: PublicKey,
-    attestation_quote_bytes: Vec<u8>,
-    tee_type: TeeType,
     api_client: TeeApiClient,
 }
 
 impl fmt::Debug for TeeProver {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("TeeProver")
             .field("config", &self.config)
-            .field("public_key", &self.public_key)
-            .field("attestation_quote_bytes", &self.attestation_quote_bytes)
-            .field("tee_type", &self.tee_type)
+            .field("public_key", &self.config.public_key)
             .finish()
     }
 }
@@ -101,7 +80,7 @@ impl TeeProver {
                 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 = self.signing_key.sign_ecdsa(msg_to_sign);
+                let signature = self.config.signing_key.sign_ecdsa(msg_to_sign);
                 observer.observe();
                 Ok((signature, batch_number, verification_result.value_hash))
             }
@@ -112,16 +91,16 @@ impl TeeProver {
     }
 
     async fn step(&self) -> Result<Option<L1BatchNumber>, TeeProverError> {
-        match self.api_client.get_job(self.tee_type).await? {
+        match self.api_client.get_job(self.config.tee_type).await? {
             Some(job) => {
                 let (signature, batch_number, root_hash) = self.verify(*job)?;
                 self.api_client
                     .submit_proof(
                         batch_number,
                         signature,
-                        &self.public_key,
+                        &self.config.public_key,
                         root_hash,
-                        self.tee_type,
+                        self.config.tee_type,
                     )
                     .await?;
                 Ok(Some(batch_number))
@@ -134,30 +113,6 @@ impl TeeProver {
     }
 }
 
-/// TEE prover configuration options.
-#[derive(Debug, Clone)]
-pub struct TeeProverConfig {
-    /// Number of retries for retriable errors before giving up on recovery (i.e., returning an error
-    /// from [`Self::run()`]).
-    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`].
-    pub initial_retry_backoff: Duration,
-    pub retry_backoff_multiplier: f32,
-    pub max_backoff: Duration,
-}
-
-impl Default for TeeProverConfig {
-    fn default() -> Self {
-        Self {
-            max_retries: 10,
-            initial_retry_backoff: Duration::from_secs(1),
-            retry_backoff_multiplier: 2.0,
-            max_backoff: Duration::from_secs(128),
-        }
-    }
-}
-
 #[async_trait::async_trait]
 impl Task for TeeProver {
     fn id(&self) -> TaskId {
@@ -167,8 +122,9 @@ impl Task for TeeProver {
     async fn run(self: Box<Self>, mut stop_receiver: StopReceiver) -> anyhow::Result<()> {
         tracing::info!("Starting the task {}", self.id());
 
+        let attestation_quote_bytes = std::fs::read(&self.config.attestation_quote_file_path)?;
         self.api_client
-            .register_attestation(self.attestation_quote_bytes.clone(), &self.public_key)
+            .register_attestation(attestation_quote_bytes, &self.config.public_key)
             .await?;
 
         let mut retries = 1;

etc/nix/container-tee_prover.nix

@@ -28,7 +28,11 @@ nixsgxLib.mkSGXContainer {
       log_level = "error";
 
       env = {
-        TEE_API_URL.passthrough = true;
+        TEE_PROVER_API_URL.passthrough = true;
+        TEE_PROVER_MAX_RETRIES.passthrough = true;
+        TEE_PROVER_INITIAL_RETRY_BACKOFF_SECONDS.passthrough = true;
+        TEE_PROVER_RETRY_BACKOFF_MULTIPLIER.passthrough = true;
+        TEE_PROVER_MAX_BACKOFF_SECONDS.passthrough = true;
         API_PROMETHEUS_LISTENER_PORT.passthrough = true;
         API_PROMETHEUS_PUSHGATEWAY_URL.passthrough = true;
         API_PROMETHEUS_PUSH_INTERVAL_MS.passthrough = true;