diff --git a/coordinator/tributary/src/lib.rs b/coordinator/tributary/src/lib.rs
index 446af37e..1040b379 100644
--- a/coordinator/tributary/src/lib.rs
+++ b/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;
diff --git a/coordinator/tributary/src/mempool.rs b/coordinator/tributary/src/mempool.rs
new file mode 100644
index 00000000..9f647b00
--- /dev/null
+++ b/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
+ }
+}
diff --git a/coordinator/tributary/src/tests/mempool.rs b/coordinator/tributary/src/tests/mempool.rs
new file mode 100644
index 00000000..7556e13a
--- /dev/null
+++ b/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)]));
+}
diff --git a/coordinator/tributary/src/tests/mod.rs b/coordinator/tributary/src/tests/mod.rs
index 01361e90..fb808477 100644
--- a/coordinator/tributary/src/tests/mod.rs
+++ b/coordinator/tributary/src/tests/mod.rs
@@ -8,3 +8,5 @@ mod merkle;
mod block;
#[cfg(test)]
mod blockchain;
+#[cfg(test)]
+mod mempool;
diff --git a/coordinator/tributary/src/tests/transaction/mod.rs b/coordinator/tributary/src/tests/transaction/mod.rs
index eeafd7e5..1ecafa55 100644
--- a/coordinator/tributary/src/tests/transaction/mod.rs
+++ b/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),
diff --git a/coordinator/tributary/src/transaction.rs b/coordinator/tributary/src/transaction.rs
index 8a63b852..b829af83 100644
--- a/coordinator/tributary/src/transaction.rs
+++ b/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(())
}