WIP

core/lib/dal/src/tee_proof_generation_dal.rs

@@ -196,38 +196,6 @@ impl TeeProofGenerationDal<'_, '_> {
         Ok(())
     }
 
-    pub async fn insert_tee_proof_generation_job(
-        &mut self,
-        batch_number: L1BatchNumber,
-        tee_type: TeeType,
-    ) -> DalResult<()> {
-        let batch_number = i64::from(batch_number.0);
-        let query = sqlx::query!(
-            r#"
-            INSERT INTO
-            tee_proof_generation_details (
-                l1_batch_number, tee_type, status, created_at, updated_at
-            )
-            VALUES
-            ($1, $2, $3, NOW(), NOW())
-            ON CONFLICT (l1_batch_number, tee_type) DO NOTHING
-            "#,
-            batch_number,
-            tee_type.to_string(),
-            TeeProofGenerationJobStatus::Unpicked.to_string(),
-        );
-        let instrumentation = Instrumented::new("insert_tee_proof_generation_job")
-            .with_arg("l1_batch_number", &batch_number)
-            .with_arg("tee_type", &tee_type);
-        instrumentation
-            .clone()
-            .with(query)
-            .execute(self.storage)
-            .await?;
-
-        Ok(())
-    }
-
     pub async fn save_attestation(&mut self, pubkey: &[u8], attestation: &[u8]) -> DalResult<()> {
         let query = sqlx::query!(
             r#"

core/node/api_server/src/web3/tests/unstable.rs

@@ -34,9 +34,6 @@ impl HttpTest for GetTeeProofsTest {
         tee_proof_generation_dal
             .save_attestation(&pubkey, &attestation)
             .await?;
-        tee_proof_generation_dal
-            .insert_tee_proof_generation_job(batch_no, tee_type)
-            .await?;
 
         let signature = vec![0, 1, 2, 3, 4];
         let proof_vec = vec![5, 6, 7, 8, 9];

core/node/proof_data_handler/src/tee_request_processor.rs

@@ -44,9 +44,11 @@ impl TeeRequestProcessor {
     ) -> Result<Json<TeeProofGenerationDataResponse>, RequestProcessorError> {
         tracing::info!("Received request for proof generation data: {:?}", request);
 
+        println!("foobar_01");
         let mut min_batch_number: Option<L1BatchNumber> = None;
         let mut missing_range: Option<(L1BatchNumber, L1BatchNumber)> = None;
 
+        println!("foobar_02");
         let result = loop {
             let l1_batch_number = match self
                 .lock_batch_for_proving(request.tee_type, min_batch_number)
@@ -56,12 +58,17 @@ impl TeeRequestProcessor {
                 None => break Ok(Json(TeeProofGenerationDataResponse(None))),
             };
 
+            println!("foobar_03");
             match self
                 .tee_verifier_input_for_existing_batch(l1_batch_number)
                 .await
             {
-                Ok(input) => break Ok(Json(TeeProofGenerationDataResponse(Some(Box::new(input))))),
+                Ok(input) => {
+                    println!("foobar_04");
+                    break Ok(Json(TeeProofGenerationDataResponse(Some(Box::new(input)))));
+                }
                 Err(RequestProcessorError::ObjectStore(ObjectStoreError::KeyNotFound(_))) => {
+                    println!("foobar_05");
                     missing_range = match missing_range {
                         Some((start, _)) => Some((start, l1_batch_number)),
                         None => Some((l1_batch_number, l1_batch_number)),
@@ -70,11 +77,14 @@ impl TeeRequestProcessor {
                     min_batch_number = Some(min_batch_number.unwrap_or(l1_batch_number) + 1);
                 }
                 Err(err) => {
+                    println!("foobar_06");
                     self.unlock_batch(l1_batch_number, request.tee_type).await?;
                     break Err(err);
                 }
             }
+            println!("foobar_07");
         };
+        println!("foobar_08");
 
         if let Some((start, end)) = missing_range {
             tracing::warn!(

core/node/proof_data_handler/src/tests.rs

@@ -1,36 +1,73 @@
+use std::env;
+
 use axum::{
     body::Body,
     http::{self, Method, Request, StatusCode},
     response::Response,
     Router,
 };
 use serde_json::json;
+use tokio::time::{sleep, Duration};
 use tower::ServiceExt;
 use zksync_config::configs::ProofDataHandlerConfig;
+use zksync_contracts::BaseSystemContractsHashes;
 use zksync_dal::{ConnectionPool, CoreDal};
 use zksync_object_store::MockObjectStore;
 use zksync_prover_interface::api::SubmitTeeProofRequest;
-use zksync_types::{commitment::L1BatchCommitmentMode, tee_types::TeeType, L1BatchNumber};
+use zksync_types::{
+    block::L1BatchHeader,
+    commitment::{L1BatchCommitmentArtifacts, L1BatchCommitmentMode},
+    protocol_version::ProtocolSemanticVersion,
+    tee_types::TeeType,
+    L1BatchNumber, ProtocolVersionId, H256,
+};
 
 use crate::create_proof_processing_router;
 
 #[tokio::test]
 async fn request_tee_proof_inputs() {
     let batch_number = L1BatchNumber::from(1);
     let db_conn_pool = ConnectionPool::test_pool().await;
-
-    mock_tee_batch_status(db_conn_pool.clone(), batch_number).await;
+    println!("DATABASE_URL: {:?}", db_conn_pool);
 
     let app = create_proof_processing_router(
         MockObjectStore::arc(),
-        db_conn_pool,
+        db_conn_pool.clone(),
         ProofDataHandlerConfig {
             http_port: 1337,
             proof_generation_timeout_in_secs: 10,
             tee_support: true,
         },
         L1BatchCommitmentMode::Rollup,
     );
+    let test_cases = vec![
+        (json!({ "tee_type": "sgx" }), StatusCode::NO_CONTENT),
+        (
+            json!({ "tee_type": "Sgx" }),
+            StatusCode::UNPROCESSABLE_ENTITY,
+        ),
+    ];
+
+    for (body, expected_status) in test_cases {
+        let req_body = Body::from(serde_json::to_vec(&body).unwrap());
+        let response = app
+            .clone()
+            .oneshot(
+                Request::builder()
+                    .method(Method::POST)
+                    .uri("/tee/proof_inputs")
+                    .header(http::header::CONTENT_TYPE, "application/json")
+                    .body(req_body)
+                    .unwrap(),
+            )
+            .await
+            .unwrap();
+
+        assert_eq!(response.status(), expected_status);
+    }
+
+    mock_l1_batch(db_conn_pool.clone(), batch_number).await;
+
     let test_cases = vec![
         (json!({ "tee_type": "sgx" }), StatusCode::OK),
         (
@@ -39,6 +76,8 @@ async fn request_tee_proof_inputs() {
         ),
     ];
 
+    sleep(Duration::from_secs(1000)).await;
+
     for (body, expected_status) in test_cases {
         let req_body = Body::from(serde_json::to_vec(&body).unwrap());
         let response = app
@@ -64,8 +103,6 @@ async fn submit_tee_proof() {
     let batch_number = L1BatchNumber::from(1);
     let db_conn_pool = ConnectionPool::test_pool().await;
 
-    mock_tee_batch_status(db_conn_pool.clone(), batch_number).await;
-
     let tee_proof_request_str = r#"{
         "signature": "0001020304",
         "pubkey": "0506070809",
@@ -138,34 +175,80 @@ async fn submit_tee_proof() {
     assert_eq!(proof.pubkey.as_ref().unwrap(), &tee_proof_request.0.pubkey);
 }
 
-// Mock SQL db with information about the status of the TEE proof generation
-async fn mock_tee_batch_status(
+async fn mock_l1_batch(
     db_conn_pool: ConnectionPool<zksync_dal::Core>,
     batch_number: L1BatchNumber,
 ) {
     let mut proof_db_conn = db_conn_pool.connection().await.unwrap();
-    let mut proof_dal = proof_db_conn.tee_proof_generation_dal();
 
-    // there should not be any batches awaiting proof in the db yet
+    // Common values
+    let protocol_version_id = ProtocolVersionId::latest();
+    let protocol_version = ProtocolSemanticVersion {
+        minor: protocol_version_id,
+        patch: 0.into(),
+    };
+    let base_system_contracts_hashes = BaseSystemContractsHashes {
+        bootloader: H256::repeat_byte(1),
+        default_aa: H256::repeat_byte(42),
+    };
 
-    let oldest_batch_number = proof_dal.get_oldest_unpicked_batch().await.unwrap();
-    assert!(oldest_batch_number.is_none());
+    // Save protocol version
+    proof_db_conn
+        .protocol_versions_dal()
+        .save_protocol_version(
+            protocol_version,
+            0,
+            Default::default(),
+            base_system_contracts_hashes,
+            None,
+        )
+        .await
+        .unwrap();
+
+    // Insert mock L1 batch
+    let header = L1BatchHeader::new(
+        batch_number,
+        100,
+        base_system_contracts_hashes,
+        protocol_version_id,
+    );
+    assert!(proof_db_conn
+        .blocks_dal()
+        .insert_mock_l1_batch(&header)
+        .await
+        .is_ok());
 
-    // mock SQL table with relevant information about the status of TEE proof generation
+    // Save L1 batch commitment artifacts and set hash
+    let hash = H256::repeat_byte(1);
+    proof_db_conn
+        .blocks_dal()
+        .save_l1_batch_commitment_artifacts(batch_number, &L1BatchCommitmentArtifacts::default())
+        .await
+        .unwrap();
+    proof_db_conn
+        .blocks_dal()
+        .set_l1_batch_hash(batch_number, hash)
+        .await
+        .unwrap();
 
-    proof_dal
-        .insert_tee_proof_generation_job(batch_number, TeeType::Sgx)
+    // Insert proof generation details
+    proof_db_conn
+        .proof_generation_dal()
+        .insert_proof_generation_details(batch_number)
         .await
-        .expect("Failed to insert tee_proof_generation_job");
+        .unwrap();
 
-    // now, there should be one batch in the db awaiting proof
+    proof_db_conn
+        .proof_generation_dal()
+        .save_vm_runner_artifacts_metadata(batch_number, "vm_run_data_blob_url")
+        .await
+        .unwrap();
 
-    let oldest_batch_number = proof_dal
-        .get_oldest_unpicked_batch()
+    proof_db_conn
+        .proof_generation_dal()
+        .save_merkle_paths_artifacts_metadata(batch_number, "proof_gen_data_blob_url")
         .await
-        .unwrap()
         .unwrap();
-    assert_eq!(oldest_batch_number, batch_number);
 }
 
 async fn send_submit_tee_proof_request(