Add a mempool

coordinator/tributary/src/lib.rs

@@ -15,6 +15,9 @@ pub use block::*;
 mod blockchain;
 pub use blockchain::*;
 
+mod mempool;
+pub use mempool::*;
+
 #[cfg(any(test, feature = "tests"))]
 pub mod tests;
 

coordinator/tributary/src/mempool.rs

@@ -0,0 +1,85 @@
+use std::collections::{HashSet, HashMap};
+
+use ciphersuite::{Ciphersuite, Ristretto};
+
+use crate::{Signed, TransactionKind, Transaction, verify_transaction};
+
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub struct Mempool<T: Transaction> {
+  genesis: [u8; 32],
+  txs: HashMap<[u8; 32], T>,
+  next_nonces: HashMap<<Ristretto as Ciphersuite>::G, u32>,
+}
+
+impl<T: Transaction> Mempool<T> {
+  pub fn new(genesis: [u8; 32]) -> Self {
+    Mempool { genesis, txs: HashMap::new(), next_nonces: HashMap::new() }
+  }
+
+  /// Returns true if this is a valid, new transaction.
+  pub fn add(
+    &mut self,
+    blockchain_nonces: &HashMap<<Ristretto as Ciphersuite>::G, u32>,
+    tx: T,
+  ) -> bool {
+    match tx.kind() {
+      TransactionKind::Signed(Signed { signer, nonce, .. }) => {
+        // If the mempool doesn't have a nonce tracked, grab it from the blockchain
+        if !self.next_nonces.contains_key(signer) {
+          // TODO: Same commentary here as present in verify_transaction about a whitelist
+          self.next_nonces.insert(*signer, blockchain_nonces.get(signer).cloned().unwrap_or(0));
+        }
+
+        if verify_transaction(&tx, self.genesis, &mut HashSet::new(), &mut self.next_nonces)
+          .is_err()
+        {
+          return false;
+        }
+        assert_eq!(self.next_nonces[signer], nonce + 1);
+
+        self.txs.insert(tx.hash(), tx);
+        true
+      }
+      _ => false,
+    }
+  }
+
+  pub fn next_nonce(&self, signer: &<Ristretto as Ciphersuite>::G) -> Option<u32> {
+    self.next_nonces.get(signer).cloned()
+  }
+
+  /// Get transactions to include in a block.
+  pub fn block(
+    &mut self,
+    blockchain_nonces: &HashMap<<Ristretto as Ciphersuite>::G, u32>,
+  ) -> HashMap<[u8; 32], T> {
+    let mut res = HashMap::new();
+    for hash in self.txs.keys().cloned().collect::<Vec<_>>() {
+      let tx = &self.txs[&hash];
+      // Verify this hasn't gone stale
+      match tx.kind() {
+        TransactionKind::Signed(Signed { signer, nonce, .. }) => {
+          if blockchain_nonces.get(signer).cloned().unwrap_or(0) > *nonce {
+            self.txs.remove(&hash);
+            continue;
+          }
+        }
+        _ => panic!("non-signed transaction entered mempool"),
+      }
+
+      // Since this TX isn't stale, include it
+      res.insert(hash, tx.clone());
+    }
+    res
+  }
+
+  /// Remove a transaction from the mempool.
+  pub fn remove(&mut self, tx: &[u8; 32]) {
+    self.txs.remove(tx);
+  }
+
+  #[cfg(test)]
+  pub(crate) fn txs(&self) -> &HashMap<[u8; 32], T> {
+    &self.txs
+  }
+}

coordinator/tributary/src/tests/mempool.rs

