feat(tee): introduce get_tee_proofs RPC method for TEE proofs

core/lib/dal/src/tee_proof_generation_dal.rs

@@ -1,5 +1,6 @@
 use std::time::Duration;
 
+use serde::{Deserialize, Serialize};
 use strum::{Display, EnumString};
 use zksync_db_connection::{
     connection::Connection,
@@ -16,6 +17,19 @@ pub struct TeeProofGenerationDal<'a, 'c> {
     pub(crate) storage: &'a mut Connection<'c, Core>,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct TeeProof {
+    // signature generated within the TEE enclave, using the privkey corresponding to the pubkey
+    pub signature: Option<Vec<u8>>,
+    // pubkey used for signature verification; each key pair is attested by the TEE attestation
+    // stored in the db
+    pub pubkey: Option<Vec<u8>>,
+    // data that was signed
+    pub proof: Option<Vec<u8>>,
+    // type of TEE used for attestation
+    pub tee_type: Option<TeeType>,
+}
+
 #[derive(Debug, EnumString, Display)]
 enum TeeProofGenerationJobStatus {
     #[strum(serialize = "ready_to_be_proven")]
@@ -34,7 +48,7 @@ impl TeeProofGenerationDal<'_, '_> {
         processing_timeout: Duration,
     ) -> DalResult<Option<L1BatchNumber>> {
         let processing_timeout = pg_interval_from_duration(processing_timeout);
-        let result: Option<L1BatchNumber> = sqlx::query!(
+        let query = sqlx::query!(
             r#"
             UPDATE tee_proof_generation_details
             SET
@@ -68,13 +82,15 @@ impl TeeProofGenerationDal<'_, '_> {
                 tee_proof_generation_details.l1_batch_number
             "#,
             &processing_timeout,
-        )
-        .fetch_optional(self.storage.conn())
-        .await
-        .unwrap()
-        .map(|row| L1BatchNumber(row.l1_batch_number as u32));
+        );
+        let batch_number = Instrumented::new("get_next_block_to_be_proven")
+            .with_arg("processing_timeout", &processing_timeout)
+            .with(query)
+            .fetch_optional(self.storage)
+            .await?
+            .map(|row| L1BatchNumber(row.l1_batch_number as u32));
 
-        Ok(result)
+        Ok(batch_number)
     }
 
     pub async fn save_proof_artifacts_metadata(
@@ -149,7 +165,7 @@ impl TeeProofGenerationDal<'_, '_> {
     }
 
     pub async fn get_oldest_unpicked_batch(&mut self) -> DalResult<Option<L1BatchNumber>> {
-        let result: Option<L1BatchNumber> = sqlx::query!(
+        let query = sqlx::query!(
             r#"
             SELECT
                 proofs.l1_batch_number
@@ -164,13 +180,14 @@ impl TeeProofGenerationDal<'_, '_> {
             LIMIT
                 1
             "#,
-        )
-        .fetch_optional(self.storage.conn())
-        .await
-        .unwrap()
-        .map(|row| L1BatchNumber(row.l1_batch_number as u32));
+        );
+        let batch_number = Instrumented::new("get_oldest_unpicked_batch")
+            .with(query)
+            .fetch_optional(self.storage)
+            .await?
+            .map(|row| L1BatchNumber(row.l1_batch_number as u32));
 
-        Ok(result)
+        Ok(batch_number)
     }
 
     pub async fn save_attestation(&mut self, pubkey: &[u8], attestation: &[u8]) -> DalResult<()> {
@@ -202,4 +219,39 @@ impl TeeProofGenerationDal<'_, '_> {
 
         Ok(())
     }
+
+    pub async fn get_tee_proof(
+        &mut self,
+        block_number: L1BatchNumber,
+    ) -> DalResult<Option<TeeProof>> {
+        let query = sqlx::query!(
+            r#"
+            SELECT
+                signature,
+                pubkey,
+                proof,
+                tee_type
+            FROM
+                tee_proof_generation_details
+            WHERE
+                l1_batch_number = $1
+            "#,
+            i64::from(block_number.0)
+        );
+        let proof = Instrumented::new("get_tee_proof")
+            .with_arg("l1_batch_number", &block_number)
+            .with(query)
+            .fetch_optional(self.storage)
+            .await?
+            .map(|row| TeeProof {
+                signature: row.signature,
+                pubkey: row.pubkey,
+                proof: row.proof,
+                tee_type: row
+                    .tee_type
+                    .map(|tt| tt.parse::<TeeType>().expect("Invalid TEE type")),
+            });
+
+        Ok(proof)
+    }
 }

core/node/proof_data_handler/src/lib.rs

@@ -1,7 +1,11 @@
 use std::{net::SocketAddr, sync::Arc};
 
 use anyhow::Context as _;
-use axum::{extract::Path, routing::post, Json, Router};
+use axum::{
+    extract::Path,
+    routing::{get, post},
+    Json, Router,
+};
 use request_processor::RequestProcessor;
 use tee_request_processor::TeeRequestProcessor;
 use tokio::sync::watch;
@@ -91,6 +95,7 @@ fn create_proof_processing_router(
             TeeRequestProcessor::new(blob_store, connection_pool, config.clone());
         let submit_tee_proof_processor = get_tee_proof_gen_processor.clone();
         let register_tee_attestation_processor = get_tee_proof_gen_processor.clone();
+        let get_tee_proof_processor = get_tee_proof_gen_processor.clone();
 
         router = router.route(
             "/tee/proof_inputs",
@@ -121,6 +126,15 @@ fn create_proof_processing_router(
                         .await
                 },
             ),
+        )
+        .route("/tee/get_proof/:l1_batch_number",
+            get(
+                move |l1_batch_number: Path<u32>| async move {
+                    get_tee_proof_processor
+                        .get_proof(l1_batch_number)
+                        .await
+                },
+            )
         );
     }
 

core/node/proof_data_handler/src/tee_request_processor.rs

@@ -2,7 +2,7 @@ use std::sync::Arc;
 
 use axum::{extract::Path, Json};
 use zksync_config::configs::ProofDataHandlerConfig;
-use zksync_dal::{ConnectionPool, Core, CoreDal};
+use zksync_dal::{tee_proof_generation_dal::TeeProof, ConnectionPool, Core, CoreDal};
 use zksync_object_store::ObjectStore;
 use zksync_prover_interface::{
     api::{
@@ -118,4 +118,23 @@ impl TeeRequestProcessor {
 
         Ok(Json(RegisterTeeAttestationResponse::Success))
     }
+
+    pub(crate) async fn get_proof(
+        &self,
+        Path(l1_batch_number): Path<u32>,
+    ) -> Result<Json<TeeProof>, RequestProcessorError> {
+        let mut connection = self
+            .pool
+            .connection()
+            .await
+            .map_err(RequestProcessorError::Dal)?;
+        let mut dal = connection.tee_proof_generation_dal();
+        let l1_batch_number = L1BatchNumber(l1_batch_number);
+        let tee_proof = dal
+            .get_tee_proof(l1_batch_number)
+            .await
+            .map_err(RequestProcessorError::Dal)?;
+
+        Ok(Json(tee_proof.unwrap()))
+    }
 }