fix(tee-prover): mitigate panic on redeployments

core/bin/zksync_tee_prover/Cargo.toml

@@ -14,8 +14,9 @@ publish = false
 [dependencies]
 anyhow.workspace = true
 async-trait.workspace = true
+envy.workspace = true
 reqwest.workspace = true
-secp256k1.workspace = true
+secp256k1 = { workspace = true, features = ["serde"] }
 serde = { workspace = true, features = ["derive"] }
 thiserror.workspace = true
 tokio = { workspace = true, features = ["full"] }

core/bin/zksync_tee_prover/src/config.rs

@@ -1,17 +1,16 @@
 use std::{path::PathBuf, time::Duration};
 
-use secp256k1::{PublicKey, Secp256k1, SecretKey};
+use secp256k1::SecretKey;
+use serde::Deserialize;
 use url::Url;
 use zksync_env_config::FromEnv;
 use zksync_types::tee_types::TeeType;
 
 /// Configuration for the TEE prover.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Deserialize)]
 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.
@@ -36,36 +35,16 @@ impl FromEnv for TeeProverConfig {
     /// Example usage of environment variables for tests:
     /// ```
     /// 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_ATTESTATION_QUOTE_FILE_PATH="/tmp/test"  # run `echo test > /tmp/test` beforehand
+    /// export TEE_PROVER_TEE_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_INITIAL_RETRY_BACKOFF=1
     /// export TEE_PROVER_RETRY_BACKOFF_MULTIPLIER=2.0
-    /// export TEE_PROVER_MAX_BACKOFF_SECONDS=128
+    /// export TEE_PROVER_MAX_BACKOFF=128
     /// ```
     fn from_env() -> anyhow::Result<Self> {
-        let signing_key = std::env::var("TEE_PROVER_SIGNING_KEY")?.parse()?;
-        Ok(Self {
-            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()?,
-            ),
-        })
+        let config: Self = envy::prefixed("TEE_PROVER_").from_env()?;
+        Ok(config)
     }
 }

core/bin/zksync_tee_prover/src/tee_prover.rs

@@ -1,6 +1,6 @@
 use std::fmt;
 
-use secp256k1::{ecdsa::Signature, Message};
+use secp256k1::{ecdsa::Signature, Message, PublicKey, Secp256k1};
 use zksync_basic_types::H256;
 use zksync_node_framework::{
     service::StopReceiver,
@@ -62,7 +62,6 @@ 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.config.public_key)
             .finish()
     }
 }
@@ -90,15 +89,15 @@ impl TeeProver {
         }
     }
 
-    async fn step(&self) -> Result<Option<L1BatchNumber>, TeeProverError> {
+    async fn step(&self, public_key: &PublicKey) -> Result<Option<L1BatchNumber>, TeeProverError> {
         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.config.public_key,
+                        public_key,
                         root_hash,
                         self.config.tee_type,
                     )
@@ -122,25 +121,27 @@ 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)?;
+        let config = &self.config;
+        let attestation_quote_bytes = std::fs::read(&config.attestation_quote_file_path)?;
+        let public_key = config.signing_key.public_key(&Secp256k1::new());
         self.api_client
-            .register_attestation(attestation_quote_bytes, &self.config.public_key)
+            .register_attestation(attestation_quote_bytes, &public_key)
             .await?;
 
         let mut retries = 1;
-        let mut backoff = self.config.initial_retry_backoff;
+        let mut backoff = config.initial_retry_backoff;
         let mut observer = METRICS.job_waiting_time.start();
 
         loop {
             if *stop_receiver.0.borrow() {
                 tracing::info!("Stop signal received, shutting down TEE Prover component");
                 return Ok(());
             }
-            let result = self.step().await;
+            let result = self.step(&public_key).await;
             let need_to_sleep = match result {
                 Ok(batch_number) => {
                     retries = 1;
-                    backoff = self.config.initial_retry_backoff;
+                    backoff = config.initial_retry_backoff;
                     if let Some(batch_number) = batch_number {
                         observer.observe();
                         observer = METRICS.job_waiting_time.start();
@@ -154,14 +155,14 @@ impl Task for TeeProver {
                 }
                 Err(err) => {
                     METRICS.network_errors_counter.inc_by(1);
-                    if !err.is_retriable() || retries > self.config.max_retries {
+                    if !err.is_retriable() || retries > config.max_retries {
                         return Err(err.into());
                     }
-                    tracing::warn!(%err, "Failed TEE prover step function {retries}/{}, retrying in {} milliseconds.", self.config.max_retries, backoff.as_millis());
+                    tracing::warn!(%err, "Failed TEE prover step function {retries}/{}, retrying in {} milliseconds.", config.max_retries, backoff.as_millis());
                     retries += 1;
                     backoff = std::cmp::min(
-                        backoff.mul_f32(self.config.retry_backoff_multiplier),
-                        self.config.max_backoff,
+                        backoff.mul_f32(config.retry_backoff_multiplier),
+                        config.max_backoff,
                     );
                     true
                 }

core/node/external_proof_integration_api/src/lib.rs

@@ -4,8 +4,6 @@ mod middleware;
 mod processor;
 mod types;
 
-pub use crate::processor::Processor;
-
 use std::net::SocketAddr;
 
 use anyhow::Context;
@@ -20,6 +18,7 @@ use tokio::sync::watch;
 use types::{ExternalProof, ProofGenerationDataResponse};
 use zksync_basic_types::L1BatchNumber;
 
+pub use crate::processor::Processor;
 use crate::{
     metrics::{CallOutcome, Method},
     middleware::MetricsMiddleware,