Skip to main content

urma_core/multipart/
consistency.rs

1use crate::error::{Context, Error, ensure};
2use crate::multipart::{
3    ChildReference, DataPart, Geometry, LeafManifest, MultipartRecord, RootManifest,
4};
5use bitcoin::{Txid, hashes::Hash};
6use sha2::{Digest, Sha256};
7
8pub struct MultipartConsistency {
9    parts: u32,
10    leaves: u16,
11    length: u64,
12    last_part_length: usize,
13    last_leaf_entries: usize,
14    payload_hash: Sha256,
15    group_hash: Sha256,
16    groups: Vec<[u8; 32]>,
17    leaves_hash: Sha256,
18}
19
20impl Default for MultipartConsistency {
21    fn default() -> Self {
22        Self::new()
23    }
24}
25
26impl MultipartConsistency {
27    pub fn new() -> Self {
28        Self {
29            parts: 0,
30            leaves: 0,
31            length: 0,
32            last_part_length: 0,
33            last_leaf_entries: 0,
34            payload_hash: Sha256::new(),
35            group_hash: Sha256::new(),
36            groups: Vec::new(),
37            leaves_hash: Sha256::new(),
38        }
39    }
40
41    pub fn accept(&mut self, bytes: &[u8], txid: Txid) -> Result<(), Error> {
42        let record = MultipartRecord::decode(bytes)?;
43        let reference = ChildReference {
44            txid,
45            record_hash: Sha256::digest(bytes).into(),
46        };
47        match record {
48            MultipartRecord::Data(part) => self.accept_data(part, reference),
49            MultipartRecord::Leaf(leaf) => self.accept_leaf(leaf, reference),
50            MultipartRecord::Root(root) => self.accept_root(root),
51        }
52    }
53
54    fn accept_data(&mut self, part: DataPart, reference: ChildReference) -> Result<(), Error> {
55        ensure!(
56            part.index == self.parts,
57            "noncanonical data publication order"
58        );
59        if self.parts > 0 {
60            ensure!(
61                self.last_part_length == Geometry::DATA_BYTES,
62                "short nonfinal data part"
63            );
64        }
65        self.payload_hash.update(&part.payload);
66        self.group_hash.update(reference.txid.to_byte_array());
67        self.group_hash.update(reference.record_hash);
68        self.last_part_length = part.payload.len();
69        self.length = self
70            .length
71            .checked_add(u64::try_from(part.payload.len())?)
72            .context("payload length overflow")?;
73        self.parts = self.parts.checked_add(1).context("part count overflow")?;
74        if self.parts.is_multiple_of(u32::from(Geometry::FANOUT)) {
75            self.groups.push(self.group_hash.finalize_reset().into());
76        }
77        Ok(())
78    }
79
80    fn accept_leaf(&mut self, leaf: LeafManifest, reference: ChildReference) -> Result<(), Error> {
81        ensure!(
82            self.parts > 0 && leaf.index == self.leaves,
83            "noncanonical leaf publication order"
84        );
85        if self.leaves > 0 {
86            ensure!(
87                self.last_leaf_entries == usize::from(Geometry::FANOUT),
88                "short nonfinal leaf"
89            );
90        }
91        if self.leaves == 0 && !self.parts.is_multiple_of(u32::from(Geometry::FANOUT)) {
92            self.groups.push(self.group_hash.finalize_reset().into());
93        }
94        ensure!(
95            references_hash(&leaf.entries)
96                == *self
97                    .groups
98                    .get(usize::from(self.leaves))
99                    .context("leaf references exceed data groups")?,
100            "leaf references differ from signed data parts"
101        );
102        self.leaves_hash.update(reference.txid.to_byte_array());
103        self.leaves_hash.update(reference.record_hash);
104        self.last_leaf_entries = leaf.entries.len();
105        self.leaves = self.leaves.checked_add(1).context("leaf count overflow")?;
106        Ok(())
107    }
108
109    fn accept_root(&self, root: RootManifest) -> Result<(), Error> {
110        ensure!(
111            root.payload_hash == <[u8; 32]>::from(self.payload_hash.clone().finalize()),
112            "root payload hash mismatch"
113        );
114        ensure!(
115            references_hash(&root.entries) == <[u8; 32]>::from(self.leaves_hash.clone().finalize()),
116            "root references differ from signed leaves"
117        );
118        let geometry = Geometry::new(root.length)?;
119        ensure!(
120            self.length == root.length
121                && self.parts == geometry.parts()
122                && self.leaves == geometry.leaves(),
123            "publication geometry mismatch"
124        );
125        ensure!(
126            self.last_leaf_entries == usize::from(geometry.leaf_entries(self.leaves - 1)?),
127            "final leaf size mismatch"
128        );
129        Ok(())
130    }
131
132    pub fn validate_count(&self, count: u32) -> Result<(), Error> {
133        ensure!(
134            Geometry::new(self.length)?.nodes() == count,
135            "publication root or geometry missing"
136        );
137        Ok(())
138    }
139}
140
141fn references_hash(entries: &[ChildReference]) -> [u8; 32] {
142    let mut hash = Sha256::new();
143    for entry in entries {
144        hash.update(entry.txid.to_byte_array());
145        hash.update(entry.record_hash);
146    }
147    hash.finalize().into()
148}