1use super::{
2 ChildReference, Geometry, ManifestInventory, MultipartRecord, RecordRequest, RootManifest,
3 VerifiedRecord,
4};
5use crate::error::Error as ServiceError;
6use bitcoin::{Txid, XOnlyPublicKey, hashes::Hash};
7use sha2::{Digest, Sha256};
8use std::{
9 fmt,
10 fs::File,
11 io::{Read, Seek, SeekFrom, Write},
12 path::Path,
13};
14use urma_core::error::Error;
15
16#[derive(Debug)]
17pub enum FetchError {
18 Unavailable,
19 Rejected(Error),
20 Source(ServiceError),
21}
22
23pub trait MultipartSource {
24 fn fetch(&mut self, request: &RecordRequest) -> Result<VerifiedRecord, FetchError>;
25
26 fn prefetch(&mut self, _requests: &[RecordRequest]) {}
27}
28
29#[derive(Clone, Copy, Debug)]
30pub struct RecoveryLimits {
31 pub max_payload_bytes: u64,
32 pub max_nodes: u32,
33}
34
35#[derive(Debug)]
36pub enum RecoveryError {
37 InvalidObject(Error),
38 InvalidCandidate { txid: Txid, cause: Error },
39 Incomplete { txid: Txid },
40 Source { txid: Txid, cause: ServiceError },
41 Capacity(String),
42 Storage(std::io::Error),
43}
44
45impl fmt::Display for RecoveryError {
46 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
47 match self {
48 Self::InvalidObject(cause) => write!(formatter, "invalid multipart object: {cause}"),
49 Self::InvalidCandidate { txid, cause } => {
50 write!(formatter, "invalid candidate {txid}: {cause}")
51 }
52 Self::Incomplete { txid } => write!(formatter, "incomplete object; missing {txid}"),
53 Self::Source { txid, cause } => write!(formatter, "source failure for {txid}: {cause}"),
54 Self::Capacity(message) => write!(formatter, "multipart capacity: {message}"),
55 Self::Storage(cause) => write!(formatter, "multipart storage: {cause}"),
56 }
57 }
58}
59impl std::error::Error for RecoveryError {}
60impl From<Error> for RecoveryError {
61 fn from(cause: Error) -> Self {
62 Self::InvalidObject(cause)
63 }
64}
65impl From<std::io::Error> for RecoveryError {
66 fn from(cause: std::io::Error) -> Self {
67 Self::Storage(cause)
68 }
69}
70
71#[derive(Debug)]
72pub struct RecoveredObject {
73 root: Txid,
74 author: XOnlyPublicKey,
75 manifest: RootManifest,
76 payload: File,
77}
78
79impl RecoveredObject {
80 pub fn root(&self) -> Txid {
81 self.root
82 }
83 pub fn author(&self) -> XOnlyPublicKey {
84 self.author
85 }
86 pub fn manifest(&self) -> &RootManifest {
87 &self.manifest
88 }
89}
90impl Read for RecoveredObject {
91 fn read(&mut self, bytes: &mut [u8]) -> std::io::Result<usize> {
92 self.payload.read(bytes)
93 }
94}
95impl Seek for RecoveredObject {
96 fn seek(&mut self, position: SeekFrom) -> std::io::Result<u64> {
97 self.payload.seek(position)
98 }
99}
100
101fn fetch<S: MultipartSource>(
102 source: &mut S,
103 reference: ChildReference,
104 author: XOnlyPublicKey,
105) -> Result<VerifiedRecord, RecoveryError> {
106 let request = RecordRequest { reference };
107 let txid = reference.txid;
108 let verified = source.fetch(&request).map_err(|error| match error {
109 FetchError::Unavailable => RecoveryError::Incomplete { txid },
110 FetchError::Rejected(cause) => RecoveryError::InvalidCandidate { txid, cause },
111 FetchError::Source(cause) => RecoveryError::Source { txid, cause },
112 })?;
113 request
114 .check(&verified)
115 .map_err(|cause| RecoveryError::InvalidCandidate { txid, cause })?;
116 if verified.author() != author {
117 return Err(Error::Invalid("referenced child author differs from root".into()).into());
118 }
119 Ok(verified)
120}
121
122fn inventory_error(cause: Error) -> RecoveryError {
123 match cause {
124 Error::Allocation(error) => RecoveryError::Capacity(error.to_string()),
125 error => RecoveryError::InvalidObject(error),
126 }
127}
128
129fn inventory<S: MultipartSource>(
130 root: &VerifiedRecord,
131 manifest: &RootManifest,
132 source: &mut S,
133 file: &mut File,
134) -> Result<(), RecoveryError> {
135 let mut inventory = ManifestInventory::new(root.txid(), manifest).map_err(inventory_error)?;
136 let mut received = 0u64;
137 tracing::info!(target: "urma_ui", phase = "Receiving manifests", done = received, total = manifest.entries.len());
138 for batch in manifest.entries.chunks(8) {
139 let requests: Vec<_> = batch
140 .iter()
141 .map(|reference| RecordRequest {
142 reference: *reference,
143 })
144 .collect();
145 source.prefetch(&requests);
146 for entry in batch {
147 let verified = fetch(source, *entry, root.author())?;
148 let MultipartRecord::Leaf(leaf) = verified.decode()? else {
149 return Err(Error::Invalid("root must reference leaf manifests".into()).into());
150 };
151 inventory.accept_leaf(&leaf).map_err(inventory_error)?;
152 for reference in leaf.entries {
153 file.write_all(reference.txid.as_byte_array())?;
154 file.write_all(&reference.record_hash)?;
155 }
156 received += 1;
157 tracing::info!(target: "urma_ui", phase = "Receiving manifests", done = received, total = manifest.entries.len());
158 }
159 }
160 file.rewind()?;
161 Ok(())
162}
163
164fn read_requests(file: &mut File, count: u32) -> Result<Vec<RecordRequest>, RecoveryError> {
165 let mut requests = Vec::new();
166 for _ in 0..count {
167 let mut txid = [0; 32];
168 let mut record_hash = [0; 32];
169 file.read_exact(&mut txid)?;
170 file.read_exact(&mut record_hash)?;
171 requests.push(RecordRequest {
172 reference: ChildReference {
173 txid: Txid::from_byte_array(txid),
174 record_hash,
175 },
176 });
177 }
178 Ok(requests)
179}
180
181fn reconstruct_parts<S: MultipartSource>(
182 source: &mut S,
183 author: XOnlyPublicKey,
184 manifest: &RootManifest,
185 geometry: Geometry,
186 inventory: &mut File,
187 payload: &mut File,
188) -> Result<(), RecoveryError> {
189 let mut digest = Sha256::new();
190 let mut length = 0u64;
191 tracing::info!(target: "urma_ui", phase = "Receiving payload bytes", done = length, total = manifest.length);
192 for start in (0..geometry.parts()).step_by(8) {
193 let count = (geometry.parts() - start).min(8);
194 let requests = read_requests(inventory, count)?;
195 source.prefetch(&requests);
196 for (offset, request) in requests.into_iter().enumerate() {
197 let index = start + u32::try_from(offset).map_err(Error::from)?;
198 let verified = fetch(source, request.reference, author)?;
199 let MultipartRecord::Data(part) = verified.decode()? else {
200 return Err(Error::Invalid("leaf must reference data parts".into()).into());
201 };
202 if part.index != index || part.payload.len() != geometry.part_length(index)? {
203 return Err(
204 Error::Invalid("data position or length disagrees with root".into()).into(),
205 );
206 }
207 length = length
208 .checked_add(u64::try_from(part.payload.len()).map_err(Error::from)?)
209 .ok_or_else(|| Error::Invalid("reconstruction length overflow".into()))?;
210 digest.update(&part.payload);
211 payload.write_all(&part.payload)?;
212 tracing::info!(target: "urma_ui", phase = "Receiving payload bytes", done = length, total = manifest.length);
213 }
214 }
215 let actual_hash: [u8; 32] = digest.finalize().into();
216 tracing::info!(target: "urma_ui", phase = "Verifying complete payload digest");
217 if length != manifest.length || actual_hash != manifest.payload_hash {
218 return Err(Error::Invalid("complete payload length or digest mismatch".into()).into());
219 }
220 payload.rewind()?;
221 Ok(())
222}
223
224pub fn reconstruct<S: MultipartSource>(
225 root: &VerifiedRecord,
226 source: &mut S,
227 limits: RecoveryLimits,
228 scratch_directory: &Path,
229) -> Result<RecoveredObject, RecoveryError> {
230 let MultipartRecord::Root(manifest) = root.decode()? else {
231 return Err(Error::Invalid("recovery locator must name a root manifest".into()).into());
232 };
233 let geometry = manifest.geometry()?;
234 if geometry.length() > limits.max_payload_bytes || geometry.nodes() > limits.max_nodes {
235 return Err(RecoveryError::Capacity(
236 "object exceeds caller's byte/node budget".into(),
237 ));
238 }
239 let mut references = tempfile::tempfile_in(scratch_directory)?;
240 inventory(root, &manifest, source, &mut references)?;
241 let mut payload = tempfile::tempfile_in(scratch_directory)?;
242 reconstruct_parts(
243 source,
244 root.author(),
245 &manifest,
246 geometry,
247 &mut references,
248 &mut payload,
249 )?;
250 Ok(RecoveredObject {
251 root: root.txid(),
252 author: root.author(),
253 manifest,
254 payload,
255 })
256}