use crate::{
journal::contiguous::{variable, Contiguous},
merkle::{
self, compact, Family, Location, Proof, MAX_PINNED_NODES, MAX_PROOF_DIGESTS_PER_ELEMENT,
},
qmdb::{self, sync::compact::Target, Error},
Context,
};
use commonware_codec::{Decode as _, EncodeSize, Read, Write};
use commonware_cryptography::{Digest, Hasher};
use commonware_parallel::Strategy;
use commonware_utils::sync::RwLock;
use std::sync::atomic::{AtomicBool, Ordering};
#[derive(Clone)]
pub(crate) struct Witness<F: Family, D: Digest> {
pub(crate) op_bytes: Vec<u8>,
pub(crate) proof: Proof<F, D>,
pub(crate) pinned_nodes: Vec<D>,
}
impl<F: Family, D: Digest> EncodeSize for Witness<F, D> {
fn encode_size(&self) -> usize {
self.op_bytes.encode_size() + self.proof.encode_size() + self.pinned_nodes.encode_size()
}
}
impl<F: Family, D: Digest> Write for Witness<F, D> {
fn write(&self, buf: &mut impl bytes::BufMut) {
self.op_bytes.write(buf);
self.proof.write(buf);
self.pinned_nodes.write(buf);
}
}
impl<F: Family, D: Digest> Read for Witness<F, D> {
type Cfg = ();
fn read_cfg(buf: &mut impl bytes::Buf, _: &()) -> Result<Self, commonware_codec::Error> {
let op_bytes = Vec::<u8>::read_cfg(buf, &((..).into(), ()))?;
let proof = Proof::<F, D>::read_cfg(buf, &MAX_PROOF_DIGESTS_PER_ELEMENT)?;
let pinned_nodes = Vec::<D>::read_cfg(buf, &((..=MAX_PINNED_NODES).into(), ()))?;
Ok(Self {
op_bytes,
proof,
pinned_nodes,
})
}
}
#[cfg(feature = "arbitrary")]
impl<F: Family, D: Digest> arbitrary::Arbitrary<'_> for Witness<F, D>
where
D: for<'a> arbitrary::Arbitrary<'a>,
{
fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
Ok(Self {
op_bytes: u.arbitrary()?,
proof: u.arbitrary()?,
pinned_nodes: u.arbitrary()?,
})
}
}
#[derive(Clone)]
pub(crate) struct VerifiedWitness<F: Family, D: Digest> {
pub(crate) witness: Witness<F, D>,
pub(crate) root: D,
}
impl<F: Family, D: Digest> VerifiedWitness<F, D> {
pub(crate) const fn leaf_count(&self) -> Location<F> {
self.witness.proof.leaves
}
pub(crate) const fn target(&self) -> Target<F, D> {
Target {
root: self.root,
leaf_count: self.leaf_count(),
}
}
}
pub(crate) type Journal<E, F, D> = variable::Journal<E, Witness<F, D>>;
#[derive(Clone, Copy)]
enum Durability {
Commit,
Sync,
}
pub(crate) struct Store<E: Context, F: Family, D: Digest> {
journal: Journal<E, F, D>,
tip_witness: RwLock<VerifiedWitness<F, D>>,
import_pending: AtomicBool,
}
impl<E: Context, F: Family, D: Digest> Store<E, F, D> {
pub(crate) const fn new(journal: Journal<E, F, D>, witness: VerifiedWitness<F, D>) -> Self {
Self {
journal,
tip_witness: RwLock::new(witness),
import_pending: AtomicBool::new(false),
}
}
pub(crate) const fn from_import(
journal: Journal<E, F, D>,
witness: VerifiedWitness<F, D>,
) -> Self {
Self {
journal,
tip_witness: RwLock::new(witness),
import_pending: AtomicBool::new(true),
}
}
pub(crate) fn with<R>(&self, f: impl FnOnce(&VerifiedWitness<F, D>) -> R) -> R {
f(&self.tip_witness.read())
}
pub(crate) fn replace(&self, witness: VerifiedWitness<F, D>) {
*self.tip_witness.write() = witness;
}
pub(crate) async fn commit<H, S>(
&mut self,
merkle: &compact::Merkle<F, D, S>,
inactivity_floor_loc: Location<F>,
last_commit_op_bytes: impl FnOnce() -> Vec<u8>,
) -> Result<(), Error<F>>
where
H: Hasher<Digest = D>,
S: Strategy,
{
self.persist::<H, S>(
merkle,
inactivity_floor_loc,
last_commit_op_bytes,
Durability::Commit,
)
.await
}
pub(crate) async fn sync<H, S>(
&mut self,
merkle: &compact::Merkle<F, D, S>,
inactivity_floor_loc: Location<F>,
last_commit_op_bytes: impl FnOnce() -> Vec<u8>,
) -> Result<(), Error<F>>
where
H: Hasher<Digest = D>,
S: Strategy,
{
self.persist::<H, S>(
merkle,
inactivity_floor_loc,
last_commit_op_bytes,
Durability::Sync,
)
.await
}
async fn persist<H, S>(
&mut self,
merkle: &compact::Merkle<F, D, S>,
inactivity_floor_loc: Location<F>,
last_commit_op_bytes: impl FnOnce() -> Vec<u8>,
durability: Durability,
) -> Result<(), Error<F>>
where
H: Hasher<Digest = D>,
S: Strategy,
{
let Some(verified) = self
.stage::<H, S>(merkle, inactivity_floor_loc, last_commit_op_bytes)
.await?
else {
return Ok(());
};
self.journal.append(&verified.witness).await?;
match durability {
Durability::Commit => self.journal.commit().await?,
Durability::Sync => self.journal.sync().await?,
}
self.import_pending.store(false, Ordering::Relaxed);
merkle.prune_to_frontier();
self.replace(verified);
Ok(())
}
async fn stage<H, S>(
&mut self,
merkle: &compact::Merkle<F, D, S>,
inactivity_floor_loc: Location<F>,
last_commit_op_bytes: impl FnOnce() -> Vec<u8>,
) -> Result<Option<VerifiedWitness<F, D>>, Error<F>>
where
H: Hasher<Digest = D>,
S: Strategy,
{
let cached_leaves = self.with(|w| w.leaf_count());
let verified = if cached_leaves == merkle.leaves() {
if !self.import_pending.load(Ordering::Relaxed) {
return Ok(None);
}
self.with(|w| w.clone())
} else if cached_leaves > merkle.leaves() {
return Err(Error::DataCorrupted("witness ahead of in-memory state"));
} else {
build_witness::<F, H, S>(merkle, inactivity_floor_loc, last_commit_op_bytes())?
};
if self.import_pending.load(Ordering::Relaxed) {
self.clear_for_import().await?;
}
Ok(Some(verified))
}
pub(crate) async fn rewind<H, S, Op>(
&mut self,
merkle: &compact::Merkle<F, D, S>,
target: Location<F>,
commit_codec_config: &Op::Cfg,
last_commit_floor: impl FnOnce(&Op) -> Option<Location<F>>,
) -> Result<Op, Error<F>>
where
H: Hasher<Digest = D>,
S: Strategy,
Op: Read,
{
self.check_import_persisted()?;
let (pos, entry) = self
.position_of(target)
.await?
.ok_or(Error::Merkle(merkle::Error::RewindBeyondHistory))?;
let (witness, op) = rebuild_and_verify::<F, D, H, S, Op>(
entry,
merkle,
commit_codec_config,
last_commit_floor,
)?;
self.journal.rewind(pos + 1).await?;
self.journal.sync().await?;
self.replace(witness);
Ok(op)
}
pub(crate) async fn prune(&mut self, pruning_boundary: Location<F>) -> Result<(), Error<F>> {
self.check_import_persisted()?;
let bounds = self.journal.bounds();
if bounds.is_empty() {
return Ok(());
}
let pos = Self::first_at_or_above(&self.journal, pruning_boundary)
.await?
.min(bounds.end - 1);
self.journal.prune(pos).await?;
self.journal.sync().await?;
Ok(())
}
pub(crate) fn import_pending(&self) -> bool {
self.import_pending.load(Ordering::Relaxed)
}
fn check_import_persisted(&self) -> Result<(), Error<F>> {
if self.import_pending.load(Ordering::Relaxed) {
return Err(Error::DataCorrupted("compact-sync import not persisted"));
}
Ok(())
}
async fn position_of(
&self,
target: Location<F>,
) -> Result<Option<(u64, Witness<F, D>)>, Error<F>> {
let pos = Self::first_at_or_above(&self.journal, target).await?;
if pos >= self.journal.bounds().end {
return Ok(None);
}
let entry = self.journal.read(pos).await?;
Ok((entry.proof.leaves == target).then_some((pos, entry)))
}
async fn first_at_or_above(
reader: &impl Contiguous<Item = Witness<F, D>>,
leaf_count: Location<F>,
) -> Result<u64, Error<F>> {
let bounds = reader.bounds();
let (mut lo, mut hi) = (bounds.start, bounds.end);
while lo < hi {
let mid = lo + (hi - lo) / 2;
if reader.read(mid).await?.proof.leaves < leaf_count {
lo = mid + 1;
} else {
hi = mid;
}
}
Ok(lo)
}
async fn clear_for_import(&mut self) -> Result<(), Error<F>> {
let size = self.journal.size();
self.journal.clear_to_size(size.max(1)).await?;
Ok(())
}
pub(crate) async fn destroy(self) -> Result<(), Error<F>> {
self.journal.destroy().await?;
Ok(())
}
}
fn build_witness<F, H, S>(
merkle: &compact::Merkle<F, H::Digest, S>,
inactivity_floor_loc: Location<F>,
last_commit_op_bytes: Vec<u8>,
) -> Result<VerifiedWitness<F, H::Digest>, Error<F>>
where
F: Family,
H: Hasher,
S: Strategy,
{
let hasher = qmdb::hasher::<H>();
merkle.with_mem(|mem| {
let leaf_count = mem.leaves();
let last_commit_loc = Location::new(*leaf_count - 1);
let inactive_peaks =
F::inactive_peaks(F::location_to_position(leaf_count), inactivity_floor_loc);
let root = mem.root(&hasher, inactive_peaks)?;
let pinned_nodes = F::nodes_to_pin(leaf_count)
.map(|pos| *mem.get_node_unchecked(pos))
.collect::<Vec<_>>();
let proof = mem.proof(&hasher, last_commit_loc, inactive_peaks)?;
Ok(VerifiedWitness {
witness: Witness {
op_bytes: last_commit_op_bytes,
proof,
pinned_nodes,
},
root,
})
})
}
pub(crate) fn validate_inactivity_floor<F: Family>(
inactivity_floor_loc: Location<F>,
last_commit_loc: Location<F>,
) -> Result<(), Error<F>> {
if inactivity_floor_loc > last_commit_loc {
return Err(Error::DataCorrupted("invalid compact witness"));
}
Ok(())
}
async fn load_tip<E, F, H, S, Op>(
journal: &Journal<E, F, H::Digest>,
merkle: &compact::Merkle<F, H::Digest, S>,
commit_codec_config: &Op::Cfg,
last_commit_floor: impl FnOnce(&Op) -> Option<Location<F>>,
) -> Result<(VerifiedWitness<F, H::Digest>, Op), Error<F>>
where
E: Context,
F: Family,
H: Hasher,
S: Strategy,
Op: Read,
{
let size = journal.size();
if size == 0 {
return Err(Error::DataCorrupted("missing compact witness"));
}
let entry = journal.read(size - 1).await?;
rebuild_and_verify::<F, H::Digest, H, S, Op>(
entry,
merkle,
commit_codec_config,
last_commit_floor,
)
}
fn rebuild_and_verify<F, D, H, S, Op>(
witness: Witness<F, D>,
merkle: &compact::Merkle<F, D, S>,
commit_codec_config: &Op::Cfg,
last_commit_floor: impl FnOnce(&Op) -> Option<Location<F>>,
) -> Result<(VerifiedWitness<F, D>, Op), Error<F>>
where
F: Family,
D: Digest,
H: Hasher<Digest = D>,
S: Strategy,
Op: Read,
{
let leaf_count = witness.proof.leaves;
if leaf_count == 0 {
return Err(Error::DataCorrupted("invalid compact witness"));
}
let last_commit_loc = Location::new(*leaf_count - 1);
let last_commit_op = Op::decode_cfg(witness.op_bytes.as_ref(), commit_codec_config)
.map_err(|_| Error::DataCorrupted("invalid commit operation"))?;
let inactivity_floor_loc = last_commit_floor(&last_commit_op)
.ok_or(Error::DataCorrupted("last operation was not a commit"))?;
validate_inactivity_floor(inactivity_floor_loc, last_commit_loc)?;
merkle
.reset_to(leaf_count, witness.pinned_nodes.clone())
.map_err(|_| Error::DataCorrupted("invalid compact witness"))?;
let inactive_peaks =
F::inactive_peaks(F::location_to_position(leaf_count), inactivity_floor_loc);
let hasher = qmdb::hasher::<H>();
let root = merkle
.root(&hasher, inactive_peaks)
.map_err(|_| Error::DataCorrupted("failed to compute compact witness root"))?;
if !witness.proof.verify_range_inclusion(
&hasher,
&[witness.op_bytes.as_slice()],
last_commit_loc,
&root,
) {
return Err(Error::DataCorrupted("invalid compact witness"));
}
Ok((VerifiedWitness { witness, root }, last_commit_op))
}
pub(crate) async fn init<E, F, H, S, Op>(
mut journal: Journal<E, F, H::Digest>,
merkle: &mut compact::Merkle<F, H::Digest, S>,
commit_codec_config: &Op::Cfg,
initial_commit_op_bytes: Vec<u8>,
last_commit_floor: impl FnOnce(&Op) -> Option<Location<F>>,
) -> Result<(Store<E, F, H::Digest>, Op), Error<F>>
where
E: Context,
F: Family,
H: Hasher,
S: Strategy,
Op: Read,
{
if journal.size() == 0 {
bootstrap_initial_commit::<E, F, H, S>(&mut journal, merkle, initial_commit_op_bytes)
.await?;
}
let (witness, op) =
load_tip::<E, F, H, S, Op>(&journal, merkle, commit_codec_config, last_commit_floor)
.await?;
Ok((Store::new(journal, witness), op))
}
async fn bootstrap_initial_commit<E, F, H, S>(
journal: &mut Journal<E, F, H::Digest>,
merkle: &mut compact::Merkle<F, H::Digest, S>,
last_commit_op_bytes: Vec<u8>,
) -> Result<(), Error<F>>
where
E: Context,
F: Family,
H: Hasher,
S: Strategy,
{
let hasher = qmdb::hasher::<H>();
let batch = {
let batch = merkle.new_batch().add(&hasher, &last_commit_op_bytes);
merkle.with_mem(|mem| batch.merkleize(mem, &hasher))
};
merkle.apply_batch(&batch)?;
let verified = build_witness::<F, H, S>(merkle, Location::new(0), last_commit_op_bytes)?;
journal.append(&verified.witness).await?;
journal.sync().await?;
Ok(())
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
#[cfg(feature = "arbitrary")]
mod conformance {
use super::*;
use crate::merkle::{mmb, mmr};
use commonware_codec::conformance::CodecConformance;
use commonware_cryptography::sha256;
commonware_conformance::conformance_tests! {
CodecConformance<Witness<mmr::Family, sha256::Digest>>,
CodecConformance<Witness<mmb::Family, sha256::Digest>>,
}
}
pub(crate) async fn corrupt_entry<E, F, D>(
journal: &mut Journal<E, F, D>,
pos: u64,
f: impl FnOnce(&mut Witness<F, D>),
) where
E: Context,
F: Family,
D: Digest,
{
let mut entries = Vec::new();
{
for p in pos..journal.bounds().end {
entries.push(journal.read(p).await.unwrap());
}
}
f(&mut entries[0]);
journal.rewind(pos).await.unwrap();
for entry in &entries {
journal.append(entry).await.unwrap();
}
journal.sync().await.unwrap();
}
pub(crate) async fn tip<E, F, D>(journal: &Journal<E, F, D>) -> (Vec<u8>, Proof<F, D>, Vec<D>)
where
E: Context,
F: Family,
D: Digest,
{
let size = journal.size();
let entry = journal.read(size - 1).await.unwrap();
(entry.op_bytes, entry.proof, entry.pinned_nodes)
}
pub(crate) async fn append_unsynced<E, F, D>(
journal: &mut Journal<E, F, D>,
op_bytes: Vec<u8>,
proof: Proof<F, D>,
pinned_nodes: Vec<D>,
) where
E: Context,
F: Family,
D: Digest,
{
journal
.append(&Witness {
op_bytes,
proof,
pinned_nodes,
})
.await
.unwrap();
}
pub(crate) async fn overwrite_tip<E, F, D>(
journal: &mut Journal<E, F, D>,
op_bytes: Vec<u8>,
proof: Proof<F, D>,
pinned_nodes: Vec<D>,
) where
E: Context,
F: Family,
D: Digest,
{
let size = journal.size();
journal.rewind(size - 1).await.unwrap();
journal
.append(&Witness {
op_bytes,
proof,
pinned_nodes,
})
.await
.unwrap();
journal.sync().await.unwrap();
}
}