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

## What ❔

With this commit, the TEE verifier uses `WitnessStorageState` of
`VMRunWitnessInputData` to initialize the storage. This requires waiting
for the BasicWitnessInputProducer to complete and therefore the TEE
verifier input producer can be removed. The input for the TEE verifier
is now collected in the `proof_data_handler`, which enables to remove
the whole job queue for the TEE verifier input producer.

## Why ❔

Previously the storage for VM re-execution was initialized just from
`WitnessInputMerklePaths`. This although misses the storage values for
slots, which are only read/written to by rolled back transactions. This
led to failed verification of blocks, which would normally pass.

## Checklist

<!-- Check your PR fulfills the following items. -->
<!-- For draft PRs check the boxes as you complete them. -->

- [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 `zk_supervisor fmt` and `zk_supervisor
lint`.

---------

Signed-off-by: Harald Hoyer <[email protected]>
Co-authored-by: Patrick Beza <[email protected]>

Cargo.toml

@@ -35,7 +35,6 @@ members = [
   "core/node/consensus",
   "core/node/contract_verification_server",
   "core/node/api_server",
-  "core/node/tee_verifier_input_producer",
   "core/node/base_token_adjuster",
   "core/node/external_proof_integration_api",
   "core/node/logs_bloom_backfill",
@@ -309,6 +308,5 @@ zksync_node_storage_init = { version = "0.1.0", path = "core/node/node_storage_i
 zksync_node_consensus = { version = "0.1.0", path = "core/node/consensus" }
 zksync_contract_verification_server = { version = "0.1.0", path = "core/node/contract_verification_server" }
 zksync_node_api_server = { version = "0.1.0", path = "core/node/api_server" }
-zksync_tee_verifier_input_producer = { version = "0.1.0", path = "core/node/tee_verifier_input_producer" }
 zksync_base_token_adjuster = { version = "0.1.0", path = "core/node/base_token_adjuster" }
 zksync_logs_bloom_backfill = { version = "0.1.0", path = "core/node/logs_bloom_backfill" }

core/bin/zksync_server/src/main.rs

@@ -45,7 +45,7 @@ struct Cli {
     /// Comma-separated list of components to launch.
     #[arg(
         long,
-        default_value = "api,tree,eth,state_keeper,housekeeper,tee_verifier_input_producer,commitment_generator,da_dispatcher,vm_runner_protective_reads"
+        default_value = "api,tree,eth,state_keeper,housekeeper,commitment_generator,da_dispatcher,vm_runner_protective_reads"
     )]
     components: ComponentsToRun,
     /// Path to the yaml config. If set, it will be used instead of env vars.

core/bin/zksync_server/src/node_builder.rs

@@ -55,7 +55,6 @@ use zksync_node_framework::{
             main_batch_executor::MainBatchExecutorLayer, mempool_io::MempoolIOLayer,
             output_handler::OutputHandlerLayer, RocksdbStorageOptions, StateKeeperLayer,
         },
-        tee_verifier_input_producer::TeeVerifierInputProducerLayer,
         vm_runner::{
             bwip::BasicWitnessInputProducerLayer, playground::VmPlaygroundLayer,
             protective_reads::ProtectiveReadsWriterLayer,
@@ -288,6 +287,7 @@ impl MainNodeBuilder {
         self.node.add_layer(ProofDataHandlerLayer::new(
             try_load_config!(self.configs.proof_data_handler_config),
             self.genesis_config.l1_batch_commit_data_generator_mode,
+            self.genesis_config.l2_chain_id,
         ));
         Ok(self)
     }
@@ -493,14 +493,6 @@ impl MainNodeBuilder {
         Ok(self)
     }
 
-    fn add_tee_verifier_input_producer_layer(mut self) -> anyhow::Result<Self> {
-        self.node.add_layer(TeeVerifierInputProducerLayer::new(
-            self.genesis_config.l2_chain_id,
-        ));
-
-        Ok(self)
-    }
-
     fn add_da_client_layer(mut self) -> anyhow::Result<Self> {
         let Some(da_client_config) = self.configs.da_client_config.clone() else {
             tracing::warn!("No config for DA client, using the NoDA client");
@@ -727,9 +719,6 @@ impl MainNodeBuilder {
                 Component::EthTxManager => {
                     self = self.add_eth_tx_manager_layer()?;
                 }
-                Component::TeeVerifierInputProducer => {
-                    self = self.add_tee_verifier_input_producer_layer()?;
-                }
                 Component::Housekeeper => {
                     self = self
                         .add_house_keeper_layer()?

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