fix(tee-prover): simplify TeeProofGenerationDataResponse

This PR addresses Alex's code review comment: #3017 (comment) and partially addresses this one too: #3017 (comment)
matter-labs · Oct 11, 2024 · 7698c20 · 7698c20
1 parent 3d77b90
commit 7698c20
Show file tree

Hide file tree

Showing 6 changed files with 23 additions and 11 deletions.
diff --git a/core/bin/zksync_tee_prover/src/api_client.rs b/core/bin/zksync_tee_prover/src/api_client.rs
@@ -77,12 +77,12 @@ impl TeeApiClient {
     pub async fn get_job(
         &self,
         tee_type: TeeType,
-    ) -> Result<Option<Box<TeeVerifierInput>>, TeeProverError> {
+    ) -> Result<TeeProofGenerationDataResponse, TeeProverError> {
         let request = TeeProofGenerationDataRequest { tee_type };
         let response = self
             .post::<_, TeeProofGenerationDataResponse, _>("/tee/proof_inputs", request)
             .await?;
-        Ok(response.0)
+        Ok(response)
     }
 
     /// Submits the successfully verified proof to the TEE prover interface API.

diff --git a/core/bin/zksync_tee_prover/src/config.rs b/core/bin/zksync_tee_prover/src/config.rs
@@ -47,7 +47,7 @@ impl FromEnv for TeeProverConfig {
     /// export TEE_PROVER_SIGNING_KEY="b50b38c8d396c88728fc032ece558ebda96907a0b1a9340289715eef7bf29deb"
     /// 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_API_URL="http://127.0.0.1:3421"
     /// export TEE_PROVER_MAX_RETRIES=10
     /// export TEE_PROVER_INITIAL_RETRY_BACKOFF_SEC=1
     /// export TEE_PROVER_RETRY_BACKOFF_MULTIPLIER=2.0

diff --git 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::{
@@ -8,7 +9,7 @@ use zksync_node_framework::{
     wiring_layer::{WiringError, WiringLayer},
     IntoContext,
 };
-use zksync_prover_interface::inputs::TeeVerifierInput;
+use zksync_prover_interface::{api::TeeProofGenerationDataResponse, inputs::TeeVerifierInput};
 use zksync_tee_verifier::Verify;
 use zksync_types::L1BatchNumber;
 
@@ -91,7 +92,7 @@ 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) => {
+            TeeProofGenerationDataResponse::VerifierInputReady(job) => {
                 let (signature, batch_number, root_hash) = self.verify(*job)?;
                 self.api_client
                     .submit_proof(
@@ -104,7 +105,7 @@ impl TeeProver {
                     .await?;
                 Ok(Some(batch_number))
             }
-            None => {
+            TeeProofGenerationDataResponse::VerifierInputNotReady => {
                 tracing::trace!("There are currently no pending batches to be proven");
                 Ok(None)
             }

diff --git a/core/lib/prover_interface/src/api.rs b/core/lib/prover_interface/src/api.rs
@@ -31,7 +31,10 @@ pub enum ProofGenerationDataResponse {
 }
 
 #[derive(Debug, Serialize, Deserialize)]
-pub struct TeeProofGenerationDataResponse(pub Option<Box<TeeVerifierInput>>);
+pub enum TeeProofGenerationDataResponse {
+    VerifierInputReady(Box<TeeVerifierInput>),
+    VerifierInputNotReady,
+}
 
 #[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
@@ -109,8 +109,14 @@ 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(),
+                        Ok(data) => match data {
+                            Json(TeeProofGenerationDataResponse::VerifierInputReady(input)) => {
+                                (StatusCode::OK, Json(TeeProofGenerationDataResponse::VerifierInputReady(input))).into_response()
+                            }
+                            Json(TeeProofGenerationDataResponse::VerifierInputNotReady) => {
+                                (StatusCode::NO_CONTENT, Json(TeeProofGenerationDataResponse::VerifierInputNotReady)).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
@@ -56,15 +56,17 @@ impl TeeRequestProcessor {
                 .await?
             {
                 Some(number) => number,
-                None => break Ok(Json(TeeProofGenerationDataResponse(None))),
+                None => break Ok(Json(TeeProofGenerationDataResponse::VerifierInputNotReady)),
             };
 
             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::VerifierInputReady(
+                        Box::new(input),
+                    )));
                 }
                 Err(RequestProcessorError::ObjectStore(ObjectStoreError::KeyNotFound(_))) => {
                     missing_range = match missing_range {