apollo_proof_manager: add nonce to key

einat-starkware · einat-starkware · commit 6cad70bb27b9 · 2026-02-22T15:13:52.000+02:00
diff --git a/crates/apollo_gateway/src/gateway.rs b/crates/apollo_gateway/src/gateway.rs
@@ -29,6 +29,7 @@ use apollo_transaction_converter::{
 };
 use async_trait::async_trait;
 use blockifier::state::contract_class_manager::ContractClassManager;
+use starknet_api::core::{ContractAddress, Nonce};
 use starknet_api::executable_transaction::AccountTransaction;
 use starknet_api::rpc_transaction::{
     InternalRpcTransaction,
@@ -205,13 +206,18 @@ impl<
             .inspect_err(|e| metric_counters.record_add_tx_failure(e))?;
 
         let mut proof_archive_handle = None;
-        if let Some((proof_facts, proof)) = proof_data {
+        if let Some((proof_facts, tx_nonce, sender_address, proof)) = proof_data {
             let tx_hash = internal_tx.tx_hash;
             // Proof is verified during conversion to internal tx. It is stored here, after
             // validation, to avoid storing proofs for rejected transactions.
             let store_result = self
                 .transaction_converter
-                .store_proof_in_proof_manager(proof_facts.clone(), proof.clone())
+                .store_proof_in_proof_manager(
+                    proof_facts.clone(),
+                    proof.clone(),
+                    tx_nonce,
+                    sender_address,
+                )
                 .await;
             match store_result {
                 Ok(proof_manager_store_duration) => {
@@ -325,7 +331,11 @@ impl<
         tx: RpcTransaction,
         tx_signature: &TransactionSignature,
     ) -> Result<
-        (InternalRpcTransaction, AccountTransaction, Option<(ProofFacts, Proof)>),
+        (
+            InternalRpcTransaction,
+            AccountTransaction,
+            Option<(ProofFacts, Nonce, ContractAddress, Proof)>,
+        ),
         StarknetError,
     > {
         let (internal_tx, verification_handle) =
@@ -356,7 +366,7 @@ impl<
         &self,
         verification_handle: Option<VerificationHandle>,
         tx_signature: &TransactionSignature,
-    ) -> Result<Option<(ProofFacts, Proof)>, StarknetError> {
+    ) -> Result<Option<(ProofFacts, Nonce, ContractAddress, Proof)>, StarknetError> {
         let Some(handle) = verification_handle else {
             return Ok(None);
         };
@@ -373,7 +383,7 @@ impl<
                 transaction_converter_err_to_deprecated_gw_err(tx_signature, e)
             })?;
 
-        Ok(Some((handle.proof_facts, handle.proof)))
+        Ok(Some((handle.proof_facts, handle.nonce, handle.sender_address, handle.proof)))
     }
 }
 
diff --git a/crates/apollo_gateway/src/gateway_test.rs b/crates/apollo_gateway/src/gateway_test.rs
@@ -187,12 +187,18 @@ impl MockDependencies {
         self.mock_mempool_client.expect_validate_tx().once().with(eq(args)).return_once(|_| result);
     }
 
-    fn expect_store_proof(&mut self, proof_facts: ProofFacts, proof: Proof) {
+    fn expect_store_proof(
+        &mut self,
+        proof_facts: ProofFacts,
+        proof: Proof,
+        nonce: Nonce,
+        sender_address: ContractAddress,
+    ) {
         self.mock_transaction_converter
             .expect_store_proof_in_proof_manager()
             .once()
-            .with(eq(proof_facts), eq(proof))
-            .returning(|_, _| Ok(std::time::Duration::ZERO));
+            .with(eq(proof_facts), eq(proof), eq(nonce), eq(sender_address))
+            .returning(|_, _, _, _| Ok(std::time::Duration::ZERO));
     }
 }
 
@@ -267,6 +273,8 @@ fn setup_transaction_converter_mock(
                 Some(VerificationHandle {
                     proof_facts: invoke_tx.proof_facts.clone(),
                     proof: invoke_tx.proof.clone(),
+                    nonce: invoke_tx.nonce,
+                    sender_address: invoke_tx.sender_address,
                     verification_task,
                 })
             } else {
@@ -346,8 +354,12 @@ async fn setup_mock_state(
     // 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_store_proof(invoke_tx.proof_facts.clone(), invoke_tx.proof.clone());
+            mock_dependencies.expect_store_proof(
+                invoke_tx.proof_facts.clone(),
+                invoke_tx.proof.clone(),
+                invoke_tx.nonce,
+                invoke_tx.sender_address,
+            );
         }
     }
 
diff --git a/crates/apollo_proof_manager/src/communication.rs b/crates/apollo_proof_manager/src/communication.rs
@@ -13,21 +13,27 @@ pub type RemoteProofManagerServer =
 impl ComponentRequestHandler<ProofManagerRequest, ProofManagerResponse> for ProofManager {
     async fn handle_request(&mut self, request: ProofManagerRequest) -> ProofManagerResponse {
         match request {
-            ProofManagerRequest::SetProof(proof_facts, proof) => ProofManagerResponse::SetProof(
-                self.set_proof(proof_facts, proof)
-                    .await
-                    .map_err(|e| ProofManagerError::ProofStorage(e.to_string())),
-            ),
-            ProofManagerRequest::GetProof(proof_facts) => ProofManagerResponse::GetProof(
-                self.get_proof(proof_facts)
-                    .await
-                    .map_err(|e| ProofManagerError::ProofStorage(e.to_string())),
-            ),
-            ProofManagerRequest::ContainsProof(proof_facts) => ProofManagerResponse::ContainsProof(
-                self.contains_proof(proof_facts)
-                    .await
-                    .map_err(|e| ProofManagerError::ProofStorage(e.to_string())),
-            ),
+            ProofManagerRequest::SetProof(proof_facts, nonce, sender_address, proof) => {
+                ProofManagerResponse::SetProof(
+                    self.set_proof(proof_facts, nonce, sender_address, proof)
+                        .await
+                        .map_err(|e| ProofManagerError::ProofStorage(e.to_string())),
+                )
+            }
+            ProofManagerRequest::GetProof(proof_facts, nonce, sender_address) => {
+                ProofManagerResponse::GetProof(
+                    self.get_proof(proof_facts, nonce, sender_address)
+                        .await
+                        .map_err(|e| ProofManagerError::ProofStorage(e.to_string())),
+                )
+            }
+            ProofManagerRequest::ContainsProof(proof_facts, nonce, sender_address) => {
+                ProofManagerResponse::ContainsProof(
+                    self.contains_proof(proof_facts, nonce, sender_address)
+                        .await
+                        .map_err(|e| ProofManagerError::ProofStorage(e.to_string())),
+                )
+            }
         }
     }
 }
diff --git a/crates/apollo_proof_manager/src/proof_manager.rs b/crates/apollo_proof_manager/src/proof_manager.rs
@@ -6,6 +6,7 @@ use apollo_infra::component_definitions::{default_component_start_fn, ComponentS
 use apollo_proof_manager_config::config::ProofManagerConfig;
 use async_trait::async_trait;
 use lru::LruCache;
+use starknet_api::core::{ContractAddress, Nonce};
 use starknet_api::transaction::fields::{Proof, ProofFacts};
 use starknet_types_core::felt::Felt;
 use tracing::info;
@@ -71,48 +72,64 @@ impl ProofManager {
         Self { proof_storage, cache: ProofCache::new(config.cache_size) }
     }
 
+    fn make_storage_key(facts_hash: Felt, nonce: Nonce, sender_address: ContractAddress) -> Felt {
+        // Simple addition of facts hash, nonce, and sender address.
+        let nonce_felt = nonce.0;
+        let sender_felt = *sender_address.key();
+        facts_hash + nonce_felt + sender_felt
+    }
+
     pub async fn set_proof(
         &self,
         proof_facts: ProofFacts,
+        nonce: Nonce,
+        sender_address: ContractAddress,
         proof: Proof,
     ) -> Result<(), FsProofStorageError> {
-        if self.contains_proof(proof_facts.clone()).await? {
+        let facts_hash = proof_facts.hash();
+        let storage_key = Self::make_storage_key(facts_hash, nonce, sender_address);
+        if self.contains_proof(proof_facts.clone(), nonce, sender_address).await? {
             return Ok(());
         }
-        let facts_hash = proof_facts.hash();
-        self.proof_storage.set_proof(facts_hash, proof.clone()).await?;
-        self.cache.insert(facts_hash, proof);
+        self.proof_storage.set_proof(storage_key, proof.clone()).await?;
+        self.cache.insert(storage_key, proof);
         Ok(())
     }
 
     pub async fn get_proof(
         &self,
         proof_facts: ProofFacts,
+        nonce: Nonce,
+        sender_address: ContractAddress,
     ) -> Result<Option<Proof>, FsProofStorageError> {
         let facts_hash = proof_facts.hash();
+        let storage_key = Self::make_storage_key(facts_hash, nonce, sender_address);
         // Check cache first.
-        if let Some(proof) = self.cache.get(&facts_hash) {
+        if let Some(proof) = self.cache.get(&storage_key) {
             return Ok(Some(proof));
         }
         // Fallback to filesystem.
-        let proof = self.proof_storage.get_proof(facts_hash).await?;
+        let proof = self.proof_storage.get_proof(storage_key).await?;
         if let Some(proof) = &proof {
-            self.cache.insert(facts_hash, proof.clone());
+            self.cache.insert(storage_key, proof.clone());
         }
         Ok(proof)
     }
 
     pub async fn contains_proof(
         &self,
         proof_facts: ProofFacts,
+        nonce: Nonce,
+        sender_address: ContractAddress,
     ) -> Result<bool, FsProofStorageError> {
         let facts_hash = proof_facts.hash();
+        let storage_key = Self::make_storage_key(facts_hash, nonce, sender_address);
         // Check cache first.
-        if self.cache.contains(&facts_hash) {
+        if self.cache.contains(&storage_key) {
             return Ok(true);
         }
         // Fallback to filesystem.
-        self.proof_storage.contains_proof(facts_hash).await
+        self.proof_storage.contains_proof(storage_key).await
     }
 }
 
diff --git a/crates/apollo_proof_manager_types/src/lib.rs b/crates/apollo_proof_manager_types/src/lib.rs
@@ -13,6 +13,7 @@ use async_trait::async_trait;
 #[cfg(any(feature = "testing", test))]
 use mockall::automock;
 use serde::{Deserialize, Serialize};
+use starknet_api::core::{ContractAddress, Nonce};
 use starknet_api::transaction::fields::{Proof, ProofFacts};
 use strum::EnumVariantNames;
 use strum_macros::{AsRefStr, EnumDiscriminants, EnumIter, IntoStaticStr};
@@ -36,12 +37,24 @@ pub trait ProofManagerClient: Send + Sync {
     async fn set_proof(
         &self,
         proof_facts: ProofFacts,
+        nonce: Nonce,
+        sender_address: ContractAddress,
         proof: Proof,
     ) -> ProofManagerClientResult<()>;
 
-    async fn get_proof(&self, proof_facts: ProofFacts) -> ProofManagerClientResult<Option<Proof>>;
+    async fn get_proof(
+        &self,
+        proof_facts: ProofFacts,
+        nonce: Nonce,
+        sender_address: ContractAddress,
+    ) -> ProofManagerClientResult<Option<Proof>>;
 
-    async fn contains_proof(&self, proof_facts: ProofFacts) -> ProofManagerClientResult<bool>;
+    async fn contains_proof(
+        &self,
+        proof_facts: ProofFacts,
+        nonce: Nonce,
+        sender_address: ContractAddress,
+    ) -> ProofManagerClientResult<bool>;
 }
 
 #[derive(Clone, Debug, Error, Eq, PartialEq, Serialize, Deserialize)]
@@ -69,9 +82,9 @@ pub enum ProofManagerClientError {
     strum(serialize_all = "snake_case")
 )]
 pub enum ProofManagerRequest {
-    SetProof(ProofFacts, Proof),
-    GetProof(ProofFacts),
-    ContainsProof(ProofFacts),
+    SetProof(ProofFacts, Nonce, ContractAddress, Proof),
+    GetProof(ProofFacts, Nonce, ContractAddress),
+    ContainsProof(ProofFacts, Nonce, ContractAddress),
 }
 impl_debug_for_infra_requests_and_responses!(ProofManagerRequest);
 impl_labeled_request!(ProofManagerRequest, ProofManagerRequestLabelValue);
@@ -98,9 +111,11 @@ where
     async fn set_proof(
         &self,
         proof_facts: ProofFacts,
+        nonce: Nonce,
+        sender_address: ContractAddress,
         proof: Proof,
     ) -> ProofManagerClientResult<()> {
-        let request = ProofManagerRequest::SetProof(proof_facts, proof);
+        let request = ProofManagerRequest::SetProof(proof_facts, nonce, sender_address, proof);
         handle_all_response_variants!(
             self,
             request,
@@ -112,8 +127,13 @@ where
         )
     }
 
-    async fn get_proof(&self, proof_facts: ProofFacts) -> ProofManagerClientResult<Option<Proof>> {
-        let request = ProofManagerRequest::GetProof(proof_facts);
+    async fn get_proof(
+        &self,
+        proof_facts: ProofFacts,
+        nonce: Nonce,
+        sender_address: ContractAddress,
+    ) -> ProofManagerClientResult<Option<Proof>> {
+        let request = ProofManagerRequest::GetProof(proof_facts, nonce, sender_address);
         handle_all_response_variants!(
             self,
             request,
@@ -125,8 +145,13 @@ where
         )
     }
 
-    async fn contains_proof(&self, proof_facts: ProofFacts) -> ProofManagerClientResult<bool> {
-        let request = ProofManagerRequest::ContainsProof(proof_facts);
+    async fn contains_proof(
+        &self,
+        proof_facts: ProofFacts,
+        nonce: Nonce,
+        sender_address: ContractAddress,
+    ) -> ProofManagerClientResult<bool> {
+        let request = ProofManagerRequest::ContainsProof(proof_facts, nonce, sender_address);
         handle_all_response_variants!(
             self,
             request,
diff --git a/crates/apollo_transaction_converter/src/transaction_converter.rs b/crates/apollo_transaction_converter/src/transaction_converter.rs
diff --git a/crates/apollo_transaction_converter/src/transaction_converter_test.rs b/crates/apollo_transaction_converter/src/transaction_converter_test.rs
