use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use crate::merkle::{Hash, MerkleTree};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub struct TreeHead {
pub tree_size: u64,
pub root_hash: Hash,
pub timestamp: DateTime<Utc>,
}
#[derive(Debug, Default)]
pub struct Witness {
heads: HashMap<u64, TreeHead>,
pub witness_id: String,
}
#[derive(Debug, thiserror::Error)]
pub enum WitnessError {
#[error("consistency proof failed: expected {expected:?}, got {actual:?}")]
ConsistencyFailed {
expected: Hash,
actual: Hash,
},
#[error("old tree head (size {0}) not known")]
OldHeadUnknown(u64),
}
impl Witness {
pub fn new(witness_id: &str) -> Self {
Self {
heads: HashMap::new(),
witness_id: witness_id.into(),
}
}
pub fn receive_head(&mut self, head: TreeHead, tree: &MerkleTree) -> Result<(), WitnessError> {
if let Some(old) = self.latest_head() {
if head.tree_size > old.tree_size {
tree.verify_consistency(
old.root_hash,
head.root_hash,
old.tree_size as usize,
head.tree_size as usize,
&[],
)
.map_err(|_| WitnessError::ConsistencyFailed {
expected: head.root_hash,
actual: tree.root(),
})?;
}
}
self.heads.insert(head.tree_size, head);
Ok(())
}
pub fn latest_head(&self) -> Option<&TreeHead> {
self.heads.values().max_by_key(|h| h.tree_size)
}
pub fn known_heads(&self) -> Vec<&TreeHead> {
let mut heads: Vec<&TreeHead> = self.heads.values().collect();
heads.sort_by_key(|h| h.tree_size);
heads
}
pub fn head_at(&self, tree_size: u64) -> Option<&TreeHead> {
self.heads.get(&tree_size)
}
pub fn head_count(&self) -> usize {
self.heads.len()
}
pub fn gossip_to(
&self,
peer: &mut Witness,
tree_size: u64,
tree: &MerkleTree,
) -> Result<(), WitnessError> {
let head = self
.heads
.get(&tree_size)
.ok_or(WitnessError::OldHeadUnknown(tree_size))?;
peer.receive_head(head.clone(), tree)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::entry::{ArtifactType, MerkleEntry};
fn build_tree(n: u64) -> MerkleTree {
let mut tree = MerkleTree::new();
for i in 0..n {
tree.append(MerkleEntry::new(
i,
ArtifactType::ThresholdSignature,
[i as u8; 32],
));
}
tree
}
#[test]
fn witness_starts_empty() {
let w = Witness::new("w1");
assert_eq!(w.head_count(), 0);
assert!(w.latest_head().is_none());
}
#[test]
fn receive_first_head() {
let tree = build_tree(5);
let mut w = Witness::new("w1");
let head = TreeHead {
tree_size: 5,
root_hash: tree.root(),
timestamp: Utc::now(),
};
w.receive_head(head, &tree).unwrap();
assert_eq!(w.head_count(), 1);
assert_eq!(w.latest_head().unwrap().tree_size, 5);
}
#[test]
fn receive_consistent_head() {
let mut tree = build_tree(5);
let mut w = Witness::new("w1");
let head1 = TreeHead {
tree_size: 5,
root_hash: tree.root(),
timestamp: Utc::now(),
};
w.receive_head(head1, &tree).unwrap();
tree.append(MerkleEntry::new(
5,
ArtifactType::ThresholdSignature,
[5; 32],
));
let head2 = TreeHead {
tree_size: 6,
root_hash: tree.root(),
timestamp: Utc::now(),
};
w.receive_head(head2, &tree).unwrap();
assert_eq!(w.head_count(), 2);
assert_eq!(w.latest_head().unwrap().tree_size, 6);
}
#[test]
fn reject_inconsistent_head() {
let tree = build_tree(5);
let mut w = Witness::new("w1");
let head1 = TreeHead {
tree_size: 5,
root_hash: tree.root(),
timestamp: Utc::now(),
};
w.receive_head(head1, &tree).unwrap();
let fake_root = [0xFF; 32];
let head2 = TreeHead {
tree_size: 10,
root_hash: fake_root,
timestamp: Utc::now(),
};
let result = w.receive_head(head2, &tree);
assert!(result.is_err());
}
#[test]
fn known_heads_sorted_by_size() {
let mut tree = build_tree(1);
let mut w = Witness::new("w1");
let head1 = TreeHead {
tree_size: 1,
root_hash: tree.root(),
timestamp: Utc::now(),
};
w.receive_head(head1, &tree).unwrap();
tree.append(MerkleEntry::new(
1,
ArtifactType::ThresholdSignature,
[1; 32],
));
let head2 = TreeHead {
tree_size: 2,
root_hash: tree.root(),
timestamp: Utc::now(),
};
w.receive_head(head2, &tree).unwrap();
tree.append(MerkleEntry::new(
2,
ArtifactType::ThresholdSignature,
[2; 32],
));
let head3 = TreeHead {
tree_size: 3,
root_hash: tree.root(),
timestamp: Utc::now(),
};
w.receive_head(head3, &tree).unwrap();
let heads = w.known_heads();
assert_eq!(heads.len(), 3);
assert_eq!(heads[0].tree_size, 1);
assert_eq!(heads[1].tree_size, 2);
assert_eq!(heads[2].tree_size, 3);
}
#[test]
fn gossip_to_peer() {
let tree = build_tree(5);
let mut w1 = Witness::new("w1");
let mut w2 = Witness::new("w2");
let head = TreeHead {
tree_size: 5,
root_hash: tree.root(),
timestamp: Utc::now(),
};
w1.receive_head(head, &tree).unwrap();
w1.gossip_to(&mut w2, 5, &tree).unwrap();
assert_eq!(w2.head_count(), 1);
assert_eq!(w2.head_at(5).unwrap().root_hash, tree.root());
}
#[test]
fn gossip_unknown_head_errors() {
let tree = build_tree(5);
let w1 = Witness::new("w1");
let mut w2 = Witness::new("w2");
assert!(w1.gossip_to(&mut w2, 99, &tree).is_err());
}
#[test]
fn head_at_returns_correct_head() {
let tree = build_tree(5);
let mut w = Witness::new("w1");
let head = TreeHead {
tree_size: 5,
root_hash: tree.root(),
timestamp: Utc::now(),
};
w.receive_head(head, &tree).unwrap();
assert!(w.head_at(5).is_some());
assert!(w.head_at(3).is_none());
}
}