1use std::collections::HashMap;
31use std::io::SeekFrom;
32use std::path::{Path, PathBuf};
33use std::sync::{Mutex, OnceLock};
34use std::time::Instant;
35
36use chrono::{DateTime, Duration, Utc};
37use serde::{Deserialize, Serialize};
38use sha2::{Digest, Sha256};
39use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
40use tracing::info;
41use uuid::Uuid;
42
43use vta_sdk::protocols::backup_management::chunked::{
44 ALGORITHM_CHUNKED, MAX_CHUNK_SIZE, chunk_count, chunk_range, sha256_digest_multibase,
45 sha256_from_digest_multibase,
46};
47
48use super::descriptors::{
49 DescriptorDeps, MAX_BUNDLE_TTL_SECS, bundle_ttl, enforce_kind, enforce_open_bundle_cap,
50 parse_bundle_id, require_owned, sha256_hex,
51};
52use crate::backup_bundle_store::{self, BundleKind, BundleRecord, BundleState, mint_token};
53use vti_common::auth::AuthClaims;
54use vti_common::error::AppError;
55use vti_common::store::KeyspaceHandle;
56
57pub const MAX_REPORTED_MISSING: usize = 256;
60
61pub const CHUNK_REQUESTS_PER_SECOND: f64 = 50.0;
69
70pub const CHUNK_REQUEST_BURST: f64 = 100.0;
72
73#[derive(Debug)]
77pub enum ChunkedError {
78 App(AppError),
81 NotFound,
84 TerminalState(String),
86 ChunkOutOfRange { index: u64, chunk_count: u64 },
88 DigestMismatch { expected_digest_multibase: String },
91 ChunkSizeMismatch { expected: u64, actual: u64 },
93 IncompleteUpload {
96 missing_count: u64,
97 missing_indices: Vec<u64>,
98 },
99 BundleDigestMismatch,
102 BundleTooLarge { size_bytes: u64 },
105 InvalidManifest(String),
108 RateLimited { retry_after_secs: u64 },
110}
111
112impl From<AppError> for ChunkedError {
113 fn from(e: AppError) -> Self {
114 match e {
118 AppError::NotFound(_) => Self::NotFound,
119 other => Self::App(other),
120 }
121 }
122}
123
124impl std::fmt::Display for ChunkedError {
125 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
126 match self {
127 Self::App(e) => write!(f, "{e}"),
128 Self::NotFound => write!(f, "no chunked backup bundle with that identifier"),
129 Self::TerminalState(s) => write!(f, "bundle is {s}; nothing more moves under it"),
130 Self::ChunkOutOfRange { index, chunk_count } => {
131 write!(
132 f,
133 "chunk index {index} is not below the chunk count {chunk_count}"
134 )
135 }
136 Self::DigestMismatch { .. } => {
137 write!(f, "chunk does not match the manifest digest for its index")
138 }
139 Self::ChunkSizeMismatch { expected, actual } => write!(
140 f,
141 "chunk is {actual} bytes; the manifest requires {expected} at this index"
142 ),
143 Self::IncompleteUpload { missing_count, .. } => {
144 write!(f, "{missing_count} chunk(s) have not been uploaded")
145 }
146 Self::BundleDigestMismatch => write!(
147 f,
148 "the assembled bundle does not match the committed digest; abort and upload again"
149 ),
150 Self::BundleTooLarge { size_bytes } => write!(
151 f,
152 "a {size_bytes}-byte bundle does not fit in 4096 chunks at the permitted chunk size"
153 ),
154 Self::InvalidManifest(why) => write!(f, "invalid chunk manifest: {why}"),
155 Self::RateLimited { retry_after_secs } => {
156 write!(f, "too many chunk requests; retry in {retry_after_secs}s")
157 }
158 }
159 }
160}
161
162#[derive(Debug, Clone, Serialize, Deserialize)]
165pub struct ChunkPlan {
166 pub bundle_id: Uuid,
167 pub chunk_size: u64,
168 pub chunk_count: u64,
169 pub digests: Vec<String>,
171 pub done: Vec<bool>,
174 pub expiry_ceiling: DateTime<Utc>,
176}
177
178impl ChunkPlan {
179 fn missing(&self) -> impl Iterator<Item = u64> + '_ {
180 self.done
181 .iter()
182 .enumerate()
183 .filter(|(_, d)| !**d)
184 .map(|(i, _)| i as u64)
185 }
186
187 fn remaining(&self) -> u64 {
188 self.done.iter().filter(|d| !**d).count() as u64
189 }
190}
191
192fn plan_key(id: &Uuid) -> String {
193 format!("chunks:{id}")
194}
195
196pub async fn get_plan(ks: &KeyspaceHandle, id: &Uuid) -> Result<Option<ChunkPlan>, AppError> {
198 ks.get(plan_key(id)).await
199}
200
201async fn store_plan(ks: &KeyspaceHandle, plan: &ChunkPlan) -> Result<(), AppError> {
202 ks.insert(plan_key(&plan.bundle_id), plan).await
203}
204
205pub async fn delete_plan(ks: &KeyspaceHandle, id: &Uuid) -> Result<(), AppError> {
208 ks.remove(plan_key(id)).await
209}
210
211fn chunk_lock() -> &'static tokio::sync::Mutex<()> {
217 static LOCK: OnceLock<tokio::sync::Mutex<()>> = OnceLock::new();
218 LOCK.get_or_init(|| tokio::sync::Mutex::new(()))
219}
220
221pub struct ChunkRateLimiter {
228 per_second: f64,
229 burst: f64,
230 buckets: Mutex<HashMap<String, (f64, Instant)>>,
231}
232
233impl ChunkRateLimiter {
234 pub fn new(per_second: f64, burst: f64) -> Self {
235 Self {
236 per_second,
237 burst,
238 buckets: Mutex::new(HashMap::new()),
239 }
240 }
241
242 pub fn global() -> &'static Self {
244 static GLOBAL: OnceLock<ChunkRateLimiter> = OnceLock::new();
245 GLOBAL.get_or_init(|| Self::new(CHUNK_REQUESTS_PER_SECOND, CHUNK_REQUEST_BURST))
246 }
247
248 pub fn check(&self, did: &str) -> Result<(), ChunkedError> {
250 let now = Instant::now();
251 let mut buckets = self.buckets.lock().unwrap_or_else(|p| p.into_inner());
252 let full_after = self.burst / self.per_second;
255 buckets.retain(|_, (_, last)| now.duration_since(*last).as_secs_f64() < full_after * 4.0);
256 let (tokens, last) = buckets.entry(did.to_string()).or_insert((self.burst, now));
257 let refilled =
258 (*tokens + now.duration_since(*last).as_secs_f64() * self.per_second).min(self.burst);
259 *last = now;
260 if refilled >= 1.0 {
261 *tokens = refilled - 1.0;
262 Ok(())
263 } else {
264 *tokens = refilled;
265 let wait = ((1.0 - refilled) / self.per_second).ceil().max(1.0);
266 Err(ChunkedError::RateLimited {
267 retry_after_secs: wait as u64,
268 })
269 }
270 }
271}
272
273#[derive(Debug, Clone)]
277pub struct ChunkedBundle {
278 pub bundle_id: Uuid,
279 pub chunk_size: u64,
280 pub chunk_count: u64,
281 pub digests: Vec<String>,
282 pub expected_sha256: String,
283 pub expected_size_bytes: u64,
284 pub expires_at: DateTime<Utc>,
285}
286
287pub async fn initiate_export(
296 deps: &DescriptorDeps<'_>,
297 auth: &AuthClaims,
298 password: &str,
299 include_audit: bool,
300 max_chunk_size: Option<u64>,
301) -> Result<ChunkedBundle, ChunkedError> {
302 auth.require_super_admin()?;
303 enforce_open_bundle_cap(deps.bundles_ks, &auth.did).await?;
304
305 let envelope = {
306 let config_guard = deps.config.read().await;
307 super::export_backup(
308 &deps.keyspaces,
309 deps.seed_store.as_ref(),
310 &config_guard,
311 auth,
312 password,
313 include_audit,
314 )
315 .await?
316 };
317 let bytes = serde_json::to_vec(&envelope)
318 .map_err(|e| AppError::Internal(format!("serialize backup envelope: {e}")))?;
319 stage_export(
320 deps.bundles_ks,
321 deps.blob_dir,
322 &auth.did,
323 &bytes,
324 max_chunk_size,
325 )
326 .await
327}
328
329pub async fn stage_export(
333 bundles_ks: &KeyspaceHandle,
334 blob_dir: &Path,
335 owner_did: &str,
336 bytes: &[u8],
337 max_chunk_size: Option<u64>,
338) -> Result<ChunkedBundle, ChunkedError> {
339 let chunk_size = max_chunk_size.unwrap_or(MAX_CHUNK_SIZE).min(MAX_CHUNK_SIZE);
340 let size = bytes.len() as u64;
341 let count =
342 chunk_count(size, chunk_size).ok_or(ChunkedError::BundleTooLarge { size_bytes: size })?;
343
344 let digests: Vec<String> = bytes
345 .chunks(chunk_size as usize)
346 .map(|c| sha256_digest_multibase(&Sha256::digest(c).into()))
347 .collect();
348 debug_assert_eq!(digests.len() as u64, count);
349
350 let bundle_id = Uuid::new_v4();
351 let blob_path = prepare_blob_path(blob_dir, &bundle_id).await?;
352 tokio::fs::write(&blob_path, bytes)
353 .await
354 .map_err(AppError::Io)?;
355 #[cfg(unix)]
356 super::descriptors::set_file_mode_600(&blob_path).await?;
357
358 let now = Utc::now();
359 let (_unused_token, token_hash) = mint_token()?;
364 let record = BundleRecord {
365 bundle_id,
366 kind: BundleKind::Export,
367 state: BundleState::ExportReady,
368 created_at: now,
369 expires_at: now + bundle_ttl(),
370 created_by: owner_did.to_string(),
371 algorithm: ALGORITHM_CHUNKED.into(),
372 expected_sha256: sha256_hex(bytes),
373 expected_size_bytes: size,
374 token_hash,
375 blob_path: Some(blob_path),
376 };
377 let plan = ChunkPlan {
378 bundle_id,
379 chunk_size,
380 chunk_count: count,
381 digests: digests.clone(),
382 done: vec![false; count as usize],
383 expiry_ceiling: now + Duration::seconds(MAX_BUNDLE_TTL_SECS as i64),
384 };
385 store_plan(bundles_ks, &plan).await?;
388 backup_bundle_store::store_bundle(bundles_ks, &record).await?;
389
390 info!(bundle_id = %bundle_id, size, chunks = count, "initiate-export (chunked): bundle ready");
391 Ok(ChunkedBundle {
392 bundle_id,
393 chunk_size,
394 chunk_count: count,
395 digests,
396 expected_sha256: record.expected_sha256,
397 expected_size_bytes: size,
398 expires_at: record.expires_at,
399 })
400}
401
402#[derive(Debug, Clone)]
404pub struct ServedChunk {
405 pub bundle_id: Uuid,
406 pub index: u64,
407 pub digest_multibase: String,
408 pub data: Vec<u8>,
409 pub expires_at: DateTime<Utc>,
410}
411
412pub async fn get_chunk(
419 bundles_ks: &KeyspaceHandle,
420 limiter: &ChunkRateLimiter,
421 auth: &AuthClaims,
422 bundle_id: &str,
423 index: u64,
424) -> Result<ServedChunk, ChunkedError> {
425 auth.require_super_admin()?;
426 limiter.check(&auth.did)?;
427 let id = parse_bundle_id(bundle_id)?;
428
429 let _guard = chunk_lock().lock().await;
430 let (mut record, mut plan) =
431 load_chunked(bundles_ks, &id, &auth.did, BundleKind::Export).await?;
432 match record.state {
433 BundleState::ExportReady => {}
434 BundleState::ExportAcked => return Err(ChunkedError::TerminalState("completed".into())),
435 BundleState::Aborted => return Err(ChunkedError::TerminalState("aborted".into())),
436 _ => return Err(ChunkedError::TerminalState("expired".into())),
437 }
438 let now = Utc::now();
439 if record.expires_at <= now {
440 return Err(ChunkedError::TerminalState("expired".into()));
441 }
442 let (start, end) = chunk_range(record.expected_size_bytes, plan.chunk_size, index).ok_or(
443 ChunkedError::ChunkOutOfRange {
444 index,
445 chunk_count: plan.chunk_count,
446 },
447 )?;
448
449 let path = record
450 .blob_path
451 .clone()
452 .ok_or_else(|| ChunkedError::TerminalState("expired".into()))?;
453 let data = read_range(&path, start, end).await?;
454
455 let expected = &plan.digests[index as usize];
459 if sha256_from_digest_multibase(expected) != Some(Sha256::digest(&data).into()) {
460 return Err(ChunkedError::App(AppError::Internal(format!(
461 "staged bytes for chunk {index} of bundle {id} no longer match the manifest"
462 ))));
463 }
464
465 plan.done[index as usize] = true;
466 extend_expiry(&mut record, &plan, now);
467 store_plan(bundles_ks, &plan).await?;
468 backup_bundle_store::store_bundle(bundles_ks, &record).await?;
469
470 Ok(ServedChunk {
471 bundle_id: id,
472 index,
473 digest_multibase: expected.clone(),
474 data,
475 expires_at: record.expires_at,
476 })
477}
478
479pub async fn all_served(ks: &KeyspaceHandle, id: &Uuid) -> Result<Option<bool>, AppError> {
483 Ok(get_plan(ks, id).await?.map(|p| p.remaining() == 0))
484}
485
486pub async fn initiate_import(
494 bundles_ks: &KeyspaceHandle,
495 auth: &AuthClaims,
496 expected_sha256: &str,
497 expected_size_bytes: u64,
498 chunk_size: u64,
499 declared_chunk_count: u64,
500 digests: Vec<String>,
501) -> Result<ChunkedBundle, ChunkedError> {
502 auth.require_super_admin()?;
503 enforce_open_bundle_cap(bundles_ks, &auth.did).await?;
504
505 if expected_sha256.len() != 64
506 || !expected_sha256
507 .chars()
508 .all(|c| c.is_ascii_digit() || ('a'..='f').contains(&c))
509 {
510 return Err(ChunkedError::InvalidManifest(
511 "expectedSha256 must be 64 lowercase hex characters".into(),
512 ));
513 }
514 let count = chunk_count(expected_size_bytes, chunk_size).ok_or_else(|| {
515 ChunkedError::InvalidManifest(format!(
516 "{expected_size_bytes} bytes cannot be divided into at most 4096 chunks of {chunk_size}"
517 ))
518 })?;
519 if count != declared_chunk_count {
520 return Err(ChunkedError::InvalidManifest(format!(
521 "chunkCount is {declared_chunk_count}; {expected_size_bytes} bytes at {chunk_size} is {count}"
522 )));
523 }
524 if digests.len() as u64 != count {
525 return Err(ChunkedError::InvalidManifest(format!(
526 "chunkDigests has {} entries; the manifest has {count} chunks",
527 digests.len()
528 )));
529 }
530 if let Some(bad) = digests
531 .iter()
532 .position(|d| sha256_from_digest_multibase(d).is_none())
533 {
534 return Err(ChunkedError::InvalidManifest(format!(
535 "chunkDigests[{bad}] is not a sha2-256 multihash this agent can verify"
536 )));
537 }
538
539 let bundle_id = Uuid::new_v4();
540 let now = Utc::now();
541 let (_unused_token, token_hash) = mint_token()?;
542 let record = BundleRecord {
543 bundle_id,
544 kind: BundleKind::Import,
545 state: BundleState::ImportPending,
546 created_at: now,
547 expires_at: now + bundle_ttl(),
548 created_by: auth.did.clone(),
549 algorithm: ALGORITHM_CHUNKED.into(),
550 expected_sha256: expected_sha256.to_string(),
551 expected_size_bytes,
552 token_hash,
553 blob_path: None,
555 };
556 let plan = ChunkPlan {
557 bundle_id,
558 chunk_size,
559 chunk_count: count,
560 digests: digests.clone(),
561 done: vec![false; count as usize],
562 expiry_ceiling: now + Duration::seconds(MAX_BUNDLE_TTL_SECS as i64),
563 };
564 store_plan(bundles_ks, &plan).await?;
565 backup_bundle_store::store_bundle(bundles_ks, &record).await?;
566
567 info!(bundle_id = %bundle_id, chunks = count, "initiate-import (chunked): slot ready");
568 Ok(ChunkedBundle {
569 bundle_id,
570 chunk_size,
571 chunk_count: count,
572 digests,
573 expected_sha256: record.expected_sha256,
574 expected_size_bytes,
575 expires_at: record.expires_at,
576 })
577}
578
579#[derive(Debug, Clone, Copy)]
581pub struct PutOutcome {
582 pub stored: bool,
583 pub remaining_count: u64,
584 pub expires_at: DateTime<Utc>,
585}
586
587#[derive(Debug, Clone, Copy)]
589pub struct ChunkWrite<'a> {
590 pub bundle_id: &'a str,
591 pub index: u64,
592 pub digest_multibase: &'a str,
594 pub data: &'a [u8],
596}
597
598pub async fn put_chunk(
606 bundles_ks: &KeyspaceHandle,
607 blob_dir: &Path,
608 limiter: &ChunkRateLimiter,
609 auth: &AuthClaims,
610 write: ChunkWrite<'_>,
611) -> Result<PutOutcome, ChunkedError> {
612 let ChunkWrite {
613 bundle_id,
614 index,
615 digest_multibase,
616 data,
617 } = write;
618 auth.require_super_admin()?;
619 limiter.check(&auth.did)?;
620 let id = parse_bundle_id(bundle_id)?;
621
622 let _guard = chunk_lock().lock().await;
623 let (mut record, mut plan) =
624 load_chunked(bundles_ks, &id, &auth.did, BundleKind::Import).await?;
625 match record.state {
626 BundleState::ImportPending => {}
627 BundleState::ImportReceived | BundleState::ImportPreviewed => {}
631 BundleState::ImportCommitted => {
632 return Err(ChunkedError::TerminalState("committed".into()));
633 }
634 BundleState::Aborted => return Err(ChunkedError::TerminalState("aborted".into())),
635 _ => return Err(ChunkedError::TerminalState("expired".into())),
636 }
637 let now = Utc::now();
638 if record.expires_at <= now {
639 return Err(ChunkedError::TerminalState("expired".into()));
640 }
641 let (start, end) = chunk_range(record.expected_size_bytes, plan.chunk_size, index).ok_or(
642 ChunkedError::ChunkOutOfRange {
643 index,
644 chunk_count: plan.chunk_count,
645 },
646 )?;
647 if data.len() as u64 != end - start {
648 return Err(ChunkedError::ChunkSizeMismatch {
649 expected: end - start,
650 actual: data.len() as u64,
651 });
652 }
653 let expected = plan.digests[index as usize].clone();
654 let committed = sha256_from_digest_multibase(&expected);
655 let restated = sha256_from_digest_multibase(digest_multibase);
656 let actual: [u8; 32] = Sha256::digest(data).into();
657 if committed.is_none() || restated != committed || committed != Some(actual) {
658 return Err(ChunkedError::DigestMismatch {
659 expected_digest_multibase: expected,
660 });
661 }
662
663 let stored = if plan.done[index as usize] {
664 false
665 } else {
666 let path = match record.blob_path.clone() {
667 Some(p) => p,
668 None => {
669 let p = prepare_blob_path(blob_dir, &id).await?;
670 record.blob_path = Some(p.clone());
671 p
672 }
673 };
674 write_range(&path, record.expected_size_bytes, start, data).await?;
675 plan.done[index as usize] = true;
676 true
677 };
678
679 extend_expiry(&mut record, &plan, now);
680 store_plan(bundles_ks, &plan).await?;
681 backup_bundle_store::store_bundle(bundles_ks, &record).await?;
682
683 Ok(PutOutcome {
684 stored,
685 remaining_count: plan.remaining(),
686 expires_at: record.expires_at,
687 })
688}
689
690pub async fn finalize_precheck(
699 bundles_ks: &KeyspaceHandle,
700 auth: &AuthClaims,
701 bundle_id: &str,
702) -> Result<(), ChunkedError> {
703 auth.require_super_admin()?;
704 let id = parse_bundle_id(bundle_id)?;
705 let _guard = chunk_lock().lock().await;
706 let record = require_owned(bundles_ks, &id, &auth.did).await?;
707 enforce_kind(&record, BundleKind::Import)?;
708 if record.algorithm != ALGORITHM_CHUNKED || record.state != BundleState::ImportPending {
709 return Ok(());
710 }
711 let plan = get_plan(bundles_ks, &id)
712 .await?
713 .ok_or(ChunkedError::NotFound)?;
714
715 let missing_count = plan.remaining();
716 if missing_count > 0 {
717 return Err(ChunkedError::IncompleteUpload {
718 missing_count,
719 missing_indices: plan.missing().take(MAX_REPORTED_MISSING).collect(),
720 });
721 }
722
723 let path = record
724 .blob_path
725 .clone()
726 .ok_or_else(|| AppError::Internal(format!("chunked bundle {id} has no staging file")))?;
727 if !assembled_matches(&path, &record).await? {
728 return Err(ChunkedError::BundleDigestMismatch);
729 }
730
731 let mut record = record;
732 record.state = BundleState::ImportReceived;
733 backup_bundle_store::store_bundle(bundles_ks, &record).await?;
734 info!(bundle_id = %id, "finalize-import (chunked): upload assembled and verified");
735 Ok(())
736}
737
738async fn load_chunked(
748 ks: &KeyspaceHandle,
749 id: &Uuid,
750 caller: &str,
751 kind: BundleKind,
752) -> Result<(BundleRecord, ChunkPlan), ChunkedError> {
753 let record = require_owned(ks, id, caller).await?;
754 enforce_kind(&record, kind)?;
755 if record.algorithm != ALGORITHM_CHUNKED {
756 return Err(ChunkedError::NotFound);
757 }
758 match record.state {
759 BundleState::ExportAcked => return Err(ChunkedError::TerminalState("completed".into())),
760 BundleState::ImportCommitted => {
761 return Err(ChunkedError::TerminalState("committed".into()));
762 }
763 BundleState::Aborted => return Err(ChunkedError::TerminalState("aborted".into())),
764 BundleState::Expired | BundleState::ExportDownloaded => {
765 return Err(ChunkedError::TerminalState("expired".into()));
766 }
767 _ => {}
768 }
769 let plan = get_plan(ks, id).await?.ok_or(ChunkedError::NotFound)?;
770 Ok((record, plan))
771}
772
773fn extend_expiry(record: &mut BundleRecord, plan: &ChunkPlan, now: DateTime<Utc>) {
776 let proposed = (now + bundle_ttl()).min(plan.expiry_ceiling);
777 if proposed > record.expires_at {
778 record.expires_at = proposed;
779 }
780}
781
782async fn prepare_blob_path(blob_dir: &Path, id: &Uuid) -> Result<PathBuf, AppError> {
783 tokio::fs::create_dir_all(blob_dir)
784 .await
785 .map_err(AppError::Io)?;
786 #[cfg(unix)]
787 super::descriptors::set_dir_mode_700(blob_dir).await?;
788 Ok(blob_dir.join(format!("{id}.vtabak")))
789}
790
791async fn read_range(path: &Path, start: u64, end: u64) -> Result<Vec<u8>, ChunkedError> {
792 let mut file = match tokio::fs::File::open(path).await {
793 Ok(f) => f,
794 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
795 return Err(ChunkedError::TerminalState("expired".into()));
796 }
797 Err(e) => return Err(AppError::Io(e).into()),
798 };
799 file.seek(SeekFrom::Start(start))
800 .await
801 .map_err(AppError::Io)?;
802 let mut data = vec![0u8; (end - start) as usize];
803 file.read_exact(&mut data).await.map_err(AppError::Io)?;
804 Ok(data)
805}
806
807async fn write_range(
810 path: &Path,
811 total_size: u64,
812 offset: u64,
813 data: &[u8],
814) -> Result<(), AppError> {
815 let existed = tokio::fs::try_exists(path).await.map_err(AppError::Io)?;
816 let mut file = tokio::fs::OpenOptions::new()
817 .create(true)
818 .write(true)
819 .truncate(false)
820 .open(path)
821 .await
822 .map_err(AppError::Io)?;
823 if !existed {
824 #[cfg(unix)]
825 super::descriptors::set_file_mode_600(path).await?;
826 file.set_len(total_size).await.map_err(AppError::Io)?;
827 }
828 file.seek(SeekFrom::Start(offset))
829 .await
830 .map_err(AppError::Io)?;
831 file.write_all(data).await.map_err(AppError::Io)?;
832 file.sync_data().await.map_err(AppError::Io)?;
833 Ok(())
834}
835
836async fn assembled_matches(path: &Path, record: &BundleRecord) -> Result<bool, AppError> {
839 let mut file = tokio::fs::File::open(path).await.map_err(AppError::Io)?;
840 let len = file.metadata().await.map_err(AppError::Io)?.len();
841 if len != record.expected_size_bytes {
842 return Ok(false);
843 }
844 let mut hasher = Sha256::new();
845 let mut buf = vec![0u8; MAX_CHUNK_SIZE as usize];
846 loop {
847 let n = file.read(&mut buf).await.map_err(AppError::Io)?;
848 if n == 0 {
849 break;
850 }
851 hasher.update(&buf[..n]);
852 }
853 let digest = hasher.finalize();
854 let mut hex = String::with_capacity(64);
855 for b in digest {
856 hex.push_str(&format!("{b:02x}"));
857 }
858 Ok(hex == record.expected_sha256)
859}
860
861#[cfg(test)]
862mod tests {
863 use super::*;
864 use vti_common::acl::Role;
865 use vti_common::config::StoreConfig as VtiStoreConfig;
866 use vti_common::store::Store;
867
868 const OWNER: &str = "did:key:z6MkChunkOwner";
869 const MIN: u64 = vta_sdk::protocols::backup_management::chunked::MIN_CHUNK_SIZE;
870
871 fn admin(did: &str) -> AuthClaims {
872 AuthClaims {
873 did: did.into(),
874 role: Role::Admin,
875 allowed_contexts: Vec::new(),
876 session_id: "test-session".into(),
877 access_expires_at: 0,
878 issued_at: 0,
879 amr: Vec::new(),
880 acr: String::new(),
881 }
882 }
883
884 struct Env {
885 _dir: tempfile::TempDir,
886 ks: KeyspaceHandle,
887 blob_dir: PathBuf,
888 }
889
890 fn env() -> Env {
891 let dir = tempfile::tempdir().unwrap();
892 let store = Store::open(&VtiStoreConfig {
893 data_dir: dir.path().join("store"),
894 })
895 .unwrap();
896 let ks = store.keyspace(crate::BACKUP_BUNDLES_TEST).unwrap();
897 let blob_dir = dir.path().join("backups");
898 Env {
899 _dir: dir,
900 ks,
901 blob_dir,
902 }
903 }
904
905 fn unlimited() -> ChunkRateLimiter {
906 ChunkRateLimiter::new(1.0e9, 1.0e9)
907 }
908
909 fn bundle_bytes() -> Vec<u8> {
912 let mut v = Vec::new();
913 v.extend(std::iter::repeat_n(1u8, MIN as usize));
914 v.extend(std::iter::repeat_n(2u8, MIN as usize));
915 v.extend_from_slice(b"backup-tail!");
916 v
917 }
918
919 fn digest(bytes: &[u8]) -> String {
920 sha256_digest_multibase(&Sha256::digest(bytes).into())
921 }
922
923 #[tokio::test]
924 async fn export_round_trips_chunk_by_chunk_and_is_non_consuming() {
925 let e = env();
926 let auth = admin(OWNER);
927 let bytes = bundle_bytes();
928 let bundle = stage_export(&e.ks, &e.blob_dir, OWNER, &bytes, Some(MIN))
929 .await
930 .unwrap();
931 assert_eq!(bundle.chunk_count, 3);
932
933 let mut assembled = Vec::new();
934 for i in 0..bundle.chunk_count {
935 let c = get_chunk(&e.ks, &unlimited(), &auth, &bundle.bundle_id.to_string(), i)
936 .await
937 .unwrap();
938 assert_eq!(c.digest_multibase, bundle.digests[i as usize]);
939 assert_eq!(digest(&c.data), bundle.digests[i as usize]);
940 assembled.extend(c.data);
941 }
942 assert_eq!(assembled, bytes);
943 assert_eq!(sha256_hex(&assembled), bundle.expected_sha256);
944
945 let again = get_chunk(&e.ks, &unlimited(), &auth, &bundle.bundle_id.to_string(), 2)
947 .await
948 .unwrap();
949 assert_eq!(again.data, b"backup-tail!");
950 assert_eq!(
951 all_served(&e.ks, &bundle.bundle_id).await.unwrap(),
952 Some(true)
953 );
954 }
955
956 #[tokio::test]
957 async fn an_index_past_the_last_chunk_is_out_of_range() {
958 let e = env();
959 let bundle = stage_export(&e.ks, &e.blob_dir, OWNER, &bundle_bytes(), Some(MIN))
960 .await
961 .unwrap();
962 let err = get_chunk(
963 &e.ks,
964 &unlimited(),
965 &admin(OWNER),
966 &bundle.bundle_id.to_string(),
967 3,
968 )
969 .await
970 .unwrap_err();
971 assert!(
972 matches!(
973 err,
974 ChunkedError::ChunkOutOfRange {
975 index: 3,
976 chunk_count: 3
977 }
978 ),
979 "{err:?}"
980 );
981 }
982
983 #[tokio::test]
984 async fn another_operators_bundle_is_not_found() {
985 let e = env();
986 let bundle = stage_export(&e.ks, &e.blob_dir, OWNER, &bundle_bytes(), Some(MIN))
987 .await
988 .unwrap();
989 let err = get_chunk(
990 &e.ks,
991 &unlimited(),
992 &admin("did:key:z6MkSomeoneElse"),
993 &bundle.bundle_id.to_string(),
994 0,
995 )
996 .await
997 .unwrap_err();
998 assert!(matches!(err, ChunkedError::NotFound), "{err:?}");
999 }
1000
1001 #[tokio::test]
1002 async fn serving_a_chunk_slides_the_expiry_but_never_past_the_ceiling() {
1003 let e = env();
1004 let auth = admin(OWNER);
1005 let bundle = stage_export(&e.ks, &e.blob_dir, OWNER, &bundle_bytes(), Some(MIN))
1006 .await
1007 .unwrap();
1008 let id = bundle.bundle_id;
1009
1010 let now = Utc::now();
1014 let mut record = backup_bundle_store::get_bundle(&e.ks, &id)
1015 .await
1016 .unwrap()
1017 .unwrap();
1018 record.expires_at = now + Duration::seconds(10);
1019 backup_bundle_store::store_bundle(&e.ks, &record)
1020 .await
1021 .unwrap();
1022 let mut plan = get_plan(&e.ks, &id).await.unwrap().unwrap();
1023 plan.expiry_ceiling = now + Duration::seconds(20);
1024 store_plan(&e.ks, &plan).await.unwrap();
1025
1026 let served = get_chunk(&e.ks, &unlimited(), &auth, &id.to_string(), 0)
1027 .await
1028 .unwrap();
1029 assert_eq!(served.expires_at, plan.expiry_ceiling);
1030
1031 let mut plan = get_plan(&e.ks, &id).await.unwrap().unwrap();
1033 plan.expiry_ceiling = now + Duration::hours(1);
1034 store_plan(&e.ks, &plan).await.unwrap();
1035 let served = get_chunk(&e.ks, &unlimited(), &auth, &id.to_string(), 1)
1036 .await
1037 .unwrap();
1038 assert!(served.expires_at > now + Duration::seconds(20));
1039 assert!(served.expires_at <= plan.expiry_ceiling);
1040 }
1041
1042 #[tokio::test]
1043 async fn an_expired_bundle_serves_nothing() {
1044 let e = env();
1045 let bundle = stage_export(&e.ks, &e.blob_dir, OWNER, &bundle_bytes(), Some(MIN))
1046 .await
1047 .unwrap();
1048 let mut record = backup_bundle_store::get_bundle(&e.ks, &bundle.bundle_id)
1049 .await
1050 .unwrap()
1051 .unwrap();
1052 record.expires_at = Utc::now() - Duration::seconds(1);
1053 backup_bundle_store::store_bundle(&e.ks, &record)
1054 .await
1055 .unwrap();
1056 let err = get_chunk(
1057 &e.ks,
1058 &unlimited(),
1059 &admin(OWNER),
1060 &bundle.bundle_id.to_string(),
1061 0,
1062 )
1063 .await
1064 .unwrap_err();
1065 assert!(matches!(err, ChunkedError::TerminalState(_)), "{err:?}");
1066 }
1067
1068 #[tokio::test]
1069 async fn a_bundle_too_large_for_the_chunk_bound_is_refused() {
1070 let e = env();
1071 let bytes = vec![0u8; (MIN * 4096 + 1) as usize];
1073 let err = stage_export(&e.ks, &e.blob_dir, OWNER, &bytes, Some(MIN))
1074 .await
1075 .unwrap_err();
1076 assert!(
1077 matches!(err, ChunkedError::BundleTooLarge { .. }),
1078 "{err:?}"
1079 );
1080 }
1081
1082 async fn open_import(e: &Env, bytes: &[u8]) -> ChunkedBundle {
1083 let digests: Vec<String> = bytes.chunks(MIN as usize).map(digest).collect();
1084 initiate_import(
1085 &e.ks,
1086 &admin(OWNER),
1087 &sha256_hex(bytes),
1088 bytes.len() as u64,
1089 MIN,
1090 digests.len() as u64,
1091 digests,
1092 )
1093 .await
1094 .unwrap()
1095 }
1096
1097 #[tokio::test]
1098 async fn import_resumes_after_a_missing_chunk_and_assembles() {
1099 let e = env();
1100 let auth = admin(OWNER);
1101 let bytes = bundle_bytes();
1102 let slot = open_import(&e, &bytes).await;
1103 let id = slot.bundle_id.to_string();
1104 let chunks: Vec<&[u8]> = bytes.chunks(MIN as usize).collect();
1105
1106 for i in [0usize, 2] {
1108 let out = put_chunk(
1109 &e.ks,
1110 &e.blob_dir,
1111 &unlimited(),
1112 &auth,
1113 ChunkWrite {
1114 bundle_id: &id,
1115 index: i as u64,
1116 digest_multibase: &slot.digests[i],
1117 data: chunks[i],
1118 },
1119 )
1120 .await
1121 .unwrap();
1122 assert!(out.stored);
1123 }
1124
1125 let err = finalize_precheck(&e.ks, &auth, &id).await.unwrap_err();
1127 match err {
1128 ChunkedError::IncompleteUpload {
1129 missing_count,
1130 missing_indices,
1131 } => {
1132 assert_eq!(missing_count, 1);
1133 assert_eq!(missing_indices, vec![1]);
1134 }
1135 other => panic!("expected IncompleteUpload, got {other:?}"),
1136 }
1137
1138 let out = put_chunk(
1140 &e.ks,
1141 &e.blob_dir,
1142 &unlimited(),
1143 &auth,
1144 ChunkWrite {
1145 bundle_id: &id,
1146 index: 1,
1147 digest_multibase: &slot.digests[1],
1148 data: chunks[1],
1149 },
1150 )
1151 .await
1152 .unwrap();
1153 assert!(out.stored);
1154 assert_eq!(out.remaining_count, 0);
1155
1156 finalize_precheck(&e.ks, &auth, &id).await.unwrap();
1157 let record = backup_bundle_store::get_bundle(&e.ks, &slot.bundle_id)
1158 .await
1159 .unwrap()
1160 .unwrap();
1161 assert_eq!(record.state, BundleState::ImportReceived);
1162 let staged = tokio::fs::read(record.blob_path.unwrap()).await.unwrap();
1163 assert_eq!(staged, bytes);
1164 }
1165
1166 #[tokio::test]
1167 async fn an_identical_re_put_is_stored_false_and_a_mismatched_one_is_refused() {
1168 let e = env();
1169 let auth = admin(OWNER);
1170 let bytes = bundle_bytes();
1171 let slot = open_import(&e, &bytes).await;
1172 let id = slot.bundle_id.to_string();
1173 let chunk0 = &bytes[..MIN as usize];
1174
1175 let first = put_chunk(
1176 &e.ks,
1177 &e.blob_dir,
1178 &unlimited(),
1179 &auth,
1180 ChunkWrite {
1181 bundle_id: &id,
1182 index: 0,
1183 digest_multibase: &slot.digests[0],
1184 data: chunk0,
1185 },
1186 )
1187 .await
1188 .unwrap();
1189 assert!(first.stored);
1190 let repeat = put_chunk(
1191 &e.ks,
1192 &e.blob_dir,
1193 &unlimited(),
1194 &auth,
1195 ChunkWrite {
1196 bundle_id: &id,
1197 index: 0,
1198 digest_multibase: &slot.digests[0],
1199 data: chunk0,
1200 },
1201 )
1202 .await
1203 .unwrap();
1204 assert!(!repeat.stored, "an identical re-put stores nothing new");
1205
1206 let tampered = vec![9u8; MIN as usize];
1208 let err = put_chunk(
1209 &e.ks,
1210 &e.blob_dir,
1211 &unlimited(),
1212 &auth,
1213 ChunkWrite {
1214 bundle_id: &id,
1215 index: 0,
1216 digest_multibase: &digest(&tampered),
1217 data: &tampered,
1218 },
1219 )
1220 .await
1221 .unwrap_err();
1222 assert!(
1223 matches!(err, ChunkedError::DigestMismatch { .. }),
1224 "{err:?}"
1225 );
1226
1227 let err = put_chunk(
1229 &e.ks,
1230 &e.blob_dir,
1231 &unlimited(),
1232 &auth,
1233 ChunkWrite {
1234 bundle_id: &id,
1235 index: 1,
1236 digest_multibase: &slot.digests[1],
1237 data: &tampered,
1238 },
1239 )
1240 .await
1241 .unwrap_err();
1242 assert!(
1243 matches!(err, ChunkedError::DigestMismatch { .. }),
1244 "{err:?}"
1245 );
1246
1247 let record = backup_bundle_store::get_bundle(&e.ks, &slot.bundle_id)
1248 .await
1249 .unwrap()
1250 .unwrap();
1251 let staged = tokio::fs::read(record.blob_path.unwrap()).await.unwrap();
1252 assert_eq!(&staged[..MIN as usize], chunk0);
1253 }
1254
1255 #[tokio::test]
1256 async fn a_wrong_length_chunk_is_refused() {
1257 let e = env();
1258 let bytes = bundle_bytes();
1259 let slot = open_import(&e, &bytes).await;
1260 let short = &bytes[..10];
1261 let err = put_chunk(
1262 &e.ks,
1263 &e.blob_dir,
1264 &unlimited(),
1265 &admin(OWNER),
1266 ChunkWrite {
1267 bundle_id: &slot.bundle_id.to_string(),
1268 index: 0,
1269 digest_multibase: &digest(short),
1270 data: short,
1271 },
1272 )
1273 .await
1274 .unwrap_err();
1275 assert!(
1276 matches!(err, ChunkedError::ChunkSizeMismatch { expected, actual: 10 } if expected == MIN),
1277 "{err:?}"
1278 );
1279 }
1280
1281 #[tokio::test]
1282 async fn a_manifest_whose_count_does_not_follow_is_refused() {
1283 let e = env();
1284 let bytes = bundle_bytes();
1285 let digests: Vec<String> = bytes.chunks(MIN as usize).map(digest).collect();
1286 let err = initiate_import(
1287 &e.ks,
1288 &admin(OWNER),
1289 &sha256_hex(&bytes),
1290 bytes.len() as u64,
1291 MIN,
1292 2,
1293 digests,
1294 )
1295 .await
1296 .unwrap_err();
1297 assert!(matches!(err, ChunkedError::InvalidManifest(_)), "{err:?}");
1298 }
1299
1300 #[tokio::test]
1301 async fn chunks_that_verify_but_do_not_assemble_to_the_bundle_are_caught() {
1302 let e = env();
1303 let auth = admin(OWNER);
1304 let bytes = bundle_bytes();
1305 let digests: Vec<String> = bytes.chunks(MIN as usize).map(digest).collect();
1308 let slot = initiate_import(
1309 &e.ks,
1310 &auth,
1311 &sha256_hex(b"something else"),
1312 bytes.len() as u64,
1313 MIN,
1314 3,
1315 digests,
1316 )
1317 .await
1318 .unwrap();
1319 let id = slot.bundle_id.to_string();
1320 for (i, c) in bytes.chunks(MIN as usize).enumerate() {
1321 put_chunk(
1322 &e.ks,
1323 &e.blob_dir,
1324 &unlimited(),
1325 &auth,
1326 ChunkWrite {
1327 bundle_id: &id,
1328 index: i as u64,
1329 digest_multibase: &slot.digests[i],
1330 data: c,
1331 },
1332 )
1333 .await
1334 .unwrap();
1335 }
1336 let err = finalize_precheck(&e.ks, &auth, &id).await.unwrap_err();
1337 assert!(matches!(err, ChunkedError::BundleDigestMismatch), "{err:?}");
1338 }
1339
1340 #[test]
1341 fn the_limiter_refuses_past_the_burst_and_names_a_wait() {
1342 let limiter = ChunkRateLimiter::new(1.0, 2.0);
1343 limiter.check(OWNER).unwrap();
1344 limiter.check(OWNER).unwrap();
1345 match limiter.check(OWNER) {
1346 Err(ChunkedError::RateLimited { retry_after_secs }) => assert!(retry_after_secs >= 1),
1347 other => panic!("expected RateLimited, got {other:?}"),
1348 }
1349 limiter.check("did:key:z6MkAnotherOperator").unwrap();
1351 }
1352}