fix(tee_verifier): correctly initialize storage for re-execution

core/bin/zksync_tee_prover/src/api_client.rs

@@ -1,13 +1,10 @@
-use reqwest::Client;
+use reqwest::{Client, Response, StatusCode};
 use secp256k1::{ecdsa::Signature, PublicKey};
-use serde::{de::DeserializeOwned, Serialize};
+use serde::Serialize;
 use url::Url;
 use zksync_basic_types::H256;
 use zksync_prover_interface::{
-    api::{
-        RegisterTeeAttestationRequest, RegisterTeeAttestationResponse, SubmitTeeProofRequest,
-        SubmitTeeProofResponse, TeeProofGenerationDataRequest, TeeProofGenerationDataResponse,
-    },
+    api::{RegisterTeeAttestationRequest, SubmitTeeProofRequest, TeeProofGenerationDataRequest},
     inputs::TeeVerifierInput,
     outputs::L1BatchTeeProofForL1,
 };
@@ -31,10 +28,9 @@ impl TeeApiClient {
         }
     }
 
-    async fn post<Req, Resp, S>(&self, endpoint: S, request: Req) -> Result<Resp, reqwest::Error>
+    async fn post<Req, S>(&self, endpoint: S, request: Req) -> Result<Response, reqwest::Error>
     where
         Req: Serialize + std::fmt::Debug,
