consensus: Implement state transitions for block sync task.

src/consensus/state.rs

@@ -23,7 +23,10 @@ use crate::{
         schnorr::{SchnorrPublic, SchnorrSecret},
     },
     net,
-    node::{Client, State},
+    node::{
+        state::{state_transition, StateUpdate},
+        Client, MemoryState, State,
+    },
     util::serial::{serialize, Encodable, SerialDecodable, SerialEncodable},
     Result,
 };
@@ -807,4 +810,52 @@ impl ValidatorState {
         self.consensus = consensus;
         Ok(())
     }
+
+    // ==========================
+    // State transition functions
+    // ==========================
+
+    /// Validate state transitions for given transactions and state and
+    /// return a vector of [`StateUpdate`]
+    pub fn validate_state_transitions(state: MemoryState, txs: &[Tx]) -> Result<Vec<StateUpdate>> {
+        let mut ret = vec![];
+        let mut st = state.clone();
+
+        for (i, tx) in txs.iter().enumerate() {
+            let update = match state_transition(&st, tx.0.clone()) {
+                Ok(v) => v,
+                Err(e) => {
+                    warn!("validate_state_transition(): Failed for tx {}: {}", i, e);
+                    return Err(e.into())
+                }
+            };
+            st.apply(update.clone());
+            ret.push(update);
+        }
+
+        Ok(ret)
+    }
+
+    /// Apply a vector of [`StateUpdate`] to the canonical state.
+    pub async fn update_canon_state(
+        &self,
+        updates: Vec<StateUpdate>,
+        notify: Option<async_channel::Sender<(PublicKey, u64)>>,
+    ) -> Result<()> {
+        let secret_keys: Vec<SecretKey> =
+            self.client.get_keypairs().await?.iter().map(|x| x.secret).collect();
+
+        debug!("update_canon_state(): Acquiring state machine lock");
+        let mut state = self.state_machine.lock().await;
+        for update in updates {
+            state
+                .apply(update, secret_keys.clone(), notify.clone(), self.client.wallet.clone())
+                .await?;
+        }
+        drop(state);
+        debug!("update_canon_state(): Dropped state machine lock");
+
+        debug!("update_canon_state(): Successfully applied state updates");
+        Ok(())
+    }
 }

src/consensus/task/block_sync.rs

@@ -1,11 +1,13 @@
 use crate::{
     consensus::{
         block::{BlockOrder, BlockResponse},
-        ValidatorStatePtr,
+        ValidatorState, ValidatorStatePtr,
     },
-    net, Result,
+    net,
+    node::MemoryState,
+    Result,
 };
-use log::{info, warn};
+use log::{debug, info, warn};
 
 /// async task used for block syncing.
 pub async fn block_sync_task(p2p: net::P2pPtr, state: ValidatorStatePtr) -> Result<()> {
@@ -35,8 +37,30 @@ pub async fn block_sync_task(p2p: net::P2pPtr, state: ValidatorStatePtr) -> Resu
             let order = BlockOrder { sl: last.0, block: last.1 };
             channel.send(order).await?;
 
-            // Node stores response data. Extra validations can be added here.
+            // Node stores response data.
             let resp = response_sub.receive().await?;
+
+            // Verify state transitions for all blocks and their respective transactions.
+            debug!("block_sync_task(): Starting state transition validations");
+            let mut canon_updates = vec![];
+            let canon_state_clone = state.read().await.state_machine.lock().await.clone();
+            let mut mem_state = MemoryState::new(canon_state_clone);
+            for block in &resp.blocks {
+                let mut state_updates =
+                    ValidatorState::validate_state_transitions(mem_state.clone(), &block.txs)?;
+
+                for update in &state_updates {
+                    mem_state.apply(update.clone());
+                }
+
+                canon_updates.append(&mut state_updates);
+            }
+            debug!("block_sync_task(): All state transitions passed");
+
+            debug!("block_sync_task(): Updating canon state");
+            state.write().await.update_canon_state(canon_updates, None).await?;
+
+            debug!("block_sync_task(): Appending blocks to ledger");
             state.write().await.blockchain.add(&resp.blocks)?;
 
             let last_received = state.read().await.blockchain.last()?.unwrap();

src/crypto/proof.rs

@@ -15,7 +15,7 @@ use crate::{
     Result,
 };
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct VerifyingKey {
     pub params: Params<vesta::Affine>,
     pub vk: plonk::VerifyingKey<vesta::Affine>,
@@ -29,7 +29,7 @@ impl VerifyingKey {
     }
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct ProvingKey {
     pub params: Params<vesta::Affine>,
     pub pk: plonk::ProvingKey<vesta::Affine>,

src/node/client.rs

@@ -8,7 +8,7 @@ use crate::{
     crypto::{
         address::Address,
         coin::Coin,
-        keypair::{Keypair, PublicKey, SecretKey},
+        keypair::{Keypair, PublicKey},
         merkle_node::MerkleNode,
         proof::ProvingKey,
         types::DrkTokenId,
@@ -196,31 +196,6 @@ impl Client {
         Err(ClientFailed::NotEnoughValue(amount))
     }
 
-    // TODO: Should this function run on finalized blocks and iterate over its transactions?
-    async fn update_state(
-        secret_keys: Vec<SecretKey>,
-        tx: Transaction,
-        state: Arc<Mutex<State>>,
-        wallet: WalletPtr,
-        notify: Option<async_channel::Sender<(PublicKey, u64)>>,
-    ) -> Result<()> {
-        debug!("update_state(): Begin state update");
-        debug!("update_state(): Acquiring state lock");
-        let update;
-        {
-            let state = &*state.lock().await;
-            update = state_transition(state, tx)?;
-        }
-
-        debug!("update_state(): Trying to apply the new state");
-        let mut state = state.lock().await;
-        state.apply(update, secret_keys, notify, wallet).await?;
-        drop(state);
-        debug!("update_state(): Successfully updated state");
-
-        Ok(())
-    }
-
     pub async fn init_db(&self) -> Result<()> {
         self.wallet.init_db().await
     }

src/node/memorystate.rs

@@ -7,6 +7,7 @@ use crate::crypto::{
 };
 
 /// In-memory state extension for state transition validations
+#[derive(Clone)]
 pub struct MemoryState {
     /// Canonical state
     pub canon: State,
@@ -45,7 +46,16 @@ impl ProgramState for MemoryState {
 }
 
 impl MemoryState {
-    pub async fn apply(&mut self, update: StateUpdate) {
+    pub fn new(canon_state: State) -> Self {
+        Self {
+            canon: canon_state.clone(),
+            tree: canon_state.tree,
+            merkle_roots: vec![],
+            nullifiers: vec![],
+        }
+    }
+
+    pub fn apply(&mut self, update: StateUpdate) {
         debug!(target: "state_apply", "(in-memory) Extend nullifier set");
         let mut nfs = update.nullifiers.clone();
         self.nullifiers.append(&mut nfs);

src/node/state.rs

@@ -37,6 +37,7 @@ pub trait ProgramState {
 
 /// A struct representing a state update.
 /// This gets applied on top of an existing state.
+#[derive(Clone)]
 pub struct StateUpdate {
     /// All nullifiers in a transaction
     pub nullifiers: Vec<Nullifier>,
@@ -109,6 +110,7 @@ pub fn state_transition<S: ProgramState>(state: &S, tx: Transaction) -> VerifyRe
 }
 
 /// Struct holding the state which we can apply a [`StateUpdate`] onto.
+#[derive(Clone)]
 pub struct State {
     /// The entire Merkle tree state
     pub tree: BridgeTree<MerkleNode, 32>,