feat(tee): TEE Prover Gateway

Cargo.toml

@@ -10,6 +10,7 @@ members = [
     "core/bin/verified_sources_fetcher",
     "core/bin/zksync_server",
     "core/bin/genesis_generator",
+    "core/bin/zksync_tee_prover",
     # Node services
     "core/node/node_framework",
     "core/node/proof_data_handler",

core/bin/zksync_tee_prover/Cargo.toml

@@ -0,0 +1,29 @@
+[package]
+name = "zksync_tee_prover"
+version.workspace = true
+edition.workspace = true
+authors.workspace = true
+homepage.workspace = true
+repository.workspace = true
+license.workspace = true
+keywords.workspace = true
+categories.workspace = true
+
+[dependencies]
+anyhow.workspace = true
+async-trait.workspace = true
+reqwest.workspace = true
+secp256k1.workspace = true
+serde = { workspace = true, features = ["derive"] }
+thiserror.workspace = true
+tokio = { workspace = true, features = ["full"] }
+tracing.workspace = true
+url.workspace = true
+zksync_basic_types.workspace = true
+zksync_config.workspace = true
+zksync_env_config.workspace = true
+zksync_node_framework.workspace = true
+zksync_prover_interface.workspace = true
+zksync_tee_verifier.workspace = true
+zksync_types.workspace = true
+zksync_vlog.workspace = true

core/bin/zksync_tee_prover/src/api_client.rs

@@ -0,0 +1,111 @@
+use reqwest::Client;
+use secp256k1::{ecdsa::Signature, PublicKey};
+use serde::{de::DeserializeOwned, Serialize};
+use url::Url;
+use zksync_basic_types::H256;
+use zksync_prover_interface::{
+    api::{
+        RegisterTeeAttestationRequest, RegisterTeeAttestationResponse, SubmitTeeProofRequest,
+        SubmitTeeProofResponse, TeeProofGenerationDataRequest, TeeProofGenerationDataResponse,
+    },
+    inputs::TeeVerifierInput,
+    outputs::L1BatchTeeProofForL1,
+};
+use zksync_types::{tee_types::TeeType, L1BatchNumber};
+
+use crate::error::TeeProverError;
+
+/// Implementation of the API client for the proof data handler, run by
+/// [`zksync_proof_data_handler::run_server`].
+#[derive(Debug)]
+pub(crate) struct TeeApiClient {
+    api_base_url: Url,
+    http_client: Client,
+}
+
+impl TeeApiClient {
+    pub fn new(api_base_url: Url) -> Self {
+        TeeApiClient {
+            api_base_url,
+            http_client: Client::new(),
+        }
+    }
+
+    async fn post<Req, Resp, S>(&self, endpoint: S, request: Req) -> Result<Resp, reqwest::Error>
+    where
+        Req: Serialize + std::fmt::Debug,
+        Resp: DeserializeOwned,
+        S: AsRef<str>,
+    {
+        let url = self.api_base_url.join(endpoint.as_ref()).unwrap();
+
+        tracing::trace!("Sending POST request to {}: {:?}", url, request);
+
+        self.http_client
+            .post(url)
+            .json(&request)
+            .send()
+            .await?
+            .error_for_status()?
+            .json::<Resp>()
+            .await
+    }
+
+    /// Registers the attestation quote with the TEE prover interface API, effectively proving that
+    /// the private key associated with the given public key was used to sign the root hash within a
+    /// trusted execution environment.
+    pub async fn register_attestation(
+        &self,
+        attestation_quote_bytes: Vec<u8>,
+        public_key: &PublicKey,
+    ) -> Result<(), TeeProverError> {
+        let request = RegisterTeeAttestationRequest {
+            attestation: attestation_quote_bytes,
+            pubkey: public_key.serialize().to_vec(),
+        };
+        self.post::<_, RegisterTeeAttestationResponse, _>("/tee/register_attestation", request)
+            .await?;
+        tracing::info!(
+            "Attestation quote was successfully registered for the public key {}",
+            public_key
+        );
+        Ok(())
+    }
+
+    /// Fetches the next job for the TEE prover to process, verifying and signing it if the
+    /// verification is successful.
+    pub async fn get_job(&self) -> Result<Option<Box<TeeVerifierInput>>, TeeProverError> {
+        let request = TeeProofGenerationDataRequest {};
+        let response = self
+            .post::<_, TeeProofGenerationDataResponse, _>("/tee/proof_inputs", request)
+            .await?;
+        Ok(response.0)
+    }
+
+    /// Submits the successfully verified proof to the TEE prover interface API.
+    pub async fn submit_proof(
+        &self,
+        batch_number: L1BatchNumber,
+        signature: Signature,
+        pubkey: &PublicKey,
+        root_hash: H256,
+        tee_type: TeeType,
+    ) -> Result<(), TeeProverError> {
+        let request = SubmitTeeProofRequest(Box::new(L1BatchTeeProofForL1 {
+            signature: signature.serialize_compact().into(),
+            pubkey: pubkey.serialize().into(),
+            proof: root_hash.as_bytes().into(),
+            tee_type,
+        }));
+        self.post::<_, SubmitTeeProofResponse, _>(
+            format!("/tee/submit_proofs/{batch_number}").as_str(),
+            request,
+        )
+        .await?;
+        tracing::info!(
+            "Proof submitted successfully for batch number {}",
+            batch_number
+        );
+        Ok(())
+    }
+}

core/bin/zksync_tee_prover/src/config.rs

@@ -0,0 +1,39 @@
+use std::path::PathBuf;
+
+use secp256k1::SecretKey;
+use url::Url;
+use zksync_env_config::FromEnv;
+use zksync_types::tee_types::TeeType;
+
+/// Configuration for the TEE prover.
+#[derive(Debug)]
+pub(crate) struct TeeProverConfig {
+    /// The private key used to sign the proofs.
+    pub signing_key: SecretKey,
+    /// The path to the file containing the TEE quote.
+    pub attestation_quote_file_path: PathBuf,
+    /// Attestation quote file.
+    pub tee_type: TeeType,
+    /// TEE proof data handler API.
+    pub api_url: Url,
+}
+
+impl FromEnv for TeeProverConfig {
+    /// Constructs the TEE Prover configuration from environment variables.
+    ///
+    /// Example usage of environment variables for tests:
+    /// ```
+    /// export TEE_SIGNING_KEY="b50b38c8d396c88728fc032ece558ebda96907a0b1a9340289715eef7bf29deb"
+    /// export TEE_QUOTE_FILE="/tmp/test"  # run `echo test > /tmp/test` beforehand
+    /// export TEE_TYPE="sgx"
+    /// export TEE_API_URL="http://127.0.0.1:3320"
+    /// ```
+    fn from_env() -> anyhow::Result<Self> {
+        Ok(Self {
+            signing_key: std::env::var("TEE_SIGNING_KEY")?.parse()?,
+            attestation_quote_file_path: std::env::var("TEE_QUOTE_FILE")?.parse()?,
+            tee_type: std::env::var("TEE_TYPE")?.parse()?,
+            api_url: std::env::var("TEE_API_URL")?.parse()?,
+        })
+    }
+}

core/bin/zksync_tee_prover/src/error.rs

@@ -0,0 +1,47 @@
+use std::{error::Error as StdError, io};
+
+use reqwest::StatusCode;
+
+#[derive(Debug, thiserror::Error)]
+pub(crate) enum TeeProverError {
+    #[error(transparent)]
+    Request(#[from] reqwest::Error),
+    #[error(transparent)]
+    Verification(anyhow::Error),
+}
+
+impl TeeProverError {
+    pub fn is_transient(&self) -> bool {
+        match self {
+            Self::Request(err) => is_transient_http_error(err),
+            _ => false,
+        }
+    }
+}
+
+fn is_transient_http_error(err: &reqwest::Error) -> bool {
+    err.is_timeout()
+            || err.is_connect()
+            // Not all request errors are logically transient, but a significant part of them are (e.g.,
+            // `hyper` protocol-level errors), and it's safer to consider an error transient.
+            || err.is_request()
+            || has_transient_io_source(err)
+            || err.status() == Some(StatusCode::BAD_GATEWAY)
+            || err.status() == Some(StatusCode::SERVICE_UNAVAILABLE)
+}
+
+fn has_transient_io_source(err: &(dyn StdError + 'static)) -> bool {
+    // We treat any I/O errors as transient. This isn't always true, but frequently occurring I/O errors
+    // (e.g., "connection reset by peer") *are* transient, and treating an error as transient is a safer option,
+    // even if it can lead to unnecessary retries.
+    get_source::<io::Error>(err).is_some()
+}
+
+fn get_source<'a, T: StdError + 'static>(mut err: &'a (dyn StdError + 'static)) -> Option<&'a T> {
+    loop {
+        if let Some(err) = err.downcast_ref::<T>() {
+            return Some(err);
+        }
+        err = err.source()?;
+    }
+}