1use std::collections::HashMap;
35use std::io::SeekFrom;
36use std::path::{Path, PathBuf};
37use std::sync::{Mutex, OnceLock};
38use std::time::Instant;
39
40use chrono::{DateTime, Duration, Utc};
41use serde::{Deserialize, Serialize};
42use sha2::{Digest, Sha256};
43use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
44use tracing::info;
45use uuid::Uuid;
46
47use vta_sdk::protocols::backup_management::chunked::{
48 ALGORITHM_CHUNKED, MAX_CHUNK_SIZE, chunk_count, chunk_range, sha256_digest_multibase,
49 sha256_from_digest_multibase,
50};
51
52use super::bundle_store::{
53 self as backup_bundle_store, BundleKind, BundleRecord, BundleState, mint_token,
54};
55use super::{
56 MAX_BUNDLE_TTL_SECS, bundle_ttl, enforce_kind, enforce_open_bundle_cap, parse_bundle_id,
57 require_owned, sha256_hex,
58};
59use crate::auth::AuthClaims;
60use crate::error::AppError;
61use crate::store::KeyspaceHandle;
62
63pub const MAX_REPORTED_MISSING: usize = 256;
66
67pub const CHUNK_REQUESTS_PER_SECOND: f64 = 50.0;
75
76pub const CHUNK_REQUEST_BURST: f64 = 100.0;
78
79#[derive(Debug)]
83#[non_exhaustive]
84pub enum ChunkedError {
85 App(AppError),
88 NotFound,
91 TerminalState(String),
93 ChunkOutOfRange { index: u64, chunk_count: u64 },
95 DigestMismatch { expected_digest_multibase: String },
98 ChunkSizeMismatch { expected: u64, actual: u64 },
100 IncompleteUpload {
103 missing_count: u64,
104 missing_indices: Vec<u64>,
105 },
106 BundleDigestMismatch,
109 BundleTooLarge { size_bytes: u64 },
112 InvalidManifest(String),
115 RateLimited { retry_after_secs: u64 },
117}
118
119impl From<AppError> for ChunkedError {
120 fn from(e: AppError) -> Self {
121 match e {
125 AppError::NotFound(_) => Self::NotFound,
126 other => Self::App(other),
127 }
128 }
129}
130
131impl std::fmt::Display for ChunkedError {
132 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
133 match self {
134 Self::App(e) => write!(f, "{e}"),
135 Self::NotFound => write!(f, "no chunked backup bundle with that identifier"),
136 Self::TerminalState(s) => write!(f, "bundle is {s}; nothing more moves under it"),
137 Self::ChunkOutOfRange { index, chunk_count } => {
138 write!(
139 f,
140 "chunk index {index} is not below the chunk count {chunk_count}"
141 )
142 }
143 Self::DigestMismatch { .. } => {
144 write!(f, "chunk does not match the manifest digest for its index")
145 }
146 Self::ChunkSizeMismatch { expected, actual } => write!(
147 f,
148 "chunk is {actual} bytes; the manifest requires {expected} at this index"
149 ),
150 Self::IncompleteUpload { missing_count, .. } => {
151 write!(f, "{missing_count} chunk(s) have not been uploaded")
152 }
153 Self::BundleDigestMismatch => write!(
154 f,
155 "the assembled bundle does not match the committed digest; abort and upload again"
156 ),
157 Self::BundleTooLarge { size_bytes } => write!(
158 f,
159 "a {size_bytes}-byte bundle does not fit in 4096 chunks at the permitted chunk size"
160 ),
161 Self::InvalidManifest(why) => write!(f, "invalid chunk manifest: {why}"),
162 Self::RateLimited { retry_after_secs } => {
163 write!(f, "too many chunk requests; retry in {retry_after_secs}s")
164 }
165 }
166 }
167}
168
169#[derive(Debug, Clone, Serialize, Deserialize)]
172pub struct ChunkPlan {
173 pub bundle_id: Uuid,
174 pub chunk_size: u64,
175 pub chunk_count: u64,
176 pub digests: Vec<String>,
178 pub done: Vec<bool>,
181 pub expiry_ceiling: DateTime<Utc>,
183 #[serde(default, skip_serializing_if = "Option::is_none")]
187 pub idle_ttl_secs: Option<u64>,
188}
189
190#[derive(Debug, Clone, Copy, PartialEq, Eq)]
198pub struct TransferTerms {
199 pub idle_ttl: Duration,
201 pub lifetime: Duration,
203}
204
205impl TransferTerms {
206 pub fn backup() -> Self {
208 Self {
209 idle_ttl: bundle_ttl(),
210 lifetime: Duration::seconds(MAX_BUNDLE_TTL_SECS as i64),
211 }
212 }
213
214 pub fn room_file() -> Self {
216 Self {
217 idle_ttl: Duration::minutes(15),
218 lifetime: Duration::hours(24),
219 }
220 }
221}
222
223impl ChunkPlan {
224 fn missing(&self) -> impl Iterator<Item = u64> + '_ {
225 self.done
226 .iter()
227 .enumerate()
228 .filter(|(_, d)| !**d)
229 .map(|(i, _)| i as u64)
230 }
231
232 fn remaining(&self) -> u64 {
233 self.done.iter().filter(|d| !**d).count() as u64
234 }
235}
236
237pub(super) fn plan_key(id: &Uuid) -> String {
238 format!("chunks:{id}")
239}
240
241pub async fn get_plan(ks: &KeyspaceHandle, id: &Uuid) -> Result<Option<ChunkPlan>, AppError> {
243 ks.get(plan_key(id)).await
244}
245
246async fn store_plan(ks: &KeyspaceHandle, plan: &ChunkPlan) -> Result<(), AppError> {
247 ks.insert(plan_key(&plan.bundle_id), plan).await
248}
249
250pub async fn delete_plan(ks: &KeyspaceHandle, id: &Uuid) -> Result<(), AppError> {
253 ks.remove(plan_key(id)).await
254}
255
256pub(super) fn chunk_lock() -> &'static tokio::sync::Mutex<()> {
262 static LOCK: OnceLock<tokio::sync::Mutex<()>> = OnceLock::new();
263 LOCK.get_or_init(|| tokio::sync::Mutex::new(()))
264}
265
266pub struct ChunkRateLimiter {
273 per_second: f64,
274 burst: f64,
275 buckets: Mutex<HashMap<String, (f64, Instant)>>,
276}
277
278impl ChunkRateLimiter {
279 pub fn new(per_second: f64, burst: f64) -> Self {
280 Self {
281 per_second,
282 burst,
283 buckets: Mutex::new(HashMap::new()),
284 }
285 }
286
287 pub fn global() -> &'static Self {
289 static GLOBAL: OnceLock<ChunkRateLimiter> = OnceLock::new();
290 GLOBAL.get_or_init(|| Self::new(CHUNK_REQUESTS_PER_SECOND, CHUNK_REQUEST_BURST))
291 }
292
293 pub fn check(&self, did: &str) -> Result<(), ChunkedError> {
295 let now = Instant::now();
296 let mut buckets = self.buckets.lock().unwrap_or_else(|p| p.into_inner());
297 let full_after = self.burst / self.per_second;
300 buckets.retain(|_, (_, last)| now.duration_since(*last).as_secs_f64() < full_after * 4.0);
301 let (tokens, last) = buckets.entry(did.to_string()).or_insert((self.burst, now));
302 let refilled =
303 (*tokens + now.duration_since(*last).as_secs_f64() * self.per_second).min(self.burst);
304 *last = now;
305 if refilled >= 1.0 {
306 *tokens = refilled - 1.0;
307 Ok(())
308 } else {
309 *tokens = refilled;
310 let wait = ((1.0 - refilled) / self.per_second).ceil().max(1.0);
311 Err(ChunkedError::RateLimited {
312 retry_after_secs: wait as u64,
313 })
314 }
315 }
316}
317
318#[derive(Debug, Clone)]
322pub struct ChunkedBundle {
323 pub bundle_id: Uuid,
324 pub chunk_size: u64,
325 pub chunk_count: u64,
326 pub digests: Vec<String>,
327 pub expected_sha256: String,
328 pub expected_size_bytes: u64,
329 pub expires_at: DateTime<Utc>,
330}
331
332pub async fn stage_export(
342 bundles_ks: &KeyspaceHandle,
343 blob_dir: &Path,
344 owner_did: &str,
345 bytes: &[u8],
346 max_chunk_size: Option<u64>,
347) -> Result<ChunkedBundle, ChunkedError> {
348 let chunk_size = max_chunk_size.unwrap_or(MAX_CHUNK_SIZE).min(MAX_CHUNK_SIZE);
349 let size = bytes.len() as u64;
350 let count =
351 chunk_count(size, chunk_size).ok_or(ChunkedError::BundleTooLarge { size_bytes: size })?;
352
353 let digests: Vec<String> = bytes
354 .chunks(chunk_size as usize)
355 .map(|c| sha256_digest_multibase(&Sha256::digest(c).into()))
356 .collect();
357 debug_assert_eq!(digests.len() as u64, count);
358
359 let bundle_id = Uuid::new_v4();
360 let blob_path = prepare_blob_path(blob_dir, &bundle_id).await?;
361 tokio::fs::write(&blob_path, bytes)
362 .await
363 .map_err(AppError::Io)?;
364 #[cfg(unix)]
365 super::set_file_mode_600(&blob_path).await?;
366
367 let now = Utc::now();
368 let (_unused_token, token_hash) = mint_token()?;
373 let record = BundleRecord {
374 bundle_id,
375 kind: BundleKind::Export,
376 state: BundleState::ExportReady,
377 created_at: now,
378 expires_at: now + bundle_ttl(),
379 created_by: owner_did.to_string(),
380 algorithm: ALGORITHM_CHUNKED.into(),
381 expected_sha256: sha256_hex(bytes),
382 expected_size_bytes: size,
383 token_hash,
384 blob_path: Some(blob_path),
385 };
386 let plan = ChunkPlan {
387 bundle_id,
388 chunk_size,
389 chunk_count: count,
390 digests: digests.clone(),
391 done: vec![false; count as usize],
392 expiry_ceiling: now + Duration::seconds(MAX_BUNDLE_TTL_SECS as i64),
393 idle_ttl_secs: None,
394 };
395 store_plan(bundles_ks, &plan).await?;
398 backup_bundle_store::store_bundle(bundles_ks, &record).await?;
399
400 info!(bundle_id = %bundle_id, size, chunks = count, "initiate-export (chunked): bundle ready");
401 Ok(ChunkedBundle {
402 bundle_id,
403 chunk_size,
404 chunk_count: count,
405 digests,
406 expected_sha256: record.expected_sha256,
407 expected_size_bytes: size,
408 expires_at: record.expires_at,
409 })
410}
411
412pub async fn check_initiate(ks: &KeyspaceHandle, auth: &AuthClaims) -> Result<(), ChunkedError> {
416 auth.require_super_admin()?;
417 enforce_open_bundle_cap(ks, &auth.did).await?;
418 Ok(())
419}
420
421pub async fn complete_export(
428 ks: &KeyspaceHandle,
429 auth: &AuthClaims,
430 bundle_id: &str,
431) -> Result<bool, ChunkedError> {
432 auth.require_super_admin()?;
433 let id = parse_bundle_id(bundle_id)?;
434 let _guard = chunk_lock().lock().await;
435 let (mut record, plan) = load_chunked(ks, &id, &auth.did, BundleKind::Export).await?;
436 if record.state != BundleState::ExportReady {
437 return Err(ChunkedError::TerminalState("expired".into()));
438 }
439 let downloaded = plan.remaining() == 0;
440 if let Some(path) = record.blob_path.take()
441 && let Err(e) = tokio::fs::remove_file(&path).await
442 && e.kind() != std::io::ErrorKind::NotFound
443 {
444 tracing::warn!(
445 bundle_id = %id,
446 path = %path.display(),
447 error = %e,
448 "complete-export (chunked): failed to delete staged bytes; sweeper will retry"
449 );
450 record.blob_path = Some(path);
451 }
452 record.state = BundleState::ExportAcked;
453 backup_bundle_store::store_bundle(ks, &record).await?;
454 delete_plan(ks, &id).await?;
455 info!(bundle_id = %id, downloaded, "complete-export (chunked): released");
456 Ok(downloaded)
457}
458
459#[derive(Debug, Clone)]
461pub struct ServedChunk {
462 pub bundle_id: Uuid,
463 pub index: u64,
464 pub digest_multibase: String,
465 pub data: Vec<u8>,
466 pub expires_at: DateTime<Utc>,
467}
468
469pub async fn get_chunk(
476 bundles_ks: &KeyspaceHandle,
477 limiter: &ChunkRateLimiter,
478 auth: &AuthClaims,
479 bundle_id: &str,
480 index: u64,
481) -> Result<ServedChunk, ChunkedError> {
482 auth.require_super_admin()?;
483 limiter.check(&auth.did)?;
484 let id = parse_bundle_id(bundle_id)?;
485
486 let _guard = chunk_lock().lock().await;
487 let (mut record, mut plan) =
488 load_chunked(bundles_ks, &id, &auth.did, BundleKind::Export).await?;
489 match record.state {
490 BundleState::ExportReady => {}
491 BundleState::ExportAcked => return Err(ChunkedError::TerminalState("completed".into())),
492 BundleState::Aborted => return Err(ChunkedError::TerminalState("aborted".into())),
493 _ => return Err(ChunkedError::TerminalState("expired".into())),
494 }
495 let now = Utc::now();
496 if record.expires_at <= now {
497 return Err(ChunkedError::TerminalState("expired".into()));
498 }
499 let (start, end) = chunk_range(record.expected_size_bytes, plan.chunk_size, index).ok_or(
500 ChunkedError::ChunkOutOfRange {
501 index,
502 chunk_count: plan.chunk_count,
503 },
504 )?;
505
506 let path = record
507 .blob_path
508 .clone()
509 .ok_or_else(|| ChunkedError::TerminalState("expired".into()))?;
510 let data = read_range(&path, start, end).await?;
511
512 let expected = &plan.digests[index as usize];
516 if sha256_from_digest_multibase(expected) != Some(Sha256::digest(&data).into()) {
517 return Err(ChunkedError::App(AppError::Internal(format!(
518 "staged bytes for chunk {index} of bundle {id} no longer match the manifest"
519 ))));
520 }
521
522 plan.done[index as usize] = true;
523 extend_expiry(&mut record, &plan, now);
524 store_plan(bundles_ks, &plan).await?;
525 backup_bundle_store::store_bundle(bundles_ks, &record).await?;
526
527 Ok(ServedChunk {
528 bundle_id: id,
529 index,
530 digest_multibase: expected.clone(),
531 data,
532 expires_at: record.expires_at,
533 })
534}
535
536pub async fn all_served(ks: &KeyspaceHandle, id: &Uuid) -> Result<Option<bool>, AppError> {
540 Ok(get_plan(ks, id).await?.map(|p| p.remaining() == 0))
541}
542
543pub async fn initiate_import(
551 bundles_ks: &KeyspaceHandle,
552 auth: &AuthClaims,
553 expected_sha256: &str,
554 expected_size_bytes: u64,
555 chunk_size: u64,
556 declared_chunk_count: u64,
557 digests: Vec<String>,
558) -> Result<ChunkedBundle, ChunkedError> {
559 auth.require_super_admin()?;
560 initiate_import_for(
561 bundles_ks,
562 &auth.did,
563 expected_sha256,
564 expected_size_bytes,
565 chunk_size,
566 declared_chunk_count,
567 digests,
568 )
569 .await
570}
571
572pub async fn initiate_import_for(
576 bundles_ks: &KeyspaceHandle,
577 owner: &str,
578 expected_sha256: &str,
579 expected_size_bytes: u64,
580 chunk_size: u64,
581 declared_chunk_count: u64,
582 digests: Vec<String>,
583) -> Result<ChunkedBundle, ChunkedError> {
584 initiate_import_with(
585 bundles_ks,
586 owner,
587 expected_sha256,
588 expected_size_bytes,
589 chunk_size,
590 declared_chunk_count,
591 digests,
592 TransferTerms::backup(),
593 )
594 .await
595}
596
597#[allow(clippy::too_many_arguments)]
599pub async fn initiate_import_with(
600 bundles_ks: &KeyspaceHandle,
601 owner: &str,
602 expected_sha256: &str,
603 expected_size_bytes: u64,
604 chunk_size: u64,
605 declared_chunk_count: u64,
606 digests: Vec<String>,
607 terms: TransferTerms,
608) -> Result<ChunkedBundle, ChunkedError> {
609 enforce_open_bundle_cap(bundles_ks, owner).await?;
610
611 if expected_sha256.len() != 64
612 || !expected_sha256
613 .chars()
614 .all(|c| c.is_ascii_digit() || ('a'..='f').contains(&c))
615 {
616 return Err(ChunkedError::InvalidManifest(
617 "expectedSha256 must be 64 lowercase hex characters".into(),
618 ));
619 }
620 let count = chunk_count(expected_size_bytes, chunk_size).ok_or_else(|| {
621 ChunkedError::InvalidManifest(format!(
622 "{expected_size_bytes} bytes cannot be divided into at most 4096 chunks of {chunk_size}"
623 ))
624 })?;
625 if count != declared_chunk_count {
626 return Err(ChunkedError::InvalidManifest(format!(
627 "chunkCount is {declared_chunk_count}; {expected_size_bytes} bytes at {chunk_size} is {count}"
628 )));
629 }
630 if digests.len() as u64 != count {
631 return Err(ChunkedError::InvalidManifest(format!(
632 "chunkDigests has {} entries; the manifest has {count} chunks",
633 digests.len()
634 )));
635 }
636 if let Some(bad) = digests
637 .iter()
638 .position(|d| sha256_from_digest_multibase(d).is_none())
639 {
640 return Err(ChunkedError::InvalidManifest(format!(
641 "chunkDigests[{bad}] is not a sha2-256 multihash this agent can verify"
642 )));
643 }
644
645 let bundle_id = Uuid::new_v4();
646 let now = Utc::now();
647 let (_unused_token, token_hash) = mint_token()?;
648 let record = BundleRecord {
649 bundle_id,
650 kind: BundleKind::Import,
651 state: BundleState::ImportPending,
652 created_at: now,
653 expires_at: now + terms.idle_ttl,
654 created_by: owner.to_string(),
655 algorithm: ALGORITHM_CHUNKED.into(),
656 expected_sha256: expected_sha256.to_string(),
657 expected_size_bytes,
658 token_hash,
659 blob_path: None,
661 };
662 let plan = ChunkPlan {
663 bundle_id,
664 chunk_size,
665 chunk_count: count,
666 digests: digests.clone(),
667 done: vec![false; count as usize],
668 expiry_ceiling: now + terms.lifetime,
669 idle_ttl_secs: (terms != TransferTerms::backup())
670 .then(|| terms.idle_ttl.num_seconds().max(1) as u64),
671 };
672 store_plan(bundles_ks, &plan).await?;
673 backup_bundle_store::store_bundle(bundles_ks, &record).await?;
674
675 info!(bundle_id = %bundle_id, chunks = count, "initiate-import (chunked): slot ready");
676 Ok(ChunkedBundle {
677 bundle_id,
678 chunk_size,
679 chunk_count: count,
680 digests,
681 expected_sha256: record.expected_sha256,
682 expected_size_bytes,
683 expires_at: record.expires_at,
684 })
685}
686
687#[derive(Debug, Clone, Copy)]
689pub struct PutOutcome {
690 pub stored: bool,
691 pub remaining_count: u64,
692 pub expires_at: DateTime<Utc>,
693}
694
695#[derive(Debug, Clone, Copy)]
697pub struct ChunkWrite<'a> {
698 pub bundle_id: &'a str,
699 pub index: u64,
700 pub digest_multibase: &'a str,
702 pub data: &'a [u8],
704}
705
706pub async fn put_chunk(
714 bundles_ks: &KeyspaceHandle,
715 blob_dir: &Path,
716 limiter: &ChunkRateLimiter,
717 auth: &AuthClaims,
718 write: ChunkWrite<'_>,
719) -> Result<PutOutcome, ChunkedError> {
720 auth.require_super_admin()?;
721 put_chunk_for(bundles_ks, blob_dir, limiter, &auth.did, write).await
722}
723
724pub async fn put_chunk_for(
727 bundles_ks: &KeyspaceHandle,
728 blob_dir: &Path,
729 limiter: &ChunkRateLimiter,
730 owner: &str,
731 write: ChunkWrite<'_>,
732) -> Result<PutOutcome, ChunkedError> {
733 let ChunkWrite {
734 bundle_id,
735 index,
736 digest_multibase,
737 data,
738 } = write;
739 limiter.check(owner)?;
740 let id = parse_bundle_id(bundle_id)?;
741
742 let _guard = chunk_lock().lock().await;
743 let (mut record, mut plan) = load_chunked(bundles_ks, &id, owner, BundleKind::Import).await?;
744 match record.state {
745 BundleState::ImportPending => {}
746 BundleState::ImportReceived | BundleState::ImportPreviewed => {}
750 BundleState::ImportCommitted => {
751 return Err(ChunkedError::TerminalState("committed".into()));
752 }
753 BundleState::Aborted => return Err(ChunkedError::TerminalState("aborted".into())),
754 _ => return Err(ChunkedError::TerminalState("expired".into())),
755 }
756 let now = Utc::now();
757 if record.expires_at <= now {
758 return Err(ChunkedError::TerminalState("expired".into()));
759 }
760 let (start, end) = chunk_range(record.expected_size_bytes, plan.chunk_size, index).ok_or(
761 ChunkedError::ChunkOutOfRange {
762 index,
763 chunk_count: plan.chunk_count,
764 },
765 )?;
766 if data.len() as u64 != end - start {
767 return Err(ChunkedError::ChunkSizeMismatch {
768 expected: end - start,
769 actual: data.len() as u64,
770 });
771 }
772 let expected = plan.digests[index as usize].clone();
773 let committed = sha256_from_digest_multibase(&expected);
774 let restated = sha256_from_digest_multibase(digest_multibase);
775 let actual: [u8; 32] = Sha256::digest(data).into();
776 if committed.is_none() || restated != committed || committed != Some(actual) {
777 return Err(ChunkedError::DigestMismatch {
778 expected_digest_multibase: expected,
779 });
780 }
781
782 let stored = if plan.done[index as usize] {
783 false
784 } else {
785 let path = match record.blob_path.clone() {
786 Some(p) => p,
787 None => {
788 let p = prepare_blob_path(blob_dir, &id).await?;
789 record.blob_path = Some(p.clone());
790 p
791 }
792 };
793 write_range(&path, record.expected_size_bytes, start, data).await?;
794 plan.done[index as usize] = true;
795 true
796 };
797
798 extend_expiry(&mut record, &plan, now);
799 store_plan(bundles_ks, &plan).await?;
800 backup_bundle_store::store_bundle(bundles_ks, &record).await?;
801
802 Ok(PutOutcome {
803 stored,
804 remaining_count: plan.remaining(),
805 expires_at: record.expires_at,
806 })
807}
808
809pub async fn finalize_precheck(
818 bundles_ks: &KeyspaceHandle,
819 auth: &AuthClaims,
820 bundle_id: &str,
821) -> Result<(), ChunkedError> {
822 auth.require_super_admin()?;
823 finalize_precheck_for(bundles_ks, &auth.did, bundle_id).await
824}
825
826pub async fn finalize_precheck_for(
829 bundles_ks: &KeyspaceHandle,
830 owner: &str,
831 bundle_id: &str,
832) -> Result<(), ChunkedError> {
833 let id = parse_bundle_id(bundle_id)?;
834 let _guard = chunk_lock().lock().await;
835 let record = require_owned(bundles_ks, &id, owner).await?;
836 enforce_kind(&record, BundleKind::Import)?;
837 if record.algorithm != ALGORITHM_CHUNKED || record.state != BundleState::ImportPending {
838 return Ok(());
839 }
840 let plan = get_plan(bundles_ks, &id)
841 .await?
842 .ok_or(ChunkedError::NotFound)?;
843
844 let missing_count = plan.remaining();
845 if missing_count > 0 {
846 return Err(ChunkedError::IncompleteUpload {
847 missing_count,
848 missing_indices: plan.missing().take(MAX_REPORTED_MISSING).collect(),
849 });
850 }
851
852 let path = record
853 .blob_path
854 .clone()
855 .ok_or_else(|| AppError::Internal(format!("chunked bundle {id} has no staging file")))?;
856 if !assembled_matches(&path, &record).await? {
857 return Err(ChunkedError::BundleDigestMismatch);
858 }
859
860 let mut record = record;
861 record.state = BundleState::ImportReceived;
862 backup_bundle_store::store_bundle(bundles_ks, &record).await?;
863 info!(bundle_id = %id, "finalize-import (chunked): upload assembled and verified");
864 Ok(())
865}
866
867async fn load_chunked(
877 ks: &KeyspaceHandle,
878 id: &Uuid,
879 caller: &str,
880 kind: BundleKind,
881) -> Result<(BundleRecord, ChunkPlan), ChunkedError> {
882 let record = require_owned(ks, id, caller).await?;
883 enforce_kind(&record, kind)?;
884 if record.algorithm != ALGORITHM_CHUNKED {
885 return Err(ChunkedError::NotFound);
886 }
887 match record.state {
888 BundleState::ExportAcked => return Err(ChunkedError::TerminalState("completed".into())),
889 BundleState::ImportCommitted => {
890 return Err(ChunkedError::TerminalState("committed".into()));
891 }
892 BundleState::Aborted => return Err(ChunkedError::TerminalState("aborted".into())),
893 BundleState::Expired | BundleState::ExportDownloaded => {
894 return Err(ChunkedError::TerminalState("expired".into()));
895 }
896 _ => {}
897 }
898 let plan = get_plan(ks, id).await?.ok_or(ChunkedError::NotFound)?;
899 Ok((record, plan))
900}
901
902fn extend_expiry(record: &mut BundleRecord, plan: &ChunkPlan, now: DateTime<Utc>) {
905 let idle = plan
906 .idle_ttl_secs
907 .map(|s| Duration::seconds(s as i64))
908 .unwrap_or_else(bundle_ttl);
909 let proposed = (now + idle).min(plan.expiry_ceiling);
910 if proposed > record.expires_at {
911 record.expires_at = proposed;
912 }
913}
914
915async fn prepare_blob_path(blob_dir: &Path, id: &Uuid) -> Result<PathBuf, AppError> {
916 tokio::fs::create_dir_all(blob_dir)
917 .await
918 .map_err(AppError::Io)?;
919 #[cfg(unix)]
920 super::set_dir_mode_700(blob_dir).await?;
921 Ok(blob_dir.join(format!("{id}.vtabak")))
922}
923
924async fn read_range(path: &Path, start: u64, end: u64) -> Result<Vec<u8>, ChunkedError> {
925 let mut file = match tokio::fs::File::open(path).await {
926 Ok(f) => f,
927 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
928 return Err(ChunkedError::TerminalState("expired".into()));
929 }
930 Err(e) => return Err(AppError::Io(e).into()),
931 };
932 file.seek(SeekFrom::Start(start))
933 .await
934 .map_err(AppError::Io)?;
935 let mut data = vec![0u8; (end - start) as usize];
936 file.read_exact(&mut data).await.map_err(AppError::Io)?;
937 Ok(data)
938}
939
940async fn write_range(
943 path: &Path,
944 total_size: u64,
945 offset: u64,
946 data: &[u8],
947) -> Result<(), AppError> {
948 let existed = tokio::fs::try_exists(path).await.map_err(AppError::Io)?;
949 let mut file = tokio::fs::OpenOptions::new()
950 .create(true)
951 .write(true)
952 .truncate(false)
953 .open(path)
954 .await
955 .map_err(AppError::Io)?;
956 if !existed {
957 #[cfg(unix)]
958 super::set_file_mode_600(path).await?;
959 file.set_len(total_size).await.map_err(AppError::Io)?;
960 }
961 file.seek(SeekFrom::Start(offset))
962 .await
963 .map_err(AppError::Io)?;
964 file.write_all(data).await.map_err(AppError::Io)?;
965 file.sync_data().await.map_err(AppError::Io)?;
966 Ok(())
967}
968
969async fn assembled_matches(path: &Path, record: &BundleRecord) -> Result<bool, AppError> {
972 let mut file = tokio::fs::File::open(path).await.map_err(AppError::Io)?;
973 let len = file.metadata().await.map_err(AppError::Io)?.len();
974 if len != record.expected_size_bytes {
975 return Ok(false);
976 }
977 let mut hasher = Sha256::new();
978 let mut buf = vec![0u8; MAX_CHUNK_SIZE as usize];
979 loop {
980 let n = file.read(&mut buf).await.map_err(AppError::Io)?;
981 if n == 0 {
982 break;
983 }
984 hasher.update(&buf[..n]);
985 }
986 let digest = hasher.finalize();
987 let mut hex = String::with_capacity(64);
988 for b in digest {
989 hex.push_str(&format!("{b:02x}"));
990 }
991 Ok(hex == record.expected_sha256)
992}
993
994#[cfg(test)]
995mod tests {
996 use super::*;
997 use crate::acl::Role;
998 use crate::config::StoreConfig as VtiStoreConfig;
999 use crate::store::Store;
1000
1001 const OWNER: &str = "did:key:z6MkChunkOwner";
1002 const MIN: u64 = vta_sdk::protocols::backup_management::chunked::MIN_CHUNK_SIZE;
1003
1004 fn admin(did: &str) -> AuthClaims {
1005 AuthClaims {
1006 did: did.into(),
1007 role: Role::Admin,
1008 allowed_contexts: Vec::new(),
1009 session_id: "test-session".into(),
1010 access_expires_at: 0,
1011 issued_at: 0,
1012 amr: Vec::new(),
1013 acr: String::new(),
1014 }
1015 }
1016
1017 struct Env {
1018 _dir: tempfile::TempDir,
1019 ks: KeyspaceHandle,
1020 blob_dir: PathBuf,
1021 }
1022
1023 fn env() -> Env {
1024 let dir = tempfile::tempdir().unwrap();
1025 let store = Store::open(&VtiStoreConfig {
1026 data_dir: dir.path().join("store"),
1027 })
1028 .unwrap();
1029 let ks = store.keyspace("backup_bundles_test").unwrap();
1030 let blob_dir = dir.path().join("backups");
1031 Env {
1032 _dir: dir,
1033 ks,
1034 blob_dir,
1035 }
1036 }
1037
1038 fn unlimited() -> ChunkRateLimiter {
1039 ChunkRateLimiter::new(1.0e9, 1.0e9)
1040 }
1041
1042 fn bundle_bytes() -> Vec<u8> {
1045 let mut v = Vec::new();
1046 v.extend(std::iter::repeat_n(1u8, MIN as usize));
1047 v.extend(std::iter::repeat_n(2u8, MIN as usize));
1048 v.extend_from_slice(b"backup-tail!");
1049 v
1050 }
1051
1052 fn digest(bytes: &[u8]) -> String {
1053 sha256_digest_multibase(&Sha256::digest(bytes).into())
1054 }
1055
1056 #[tokio::test]
1057 async fn export_round_trips_chunk_by_chunk_and_is_non_consuming() {
1058 let e = env();
1059 let auth = admin(OWNER);
1060 let bytes = bundle_bytes();
1061 let bundle = stage_export(&e.ks, &e.blob_dir, OWNER, &bytes, Some(MIN))
1062 .await
1063 .unwrap();
1064 assert_eq!(bundle.chunk_count, 3);
1065
1066 let mut assembled = Vec::new();
1067 for i in 0..bundle.chunk_count {
1068 let c = get_chunk(&e.ks, &unlimited(), &auth, &bundle.bundle_id.to_string(), i)
1069 .await
1070 .unwrap();
1071 assert_eq!(c.digest_multibase, bundle.digests[i as usize]);
1072 assert_eq!(digest(&c.data), bundle.digests[i as usize]);
1073 assembled.extend(c.data);
1074 }
1075 assert_eq!(assembled, bytes);
1076 assert_eq!(sha256_hex(&assembled), bundle.expected_sha256);
1077
1078 let again = get_chunk(&e.ks, &unlimited(), &auth, &bundle.bundle_id.to_string(), 2)
1080 .await
1081 .unwrap();
1082 assert_eq!(again.data, b"backup-tail!");
1083 assert_eq!(
1084 all_served(&e.ks, &bundle.bundle_id).await.unwrap(),
1085 Some(true)
1086 );
1087 }
1088
1089 #[tokio::test]
1090 async fn an_index_past_the_last_chunk_is_out_of_range() {
1091 let e = env();
1092 let bundle = stage_export(&e.ks, &e.blob_dir, OWNER, &bundle_bytes(), Some(MIN))
1093 .await
1094 .unwrap();
1095 let err = get_chunk(
1096 &e.ks,
1097 &unlimited(),
1098 &admin(OWNER),
1099 &bundle.bundle_id.to_string(),
1100 3,
1101 )
1102 .await
1103 .unwrap_err();
1104 assert!(
1105 matches!(
1106 err,
1107 ChunkedError::ChunkOutOfRange {
1108 index: 3,
1109 chunk_count: 3
1110 }
1111 ),
1112 "{err:?}"
1113 );
1114 }
1115
1116 #[tokio::test]
1117 async fn another_operators_bundle_is_not_found() {
1118 let e = env();
1119 let bundle = stage_export(&e.ks, &e.blob_dir, OWNER, &bundle_bytes(), Some(MIN))
1120 .await
1121 .unwrap();
1122 let err = get_chunk(
1123 &e.ks,
1124 &unlimited(),
1125 &admin("did:key:z6MkSomeoneElse"),
1126 &bundle.bundle_id.to_string(),
1127 0,
1128 )
1129 .await
1130 .unwrap_err();
1131 assert!(matches!(err, ChunkedError::NotFound), "{err:?}");
1132 }
1133
1134 #[tokio::test]
1135 async fn serving_a_chunk_slides_the_expiry_but_never_past_the_ceiling() {
1136 let e = env();
1137 let auth = admin(OWNER);
1138 let bundle = stage_export(&e.ks, &e.blob_dir, OWNER, &bundle_bytes(), Some(MIN))
1139 .await
1140 .unwrap();
1141 let id = bundle.bundle_id;
1142
1143 let now = Utc::now();
1147 let mut record = backup_bundle_store::get_bundle(&e.ks, &id)
1148 .await
1149 .unwrap()
1150 .unwrap();
1151 record.expires_at = now + Duration::seconds(10);
1152 backup_bundle_store::store_bundle(&e.ks, &record)
1153 .await
1154 .unwrap();
1155 let mut plan = get_plan(&e.ks, &id).await.unwrap().unwrap();
1156 plan.expiry_ceiling = now + Duration::seconds(20);
1157 store_plan(&e.ks, &plan).await.unwrap();
1158
1159 let served = get_chunk(&e.ks, &unlimited(), &auth, &id.to_string(), 0)
1160 .await
1161 .unwrap();
1162 assert_eq!(served.expires_at, plan.expiry_ceiling);
1163
1164 let mut plan = get_plan(&e.ks, &id).await.unwrap().unwrap();
1166 plan.expiry_ceiling = now + Duration::hours(1);
1167 store_plan(&e.ks, &plan).await.unwrap();
1168 let served = get_chunk(&e.ks, &unlimited(), &auth, &id.to_string(), 1)
1169 .await
1170 .unwrap();
1171 assert!(served.expires_at > now + Duration::seconds(20));
1172 assert!(served.expires_at <= plan.expiry_ceiling);
1173 }
1174
1175 #[tokio::test]
1176 async fn an_expired_bundle_serves_nothing() {
1177 let e = env();
1178 let bundle = stage_export(&e.ks, &e.blob_dir, OWNER, &bundle_bytes(), Some(MIN))
1179 .await
1180 .unwrap();
1181 let mut record = backup_bundle_store::get_bundle(&e.ks, &bundle.bundle_id)
1182 .await
1183 .unwrap()
1184 .unwrap();
1185 record.expires_at = Utc::now() - Duration::seconds(1);
1186 backup_bundle_store::store_bundle(&e.ks, &record)
1187 .await
1188 .unwrap();
1189 let err = get_chunk(
1190 &e.ks,
1191 &unlimited(),
1192 &admin(OWNER),
1193 &bundle.bundle_id.to_string(),
1194 0,
1195 )
1196 .await
1197 .unwrap_err();
1198 assert!(matches!(err, ChunkedError::TerminalState(_)), "{err:?}");
1199 }
1200
1201 #[tokio::test]
1202 async fn a_bundle_too_large_for_the_chunk_bound_is_refused() {
1203 let e = env();
1204 let bytes = vec![0u8; (MIN * 4096 + 1) as usize];
1206 let err = stage_export(&e.ks, &e.blob_dir, OWNER, &bytes, Some(MIN))
1207 .await
1208 .unwrap_err();
1209 assert!(
1210 matches!(err, ChunkedError::BundleTooLarge { .. }),
1211 "{err:?}"
1212 );
1213 }
1214
1215 async fn open_import(e: &Env, bytes: &[u8]) -> ChunkedBundle {
1216 let digests: Vec<String> = bytes.chunks(MIN as usize).map(digest).collect();
1217 initiate_import(
1218 &e.ks,
1219 &admin(OWNER),
1220 &sha256_hex(bytes),
1221 bytes.len() as u64,
1222 MIN,
1223 digests.len() as u64,
1224 digests,
1225 )
1226 .await
1227 .unwrap()
1228 }
1229
1230 #[tokio::test]
1231 async fn import_resumes_after_a_missing_chunk_and_assembles() {
1232 let e = env();
1233 let auth = admin(OWNER);
1234 let bytes = bundle_bytes();
1235 let slot = open_import(&e, &bytes).await;
1236 let id = slot.bundle_id.to_string();
1237 let chunks: Vec<&[u8]> = bytes.chunks(MIN as usize).collect();
1238
1239 for i in [0usize, 2] {
1241 let out = put_chunk(
1242 &e.ks,
1243 &e.blob_dir,
1244 &unlimited(),
1245 &auth,
1246 ChunkWrite {
1247 bundle_id: &id,
1248 index: i as u64,
1249 digest_multibase: &slot.digests[i],
1250 data: chunks[i],
1251 },
1252 )
1253 .await
1254 .unwrap();
1255 assert!(out.stored);
1256 }
1257
1258 let err = finalize_precheck(&e.ks, &auth, &id).await.unwrap_err();
1260 match err {
1261 ChunkedError::IncompleteUpload {
1262 missing_count,
1263 missing_indices,
1264 } => {
1265 assert_eq!(missing_count, 1);
1266 assert_eq!(missing_indices, vec![1]);
1267 }
1268 other => panic!("expected IncompleteUpload, got {other:?}"),
1269 }
1270
1271 let out = put_chunk(
1273 &e.ks,
1274 &e.blob_dir,
1275 &unlimited(),
1276 &auth,
1277 ChunkWrite {
1278 bundle_id: &id,
1279 index: 1,
1280 digest_multibase: &slot.digests[1],
1281 data: chunks[1],
1282 },
1283 )
1284 .await
1285 .unwrap();
1286 assert!(out.stored);
1287 assert_eq!(out.remaining_count, 0);
1288
1289 finalize_precheck(&e.ks, &auth, &id).await.unwrap();
1290 let record = backup_bundle_store::get_bundle(&e.ks, &slot.bundle_id)
1291 .await
1292 .unwrap()
1293 .unwrap();
1294 assert_eq!(record.state, BundleState::ImportReceived);
1295 let staged = tokio::fs::read(record.blob_path.unwrap()).await.unwrap();
1296 assert_eq!(staged, bytes);
1297 }
1298
1299 #[tokio::test]
1300 async fn an_identical_re_put_is_stored_false_and_a_mismatched_one_is_refused() {
1301 let e = env();
1302 let auth = admin(OWNER);
1303 let bytes = bundle_bytes();
1304 let slot = open_import(&e, &bytes).await;
1305 let id = slot.bundle_id.to_string();
1306 let chunk0 = &bytes[..MIN as usize];
1307
1308 let first = put_chunk(
1309 &e.ks,
1310 &e.blob_dir,
1311 &unlimited(),
1312 &auth,
1313 ChunkWrite {
1314 bundle_id: &id,
1315 index: 0,
1316 digest_multibase: &slot.digests[0],
1317 data: chunk0,
1318 },
1319 )
1320 .await
1321 .unwrap();
1322 assert!(first.stored);
1323 let repeat = put_chunk(
1324 &e.ks,
1325 &e.blob_dir,
1326 &unlimited(),
1327 &auth,
1328 ChunkWrite {
1329 bundle_id: &id,
1330 index: 0,
1331 digest_multibase: &slot.digests[0],
1332 data: chunk0,
1333 },
1334 )
1335 .await
1336 .unwrap();
1337 assert!(!repeat.stored, "an identical re-put stores nothing new");
1338
1339 let tampered = vec![9u8; MIN as usize];
1341 let err = put_chunk(
1342 &e.ks,
1343 &e.blob_dir,
1344 &unlimited(),
1345 &auth,
1346 ChunkWrite {
1347 bundle_id: &id,
1348 index: 0,
1349 digest_multibase: &digest(&tampered),
1350 data: &tampered,
1351 },
1352 )
1353 .await
1354 .unwrap_err();
1355 assert!(
1356 matches!(err, ChunkedError::DigestMismatch { .. }),
1357 "{err:?}"
1358 );
1359
1360 let err = put_chunk(
1362 &e.ks,
1363 &e.blob_dir,
1364 &unlimited(),
1365 &auth,
1366 ChunkWrite {
1367 bundle_id: &id,
1368 index: 1,
1369 digest_multibase: &slot.digests[1],
1370 data: &tampered,
1371 },
1372 )
1373 .await
1374 .unwrap_err();
1375 assert!(
1376 matches!(err, ChunkedError::DigestMismatch { .. }),
1377 "{err:?}"
1378 );
1379
1380 let record = backup_bundle_store::get_bundle(&e.ks, &slot.bundle_id)
1381 .await
1382 .unwrap()
1383 .unwrap();
1384 let staged = tokio::fs::read(record.blob_path.unwrap()).await.unwrap();
1385 assert_eq!(&staged[..MIN as usize], chunk0);
1386 }
1387
1388 #[tokio::test]
1389 async fn a_wrong_length_chunk_is_refused() {
1390 let e = env();
1391 let bytes = bundle_bytes();
1392 let slot = open_import(&e, &bytes).await;
1393 let short = &bytes[..10];
1394 let err = put_chunk(
1395 &e.ks,
1396 &e.blob_dir,
1397 &unlimited(),
1398 &admin(OWNER),
1399 ChunkWrite {
1400 bundle_id: &slot.bundle_id.to_string(),
1401 index: 0,
1402 digest_multibase: &digest(short),
1403 data: short,
1404 },
1405 )
1406 .await
1407 .unwrap_err();
1408 assert!(
1409 matches!(err, ChunkedError::ChunkSizeMismatch { expected, actual: 10 } if expected == MIN),
1410 "{err:?}"
1411 );
1412 }
1413
1414 #[tokio::test]
1415 async fn a_manifest_whose_count_does_not_follow_is_refused() {
1416 let e = env();
1417 let bytes = bundle_bytes();
1418 let digests: Vec<String> = bytes.chunks(MIN as usize).map(digest).collect();
1419 let err = initiate_import(
1420 &e.ks,
1421 &admin(OWNER),
1422 &sha256_hex(&bytes),
1423 bytes.len() as u64,
1424 MIN,
1425 2,
1426 digests,
1427 )
1428 .await
1429 .unwrap_err();
1430 assert!(matches!(err, ChunkedError::InvalidManifest(_)), "{err:?}");
1431 }
1432
1433 #[tokio::test]
1434 async fn chunks_that_verify_but_do_not_assemble_to_the_bundle_are_caught() {
1435 let e = env();
1436 let auth = admin(OWNER);
1437 let bytes = bundle_bytes();
1438 let digests: Vec<String> = bytes.chunks(MIN as usize).map(digest).collect();
1441 let slot = initiate_import(
1442 &e.ks,
1443 &auth,
1444 &sha256_hex(b"something else"),
1445 bytes.len() as u64,
1446 MIN,
1447 3,
1448 digests,
1449 )
1450 .await
1451 .unwrap();
1452 let id = slot.bundle_id.to_string();
1453 for (i, c) in bytes.chunks(MIN as usize).enumerate() {
1454 put_chunk(
1455 &e.ks,
1456 &e.blob_dir,
1457 &unlimited(),
1458 &auth,
1459 ChunkWrite {
1460 bundle_id: &id,
1461 index: i as u64,
1462 digest_multibase: &slot.digests[i],
1463 data: c,
1464 },
1465 )
1466 .await
1467 .unwrap();
1468 }
1469 let err = finalize_precheck(&e.ks, &auth, &id).await.unwrap_err();
1470 assert!(matches!(err, ChunkedError::BundleDigestMismatch), "{err:?}");
1471 }
1472
1473 #[test]
1474 fn the_limiter_refuses_past_the_burst_and_names_a_wait() {
1475 let limiter = ChunkRateLimiter::new(1.0, 2.0);
1476 limiter.check(OWNER).unwrap();
1477 limiter.check(OWNER).unwrap();
1478 match limiter.check(OWNER) {
1479 Err(ChunkedError::RateLimited { retry_after_secs }) => assert!(retry_after_secs >= 1),
1480 other => panic!("expected RateLimited, got {other:?}"),
1481 }
1482 limiter.check("did:key:z6MkAnotherOperator").unwrap();
1484 }
1485
1486 #[tokio::test]
1489 async fn import_terms_set_the_idle_window_and_the_lifetime() {
1490 let env = env();
1491 let bytes = bundle_bytes();
1492 let digests: Vec<String> = bytes.chunks(MIN as usize).map(digest).collect();
1493 let before = Utc::now();
1494 let slot = initiate_import_with(
1495 &env.ks,
1496 OWNER,
1497 &sha256_hex(&bytes),
1498 bytes.len() as u64,
1499 MIN,
1500 digests.len() as u64,
1501 digests.clone(),
1502 TransferTerms::room_file(),
1503 )
1504 .await
1505 .unwrap();
1506 assert!(slot.expires_at >= before + Duration::minutes(15));
1507 let plan = get_plan(&env.ks, &slot.bundle_id).await.unwrap().unwrap();
1508 assert!(plan.expiry_ceiling >= before + Duration::hours(24));
1509
1510 let out = put_chunk_for(
1511 &env.ks,
1512 &env.blob_dir,
1513 &unlimited(),
1514 OWNER,
1515 ChunkWrite {
1516 bundle_id: &slot.bundle_id.to_string(),
1517 index: 0,
1518 digest_multibase: &digests[0],
1519 data: &bytes[..MIN as usize],
1520 },
1521 )
1522 .await
1523 .unwrap();
1524 assert!(out.expires_at >= Utc::now() + Duration::minutes(14));
1525
1526 let backup = initiate_import_for(
1528 &env.ks,
1529 "did:key:z6MkOther",
1530 &sha256_hex(&bytes),
1531 bytes.len() as u64,
1532 MIN,
1533 digests.len() as u64,
1534 digests,
1535 )
1536 .await
1537 .unwrap();
1538 assert!(backup.expires_at <= Utc::now() + bundle_ttl());
1539 let plan = get_plan(&env.ks, &backup.bundle_id).await.unwrap().unwrap();
1540 assert_eq!(plan.idle_ttl_secs, None);
1541 }
1542}