From ab6e56552377ea74ccf619f940a03735c0e4d5ea Mon Sep 17 00:00:00 2001 From: Patrick Beza Date: Thu, 10 Oct 2024 18:17:58 +0200 Subject: [PATCH] fix(tee-prover): simplify TeeProofGenerationDataResponse This PR addresses Alex's code review comment: https://github.com/matter-labs/zksync-era/pull/3017#discussion_r1795237929 and partially addresses this one too: https://github.com/matter-labs/zksync-era/pull/3017#discussion_r1795240750 --- core/bin/zksync_tee_prover/src/api_client.rs | 2 +- core/bin/zksync_tee_prover/src/tee_prover.rs | 18 ++++++++++++------ core/lib/dal/src/tee_proof_generation_dal.rs | 9 ++++----- core/lib/prover_interface/src/api.rs | 2 +- core/node/proof_data_handler/src/lib.rs | 3 +-- .../src/tee_request_processor.rs | 12 ++++-------- core/node/proof_data_handler/src/tests.rs | 2 +- 7 files changed, 24 insertions(+), 24 deletions(-) diff --git a/core/bin/zksync_tee_prover/src/api_client.rs b/core/bin/zksync_tee_prover/src/api_client.rs index 13fbc1ba886..860c102d562 100644 --- a/core/bin/zksync_tee_prover/src/api_client.rs +++ b/core/bin/zksync_tee_prover/src/api_client.rs @@ -77,7 +77,7 @@ impl TeeApiClient { pub async fn get_job( &self, tee_type: TeeType, - ) -> Result>, TeeProverError> { + ) -> Result, TeeProverError> { let request = TeeProofGenerationDataRequest { tee_type }; let response = self .post::<_, TeeProofGenerationDataResponse, _>("/tee/proof_inputs", request) diff --git a/core/bin/zksync_tee_prover/src/tee_prover.rs b/core/bin/zksync_tee_prover/src/tee_prover.rs index 1511f0c88e3..7d40bd32431 100644 --- a/core/bin/zksync_tee_prover/src/tee_prover.rs +++ b/core/bin/zksync_tee_prover/src/tee_prover.rs @@ -1,5 +1,6 @@ use std::fmt; +use reqwest::StatusCode; use secp256k1::{ecdsa::Signature, Message, PublicKey, Secp256k1}; use zksync_basic_types::H256; use zksync_node_framework::{ @@ -90,8 +91,8 @@ impl TeeProver { } async fn step(&self, public_key: &PublicKey) -> Result, TeeProverError> { - match self.api_client.get_job(self.config.tee_type).await? { - Some(job) => { + match self.api_client.get_job(self.config.tee_type).await { + Ok(job) => { let (signature, batch_number, root_hash) = self.verify(*job)?; self.api_client .submit_proof( @@ -104,10 +105,15 @@ impl TeeProver { .await?; Ok(Some(batch_number)) } - None => { - tracing::trace!("There are currently no pending batches to be proven"); - Ok(None) - } + Err(err) => match err { + TeeProverError::Request(req_err) + if req_err.status() == Some(StatusCode::NOT_FOUND) => + { + tracing::trace!("There are currently no pending batches to be proven"); + Ok(None) + } + _ => Err(err), + }, } } } diff --git a/core/lib/dal/src/tee_proof_generation_dal.rs b/core/lib/dal/src/tee_proof_generation_dal.rs index 8cb961e40e3..192b81db1bf 100644 --- a/core/lib/dal/src/tee_proof_generation_dal.rs +++ b/core/lib/dal/src/tee_proof_generation_dal.rs @@ -33,7 +33,7 @@ impl TeeProofGenerationDal<'_, '_> { tee_type: TeeType, processing_timeout: Duration, min_batch_number: Option, - ) -> DalResult> { + ) -> DalResult { 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!( @@ -105,11 +105,10 @@ impl TeeProofGenerationDal<'_, '_> { .with_arg("tee_type", &tee_type) .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)); + .fetch_one(self.storage) + .await?; - Ok(batch_number) + Ok(L1BatchNumber(batch_number.l1_batch_number as u32)) } pub async fn unlock_batch( diff --git a/core/lib/prover_interface/src/api.rs b/core/lib/prover_interface/src/api.rs index 776cd3141cb..acf104cc4c6 100644 --- a/core/lib/prover_interface/src/api.rs +++ b/core/lib/prover_interface/src/api.rs @@ -31,7 +31,7 @@ pub enum ProofGenerationDataResponse { } #[derive(Debug, Serialize, Deserialize)] -pub struct TeeProofGenerationDataResponse(pub Option>); +pub struct TeeProofGenerationDataResponse(pub Box); #[derive(Debug, Serialize, Deserialize)] pub enum SubmitProofResponse { diff --git a/core/node/proof_data_handler/src/lib.rs b/core/node/proof_data_handler/src/lib.rs index 0f90d657be5..2253cf9df97 100644 --- a/core/node/proof_data_handler/src/lib.rs +++ b/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}; @@ -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(), } diff --git a/core/node/proof_data_handler/src/tee_request_processor.rs b/core/node/proof_data_handler/src/tee_request_processor.rs index 5956be17e97..8af5e08b475 100644 --- a/core/node/proof_data_handler/src/tee_request_processor.rs +++ b/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 { @@ -156,7 +152,7 @@ impl TeeRequestProcessor { &self, tee_type: TeeType, min_batch_number: Option, - ) -> Result, RequestProcessorError> { + ) -> Result { let result = self .pool .connection() diff --git a/core/node/proof_data_handler/src/tests.rs b/core/node/proof_data_handler/src/tests.rs index a10044cacd9..240940506f8 100644 --- a/core/node/proof_data_handler/src/tests.rs +++ b/core/node/proof_data_handler/src/tests.rs @@ -31,7 +31,7 @@ async fn request_tee_proof_inputs() { L2ChainId::default(), ); let test_cases = vec![ - (json!({ "tee_type": "sgx" }), StatusCode::NO_CONTENT), + (json!({ "tee_type": "sgx" }), StatusCode::NOT_FOUND), ( json!({ "tee_type": "Sgx" }), StatusCode::UNPROCESSABLE_ENTITY,