urma_runtime/multipart/
recovery.rs

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    time::{Duration, Instant},
14};
15use urma_core::error::Error;
16
17#[derive(Clone, Debug)]
18pub struct Candidate {
19    pub origin: String,
20    pub record: VerifiedRecord,
21}
22
23#[derive(Debug)]
24pub enum FetchError {
25    Unavailable,
26    Rejected { origin: String, cause: Error },
27    Source(ServiceError),
28}
29
30pub trait MultipartSource {
31    fn fetch(
32        &mut self,
33        request: &RecordRequest,
34        rejected: &[String],
35    ) -> Result<Candidate, FetchError>;
36
37    fn prefetch(&mut self, _requests: &[RecordRequest]) {}
38}
39
40#[derive(Clone, Copy, Debug)]
41pub struct RecoveryLimits {
42    pub max_payload_bytes: u64,
43    pub max_nodes: u32,
44    pub attempts: u32,
45    pub timeout: Duration,
46}
47
48impl RecoveryLimits {
49    pub const ATTEMPTS: u32 = 4;
50    pub const TIMEOUT: Duration = Duration::from_secs(24 * 60 * 60);
51
52    pub fn new(max_payload_bytes: u64, max_nodes: u32) -> Self {
53        Self {
54            max_payload_bytes,
55            max_nodes,
56            attempts: Self::ATTEMPTS,
57            timeout: Self::TIMEOUT,
58        }
59    }
60}
61
62#[derive(Debug)]
63pub enum RecoveryError {
64    InvalidObject(Error),
65    InvalidCandidate { txid: Txid, cause: Error },
66    Incomplete { txid: Txid },
67    Source { txid: Txid, cause: ServiceError },
68    Capacity(String),
69    Storage(std::io::Error),
70}
71
72impl fmt::Display for RecoveryError {
73    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
74        match self {
75            Self::InvalidObject(cause) => write!(formatter, "invalid multipart object: {cause}"),
76            Self::InvalidCandidate { txid, cause } => {
77                write!(formatter, "invalid candidate {txid}: {cause}")
78            }
79            Self::Incomplete { txid } => write!(formatter, "incomplete object; missing {txid}"),
80            Self::Source { txid, cause } => write!(formatter, "source failure for {txid}: {cause}"),
81            Self::Capacity(message) => write!(formatter, "multipart capacity: {message}"),
82            Self::Storage(cause) => write!(formatter, "multipart storage: {cause}"),
83        }
84    }
85}
86impl std::error::Error for RecoveryError {}
87impl From<Error> for RecoveryError {
88    fn from(cause: Error) -> Self {
89        Self::InvalidObject(cause)
90    }
91}
92impl From<std::io::Error> for RecoveryError {
93    fn from(cause: std::io::Error) -> Self {
94        Self::Storage(cause)
95    }
96}
97
98#[derive(Debug)]
99pub struct RecoveredObject {
100    root: Txid,
101    author: XOnlyPublicKey,
102    manifest: RootManifest,
103    payload: File,
104}
105
106impl RecoveredObject {
107    pub fn root(&self) -> Txid {
108        self.root
109    }
110    pub fn author(&self) -> XOnlyPublicKey {
111        self.author
112    }
113    pub fn manifest(&self) -> &RootManifest {
114        &self.manifest
115    }
116}
117impl Read for RecoveredObject {
118    fn read(&mut self, bytes: &mut [u8]) -> std::io::Result<usize> {
119        self.payload.read(bytes)
120    }
121}
122impl Seek for RecoveredObject {
123    fn seek(&mut self, position: SeekFrom) -> std::io::Result<u64> {
124        self.payload.seek(position)
125    }
126}
127
128struct Budget {
129    attempts: u32,
130    started: Instant,
131    timeout: Duration,
132}
133
134impl Budget {
135    fn new(limits: RecoveryLimits) -> Self {
136        Self {
137            attempts: limits.attempts.max(1),
138            started: Instant::now(),
139            timeout: limits.timeout,
140        }
141    }
142
143    fn expired(&self) -> bool {
144        self.started.elapsed() >= self.timeout
145    }
146}
147
148enum Attempted {
149    Nothing,
150    Rejected(Error),
151    Failed(ServiceError),
152}
153
154fn given_up(txid: Txid, attempted: Attempted) -> RecoveryError {
155    match attempted {
156        Attempted::Nothing => RecoveryError::Source {
157            txid,
158            cause: ServiceError::Io(std::io::Error::new(
159                std::io::ErrorKind::TimedOut,
160                "recovery deadline passed",
161            )),
162        },
163        Attempted::Rejected(cause) => RecoveryError::InvalidCandidate { txid, cause },
164        Attempted::Failed(cause) => RecoveryError::Source { txid, cause },
165    }
166}
167
168fn unavailable(txid: Txid, attempted: Attempted) -> RecoveryError {
169    match attempted {
170        Attempted::Nothing => RecoveryError::Incomplete { txid },
171        Attempted::Rejected(cause) => RecoveryError::InvalidCandidate { txid, cause },
172        Attempted::Failed(cause) => {
173            tracing::warn!(%txid, error = %cause, "record unavailable after a transport failure");
174            RecoveryError::Incomplete { txid }
175        }
176    }
177}
178
179fn accept(
180    request: &RecordRequest,
181    record: VerifiedRecord,
182    author: XOnlyPublicKey,
183) -> Result<VerifiedRecord, RecoveryError> {
184    request.check_hash(&record)?;
185    if record.author() != author {
186        return Err(Error::Invalid("referenced child author differs from root".into()).into());
187    }
188    Ok(record)
189}
190
191fn fetch<S: MultipartSource>(
192    source: &mut S,
193    reference: ChildReference,
194    author: XOnlyPublicKey,
195    budget: &Budget,
196) -> Result<VerifiedRecord, RecoveryError> {
197    let request = RecordRequest { reference };
198    let record = fetch_candidate(reference.txid, budget, |rejected| {
199        source.fetch(&request, rejected)
200    })?;
201    accept(&request, record, author)
202}
203
204pub(crate) fn retrieve(
205    txid: Txid,
206    limits: RecoveryLimits,
207    fetch: impl FnMut(&[String]) -> Result<Candidate, FetchError>,
208) -> Result<VerifiedRecord, RecoveryError> {
209    fetch_candidate(txid, &Budget::new(limits), fetch)
210}
211
212fn fetch_candidate(
213    txid: Txid,
214    budget: &Budget,
215    mut fetch: impl FnMut(&[String]) -> Result<Candidate, FetchError>,
216) -> Result<VerifiedRecord, RecoveryError> {
217    let mut rejected: Vec<String> = Vec::new();
218    let mut attempted = Attempted::Nothing;
219    for attempt in 1..=budget.attempts {
220        if budget.expired() {
221            return Err(given_up(txid, attempted));
222        }
223        let (origin, cause) = match fetch(&rejected) {
224            Ok(candidate) => {
225                if candidate.record.txid() == txid {
226                    return Ok(candidate.record);
227                }
228                let cause = Error::Invalid("candidate TXID mismatch".into());
229                tracing::warn!(%txid, origin = %candidate.origin, error = %cause, "candidate rejected");
230                (candidate.origin, cause)
231            }
232            Err(FetchError::Rejected { origin, cause }) => {
233                tracing::warn!(%txid, %origin, error = %cause, "candidate rejected");
234                (origin, cause)
235            }
236            Err(FetchError::Source(cause)) => {
237                tracing::warn!(%txid, attempt, error = %cause, "multipart fetch failed; asking again");
238                attempted = Attempted::Failed(cause);
239                continue;
240            }
241            Err(FetchError::Unavailable) => return Err(unavailable(txid, attempted)),
242        };
243        if rejected.contains(&origin) {
244            return Err(RecoveryError::InvalidCandidate { txid, cause });
245        }
246        rejected.push(origin);
247        attempted = Attempted::Rejected(cause);
248    }
249    Err(given_up(txid, attempted))
250}
251
252fn inventory_error(cause: Error) -> RecoveryError {
253    match cause {
254        Error::Allocation(error) => RecoveryError::Capacity(error.to_string()),
255        error => RecoveryError::InvalidObject(error),
256    }
257}
258
259fn inventory<S: MultipartSource>(
260    root: &VerifiedRecord,
261    manifest: &RootManifest,
262    source: &mut S,
263    budget: &Budget,
264    file: &mut File,
265) -> Result<(), RecoveryError> {
266    let mut inventory = ManifestInventory::new(root.txid(), manifest).map_err(inventory_error)?;
267    let mut received = 0u64;
268    tracing::info!(target: "urma_ui", phase = "Receiving manifests", done = received, total = manifest.entries.len());
269    for batch in manifest.entries.chunks(8) {
270        let requests: Vec<_> = batch
271            .iter()
272            .map(|reference| RecordRequest {
273                reference: *reference,
274            })
275            .collect();
276        source.prefetch(&requests);
277        for entry in batch {
278            let verified = fetch(source, *entry, root.author(), budget)?;
279            let MultipartRecord::Leaf(leaf) = verified.decode()? else {
280                return Err(Error::Invalid("root must reference leaf manifests".into()).into());
281            };
282            inventory.accept_leaf(&leaf).map_err(inventory_error)?;
283            for reference in leaf.entries {
284                file.write_all(reference.txid.as_byte_array())?;
285                file.write_all(&reference.record_hash)?;
286            }
287            received += 1;
288            tracing::info!(target: "urma_ui", phase = "Receiving manifests", done = received, total = manifest.entries.len());
289        }
290    }
291    file.rewind()?;
292    Ok(())
293}
294
295fn read_requests(file: &mut File, count: u32) -> Result<Vec<RecordRequest>, RecoveryError> {
296    let mut requests = Vec::new();
297    for _ in 0..count {
298        let mut txid = [0; 32];
299        let mut record_hash = [0; 32];
300        file.read_exact(&mut txid)?;
301        file.read_exact(&mut record_hash)?;
302        requests.push(RecordRequest {
303            reference: ChildReference {
304                txid: Txid::from_byte_array(txid),
305                record_hash,
306            },
307        });
308    }
309    Ok(requests)
310}
311
312fn reconstruct_parts<S: MultipartSource>(
313    source: &mut S,
314    author: XOnlyPublicKey,
315    manifest: &RootManifest,
316    geometry: Geometry,
317    budget: &Budget,
318    inventory: &mut File,
319    payload: &mut File,
320) -> Result<(), RecoveryError> {
321    let mut digest = Sha256::new();
322    let mut length = 0u64;
323    tracing::info!(target: "urma_ui", phase = "Receiving payload bytes", done = length, total = manifest.length);
324    for start in (0..geometry.parts()).step_by(8) {
325        let count = (geometry.parts() - start).min(8);
326        let requests = read_requests(inventory, count)?;
327        source.prefetch(&requests);
328        for (offset, request) in requests.into_iter().enumerate() {
329            let index = start + u32::try_from(offset).map_err(Error::from)?;
330            let verified = fetch(source, request.reference, author, budget)?;
331            let MultipartRecord::Data(part) = verified.decode()? else {
332                return Err(Error::Invalid("leaf must reference data parts".into()).into());
333            };
334            if part.index != index || part.payload.len() != geometry.part_length(index)? {
335                return Err(
336                    Error::Invalid("data position or length disagrees with root".into()).into(),
337                );
338            }
339            length = length
340                .checked_add(u64::try_from(part.payload.len()).map_err(Error::from)?)
341                .ok_or_else(|| Error::Invalid("reconstruction length overflow".into()))?;
342            digest.update(&part.payload);
343            payload.write_all(&part.payload)?;
344            tracing::info!(target: "urma_ui", phase = "Receiving payload bytes", done = length, total = manifest.length);
345        }
346    }
347    let actual_hash: [u8; 32] = digest.finalize().into();
348    tracing::info!(target: "urma_ui", phase = "Verifying complete payload digest");
349    if length != manifest.length || actual_hash != manifest.payload_hash {
350        return Err(Error::Invalid("complete payload length or digest mismatch".into()).into());
351    }
352    payload.rewind()?;
353    Ok(())
354}
355
356pub fn reconstruct<S: MultipartSource>(
357    root: &VerifiedRecord,
358    source: &mut S,
359    limits: RecoveryLimits,
360    scratch_directory: &Path,
361) -> Result<RecoveredObject, RecoveryError> {
362    let MultipartRecord::Root(manifest) = root.decode()? else {
363        return Err(Error::Invalid("recovery locator must name a root manifest".into()).into());
364    };
365    let geometry = manifest.geometry()?;
366    if geometry.length() > limits.max_payload_bytes || geometry.nodes() > limits.max_nodes {
367        return Err(RecoveryError::Capacity(
368            "object exceeds caller's byte/node budget".into(),
369        ));
370    }
371    let budget = Budget::new(limits);
372    let mut references = tempfile::tempfile_in(scratch_directory)?;
373    inventory(root, &manifest, source, &budget, &mut references)?;
374    let mut payload = tempfile::tempfile_in(scratch_directory)?;
375    reconstruct_parts(
376        source,
377        root.author(),
378        &manifest,
379        geometry,
380        &budget,
381        &mut references,
382        &mut payload,
383    )?;
384    Ok(RecoveredObject {
385        root: root.txid(),
386        author: root.author(),
387        manifest,
388        payload,
389    })
390}
urma_runtime/multipart/recovery.rs

urma_runtime/multipart/
recovery.rs
