fix: run the validation and the insertion of a transaction in the mempool at the same time

Signed-off-by: Eric Torreborre <etorreborre@yahoo.com>

Author: etorreborre

crates/amaru-mempool/src/strategies/dummy.rs

@@ -15,10 +15,7 @@
 use std::{collections::BTreeSet, mem, pin::Pin};
 
 use amaru_kernel::cbor;
-use amaru_ouroboros_traits::{
-    CanValidateTransactions, Mempool, MempoolSeqNo, TransactionValidationError, TxId, TxOrigin, TxRejectReason,
-    TxSubmissionMempool,
-};
+use amaru_ouroboros_traits::{Mempool, MempoolSeqNo, TxId, TxOrigin, TxRejectReason, TxSubmissionMempool};
 use parking_lot::RwLock;
 
 #[derive(Debug, Default)]
@@ -37,12 +34,6 @@ pub struct DummyMempoolInner<T> {
     transactions: Vec<T>,
 }
 
-impl<Tx: cbor::Encode<()> + Send + Sync + 'static> CanValidateTransactions<Tx> for DummyMempool<Tx> {
-    fn validate_transaction(&self, _tx: Tx) -> Result<(), TransactionValidationError> {
-        Ok(())
-    }
-}
-
 impl<Tx: cbor::Encode<()> + Send + Sync + 'static> TxSubmissionMempool<Tx> for DummyMempool<Tx> {
     fn insert(&self, tx: Tx, _tx_origin: TxOrigin) -> Result<(TxId, MempoolSeqNo), TxRejectReason> {
         let tx_id = TxId::from(&tx);

crates/amaru-mempool/src/strategies/in_memory_mempool.rs

@@ -21,8 +21,7 @@ use std::{
 
 use amaru_kernel::{cbor, to_cbor};
 use amaru_ouroboros_traits::{
-    CanValidateTransactions, MempoolSeqNo, TransactionValidationError, TxId, TxOrigin, TxRejectReason,
-    TxSubmissionMempool, mempool::Mempool,
+    MempoolSeqNo, TransactionValidationError, TxId, TxOrigin, TxRejectReason, TxSubmissionMempool, mempool::Mempool,
 };
 use tokio::sync::Notify;
 
@@ -31,39 +30,48 @@ use tokio::sync::Notify;
 /// It stores transactions in memory, indexed by their TxId and by a sequence number assigned
 /// at insertion time.
 ///
-/// The validation of the transactions are delegated to a `CanValidateTransactions` implementation.
+/// The validation of transactions is delegated to an injected validator.
 ///
 #[derive(Clone)]
 pub struct InMemoryMempool<Tx> {
     config: MempoolConfig,
     inner: Arc<parking_lot::RwLock<MempoolInner<Tx>>>,
-    tx_validator: Arc<dyn CanValidateTransactions<Tx>>,
+    tx_validator: Arc<dyn TxValidator<Tx>>,
 }
 
-impl<Tx> Default for InMemoryMempool<Tx> {
+pub trait TxValidator<Tx>: Send + Sync {
+    fn validate(&self, tx: &Tx) -> Result<(), TransactionValidationError>;
+}
+
+impl<Tx, F> TxValidator<Tx> for F
+where
+    F: Fn(&Tx) -> Result<(), TransactionValidationError> + Send + Sync,
+{
+    fn validate(&self, tx: &Tx) -> Result<(), TransactionValidationError> {
+        self(tx)
+    }
+}
+
+impl<Tx: 'static> Default for InMemoryMempool<Tx> {
     fn default() -> Self {
         Self::from_config(MempoolConfig::default())
     }
 }
 
 impl<Tx> InMemoryMempool<Tx> {
-    pub fn new(config: MempoolConfig, tx_validator: Arc<dyn CanValidateTransactions<Tx>>) -> Self {
+    pub fn new(config: MempoolConfig, tx_validator: Arc<dyn TxValidator<Tx>>) -> Self {
         InMemoryMempool { config, inner: Arc::new(parking_lot::RwLock::new(MempoolInner::default())), tx_validator }
     }
+}
 
+impl<Tx: 'static> InMemoryMempool<Tx> {
     pub fn from_config(config: MempoolConfig) -> Self {
-        Self::new(config, Arc::new(DefaultCanValidateTransactions))
+        Self::new(config, Arc::new(default_tx_validator::<Tx>))
     }
 }
 
-/// A default transactions validator.
-#[derive(Clone, Debug, Default)]
-pub struct DefaultCanValidateTransactions;
-
-impl<Tx> CanValidateTransactions<Tx> for DefaultCanValidateTransactions {
-    fn validate_transaction(&self, _tx: Tx) -> Result<(), TransactionValidationError> {
-        Ok(())
-    }
+pub fn default_tx_validator<Tx>(_tx: &Tx) -> Result<(), TransactionValidationError> {
+    Ok(())
 }
 
 #[derive(Debug)]
@@ -169,14 +177,9 @@ impl MempoolConfig {
     }
 }
 
-impl<Tx: Send + Sync + 'static> CanValidateTransactions<Tx> for InMemoryMempool<Tx> {
-    fn validate_transaction(&self, tx: Tx) -> Result<(), TransactionValidationError> {
-        self.tx_validator.validate_transaction(tx)
-    }
-}
-
 impl<Tx: Send + Sync + 'static + cbor::Encode<()> + Clone> TxSubmissionMempool<Tx> for InMemoryMempool<Tx> {
     fn insert(&self, tx: Tx, tx_origin: TxOrigin) -> Result<(TxId, MempoolSeqNo), TxRejectReason> {
+        self.tx_validator.validate(&tx).map_err(TxRejectReason::Invalid)?;
         let mut inner = self.inner.write();
         let res = inner.insert(&self.config, tx, tx_origin);
         if res.is_ok() {
@@ -258,7 +261,7 @@ impl<Tx: Send + Sync + 'static + cbor::Encode<()> + Clone> Mempool<Tx> for InMem
 
 #[cfg(test)]
 mod tests {
-    use std::{ops::Deref, slice, str::FromStr, time::Duration};
+    use std::{ops::Deref, slice, str::FromStr, sync::Arc, time::Duration};
 
     use amaru_kernel::{Peer, cbor, cbor as minicbor};
     use assertables::assert_some_eq_x;
@@ -283,6 +286,16 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn reject_invalid_transaction_on_insert() {
+        let mempool = InMemoryMempool::new(MempoolConfig::default(), Arc::new(reject_tx));
+        let tx = Tx::from_str("tx1").unwrap();
+
+        let result = mempool.insert(tx, TxOrigin::Local);
+
+        assert_eq!(result, Err(TxRejectReason::Invalid(anyhow::anyhow!("transaction rejected for testing").into())));
+    }
+
     // HELPERS
     #[derive(Debug, PartialEq, Eq, Clone, cbor::Encode, cbor::Decode)]
     struct Tx(#[n(0)] String);
@@ -300,4 +313,8 @@ mod tests {
             Ok(Tx(s.to_string()))
         }
     }
+
+    fn reject_tx(_tx: &Tx) -> Result<(), TransactionValidationError> {
+        Err(anyhow::anyhow!("transaction rejected for testing").into())
+    }
 }

crates/amaru-ouroboros-traits/src/can_validate_transactions/mock.rs

@@ -18,9 +18,6 @@ pub use has_stake_distribution::{HasStakeDistribution, PoolSummary};
 pub mod can_validate_blocks;
 pub use can_validate_blocks::{BlockValidationError, CanValidateBlocks};
 
-pub mod can_validate_transactions;
-pub use can_validate_transactions::{CanValidateTransactions, TransactionValidationError};
-
 pub mod connections;
 pub use connections::*;
 

crates/amaru-ouroboros-traits/src/can_validate_transactions/mod.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::{
+    error::Error,
     fmt,
     fmt::{Display, Formatter},
     pin::Pin,
@@ -22,8 +23,6 @@ use std::{
 use amaru_kernel::{Hash, Hasher, Peer, TransactionId, cbor, size::TRANSACTION_BODY};
 use serde::{Deserialize, Serialize};
 
-use crate::CanValidateTransactions;
-
 /// An simple mempool interface to:
 ///
 /// - Retrieve transactions to be included in a new block.
@@ -48,7 +47,7 @@ pub trait Mempool<Tx: Send + Sync + 'static>: TxSubmissionMempool<Tx> + Send + S
 
 pub type ResourceMempool<Tx> = Arc<dyn TxSubmissionMempool<Tx>>;
 
-pub trait TxSubmissionMempool<Tx: Send + Sync + 'static>: Send + Sync + CanValidateTransactions<Tx> {
+pub trait TxSubmissionMempool<Tx: Send + Sync + 'static>: Send + Sync {
     /// Insert a transaction into the mempool, specifying its origin.
     /// A TxOrigin::Local origin indicates the transaction was created on the current node,
     /// A TxOrigin::Remote(origin_peer) indicates the transaction was received from a remote peer.
@@ -97,6 +96,49 @@ pub trait TxSubmissionMempool<Tx: Send + Sync + 'static>: Send + Sync + CanValid
     fn last_seq_no(&self) -> MempoolSeqNo;
 }
 
+#[derive(Debug, thiserror::Error)]
+#[error("TransactionValidationError: {0}")]
+pub struct TransactionValidationError(#[from] anyhow::Error);
+
+impl TransactionValidationError {
+    pub fn to_anyhow(self) -> anyhow::Error {
+        self.0
+    }
+
+    pub fn downcast<T: Error + fmt::Debug + Send + Sync + 'static>(self) -> Result<T, anyhow::Error> {
+        self.0.downcast::<T>()
+    }
+
+    pub fn downcast_ref<T: Error + fmt::Debug + Send + Sync + 'static>(&self) -> Option<&T> {
+        self.0.downcast_ref::<T>()
+    }
+}
+
+impl Serialize for TransactionValidationError {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        serializer.serialize_str(&self.0.to_string())
+    }
+}
+
+impl<'de> Deserialize<'de> for TransactionValidationError {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let s = String::deserialize(deserializer)?;
+        Ok(TransactionValidationError(anyhow::anyhow!(s)))
+    }
+}
+
+impl PartialEq for TransactionValidationError {
+    fn eq(&self, other: &Self) -> bool {
+        self.0.to_string() == other.0.to_string()
+    }
+}
+
 /// Sequence number assigned to a transaction when inserted into the mempool.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Default)]
 pub struct MempoolSeqNo(pub u64);
@@ -111,14 +153,14 @@ impl MempoolSeqNo {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, thiserror::Error, Serialize, Deserialize)]
+#[derive(Debug, PartialEq, thiserror::Error, Serialize, Deserialize)]
 pub enum TxRejectReason {
     #[error("Mempool is full")]
     MempoolFull,
     #[error("Transaction is a duplicate")]
     Duplicate,
-    #[error("Transaction is invalid")]
-    Invalid,
+    #[error(transparent)]
+    Invalid(#[from] TransactionValidationError),
 }
 
 /// Origin of a transaction being inserted into the mempool: