diff --git a/core/lib/dal/src/tee_proof_generation_dal.rs b/core/lib/dal/src/tee_proof_generation_dal.rs index 655af716d2ce..c98e0529b7f7 100644 --- a/core/lib/dal/src/tee_proof_generation_dal.rs +++ b/core/lib/dal/src/tee_proof_generation_dal.rs @@ -259,6 +259,40 @@ impl TeeProofGenerationDal<'_, '_> { Ok(proofs) } + /// For testing purposes only. + 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(()) + } + + /// For testing purposes only. pub async fn get_oldest_unpicked_batch(&mut self) -> DalResult> { let query = sqlx::query!( r#" diff --git a/core/node/api_server/src/web3/tests/unstable.rs b/core/node/api_server/src/web3/tests/unstable.rs index 4c4b604ad3ed..e814081afa02 100644 --- a/core/node/api_server/src/web3/tests/unstable.rs +++ b/core/node/api_server/src/web3/tests/unstable.rs @@ -34,6 +34,9 @@ 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]; diff --git a/core/node/proof_data_handler/src/tee_request_processor.rs b/core/node/proof_data_handler/src/tee_request_processor.rs index 4c6a2fbe880d..b7e661450177 100644 --- a/core/node/proof_data_handler/src/tee_request_processor.rs +++ b/core/node/proof_data_handler/src/tee_request_processor.rs @@ -44,11 +44,9 @@ impl TeeRequestProcessor { ) -> Result, RequestProcessorError> { tracing::info!("Received request for proof generation data: {:?}", request); - println!("foobar_01"); let mut min_batch_number: Option = 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) @@ -58,17 +56,14 @@ 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) => { - 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)), @@ -77,14 +72,11 @@ 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!( diff --git a/core/node/proof_data_handler/src/tests.rs b/core/node/proof_data_handler/src/tests.rs index c5dc499f4890..27ab65f173d0 100644 --- a/core/node/proof_data_handler/src/tests.rs +++ b/core/node/proof_data_handler/src/tests.rs @@ -1,5 +1,3 @@ -use std::env; - use axum::{ body::Body, http::{self, Method, Request, StatusCode}, @@ -7,28 +5,18 @@ use axum::{ 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::{ - block::L1BatchHeader, - commitment::{L1BatchCommitmentArtifacts, L1BatchCommitmentMode}, - protocol_version::ProtocolSemanticVersion, - tee_types::TeeType, - L1BatchNumber, ProtocolVersionId, H256, -}; +use zksync_types::{commitment::L1BatchCommitmentMode, tee_types::TeeType, L1BatchNumber}; 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; - println!("DATABASE_URL: {:?}", db_conn_pool); let app = create_proof_processing_router( MockObjectStore::arc(), @@ -65,36 +53,6 @@ async fn request_tee_proof_inputs() { 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), - ( - json!({ "tee_type": "Sgx" }), - StatusCode::UNPROCESSABLE_ENTITY, - ), - ]; - - 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 - .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); - } } // Test /tee/submit_proofs endpoint using a mocked TEE proof and verify response and db state @@ -175,82 +133,6 @@ async fn submit_tee_proof() { assert_eq!(proof.pubkey.as_ref().unwrap(), &tee_proof_request.0.pubkey); } -async fn mock_l1_batch( - db_conn_pool: ConnectionPool, - batch_number: L1BatchNumber, -) { - let mut proof_db_conn = db_conn_pool.connection().await.unwrap(); - - // 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), - }; - - // 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()); - - // 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(); - - // Insert proof generation details - proof_db_conn - .proof_generation_dal() - .insert_proof_generation_details(batch_number) - .await - .unwrap(); - - proof_db_conn - .proof_generation_dal() - .save_vm_runner_artifacts_metadata(batch_number, "vm_run_data_blob_url") - .await - .unwrap(); - - proof_db_conn - .proof_generation_dal() - .save_merkle_paths_artifacts_metadata(batch_number, "proof_gen_data_blob_url") - .await - .unwrap(); -} - async fn send_submit_tee_proof_request( app: &Router, uri: &str,