Add a DB to Tributary

Adds support for reloading most of the blockchain.

Cargo.lock

@@ -10697,6 +10697,7 @@ dependencies = [
  "rand_chacha 0.3.1",
  "rand_core 0.6.4",
  "schnorr-signatures",
+ "serai-db",
  "subtle",
  "tendermint-machine",
  "thiserror",

coordinator/tributary/Cargo.toml

@@ -26,6 +26,8 @@ schnorr = { package = "schnorr-signatures", path = "../../crypto/schnorr" }
 hex = "0.4"
 log = "0.4"
 
+serai-db = { path = "../../common/db" }
+
 scale = { package = "parity-scale-codec", version = "3", features = ["derive"] }
 futures = "0.3"
 tendermint = { package = "tendermint-machine", path = "./tendermint" }

coordinator/tributary/src/blockchain.rs

@@ -1,17 +1,20 @@
 use std::collections::HashMap;
 
-use ciphersuite::{Ciphersuite, Ristretto};
+use ciphersuite::{group::GroupEncoding, Ciphersuite, Ristretto};
+
+use serai_db::{DbTxn, Db};
 
 use crate::{
-  Signed, TransactionKind, Transaction, verify_transaction, ProvidedTransactions, BlockError,
+  ReadWrite, Signed, TransactionKind, Transaction, verify_transaction, ProvidedTransactions,
-  Block, Mempool,
+  BlockError, Block, Mempool,
 };
 
 #[derive(Clone, PartialEq, Eq, Debug)]
-pub(crate) struct Blockchain<T: Transaction> {
+pub(crate) struct Blockchain<D: Db, T: Transaction> {
+  db: Option<D>,
   genesis: [u8; 32],
-  // TODO: db
+
-  block_number: u64,
+  block_number: u32,
   tip: [u8; 32],
   next_nonces: HashMap<<Ristretto as Ciphersuite>::G, u32>,
 
@@ -19,16 +22,43 @@ pub(crate) struct Blockchain<T: Transaction> {
   mempool: Mempool<T>,
 }
 
-impl<T: Transaction> Blockchain<T> {
+impl<D: Db, T: Transaction> Blockchain<D, T> {
-  pub(crate) fn new(genesis: [u8; 32], participants: &[<Ristretto as Ciphersuite>::G]) -> Self {
+  fn tip_key(&self) -> Vec<u8> {
-    // TODO: Reload block_number/tip/next_nonces/provided/mempool
+    D::key(b"tributary", b"tip", self.genesis)
+  }
+  fn block_number_key(&self) -> Vec<u8> {
+    D::key(b"tributary", b"block_number", self.genesis)
+  }
+  fn block_key(&self, hash: &[u8; 32]) -> Vec<u8> {
+    // Since block hashes incorporate their parent, and the first parent is the genesis, this is
+    // fine not incorporating the hash unless there's a hash collision
+    D::key(b"tributary", b"block", hash)
+  }
+  fn commit_key(&self, hash: &[u8; 32]) -> Vec<u8> {
+    D::key(b"tributary", b"commit", hash)
+  }
+  fn next_nonce_key(&self, signer: &<Ristretto as Ciphersuite>::G) -> Vec<u8> {
+    D::key(
+      b"tributary",
+      b"next_nonce",
+      [self.genesis.as_ref(), signer.to_bytes().as_ref()].concat(),
+    )
+  }
+
+  pub(crate) fn new(
+    db: D,
+    genesis: [u8; 32],
+    participants: &[<Ristretto as Ciphersuite>::G],
+  ) -> Self {
+    // TODO: Reload provided/mempool
 
     let mut next_nonces = HashMap::new();
     for participant in participants {
       next_nonces.insert(*participant, 0);
     }
 
-    Self {
+    let mut res = Self {
+      db: Some(db),
       genesis,
 
       block_number: 0,
@@ -37,17 +67,37 @@ impl<T: Transaction> Blockchain<T> {
 
       provided: ProvidedTransactions::new(),
       mempool: Mempool::new(genesis),
+    };
+
+    if let Some((block_number, tip)) = {
+      let db = res.db.as_ref().unwrap();
+      db.get(res.block_number_key()).map(|number| (number, db.get(res.tip_key()).unwrap()))
+    } {
+      res.block_number = u32::from_le_bytes(block_number.try_into().unwrap());
+      res.tip.copy_from_slice(&tip);
     }
+
+    for participant in participants {
+      if let Some(next_nonce) = res.db.as_ref().unwrap().get(res.next_nonce_key(participant)) {
+        res.next_nonces.insert(*participant, u32::from_le_bytes(next_nonce.try_into().unwrap()));
+      }
+    }
+
+    res
   }
 
   pub(crate) fn tip(&self) -> [u8; 32] {
     self.tip
   }
 
-  pub(crate) fn block_number(&self) -> u64 {
+  pub(crate) fn block_number(&self) -> u32 {
     self.block_number
   }
 
+  pub(crate) fn commit(&self, block: &[u8; 32]) -> Option<Vec<u8>> {
+    self.db.as_ref().unwrap().get(self.commit_key(block))
+  }
+
   pub(crate) fn add_transaction(&mut self, internal: bool, tx: T) -> bool {
     self.mempool.add(&self.next_nonces, internal, tx)
   }
@@ -87,12 +137,25 @@ impl<T: Transaction> Blockchain<T> {
   }
 
   /// Add a block.
-  pub(crate) fn add_block(&mut self, block: &Block<T>) -> Result<(), BlockError> {
+  pub(crate) fn add_block(&mut self, block: &Block<T>, commit: Vec<u8>) -> Result<(), BlockError> {
     self.verify_block(block)?;
 
     // None of the following assertions should be reachable since we verified the block
+
+    // Take it from the Option so Rust doesn't consider self as mutably borrowed thanks to the
+    // existence of the txn
+    let mut db = self.db.take().unwrap();
+    let mut txn = db.txn();
+
     self.tip = block.hash();
+    txn.put(self.tip_key(), self.tip);
+
     self.block_number += 1;
+    txn.put(self.block_number_key(), self.block_number.to_le_bytes());
+
+    txn.put(self.block_key(&self.tip), block.serialize());
+    txn.put(self.commit_key(&self.tip), commit);
+
     for tx in &block.transactions {
       match tx.kind() {
         TransactionKind::Provided => {
@@ -100,19 +163,25 @@ impl<T: Transaction> Blockchain<T> {
         }
         TransactionKind::Unsigned => {}
         TransactionKind::Signed(Signed { signer, nonce, .. }) => {
+          let next_nonce = nonce + 1;
           let prev = self
             .next_nonces
-            .insert(*signer, nonce + 1)
+            .insert(*signer, next_nonce)
             .expect("block had signed transaction from non-participant");
           if prev != *nonce {
             panic!("verified block had an invalid nonce");
           }
 
+          txn.put(self.next_nonce_key(signer), next_nonce.to_le_bytes());
+
           self.mempool.remove(&tx.hash());
         }
       }
     }
 
+    txn.commit();
+    self.db = Some(db);
+
     Ok(())
   }
 }

coordinator/tributary/src/lib.rs

@@ -14,11 +14,13 @@ use ciphersuite::{Ciphersuite, Ristretto};
 use scale::Decode;
 use futures::SinkExt;
 use ::tendermint::{
-  ext::{BlockNumber, Commit, Block as BlockTrait, Network as NetworkTrait},
+  ext::{BlockNumber, Commit, Block as BlockTrait, Network},
   SignedMessageFor, SyncedBlock, SyncedBlockSender, MessageSender, TendermintMachine,
   TendermintHandle,
 };
 
+use serai_db::Db;
+
 mod merkle;
 pub(crate) use merkle::*;
 
@@ -80,15 +82,16 @@ impl<P: P2p> P2p for Arc<P> {
   }
 }
 
-pub struct Tributary<T: Transaction, P: P2p> {
+pub struct Tributary<D: Db, T: Transaction, P: P2p> {
-  network: Network<T, P>,
+  network: TendermintNetwork<D, T, P>,
 
-  synced_block: SyncedBlockSender<Network<T, P>>,
+  synced_block: SyncedBlockSender<TendermintNetwork<D, T, P>>,
-  messages: MessageSender<Network<T, P>>,
+  messages: MessageSender<TendermintNetwork<D, T, P>>,
 }
 
-impl<T: Transaction, P: P2p> Tributary<T, P> {
+impl<D: Db, T: Transaction, P: P2p> Tributary<D, T, P> {
   pub async fn new(
+    db: D,
     genesis: [u8; 32],
     start_time: u64,
     key: Zeroizing<<Ristretto as Ciphersuite>::F>,
@@ -100,16 +103,21 @@ impl<T: Transaction, P: P2p> Tributary<T, P> {
     let signer = Arc::new(Signer::new(genesis, key));
     let validators = Arc::new(Validators::new(genesis, validators));
 
-    let mut blockchain = Blockchain::new(genesis, &validators_vec);
+    let mut blockchain = Blockchain::new(db, genesis, &validators_vec);
     let block_number = blockchain.block_number();
-    let start_time = start_time; // TODO: Get the start time from the blockchain
+
+    let start_time = if let Some(commit) = blockchain.commit(&blockchain.tip()) {
+      Commit::<Validators>::decode(&mut commit.as_ref()).unwrap().end_time
+    } else {
+      start_time
+    };
     let proposal = TendermintBlock(blockchain.build_block().serialize());
     let blockchain = Arc::new(RwLock::new(blockchain));
 
-    let network = Network { genesis, signer, validators, blockchain, p2p };
+    let network = TendermintNetwork { genesis, signer, validators, blockchain, p2p };
 
     // The genesis block is 0, so we're working on block #1
-    let block_number = BlockNumber(block_number + 1);
+    let block_number = BlockNumber((block_number + 1).into());
     let TendermintHandle { synced_block, messages, machine } =
       TendermintMachine::new(network.clone(), block_number, start_time, proposal).await;
     tokio::task::spawn(machine.run());
@@ -117,6 +125,8 @@ impl<T: Transaction, P: P2p> Tributary<T, P> {
     Self { network, synced_block, messages }
   }
 
+  // TODO: Is there a race condition with providing these? Since the same provided TX provided
+  // twice counts as two transactions
   pub fn provide_transaction(&self, tx: T) -> bool {
     self.network.blockchain.write().unwrap().provide_transaction(tx)
   }
@@ -156,7 +166,7 @@ impl<T: Transaction, P: P2p> Tributary<T, P> {
       return false;
     }
 
-    let number = BlockNumber(block_number + 1);
+    let number = BlockNumber((block_number + 1).into());
     self.synced_block.send(SyncedBlock { number, block, commit }).await.unwrap();
     true
   }
@@ -175,12 +185,14 @@ impl<T: Transaction, P: P2p> Tributary<T, P> {
       }
 
       TENDERMINT_MESSAGE => {
-        let Ok(msg) = SignedMessageFor::<Network<T, P>>::decode::<&[u8]>(&mut &msg[1 ..]) else {
+        let Ok(msg) = SignedMessageFor::<TendermintNetwork<D, T, P>>::decode::<&[u8]>(
+          &mut &msg[1 ..]
+        ) else {
           return false;
         };
 
         // If this message isn't to form consensus on the next block, ignore it
-        if msg.block().0 != (self.network.blockchain.read().unwrap().block_number() + 1) {
+        if msg.block().0 != (self.network.blockchain.read().unwrap().block_number() + 1).into() {
           return false;
         }
 

coordinator/tributary/src/tendermint.rs

@@ -23,12 +23,14 @@ use ciphersuite::{
 };
 use schnorr::SchnorrSignature;
 
+use serai_db::Db;
+
 use scale::{Encode, Decode};
 use tendermint::{
   SignedMessageFor,
   ext::{
     BlockNumber, RoundNumber, Signer as SignerTrait, SignatureScheme, Weights, Block as BlockTrait,
-    BlockError as TendermintBlockError, Commit, Network as NetworkTrait,
+    BlockError as TendermintBlockError, Commit, Network,
   },
 };
 
@@ -220,16 +222,18 @@ impl BlockTrait for TendermintBlock {
 }
 
 #[derive(Clone, Debug)]
-pub(crate) struct Network<T: Transaction, P: P2p> {
+pub(crate) struct TendermintNetwork<D: Db, T: Transaction, P: P2p> {
   pub(crate) genesis: [u8; 32],
+
   pub(crate) signer: Arc<Signer>,
   pub(crate) validators: Arc<Validators>,
-  pub(crate) blockchain: Arc<RwLock<Blockchain<T>>>,
+  pub(crate) blockchain: Arc<RwLock<Blockchain<D, T>>>,
+
   pub(crate) p2p: P,
 }
 
 #[async_trait]
-impl<T: Transaction, P: P2p> NetworkTrait for Network<T, P> {
+impl<D: Db, T: Transaction, P: P2p> Network for TendermintNetwork<D, T, P> {
   type ValidatorId = [u8; 32];
   type SignatureScheme = Arc<Validators>;
   type Weights = Arc<Validators>;
@@ -284,6 +288,7 @@ impl<T: Transaction, P: P2p> NetworkTrait for Network<T, P> {
       panic!("validators added invalid block to tributary {}", hex::encode(self.genesis));
     };
 
+    // Tendermint should only produce valid commits
     assert!(self.verify_commit(block.id(), &commit));
 
     let Ok(block) = Block::read::<&[u8]>(&mut block.0.as_ref()) else {
@@ -291,7 +296,7 @@ impl<T: Transaction, P: P2p> NetworkTrait for Network<T, P> {
     };
 
     loop {
-      let block_res = self.blockchain.write().unwrap().add_block(&block);
+      let block_res = self.blockchain.write().unwrap().add_block(&block, commit.encode());
       match block_res {
         Ok(()) => break,
         Err(BlockError::NonLocalProvided(hash)) => {
@@ -306,8 +311,6 @@ impl<T: Transaction, P: P2p> NetworkTrait for Network<T, P> {
       }
     }
 
-    // TODO: Save the commit to disk
-
     Some(TendermintBlock(self.blockchain.write().unwrap().build_block().serialize()))
   }
 }

coordinator/tributary/src/tests/blockchain.rs

@@ -5,6 +5,8 @@ use blake2::{Digest, Blake2s256};
 
 use ciphersuite::{group::ff::Field, Ciphersuite, Ristretto};
 
+use serai_db::MemDb;
+
 use crate::{
   merkle, Transaction, ProvidedTransactions, Block, Blockchain,
   tests::{ProvidedTransaction, SignedTransaction, random_provided_transaction},
@@ -19,8 +21,8 @@ fn new_genesis() -> [u8; 32] {
 fn new_blockchain<T: Transaction>(
   genesis: [u8; 32],
   participants: &[<Ristretto as Ciphersuite>::G],
-) -> Blockchain<T> {
+) -> Blockchain<MemDb, T> {
-  let blockchain = Blockchain::new(genesis, participants);
+  let blockchain = Blockchain::new(MemDb::new(), genesis, participants);
   assert_eq!(blockchain.tip(), genesis);
   assert_eq!(blockchain.block_number(), 0);
   blockchain
@@ -34,7 +36,7 @@ fn block_addition() {
   assert_eq!(block.header.parent, genesis);
   assert_eq!(block.header.transactions, [0; 32]);
   blockchain.verify_block(&block).unwrap();
-  assert!(blockchain.add_block(&block).is_ok());
+  assert!(blockchain.add_block(&block, vec![]).is_ok());
   assert_eq!(blockchain.tip(), block.hash());
   assert_eq!(blockchain.block_number(), 1);
 }
@@ -129,7 +131,8 @@ fn signed_transaction() {
   let mut blockchain = new_blockchain::<SignedTransaction>(genesis, &[signer]);
   assert_eq!(blockchain.next_nonce(signer), Some(0));
 
-  let test = |blockchain: &mut Blockchain<SignedTransaction>, mempool: Vec<SignedTransaction>| {
+  let test = |blockchain: &mut Blockchain<MemDb, SignedTransaction>,
+              mempool: Vec<SignedTransaction>| {
     let tip = blockchain.tip();
     for tx in mempool.clone() {
       let next_nonce = blockchain.next_nonce(signer).unwrap();
@@ -151,7 +154,7 @@ fn signed_transaction() {
 
     // Verify and add the block
     blockchain.verify_block(&block).unwrap();
-    assert!(blockchain.add_block(&block).is_ok());
+    assert!(blockchain.add_block(&block, vec![]).is_ok());
     assert_eq!(blockchain.tip(), block.hash());
   };
 
@@ -188,11 +191,11 @@ fn provided_transaction() {
   blockchain.verify_block(&block).unwrap();
 
   // add_block should work for verified blocks
-  assert!(blockchain.add_block(&block).is_ok());
+  assert!(blockchain.add_block(&block, vec![]).is_ok());
 
   let block = Block::new(blockchain.tip(), vec![tx], vec![]);
   // The provided transaction should no longer considered provided, causing this error
   assert!(blockchain.verify_block(&block).is_err());
   // add_block should fail for unverified provided transactions if told to add them
-  assert!(blockchain.add_block(&block).is_err());
+  assert!(blockchain.add_block(&block, vec![]).is_err());
 }