@@ -0,0 +1,67 @@
+use std::collections::HashMap;
+
+use zeroize::Zeroizing;
+use rand_core::{RngCore, OsRng};
+
+use ciphersuite::{group::ff::Field, Ciphersuite, Ristretto};
+
+use crate::{
+  Transaction, Mempool,
+  tests::{SignedTransaction, signed_transaction},
+};
+
+fn new_mempool<T: Transaction>() -> ([u8; 32], Mempool<T>) {
+  let mut genesis = [0; 32];
+  OsRng.fill_bytes(&mut genesis);
+  (genesis, Mempool::new(genesis))
+}
+
+#[test]
+fn mempool_addition() {
+  let (genesis, mut mempool) = new_mempool::<SignedTransaction>();
+
+  let key = Zeroizing::new(<Ristretto as Ciphersuite>::F::random(&mut OsRng));
+
+  let first_tx = signed_transaction(&mut OsRng, genesis, &key, 0);
+  let signer = first_tx.1.signer;
+  assert_eq!(mempool.next_nonce(&signer), None);
+
+  // Add TX 0
+  assert!(mempool.add(&HashMap::new(), first_tx.clone()));
+  assert_eq!(mempool.next_nonce(&signer), Some(1));
+
+  // Adding it again should fail
+  assert!(!mempool.add(&HashMap::new(), first_tx.clone()));
+
+  // Do the same with the next nonce
+  let second_tx = signed_transaction(&mut OsRng, genesis, &key, 1);
+  assert!(mempool.add(&HashMap::new(), second_tx.clone()));
+  assert_eq!(mempool.next_nonce(&signer), Some(2));
+  assert!(!mempool.add(&HashMap::new(), second_tx.clone()));
+
+  // If the mempool doesn't have a nonce for an account, it should successfully use the
+  // blockchain's
+  let second_key = Zeroizing::new(<Ristretto as Ciphersuite>::F::random(&mut OsRng));
+  let tx = signed_transaction(&mut OsRng, genesis, &second_key, 2);
+  let second_signer = tx.1.signer;
+  assert_eq!(mempool.next_nonce(&second_signer), None);
+  let mut blockchain_nonces = HashMap::from([(second_signer, 2)]);
+  assert!(mempool.add(&blockchain_nonces, tx.clone()));
+  assert_eq!(mempool.next_nonce(&second_signer), Some(3));
+
+  // Getting a block should work
+  let block = mempool.block(&HashMap::new());
+  assert_eq!(block, mempool.block(&blockchain_nonces));
+  assert_eq!(block.len(), 3);
+
+  // If the blockchain says an account had its nonce updated, it should cause a prune
+  blockchain_nonces.insert(signer, 1);
+  let block = mempool.block(&blockchain_nonces);
+  assert_eq!(block.len(), 2);
+  assert!(!block.contains_key(&first_tx.hash()));
+  assert_eq!(mempool.txs(), &block);
+
+  // Removing should also successfully prune
+  mempool.remove(&tx.hash());
+  assert_eq!(mempool.txs(), &HashMap::from([(second_tx.hash(), second_tx)]));
+}

coordinator/tributary/src/tests/mod.rs

@@ -8,3 +8,5 @@ mod merkle;
 mod block;
 #[cfg(test)]
 mod blockchain;
+#[cfg(test)]
+mod mempool;

coordinator/tributary/src/tests/transaction/mod.rs

@@ -25,7 +25,7 @@ mod provided;
 pub fn random_signed<R: RngCore + CryptoRng>(rng: &mut R) -> Signed {
   Signed {
     signer: <Ristretto as Ciphersuite>::G::random(&mut *rng),
-    nonce: u32::try_from(rng.next_u64() >> 32).unwrap(),
+    nonce: u32::try_from(rng.next_u64() >> 32 >> 1).unwrap(),
     signature: SchnorrSignature::<Ristretto> {
       R: <Ristretto as Ciphersuite>::G::random(&mut *rng),
       s: <Ristretto as Ciphersuite>::F::random(rng),

coordinator/tributary/src/transaction.rs

@@ -38,6 +38,9 @@ impl ReadWrite for Signed {
     let mut nonce = [0; 4];
     reader.read_exact(&mut nonce)?;
     let nonce = u32::from_le_bytes(nonce);
+    if nonce >= (u32::MAX - 1) {
+      Err(io::Error::new(io::ErrorKind::Other, "nonce exceeded limit"))?;
+    }
 
     let signature = SchnorrSignature::<Ristretto>::read(reader)?;
 
@@ -88,12 +91,15 @@ pub trait Transaction: Send + Sync + Clone + Eq + Debug + ReadWrite {
   }
 }
 
+// This will only cause mutations when the transaction is valid.
 pub(crate) fn verify_transaction<T: Transaction>(
   tx: &T,
   genesis: [u8; 32],
   locally_provided: &mut HashSet<[u8; 32]>,
   next_nonces: &mut HashMap<<Ristretto as Ciphersuite>::G, u32>,
 ) -> Result<(), TransactionError> {
+  tx.verify()?;
+
   match tx.kind() {
     TransactionKind::Provided => {
       if !locally_provided.remove(&tx.hash()) {
@@ -106,14 +112,15 @@ pub(crate) fn verify_transaction<T: Transaction>(
       if next_nonces.get(signer).cloned().unwrap_or(0) != *nonce {
         Err(TransactionError::Temporal)?;
       }
-      next_nonces.insert(*signer, nonce + 1);
 
       // TODO: Use Schnorr half-aggregation and a batch verification here
       if !signature.verify(*signer, tx.sig_hash(genesis)) {
         Err(TransactionError::Fatal)?;
       }
+
+      next_nonces.insert(*signer, nonce + 1);
     }
   }
 
-  tx.verify()
+  Ok(())
 }