apollo_gateway: add mempool validations (#10436)

Author: lev-starkware

crates/apollo_gateway/src/gateway_test.rs

@@ -32,7 +32,7 @@ use apollo_mempool_types::communication::{
     MockMempoolClient,
 };
 use apollo_mempool_types::errors::MempoolError;
-use apollo_mempool_types::mempool_types::{AccountState, AddTransactionArgs};
+use apollo_mempool_types::mempool_types::{AccountState, AddTransactionArgs, ValidationArgs};
 use apollo_metrics::metrics::HistogramValue;
 use apollo_network_types::network_types::BroadcastedMessageMetadata;
 use apollo_test_utils::{get_rng, GetTestInstance};
@@ -165,6 +165,10 @@ impl MockDependencies {
     fn expect_add_tx(&mut self, args: AddTransactionArgsWrapper, result: MempoolClientResult<()>) {
         self.mock_mempool_client.expect_add_tx().once().with(eq(args)).return_once(|_| result);
     }
+
+    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 account_contract() -> FeatureContract {
@@ -283,13 +287,16 @@ async fn setup_mock_state(
         tx: expected_internal_tx.clone(),
         account_state: AccountState { address, nonce: *input_tx.nonce() },
     };
+    let validation_args = ValidationArgs::from(&mempool_add_tx_args);
     mock_dependencies.expect_add_tx(
         AddTransactionArgsWrapper {
             args: mempool_add_tx_args,
             p2p_message_metadata: p2p_message_metadata(),
         },
         expected_mempool_result,
     );
+
+    mock_dependencies.expect_validate_tx(validation_args, Ok(()))
 }
 
 struct AddTxResults {

crates/apollo_gateway/src/stateful_transaction_validator.rs

@@ -8,6 +8,7 @@ use apollo_gateway_types::deprecated_gateway_error::{
 };
 use apollo_gateway_types::errors::GatewaySpecError;
 use apollo_mempool_types::communication::SharedMempoolClient;
+use apollo_mempool_types::mempool_types::ValidationArgs;
 use apollo_proc_macros::sequencer_latency_histogram;
 use async_trait::async_trait;
 use blockifier::blockifier::stateful_validator::{
@@ -165,6 +166,11 @@ impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidator<B> {
         runtime: tokio::runtime::Handle,
     ) -> StatefulTransactionValidatorResult<()> {
         self.validate_state_preconditions(executable_tx, account_nonce)?;
+        runtime.block_on(validate_by_mempool(
+            executable_tx,
+            account_nonce,
+            mempool_client.clone(),
+        ))?;
         self.run_validate_entry_point(executable_tx, account_nonce, mempool_client, runtime)?;
         Ok(())
     }
@@ -316,6 +322,19 @@ impl<B: BlockifierStatefulValidatorTrait> StatefulTransactionValidator<B> {
     }
 }
 
+/// Perform transaction validation by the mempool.
+async fn validate_by_mempool(
+    tx: &ExecutableTransaction,
+    account_nonce: Nonce,
+    mempool_client: SharedMempoolClient,
+) -> StatefulTransactionValidatorResult<()> {
+    let validation_args = ValidationArgs::new(tx, account_nonce);
+    mempool_client
+        .validate_tx(validation_args)
+        .await
+        .map_err(|err| mempool_client_err_to_deprecated_gw_err(&tx.signature(), err))
+}
+
 /// Check if validation of an invoke transaction should be skipped due to deploy_account not being
 /// processed yet. This feature is used to improve UX for users sending deploy_account + invoke at
 /// once.

crates/apollo_gateway/src/stateful_transaction_validator_test.rs

@@ -134,6 +134,7 @@ async fn test_extract_state_nonce_and_run_validations(
         // The mempool does not have any transactions from the sender.
         Ok(false)
     });
+    mock_mempool_client.expect_validate_tx().returning(|_| Ok(()));
     let mempool_client = Arc::new(mock_mempool_client);
     let runtime = tokio::runtime::Handle::current();
 
@@ -347,10 +348,12 @@ async fn validate_resource_bounds(
     };
 
     let result = tokio::task::spawn_blocking(move || {
+        let mut mempool_client = MockMempoolClient::new();
+        mempool_client.expect_validate_tx().returning(|_| Ok(()));
         stateful_validator.run_transaction_validations(
             &executable_tx,
             account_nonce,
-            Arc::new(MockMempoolClient::new()),
+            Arc::new(mempool_client),
             tokio::runtime::Handle::current(),
         )
     })
@@ -465,10 +468,12 @@ async fn run_transaction_validation_test(
     };
 
     let result = tokio::task::spawn_blocking(move || {
+        let mut mempool_client = MockMempoolClient::new();
+        mempool_client.expect_validate_tx().returning(|_| Ok(()));
         stateful_validator.run_transaction_validations(
             &executable_tx,
             account_nonce,
-            Arc::new(MockMempoolClient::new()),
+            Arc::new(mempool_client),
             tokio::runtime::Handle::current(),
         )
     })

crates/apollo_mempool_types/src/mempool_types.rs

@@ -6,6 +6,7 @@ use indexmap::IndexSet;
 use serde::{Deserialize, Serialize};
 use starknet_api::block::GasPrice;
 use starknet_api::core::{ContractAddress, Nonce};
+use starknet_api::executable_transaction::AccountTransaction;
 use starknet_api::rpc_transaction::InternalRpcTransaction;
 use starknet_api::transaction::fields::Tip;
 use starknet_api::transaction::TransactionHash;
@@ -40,7 +41,7 @@ impl AddTransactionArgs {
     }
 }
 
-#[derive(Clone, Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub struct ValidationArgs {
     pub address: ContractAddress,
     pub account_nonce: Nonce,
@@ -50,6 +51,19 @@ pub struct ValidationArgs {
     pub max_l2_gas_price: GasPrice,
 }
 
+impl ValidationArgs {
+    pub fn new(tx: &AccountTransaction, account_nonce: Nonce) -> Self {
+        Self {
+            address: tx.sender_address(),
+            account_nonce,
+            tx_hash: tx.tx_hash(),
+            tx_nonce: tx.nonce(),
+            tip: tx.tip(),
+            max_l2_gas_price: tx.resource_bounds().get_l2_bounds().max_price_per_unit,
+        }
+    }
+}
+
 impl From<&AddTransactionArgs> for ValidationArgs {
     fn from(args: &AddTransactionArgs) -> Self {
         Self {