apollo_gateway: only set proofs if they have passed tx validation in gateway

Author: einat-starkware

crates/apollo_class_manager_types/src/transaction_converter.rs

@@ -143,10 +143,29 @@ impl TransactionConverterTrait for TransactionConverter {
         tx: ConsensusTransaction,
     ) -> TransactionConverterResult<InternalConsensusTransaction> {
         match tx {
-            ConsensusTransaction::RpcTransaction(tx) => self
-                .convert_rpc_tx_to_internal_rpc_tx(tx)
-                .await
-                .map(InternalConsensusTransaction::RpcTransaction),
+            ConsensusTransaction::RpcTransaction(tx) => {
+                // Extract proof and proof_facts from v3 invoke tx before conversion.
+                let proof_data = match &tx {
+                    RpcTransaction::Invoke(RpcInvokeTransaction::V3(invoke_tx)) => {
+                        if !invoke_tx.proof_facts.is_empty() {
+                            Some((invoke_tx.proof_facts.clone(), invoke_tx.proof.clone()))
+                        } else {
+                            None
+                        }
+                    }
+                    _ => None,
+                };
+
+                let internal_tx = self.convert_rpc_tx_to_internal_rpc_tx(tx).await?;
+
+                // If we extracted proof data and conversion was successful, set it in the proof
+                // manager.
+                if let Some((proof_facts, proof)) = proof_data {
+                    self.proof_manager_client.set_proof(proof_facts, proof).await?;
+                }
+
+                Ok(InternalConsensusTransaction::RpcTransaction(internal_tx))
+            }
             ConsensusTransaction::L1Handler(tx) => self
                 .convert_consensus_l1_handler_to_internal_l1_handler(tx)
                 .map(InternalConsensusTransaction::L1Handler),
@@ -335,7 +354,6 @@ impl TransactionConverter {
         }
 
         self.verify_proof(proof_facts.clone(), proof.clone())?;
-        self.proof_manager_client.set_proof(proof_facts.clone(), proof.clone()).await?;
         Ok(())
     }
 

crates/apollo_class_manager_types/src/transaction_converter_test.rs

@@ -76,16 +76,13 @@ async fn test_proof_verification_called_for_invoke_v3_with_proof_facts() {
     );
 
     let mut mock_proof_manager_client = MockProofManagerClient::new();
+    // Expect contains_proof to be called and return false (proof does not exist).
+    // Expect set_proof not to be called when converting rpc tx to internal rpc tx.
     mock_proof_manager_client
         .expect_contains_proof()
         .once()
         .with(eq(proof_facts.clone()))
         .return_once(|_| Ok(false));
-    mock_proof_manager_client
-        .expect_set_proof()
-        .once()
-        .with(eq(proof_facts), eq(proof))
-        .return_once(|_, _| Ok(()));
 
     let mock_class_manager_client = MockClassManagerClient::new();
 

crates/apollo_gateway/src/gateway.rs

@@ -72,6 +72,7 @@ pub struct Gateway {
     pub mempool_client: SharedMempoolClient,
     pub transaction_converter: Arc<dyn TransactionConverterTrait>,
     pub proof_archive_writer: Arc<dyn ProofArchiveWriterTrait>,
+    pub proof_manager_client: SharedProofManagerClient,
 }
 
 impl Gateway {
@@ -82,6 +83,7 @@ impl Gateway {
         transaction_converter: Arc<dyn TransactionConverterTrait>,
         stateless_tx_validator: Arc<dyn StatelessTransactionValidatorTrait>,
         proof_archive_writer: Arc<dyn ProofArchiveWriterTrait>,
+        proof_manager_client: SharedProofManagerClient,
     ) -> Self {
         Self {
             config: Arc::new(config.clone()),
@@ -97,6 +99,7 @@ impl Gateway {
             mempool_client,
             transaction_converter,
             proof_archive_writer,
+            proof_manager_client,
         }
     }
 
@@ -169,8 +172,14 @@ impl Gateway {
             .inspect_err(|e| metric_counters.record_add_tx_failure(e))?;
 
         if let Some((proof_facts, proof)) = proof_data {
+            let proof_manager_client = self.proof_manager_client.clone();
+            if let Err(e) = proof_manager_client.set_proof(proof_facts.clone(), proof.clone()).await
+            {
+                warn!("Failed to set proof in proof manager: {}", e);
+            }
             let proof_archive_writer = self.proof_archive_writer.clone();
             tokio::spawn(async move {
+                // Write proof to both archive and proof manager.
                 if let Err(e) = proof_archive_writer.set_proof(proof_facts, proof).await {
                     error!("Failed to archive proof to GCS: {}", e);
                 }
@@ -265,7 +274,7 @@ pub fn create_gateway(
     });
     let transaction_converter = Arc::new(TransactionConverter::new(
         class_manager_client,
-        proof_manager_client,
+        proof_manager_client.clone(),
         config.chain_info.chain_id.clone(),
     ));
     let stateless_tx_validator = Arc::new(StatelessTransactionValidator {
@@ -287,6 +296,7 @@ pub fn create_gateway(
         transaction_converter,
         stateless_tx_validator,
         proof_archive_writer,
+        proof_manager_client,
     )
 }
 

crates/apollo_gateway/src/gateway_test.rs

@@ -36,6 +36,7 @@ use apollo_mempool_types::errors::MempoolError;
 use apollo_mempool_types::mempool_types::{AccountState, AddTransactionArgs, ValidationArgs};
 use apollo_metrics::metrics::HistogramValue;
 use apollo_network_types::network_types::BroadcastedMessageMetadata;
+use apollo_proof_manager_types::MockProofManagerClient;
 use apollo_test_utils::{get_rng, GetTestInstance};
 use blockifier::blockifier::config::ContractClassManagerConfig;
 use blockifier::context::ChainInfo;
@@ -53,6 +54,7 @@ use starknet_api::executable_transaction::AccountTransaction;
 use starknet_api::rpc_transaction::{
     InternalRpcTransaction,
     RpcDeclareTransaction,
+    RpcInvokeTransaction,
     RpcTransaction,
     RpcTransactionLabelValue,
 };
@@ -136,13 +138,15 @@ fn mock_dependencies() -> MockDependencies {
     let mock_transaction_converter = MockTransactionConverterTrait::new();
     let mock_stateless_transaction_validator = mock_stateless_transaction_validator();
     let mock_proof_archive_writer = MockProofArchiveWriterTrait::new();
+    let mock_proof_manager_client = MockProofManagerClient::new();
     MockDependencies {
         config,
         state_reader_factory,
         mock_mempool_client,
         mock_transaction_converter,
         mock_stateless_transaction_validator,
         mock_proof_archive_writer,
+        mock_proof_manager_client,
     }
 }
 
@@ -153,6 +157,7 @@ struct MockDependencies {
     mock_transaction_converter: MockTransactionConverterTrait,
     mock_stateless_transaction_validator: MockStatelessTransactionValidatorTrait,
     mock_proof_archive_writer: MockProofArchiveWriterTrait,
+    mock_proof_manager_client: MockProofManagerClient,
 }
 
 impl MockDependencies {
@@ -165,6 +170,7 @@ impl MockDependencies {
             Arc::new(self.mock_transaction_converter),
             Arc::new(self.mock_stateless_transaction_validator),
             Arc::new(self.mock_proof_archive_writer),
+            Arc::new(self.mock_proof_manager_client),
         )
     }
 
@@ -175,6 +181,14 @@ impl MockDependencies {
     fn expect_validate_tx(&mut self, args: ValidationArgs, result: MempoolClientResult<()>) {
         self.mock_mempool_client.expect_validate_tx().once().with(eq(args)).return_once(|_| result);
     }
+
+    fn expect_set_proof(&mut self, proof_facts: ProofFacts, proof: Proof) {
+        self.mock_proof_manager_client
+            .expect_set_proof()
+            .once()
+            .with(eq(proof_facts), eq(proof))
+            .return_once(|_, _| Ok(()));
+    }
 }
 
 fn account_contract() -> FeatureContract {
@@ -291,6 +305,14 @@ async fn setup_mock_state(
         &mut mock_dependencies.state_reader_factory.state_reader.blockifier_state_reader,
     );
 
+    // If the transaction has proof facts, expect set_proof to be called on the proof manager.
+    if let RpcTransaction::Invoke(RpcInvokeTransaction::V3(ref invoke_tx)) = input_tx {
+        if !invoke_tx.proof_facts.is_empty() {
+            mock_dependencies
+                .expect_set_proof(invoke_tx.proof_facts.clone(), invoke_tx.proof.clone());
+        }
+    }
+
     let mempool_add_tx_args = AddTransactionArgs {
         tx: expected_internal_tx.clone(),
         account_state: AccountState { address, nonce: *input_tx.nonce() },
@@ -580,6 +602,7 @@ async fn add_tx_returns_error_when_extract_state_nonce_and_run_validations_fails
         mempool_client: Arc::new(mock_dependencies.mock_mempool_client),
         transaction_converter: Arc::new(mock_dependencies.mock_transaction_converter),
         proof_archive_writer: Arc::new(mock_dependencies.mock_proof_archive_writer),
+        proof_manager_client: Arc::new(mock_dependencies.mock_proof_manager_client),
     };
 
     let result = gateway.add_tx(tx_args.get_rpc_tx(), None).await;
@@ -633,6 +656,7 @@ async fn add_tx_returns_error_when_instantiating_validator_fails(
         mempool_client: Arc::new(mock_dependencies.mock_mempool_client),
         transaction_converter: Arc::new(mock_dependencies.mock_transaction_converter),
         proof_archive_writer: Arc::new(mock_dependencies.mock_proof_archive_writer),
+        proof_manager_client: Arc::new(mock_dependencies.mock_proof_manager_client),
     };
 
     let result = gateway.add_tx(tx_args.get_rpc_tx(), None).await;

crates/apollo_gateway/src/test_utils.rs

@@ -185,5 +185,6 @@ pub fn gateway_for_benchmark(gateway_config: GatewayConfig) -> Gateway {
         Arc::new(transaction_converter),
         stateless_tx_validator,
         proof_archive_writer,
+        proof_manager_client,
     )
 }