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

@@ -82,6 +82,8 @@ pub trait TransactionConverterTrait: Send + Sync {
         &self,
         tx: InternalConsensusTransaction,
     ) -> TransactionConverterResult<ExecutableTransaction>;
+
+    fn get_proof_manager_client(&self) -> SharedProofManagerClient;
 }
 
 #[derive(Clone)]
@@ -143,10 +145,11 @@ 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) => {
+                let internal_tx =
+                    self.extract_proof_and_convert_rpc_tx_to_internal_rpc_tx(tx).await?;
+                Ok(InternalConsensusTransaction::RpcTransaction(internal_tx))
+            }
             ConsensusTransaction::L1Handler(tx) => self
                 .convert_consensus_l1_handler_to_internal_l1_handler(tx)
                 .map(InternalConsensusTransaction::L1Handler),
@@ -299,6 +302,10 @@ impl TransactionConverterTrait for TransactionConverter {
             InternalConsensusTransaction::L1Handler(tx) => Ok(ExecutableTransaction::L1Handler(tx)),
         }
     }
+
+    fn get_proof_manager_client(&self) -> SharedProofManagerClient {
+        self.proof_manager_client.clone()
+    }
 }
 
 impl TransactionConverter {
@@ -314,6 +321,27 @@ impl TransactionConverter {
         )?)
     }
 
+    async fn extract_proof_and_convert_rpc_tx_to_internal_rpc_tx(
+        &self,
+        tx: RpcTransaction,
+    ) -> TransactionConverterResult<InternalRpcTransaction> {
+        // 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 let Some((proof_facts, proof)) = proof_data {
+            self.proof_manager_client.set_proof(proof_facts, proof).await?;
+        }
+        Ok(internal_tx)
+    }
     async fn handle_proof_verification(
         &self,
         proof_facts: &ProofFacts,
@@ -335,7 +363,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

@@ -11,6 +11,7 @@ use mempool_test_utils::starknet_api_test_utils::{
 };
 use mockall::predicate::eq;
 use rstest::rstest;
+use starknet_api::consensus_transaction::ConsensusTransaction;
 use starknet_api::executable_transaction::ValidateCompiledClassHashError;
 use starknet_api::rpc_transaction::{RpcDeclareTransaction, RpcTransaction};
 use starknet_api::transaction::fields::Proof;
@@ -76,16 +77,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();
 
@@ -122,8 +120,7 @@ async fn test_proof_verification_skipped_for_invoke_v3_without_proof_facts() {
 
 #[rstest]
 #[tokio::test]
-async fn test_proof_verification_skipped_when_proof_already_exists() {
-    // Create an invoke transaction with proof_facts and proof.
+async fn test_consensus_tx_to_internal_with_proof_facts_verifies_and_sets_proof() {
     let proof_facts = proof_facts![felt!("0x1"), felt!("0x2"), felt!("0x3")];
     let proof = Proof::from(vec![1u32, 2u32, 3u32]);
     let invoke_tx = invoke_tx_client_side_proving(
@@ -132,14 +129,24 @@ async fn test_proof_verification_skipped_when_proof_already_exists() {
         proof.clone(),
     );
 
+    let consensus_tx = ConsensusTransaction::RpcTransaction(invoke_tx);
+
     let mut mock_proof_manager_client = MockProofManagerClient::new();
-    // Expect contains_proof to be called and return true (proof already exists).
+
+    // Expect contains_proof to be called during convert_rpc_tx_to_internal_rpc_tx.
     mock_proof_manager_client
         .expect_contains_proof()
         .once()
-        .with(eq(proof_facts))
-        .return_once(|_| Ok(true));
-    // Since proof already exists, expect set_proof to NOT be called.
+        .with(eq(proof_facts.clone()))
+        .return_once(|_| Ok(false));
+
+    // Expect set_proof to be called after the conversion succeeds.
+    // This is specific to convert_consensus_tx_to_internal_consensus_tx.
+    mock_proof_manager_client
+        .expect_set_proof()
+        .once()
+        .with(eq(proof_facts), eq(proof))
+        .return_once(|_, _| Ok(()));
 
     let mock_class_manager_client = MockClassManagerClient::new();
 
@@ -149,7 +156,10 @@ async fn test_proof_verification_skipped_when_proof_already_exists() {
         ChainInfo::create_for_testing().chain_id,
     );
 
-    // Convert the RPC transaction to an internal RPC transaction.
-    // This should succeed and only call contains_proof, not set_proof.
-    transaction_converter.convert_rpc_tx_to_internal_rpc_tx(invoke_tx).await.unwrap();
+    // Convert the consensus transaction to an internal consensus transaction.
+    // This should call contains_proof and set_proof.
+    transaction_converter
+        .convert_consensus_tx_to_internal_consensus_tx(consensus_tx)
+        .await
+        .unwrap();
 }

crates/apollo_gateway/src/gateway.rs

@@ -169,6 +169,11 @@ 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.transaction_converter.get_proof_manager_client();
+            if let Err(e) = proof_manager_client.set_proof(proof_facts.clone(), proof.clone()).await
+            {
+                error!("Failed to set proof in proof manager: {}", e);
+            }
             let proof_archive_writer = self.proof_archive_writer.clone();
             tokio::spawn(async move {
                 if let Err(e) = proof_archive_writer.set_proof(proof_facts, proof).await {

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, ProofManagerClient};
 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,16 @@ 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 mut mock_proof_manager_client = MockProofManagerClient::new();
+    mock_proof_manager_client.expect_contains_proof().returning(|_| Ok(false));
     MockDependencies {
         config,
         state_reader_factory,
         mock_mempool_client,
         mock_transaction_converter,
         mock_stateless_transaction_validator,
         mock_proof_archive_writer,
+        mock_proof_manager_client,
     }
 }
 
@@ -153,10 +158,18 @@ struct MockDependencies {
     mock_transaction_converter: MockTransactionConverterTrait,
     mock_stateless_transaction_validator: MockStatelessTransactionValidatorTrait,
     mock_proof_archive_writer: MockProofArchiveWriterTrait,
+    mock_proof_manager_client: MockProofManagerClient,
 }
 
 impl MockDependencies {
-    fn gateway(self) -> Gateway {
+    fn gateway(mut self) -> Gateway {
+        let proof_manager_client: Arc<dyn ProofManagerClient> =
+            Arc::new(self.mock_proof_manager_client);
+        let pmc = proof_manager_client.clone();
+        self.mock_transaction_converter
+            .expect_get_proof_manager_client()
+            .returning(move || pmc.clone());
+
         register_metrics();
         Gateway::new(
             self.config,
@@ -175,6 +188,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))
+            .returning(|_, _| Ok(()));
+    }
 }
 
 fn account_contract() -> FeatureContract {
@@ -291,6 +312,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() },