-        Resp: DeserializeOwned,
         S: AsRef<str>,
     {
         let url = self.api_base_url.join(endpoint.as_ref()).unwrap();
@@ -46,9 +42,7 @@ impl TeeApiClient {
             .json(&request)
             .send()
             .await?
-            .error_for_status()?
-            .json::<Resp>()
-            .await
+            .error_for_status()
     }
 
     /// Registers the attestation quote with the TEE prover interface API, effectively proving that
@@ -63,8 +57,7 @@ impl TeeApiClient {
             attestation: attestation_quote_bytes,
             pubkey: public_key.serialize().to_vec(),
         };
-        self.post::<_, RegisterTeeAttestationResponse, _>("/tee/register_attestation", request)
-            .await?;
+        self.post("/tee/register_attestation", request).await?;
         tracing::info!(
             "Attestation quote was successfully registered for the public key {}",
             public_key
@@ -77,12 +70,17 @@ impl TeeApiClient {
     pub async fn get_job(
         &self,
         tee_type: TeeType,
-    ) -> Result<Option<Box<TeeVerifierInput>>, TeeProverError> {
+    ) -> Result<Option<TeeVerifierInput>, TeeProverError> {
         let request = TeeProofGenerationDataRequest { tee_type };
-        let response = self
-            .post::<_, TeeProofGenerationDataResponse, _>("/tee/proof_inputs", request)
-            .await?;
-        Ok(response.0)
+        let response = self.post("/tee/proof_inputs", request).await?;
+        match response.status() {
+            StatusCode::OK => Ok(Some(response.json::<TeeVerifierInput>().await?)),
+            StatusCode::NO_CONTENT => Ok(None),
+            _ => response
+                .json::<Option<TeeVerifierInput>>()
+                .await
+                .map_err(TeeProverError::Request),
+        }
     }
 
     /// Submits the successfully verified proof to the TEE prover interface API.
@@ -101,7 +99,7 @@ impl TeeApiClient {
             tee_type,
         }));
         let observer = METRICS.proof_submitting_time.start();
-        self.post::<_, SubmitTeeProofResponse, _>(
+        self.post(
             format!("/tee/submit_proofs/{batch_number}").as_str(),
             request,
         )

core/bin/zksync_tee_prover/src/tee_prover.rs

@@ -90,9 +90,9 @@ impl TeeProver {
     }
 
     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)?;
+        match self.api_client.get_job(self.config.tee_type).await {
+            Ok(Some(job)) => {
+                let (signature, batch_number, root_hash) = self.verify(job)?;
                 self.api_client
                     .submit_proof(
                         batch_number,
@@ -104,10 +104,11 @@ impl TeeProver {
                     .await?;
                 Ok(Some(batch_number))
             }
-            None => {
+            Ok(None) => {
                 tracing::trace!("There are currently no pending batches to be proven");
                 Ok(None)
             }
+            Err(err) => Err(err),
         }
     }
 }

core/lib/dal/src/tee_proof_generation_dal.rs

@@ -36,7 +36,7 @@ impl TeeProofGenerationDal<'_, '_> {
     ) -> 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 batch_number = sqlx::query!(
+        sqlx::query!(
             r#"
             WITH upsert AS (
                 SELECT
@@ -106,10 +106,8 @@ impl TeeProofGenerationDal<'_, '_> {
         .with_arg("processing_timeout", &processing_timeout)
         .with_arg("l1_batch_number", &min_batch_number)
         .fetch_optional(self.storage)
-        .await?
-        .map(|row| L1BatchNumber(row.l1_batch_number as u32));
-
-        Ok(batch_number)
+        .await
+        .map(|record| record.map(|record| L1BatchNumber(record.l1_batch_number as u32)))
     }
 
     pub async fn unlock_batch(

core/lib/prover_interface/src/api.rs

@@ -31,7 +31,7 @@ pub enum ProofGenerationDataResponse {
 }
 
 #[derive(Debug, Serialize, Deserialize)]
-pub struct TeeProofGenerationDataResponse(pub Option<Box<TeeVerifierInput>>);
+pub struct TeeProofGenerationDataResponse(pub Box<TeeVerifierInput>);
 
 #[derive(Debug, Serialize, Deserialize)]
 pub enum SubmitProofResponse {

core/lib/tee_verifier/src/lib.rs

@@ -216,16 +216,16 @@ fn map_log_tree(
     idx: &mut u64,
 ) -> anyhow::Result<TreeInstruction> {
     let key = storage_log.key.hashed_key_u256();
-    let tree_instruction = match (storage_log.is_write(), storage_log.kind, *tree_log_entry) {
-        (true, _, TreeLogEntry::Updated { leaf_index, .. }) => {
+    let tree_instruction = match (storage_log.is_write(), *tree_log_entry) {
+        (true, TreeLogEntry::Updated { leaf_index, .. }) => {
             TreeInstruction::write(key, leaf_index, H256(storage_log.value.into()))
         }
-        (true, _, TreeLogEntry::Inserted) => {
+        (true, TreeLogEntry::Inserted) => {
             let leaf_index = *idx;
             *idx += 1;
             TreeInstruction::write(key, leaf_index, H256(storage_log.value.into()))
         }
-        (false, _, TreeLogEntry::Read { value, .. }) => {
+        (false, TreeLogEntry::Read { value, .. }) => {
             if storage_log.value != value {
                 tracing::error!(
                     ?storage_log,
@@ -238,11 +238,11 @@ fn map_log_tree(
             }
             TreeInstruction::Read(key)
         }
-        (false, _, TreeLogEntry::ReadMissingKey { .. }) => TreeInstruction::Read(key),
-        (true, _, TreeLogEntry::Read { .. })
-        | (true, _, TreeLogEntry::ReadMissingKey)
-        | (false, _, TreeLogEntry::Inserted)
-        | (false, _, TreeLogEntry::Updated { .. }) => {
+        (false, TreeLogEntry::ReadMissingKey { .. }) => TreeInstruction::Read(key),
+        (true, TreeLogEntry::Read { .. })
+        | (true, TreeLogEntry::ReadMissingKey)
+        | (false, TreeLogEntry::Inserted)
+        | (false, TreeLogEntry::Updated { .. }) => {
             tracing::error!(
                 ?storage_log,
                 ?tree_log_entry,

core/node/proof_data_handler/src/errors.rs

@@ -6,6 +6,7 @@ use zksync_dal::DalError;
 use zksync_object_store::ObjectStoreError;
 
 pub(crate) enum RequestProcessorError {
+    NoJob,
     GeneralError(String),
     ObjectStore(ObjectStoreError),
     Dal(DalError),
@@ -19,34 +20,35 @@ impl From<DalError> for RequestProcessorError {
 
 impl IntoResponse for RequestProcessorError {
     fn into_response(self) -> Response {
-        let (status_code, message) = match self {
+        match self {
             RequestProcessorError::GeneralError(err) => {
                 tracing::error!("Error: {:?}", err);
                 (
                     StatusCode::INTERNAL_SERVER_ERROR,
                     "An internal error occurred".to_owned(),
                 )
+                    .into_response()
             }
             RequestProcessorError::ObjectStore(err) => {
                 tracing::error!("GCS error: {:?}", err);
                 (
                     StatusCode::BAD_GATEWAY,
                     "Failed fetching/saving from GCS".to_owned(),
                 )
+                    .into_response()
             }
             RequestProcessorError::Dal(err) => {
                 tracing::error!("Sqlx error: {:?}", err);
-                match err.inner() {
-                    zksync_dal::SqlxError::RowNotFound => {
-                        (StatusCode::NOT_FOUND, "Non existing L1 batch".to_owned())
-                    }
-                    _ => (
-                        StatusCode::BAD_GATEWAY,
-                        "Failed fetching/saving from db".to_owned(),
-                    ),
-                }
+                (
+                    StatusCode::BAD_GATEWAY,
+                    "Failed fetching/saving from db".to_owned(),
+                )
+                    .into_response()
+            }
+            RequestProcessorError::NoJob => {
+                tracing::trace!("No job found");
+                (StatusCode::NO_CONTENT, ()).into_response()
             }
-        };
-        (status_code, message).into_response()
+        }
     }
 }

core/node/proof_data_handler/src/lib.rs

@@ -10,7 +10,7 @@ use zksync_dal::{ConnectionPool, Core};
 use zksync_object_store::ObjectStore;
 use zksync_prover_interface::api::{
     ProofGenerationDataRequest, RegisterTeeAttestationRequest, SubmitProofRequest,
-    SubmitTeeProofRequest, TeeProofGenerationDataRequest, TeeProofGenerationDataResponse,
+    SubmitTeeProofRequest, TeeProofGenerationDataRequest,
 };
 use zksync_types::{commitment::L1BatchCommitmentMode, L2ChainId};
 
@@ -31,7 +31,7 @@ pub async fn run_server(
     mut stop_receiver: watch::Receiver<bool>,
 ) -> anyhow::Result<()> {
     let bind_address = SocketAddr::from(([0, 0, 0, 0], config.http_port));
-    tracing::debug!("Starting proof data handler server on {bind_address}");
+    tracing::info!("Starting proof data handler server on {bind_address}");
     let app = create_proof_processing_router(
         blob_store,
         connection_pool,
@@ -109,7 +109,6 @@ fn create_proof_processing_router(
                         .await;
 
                     match result {
-                        Ok(Json(TeeProofGenerationDataResponse(None))) => (StatusCode::NO_CONTENT, Json("No new TeeVerifierInputs are available yet")).into_response(),
                         Ok(data) => (StatusCode::OK, data).into_response(),
                         Err(e) => e.into_response(),
                     }

core/node/proof_data_handler/src/tee_request_processor.rs

@@ -51,20 +51,16 @@ impl TeeRequestProcessor {
         let mut missing_range: Option<(L1BatchNumber, L1BatchNumber)> = None;
 
         let result = loop {
-            let l1_batch_number = match self
+            let l1_batch_number = self
                 .lock_batch_for_proving(request.tee_type, min_batch_number)
-                .await?
-            {
-                Some(number) => number,
-                None => break Ok(Json(TeeProofGenerationDataResponse(None))),
-            };
+                .await?;
 
             match self
                 .tee_verifier_input_for_existing_batch(l1_batch_number)
                 .await
             {
                 Ok(input) => {
-                    break Ok(Json(TeeProofGenerationDataResponse(Some(Box::new(input)))));
+                    break Ok(Json(TeeProofGenerationDataResponse(Box::new(input))));
                 }
                 Err(RequestProcessorError::ObjectStore(ObjectStoreError::KeyNotFound(_))) => {
                     missing_range = match missing_range {
@@ -111,7 +107,7 @@ impl TeeRequestProcessor {
 
         let mut connection = self
             .pool
-            .connection()
+            .connection_tagged("tee_request_processor")
             .await
             .map_err(RequestProcessorError::Dal)?;
 
@@ -156,19 +152,19 @@ impl TeeRequestProcessor {
         &self,
         tee_type: TeeType,
         min_batch_number: Option<L1BatchNumber>,
-    ) -> Result<Option<L1BatchNumber>, RequestProcessorError> {
-        let result = self
-            .pool
-            .connection()
+    ) -> Result<L1BatchNumber, RequestProcessorError> {
+        self.pool
+            .connection_tagged("tee_request_processor")
             .await?
             .tee_proof_generation_dal()
             .lock_batch_for_proving(
                 tee_type,
                 self.config.proof_generation_timeout(),
                 min_batch_number,
             )
-            .await?;
-        Ok(result)
+            .await
+            .map_err(RequestProcessorError::Dal)
+            .and_then(|batch_number| batch_number.ok_or(RequestProcessorError::NoJob))
     }
 
     async fn unlock_batch(
@@ -177,7 +173,7 @@ impl TeeRequestProcessor {
         tee_type: TeeType,
     ) -> Result<(), RequestProcessorError> {
         self.pool
-            .connection()
+            .connection_tagged("tee_request_processor")
             .await?
             .tee_proof_generation_dal()
             .unlock_batch(l1_batch_number, tee_type)
@@ -191,7 +187,7 @@ impl TeeRequestProcessor {
         Json(proof): Json<SubmitTeeProofRequest>,
     ) -> Result<Json<SubmitProofResponse>, RequestProcessorError> {
         let l1_batch_number = L1BatchNumber(l1_batch_number);
-        let mut connection = self.pool.connection().await?;
+        let mut connection = self.pool.connection_tagged("tee_request_processor").await?;
         let mut dal = connection.tee_proof_generation_dal();
 
         tracing::info!(
@@ -217,7 +213,7 @@ impl TeeRequestProcessor {
     ) -> Result<Json<RegisterTeeAttestationResponse>, RequestProcessorError> {
         tracing::info!("Received attestation: {:?}", payload);
 
-        let mut connection = self.pool.connection().await?;
+        let mut connection = self.pool.connection_tagged("tee_request_processor").await?;
         let mut dal = connection.tee_proof_generation_dal();
 
         dal.save_attestation(&payload.pubkey, &payload.attestation)