Restructure db tables for TEE proof generation

Make them more flexible to allow multiple proofs for each
(batch_number, tee_type) pair.

core/bin/zksync_tee_prover/src/api_client.rs

@@ -9,7 +9,7 @@ use zksync_prover_interface::{
         SubmitTeeProofResponse, TeeProofGenerationDataRequest, TeeProofGenerationDataResponse,
     },
     inputs::TeeVerifierInput,
-    outputs::L1BatchTeeProofForL1,
+    L1BatchTeeProofForL1,
 };
 use zksync_types::{tee_types::TeeType, L1BatchNumber};
 
@@ -74,24 +74,29 @@ impl TeeApiClient {
 
     /// 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 {};
+    pub async fn get_job(
+        &self,
+        tee_type: TeeType,
+    ) -> Result<Option<Box<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.
     pub async fn submit_proof(
         &self,
+        proof_id: i64,
         batch_number: L1BatchNumber,
         signature: Signature,
         pubkey: &PublicKey,
         root_hash: H256,
         tee_type: TeeType,
     ) -> Result<(), TeeProverError> {
         let request = SubmitTeeProofRequest(Box::new(L1BatchTeeProofForL1 {
+            proof_id,
             signature: signature.serialize_compact().into(),
             pubkey: pubkey.serialize().into(),
             proof: root_hash.as_bytes().into(),

core/bin/zksync_tee_prover/src/tee_prover.rs

@@ -112,11 +112,12 @@ impl TeeProver {
     }
 
     async fn step(&self) -> Result<Option<L1BatchNumber>, TeeProverError> {
-        match self.api_client.get_job().await? {
+        match self.api_client.get_job(self.tee_type).await? {
             Some(job) => {
-                let (signature, batch_number, root_hash) = self.verify(*job)?;
+                let (signature, batch_number, root_hash) = self.verify(*job.input)?;
                 self.api_client
                     .submit_proof(
+                        job.proof_id.unwrap(), // TODO unwrap() is not safe here
                         batch_number,
                         signature,
                         &self.public_key,

core/lib/dal/migrations/20240325133100_add_tee_verifier_input_producer_jobs_table.down.sql

@@ -1,4 +1,5 @@
 DROP INDEX IF EXISTS idx_tee_verifier_input_producer_jobs_status_processing_attempts;
+DROP INDEX IF EXISTS idx_tee_verifier_input_producer_jobs_l1_batch_number_status;
 
 DROP TABLE IF EXISTS tee_verifier_input_producer_jobs;
 

core/lib/dal/migrations/20240325143100_add_tee_verifier_input_producer_jobs_table.up.sql

@@ -5,7 +5,6 @@ CREATE TABLE IF NOT EXISTS tee_verifier_input_producer_jobs
     l1_batch_number       BIGINT    NOT NULL PRIMARY KEY,
     attempts              SMALLINT  NOT NULL DEFAULT 0,
     status                tee_verifier_input_producer_job_status,
-    picked_by             TEXT,
     input_blob_url        TEXT,
     error                 TEXT,
     created_at            TIMESTAMP NOT NULL,
@@ -16,3 +15,6 @@ CREATE TABLE IF NOT EXISTS tee_verifier_input_producer_jobs
 
 CREATE INDEX IF NOT EXISTS idx_tee_verifier_input_producer_jobs_status_processing_attempts
     ON tee_verifier_input_producer_jobs (status, processing_started_at, attempts);
+
+CREATE INDEX IF NOT EXISTS idx_tee_verifier_input_producer_jobs_l1_batch_number_status
+ON tee_verifier_input_producer_jobs (l1_batch_number, status);

core/lib/dal/migrations/20240523085604_add_tee_proof_generation_details_table.down.sql

@@ -1,4 +1,4 @@
-DROP TABLE IF EXISTS tee_attestations;
-DROP TABLE IF EXISTS tee_proof_generation_details;
+DROP INDEX IF EXISTS idx_proofs_number_per_batch_number_and_tee_type;
 
-DROP INDEX IF EXISTS idx_tee_proof_generation_details_status_prover_taken_at;
+DROP TABLE IF EXISTS tee_attestations;
+DROP TABLE IF EXISTS tee_proofs;

core/lib/dal/migrations/20240523085604_add_tee_proof_generation_details_table.up.sql

@@ -4,19 +4,18 @@ CREATE TABLE IF NOT EXISTS tee_attestations
     attestation             BYTEA
 );
 
-CREATE TABLE IF NOT EXISTS tee_proof_generation_details
+CREATE TABLE IF NOT EXISTS tee_proofs
 (
-    l1_batch_number         BIGINT PRIMARY KEY REFERENCES tee_verifier_input_producer_jobs (l1_batch_number) ON DELETE CASCADE,
-    status                  TEXT      NOT NULL,
+    id                      BIGSERIAL PRIMARY KEY,
+    l1_batch_number         BIGINT NOT NULL REFERENCES tee_verifier_input_producer_jobs (l1_batch_number) ON DELETE CASCADE,
+    tee_type                TEXT NOT NULL,
+    pubkey                  BYTEA REFERENCES tee_attestations (pubkey) ON DELETE CASCADE,
     signature               BYTEA,
-    pubkey                  BYTEA REFERENCES tee_attestations (pubkey) ON DELETE SET NULL,
     proof                   BYTEA,
-    tee_type                TEXT,
     created_at              TIMESTAMP NOT NULL,
-    updated_at              TIMESTAMP NOT NULL,
+    proved_at               TIMESTAMP,
     prover_taken_at         TIMESTAMP
 );
 
-CREATE INDEX IF NOT EXISTS idx_tee_proof_generation_details_status_prover_taken_at
-    ON tee_proof_generation_details (prover_taken_at)
-    WHERE status = 'picked_by_prover';
+CREATE INDEX IF NOT EXISTS idx_proofs_number_per_batch_number_and_tee_type
+    ON tee_proofs (l1_batch_number, tee_type);

core/lib/dal/src/models/mod.rs

@@ -11,6 +11,7 @@ pub mod storage_log;
 pub mod storage_oracle_info;
 pub mod storage_protocol_version;
 pub mod storage_sync;
+pub mod storage_tee_proof;
 pub mod storage_transaction;
 pub mod storage_verification_request;
 pub mod storage_witness_job_info;

core/lib/dal/src/models/storage_tee_proof.rs

@@ -0,0 +1,35 @@
+use serde::{Deserialize, Serialize}; // TODO needed?
+use zksync_types::{tee_types::TeeType, L1BatchNumber};
+
+/// Represents a proof generated within a TEE enclave
+/// TODO move it to core/lib/types/src/api/mod.rs and call TeeProof?
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StorageTeeProof {
+    // batch number for which the proof was generated
+    pub l1_batch_number: L1BatchNumber,
+    // type of TEE used for attestation
+    pub tee_type: Option<TeeType>,
+    // pubkey used for signature verification; each key pair is attested by the TEE attestation
+    // stored in the db
+    pub pubkey: Option<Vec<u8>>,
+    // signature generated within the TEE enclave, using the privkey corresponding to the pubkey
+    pub signature: Option<Vec<u8>>,
+    // data that was signed
+    pub proof: Option<Vec<u8>>,
+    // attestation quote generated within the TEE enclave
+    pub attestation: Option<Vec<u8>>,
+    // timestamp when the proof was generated
+    pub proved_at: chrono::DateTime<chrono::Utc>,
+}
+
+/// TODO rename it TeeProof once StorageTeeProof is moved to api/mod.rs?
+#[derive(Debug, Clone, sqlx::FromRow)]
+pub struct TmpStorageTeeProof {
+    #[allow(dead_code)]
+    pub id: i64,
+    pub pubkey: Option<Vec<u8>>,
+    pub signature: Option<Vec<u8>>,
+    pub proof: Option<Vec<u8>>,
+    pub proved_at: chrono::DateTime<chrono::Utc>,
+    pub attestation: Option<Vec<u8>>,
+}

core/lib/dal/src/proof_generation_dal.rs

@@ -1,4 +1,3 @@
-#![doc = include_str!("../doc/ProofGenerationDal.md")]
 use std::time::Duration;
 
 use strum::{Display, EnumString};