feat(proof-data-handler): add first processed batch option (#3112)

## What ❔

Add an option to the proof data handler to allow the first verified
batch to be set.

## Why ❔

To be able to skip some batches if we need to.

## Checklist

- [x] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [x] Tests for the changes have been added / updated.
- [ ] Documentation comments have been added / updated.
- [x] Code has been formatted via `zkstack dev fmt` and `zkstack dev
lint`.

core/lib/config/src/configs/mod.rs

@@ -22,7 +22,7 @@ pub use self::{
     genesis::GenesisConfig,
     object_store::ObjectStoreConfig,
     observability::{ObservabilityConfig, OpentelemetryConfig},
-    proof_data_handler::ProofDataHandlerConfig,
+    proof_data_handler::{ProofDataHandlerConfig, TeeConfig},
     prover_job_monitor::ProverJobMonitorConfig,
     pruning::PruningConfig,
     secrets::{DatabaseSecrets, L1Secrets, Secrets},

core/lib/config/src/configs/proof_data_handler.rs

@@ -1,12 +1,43 @@
 use std::time::Duration;
 
 use serde::Deserialize;
+use zksync_basic_types::L1BatchNumber;
+
+#[derive(Debug, Deserialize, Clone, PartialEq)]
+pub struct TeeConfig {
+    /// If true, the TEE support is enabled.
+    pub tee_support: bool,
+    /// All batches before this one are considered to be processed.
+    pub first_tee_processed_batch: L1BatchNumber,
+}
+
+impl Default for TeeConfig {
+    fn default() -> Self {
+        TeeConfig {
+            tee_support: Self::default_tee_support(),
+            first_tee_processed_batch: Self::default_first_tee_processed_batch(),
+        }
+    }
+}
+
+impl TeeConfig {
+    pub fn default_tee_support() -> bool {
+        false
+    }
+
+    pub fn default_first_tee_processed_batch() -> L1BatchNumber {
+        L1BatchNumber(0)
+    }
+}
 
 #[derive(Debug, Deserialize, Clone, PartialEq)]
 pub struct ProofDataHandlerConfig {
     pub http_port: u16,
     pub proof_generation_timeout_in_secs: u16,
-    pub tee_support: bool,
+    #[serde(skip)]
+    // ^ Filled in separately in `Self::from_env()`. We cannot use `serde(flatten)` because it
+    // doesn't work with `envy`: https://github.com/softprops/envy/issues/26
+    pub tee_config: TeeConfig,
 }
 
 impl ProofDataHandlerConfig {

core/lib/config/src/testonly.rs

@@ -677,7 +677,10 @@ impl Distribution<configs::ProofDataHandlerConfig> for EncodeDist {
         configs::ProofDataHandlerConfig {
             http_port: self.sample(rng),
             proof_generation_timeout_in_secs: self.sample(rng),
-            tee_support: self.sample(rng),
+            tee_config: configs::TeeConfig {
+                tee_support: self.sample(rng),
+                first_tee_processed_batch: L1BatchNumber(rng.gen()),
+            },
         }
     }
 }

core/lib/dal/src/tee_proof_generation_dal.rs

@@ -32,10 +32,10 @@ impl TeeProofGenerationDal<'_, '_> {
         &mut self,
         tee_type: TeeType,
         processing_timeout: Duration,
-        min_batch_number: Option<L1BatchNumber>,
+        min_batch_number: L1BatchNumber,
     ) -> DalResult<Option<L1BatchNumber>> {
         let processing_timeout = pg_interval_from_duration(processing_timeout);
-        let min_batch_number = min_batch_number.map_or(0, |num| i64::from(num.0));
+        let min_batch_number = i64::from(min_batch_number.0);
         sqlx::query!(
             r#"
             WITH upsert AS (

core/lib/env_config/src/proof_data_handler.rs

@@ -4,12 +4,18 @@ use crate::{envy_load, FromEnv};
 
 impl FromEnv for ProofDataHandlerConfig {
     fn from_env() -> anyhow::Result<Self> {
-        envy_load("proof_data_handler", "PROOF_DATA_HANDLER_")
+        Ok(Self {
+            tee_config: envy_load("proof_data_handler.tee", "PROOF_DATA_HANDLER_")?,
+            ..envy_load("proof_data_handler", "PROOF_DATA_HANDLER_")?
+        })
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use zksync_basic_types::L1BatchNumber;
+    use zksync_config::configs::TeeConfig;
+
     use super::*;
     use crate::test_utils::EnvMutex;
 
@@ -19,7 +25,10 @@ mod tests {
         ProofDataHandlerConfig {
             http_port: 3320,
             proof_generation_timeout_in_secs: 18000,
-            tee_support: true,
+            tee_config: TeeConfig {
+                tee_support: true,
+                first_tee_processed_batch: L1BatchNumber(1337),
+            },
         }
     }
 
@@ -29,6 +38,7 @@ mod tests {
             PROOF_DATA_HANDLER_PROOF_GENERATION_TIMEOUT_IN_SECS="18000"
             PROOF_DATA_HANDLER_HTTP_PORT="3320"
             PROOF_DATA_HANDLER_TEE_SUPPORT="true"
+            PROOF_DATA_HANDLER_FIRST_TEE_PROCESSED_BATCH="1337"
         "#;
         let mut lock = MUTEX.lock();
         lock.set_env(config);

core/lib/protobuf_config/src/proof_data_handler.rs

@@ -1,6 +1,7 @@
 use anyhow::Context as _;
 use zksync_config::configs;
 use zksync_protobuf::{repr::ProtoRepr, required};
+use zksync_types::L1BatchNumber;
 
 use crate::proto::prover as proto;
 
@@ -14,17 +15,24 @@ impl ProtoRepr for proto::ProofDataHandler {
             proof_generation_timeout_in_secs: required(&self.proof_generation_timeout_in_secs)
                 .and_then(|x| Ok((*x).try_into()?))
                 .context("proof_generation_timeout_in_secs")?,
-            tee_support: required(&self.tee_support)
-                .copied()
-                .context("tee_support")?,
+            tee_config: configs::TeeConfig {
+                tee_support: self
+                    .tee_support
+                    .unwrap_or_else(configs::TeeConfig::default_tee_support),
+                first_tee_processed_batch: self
+                    .first_tee_processed_batch
+                    .map(|x| L1BatchNumber(x as u32))
+                    .unwrap_or_else(configs::TeeConfig::default_first_tee_processed_batch),
+            },
         })
     }
 
     fn build(this: &Self::Type) -> Self {
         Self {
             http_port: Some(this.http_port.into()),
             proof_generation_timeout_in_secs: Some(this.proof_generation_timeout_in_secs.into()),
-            tee_support: Some(this.tee_support),
+            tee_support: Some(this.tee_config.tee_support),
+            first_tee_processed_batch: Some(this.tee_config.first_tee_processed_batch.0 as u64),
         }
     }
 }

core/lib/protobuf_config/src/proto/config/prover.proto

@@ -107,5 +107,6 @@ message WitnessVectorGenerator {
 message ProofDataHandler {
   optional uint32 http_port = 1; // required; u16
   optional uint32 proof_generation_timeout_in_secs = 2; // required; s
-  optional bool tee_support = 3; // required
+  optional bool tee_support = 3; // optional
+  optional uint64 first_tee_processed_batch = 4; // optional
 }

core/node/proof_data_handler/src/lib.rs

@@ -94,7 +94,7 @@ fn create_proof_processing_router(
             ),
         );
 
-    if config.tee_support {
+    if config.tee_config.tee_support {
         let get_tee_proof_gen_processor =
             TeeRequestProcessor::new(blob_store, connection_pool, config.clone(), l2_chain_id);
         let submit_tee_proof_processor = get_tee_proof_gen_processor.clone();

core/node/proof_data_handler/src/tee_request_processor.rs

@@ -47,7 +47,7 @@ impl TeeRequestProcessor {
     ) -> Result<Option<Json<TeeProofGenerationDataResponse>>, RequestProcessorError> {
         tracing::info!("Received request for proof generation data: {:?}", request);
 
-        let mut min_batch_number: Option<L1BatchNumber> = None;
+        let mut min_batch_number = self.config.tee_config.first_tee_processed_batch;
         let mut missing_range: Option<(L1BatchNumber, L1BatchNumber)> = None;
 
         let result = loop {
@@ -72,7 +72,7 @@ impl TeeRequestProcessor {
                         None => Some((l1_batch_number, l1_batch_number)),
                     };
                     self.unlock_batch(l1_batch_number, request.tee_type).await?;
-                    min_batch_number = Some(min_batch_number.unwrap_or(l1_batch_number) + 1);
+                    min_batch_number = l1_batch_number + 1;
                 }
                 Err(err) => {
                     self.unlock_batch(l1_batch_number, request.tee_type).await?;
@@ -156,7 +156,7 @@ impl TeeRequestProcessor {
     async fn lock_batch_for_proving(
         &self,
         tee_type: TeeType,
-        min_batch_number: Option<L1BatchNumber>,
+        min_batch_number: L1BatchNumber,
     ) -> Result<Option<L1BatchNumber>, RequestProcessorError> {
         self.pool
             .connection_tagged("tee_request_processor")

core/node/proof_data_handler/src/tests.rs

@@ -7,7 +7,7 @@ use axum::{
 use serde_json::json;
 use tower::ServiceExt;
 use zksync_basic_types::L2ChainId;
-use zksync_config::configs::ProofDataHandlerConfig;
+use zksync_config::configs::{ProofDataHandlerConfig, TeeConfig};
 use zksync_dal::{ConnectionPool, CoreDal};
 use zksync_object_store::MockObjectStore;
 use zksync_prover_interface::api::SubmitTeeProofRequest;
@@ -25,7 +25,10 @@ async fn request_tee_proof_inputs() {
         ProofDataHandlerConfig {
             http_port: 1337,
             proof_generation_timeout_in_secs: 10,
-            tee_support: true,
+            tee_config: TeeConfig {
+                tee_support: true,
+                first_tee_processed_batch: L1BatchNumber(0),
+            },
         },
         L1BatchCommitmentMode::Rollup,
         L2ChainId::default(),
@@ -80,7 +83,10 @@ async fn submit_tee_proof() {
         ProofDataHandlerConfig {
             http_port: 1337,
             proof_generation_timeout_in_secs: 10,
-            tee_support: true,
+            tee_config: TeeConfig {
+                tee_support: true,
+                first_tee_processed_batch: L1BatchNumber(0),
+            },
         },
         L1BatchCommitmentMode::Rollup,
         L2ChainId::default(),