fix(tee-prover): simplify TeeProofGenerationDataResponse

core/bin/zksync_tee_prover/src/api_client.rs

@@ -1,14 +1,13 @@
-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,
+        RegisterTeeAttestationRequest, SubmitTeeProofRequest, TeeProofGenerationDataRequest,
+        TeeProofGenerationDataResponse,
     },
-    inputs::TeeVerifierInput,
     outputs::L1BatchTeeProofForL1,
 };
 use zksync_types::{tee_types::TeeType, L1BatchNumber};
@@ -31,10 +30,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 +44,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 +59,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 +72,18 @@ 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)
+        let response = self.post("/tee/proof_inputs", request).await?;
+
+        match response.status() {
+            StatusCode::OK => Ok(response.json::<TeeProofGenerationDataResponse>().await?),
+            StatusCode::NO_CONTENT => Ok(TeeProofGenerationDataResponse::VerifierInputNotReady),
+            _ => response
+                .json::<TeeProofGenerationDataResponse>()
+                .await
+                .map_err(TeeProverError::Request),
+        }
     }
 
     /// Submits the successfully verified proof to the TEE prover interface API.
@@ -101,7 +102,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/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

core/bin/zksync_tee_prover/src/tee_prover.rs

@@ -8,7 +8,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 +91,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 +104,7 @@ impl TeeProver {
                     .await?;
                 Ok(Some(batch_number))
             }
-            None => {
+            TeeProofGenerationDataResponse::VerifierInputNotReady => {
                 tracing::trace!("There are currently no pending batches to be proven");
                 Ok(None)
             }

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 {

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(),
                     }
                 },

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 {