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 for batch in manifest.entries.chunks(8) {
137 let requests: Vec<_> = batch
138 .iter()
139 .map(|reference| RecordRequest {
140 reference: *reference,
141 })
142 .collect();
143 source.prefetch(&requests);
144 for entry in batch {
145 let verified = fetch(source, *entry, root.author())?;
146 let MultipartRecord::Leaf(leaf) = verified.decode()? else {
147 return Err(Error::Invalid("root must reference leaf manifests".into()).into());
148 };
149 inventory.accept_leaf(&leaf).map_err(inventory_error)?;
150 for reference in leaf.entries {
151 file.write_all(reference.txid.as_byte_array())?;
152 file.write_all(&reference.record_hash)?;
153 }
154 }
155 }
156 file.rewind()?;
157 Ok(())
158}
159
160fn read_requests(file: &mut File, count: u32) -> Result<Vec<RecordRequest>, RecoveryError> {
161 let mut requests = Vec::new();
162 for _ in 0..count {
163 let mut txid = [0; 32];
164 let mut record_hash = [0; 32];
165 file.read_exact(&mut txid)?;
166 file.read_exact(&mut record_hash)?;
167 requests.push(RecordRequest {
168 reference: ChildReference {
169 txid: Txid::from_byte_array(txid),
170 record_hash,
171 },
172 });
173 }
174 Ok(requests)
175}
176
177fn reconstruct_parts<S: MultipartSource>(
178 source: &mut S,
179 author: XOnlyPublicKey,
180 manifest: &RootManifest,
181 geometry: Geometry,
182 inventory: &mut File,
183 payload: &mut File,
184) -> Result<(), RecoveryError> {
185 let mut digest = Sha256::new();
186 let mut length = 0u64;
187 for start in (0..geometry.parts()).step_by(8) {
188 let count = (geometry.parts() - start).min(8);
189 let requests = read_requests(inventory, count)?;
190 source.prefetch(&requests);
191 for (offset, request) in requests.into_iter().enumerate() {
192 let index = start + u32::try_from(offset).map_err(Error::from)?;
193 let verified = fetch(source, request.reference, author)?;
194 let MultipartRecord::Data(part) = verified.decode()? else {
195 return Err(Error::Invalid("leaf must reference data parts".into()).into());
196 };
197 if part.index != index || part.payload.len() != geometry.part_length(index)? {
198 return Err(
199 Error::Invalid("data position or length disagrees with root".into()).into(),
200 );
201 }
202 length = length
203 .checked_add(u64::try_from(part.payload.len()).map_err(Error::from)?)
204 .ok_or_else(|| Error::Invalid("reconstruction length overflow".into()))?;
205 digest.update(&part.payload);
206 payload.write_all(&part.payload)?;
207 }
208 }
209 let actual_hash: [u8; 32] = digest.finalize().into();
210 if length != manifest.length || actual_hash != manifest.payload_hash {
211 return Err(Error::Invalid("complete payload length or digest mismatch".into()).into());
212 }
213 payload.rewind()?;
214 Ok(())
215}
216
217pub fn reconstruct<S: MultipartSource>(
218 root: &VerifiedRecord,
219 source: &mut S,
220 limits: RecoveryLimits,
221 scratch_directory: &Path,
222) -> Result<RecoveredObject, RecoveryError> {
223 let MultipartRecord::Root(manifest) = root.decode()? else {
224 return Err(Error::Invalid("recovery locator must name a root manifest".into()).into());
225 };
226 let geometry = manifest.geometry()?;
227 if geometry.length() > limits.max_payload_bytes || geometry.nodes() > limits.max_nodes {
228 return Err(RecoveryError::Capacity(
229 "object exceeds caller's byte/node budget".into(),
230 ));
231 }
232 let mut references = tempfile::tempfile_in(scratch_directory)?;
233 inventory(root, &manifest, source, &mut references)?;
234 let mut payload = tempfile::tempfile_in(scratch_directory)?;
235 reconstruct_parts(
236 source,
237 root.author(),
238 &manifest,
239 geometry,
240 &mut references,
241 &mut payload,
242 )?;
243 Ok(RecoveredObject {
244 root: root.txid(),
245 author: root.author(),
246 manifest,
247 payload,
248 })
249}