diff --git a/Cargo.lock b/Cargo.lock index 8f88dce7e8..541491e104 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7191,6 +7191,7 @@ dependencies = [ "itertools 0.14.0", "log", "memory", + "mockito", "parking_lot", "proptest", "prost 0.11.9", @@ -7204,6 +7205,7 @@ dependencies = [ "serde", "serde_cbor", "serde_json", + "sha2", "shard", "strum", "tap", diff --git a/Cargo.toml b/Cargo.toml index 50db29cf44..f60da8efef 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,13 +43,12 @@ staging = ["collection/staging", "storage/staging", "shard/staging"] serde_urlencoded = "0.7" sealed_test = "1.1.0" +mockito = { workspace = true } tempfile = { workspace = true } rusty-hook = "^0.11.2" nix = { workspace = true, features = ["process"] } fs-err = { workspace = true, features = ["debug", "debug_tokio", "tokio"] } # for nicer error messages -mockito = "1.7" - [dependencies] parking_lot = { workspace = true } @@ -238,6 +237,7 @@ integer-encoding = "4.1.0" itertools = "0.14.0" log = "0.4.29" memmap2 = "0.9.9" +mockito = "1.7" nix = { version = "0.31", features = ["fs", "feature"] } num-traits = "0.2.19" ordered-float = { version = "5.1.0", features = ["serde", "schemars"] } diff --git a/clippy.toml b/clippy.toml index 33e1fcba10..2f939ecd12 100644 --- a/clippy.toml +++ b/clippy.toml @@ -2,6 +2,7 @@ large-error-threshold = 256 disallowed-types = [ + # Use fs_err instead of std::fs and tokio::fs { path = "std::fs::DirEntry", replacement = "fs_err::DirEntry" }, { path = "std::fs::File", replacement = "fs_err::File" }, { path = "std::fs::OpenOptions", replacement = "fs_err::OpenOptions" }, @@ -14,6 +15,9 @@ disallowed-types = [ ] disallowed-methods = [ + { path = "tar::Archive::unpack", replacement = "common::tar_unpack", reason = "Use safe unpackers" }, + + # Use fs_err instead of std::fs and tokio::fs { path = "std::fs::DirEntry::file_type", replacement = "fs_err::DirEntry::file_type" }, { path = "std::fs::DirEntry::metadata", replacement = "fs_err::DirEntry::metadata" }, { path = "std::fs::File::create", replacement = "fs_err::File::create" }, diff --git a/lib/collection/src/collection/snapshots.rs b/lib/collection/src/collection/snapshots.rs index 73a7ffcfa1..25a7f76261 100644 --- a/lib/collection/src/collection/snapshots.rs +++ b/lib/collection/src/collection/snapshots.rs @@ -2,12 +2,13 @@ use std::collections::HashSet; use std::path::Path; use common::tar_ext::BuilderExt; -use common::tempfile_ext::MaybeTempPath; +use common::tar_unpack::tar_unpack_file; use fs_err::File; use io::file_operations::read_json; use io::storage_version::StorageVersion as _; -use segment::common::validate_snapshot_archive::open_snapshot_archive_with_validation; use segment::types::SnapshotFormat; +use segment::utils::fs::move_all; +use shard::snapshots::snapshot_data::SnapshotData; use shard::snapshots::snapshot_manifest::{RecoveryType, SnapshotManifest}; use tokio::sync::OwnedRwLockReadGuard; @@ -157,14 +158,22 @@ impl Collection { /// /// This method performs blocking IO. pub fn restore_snapshot( - snapshot_path: &Path, + snapshot_data: SnapshotData, target_dir: &Path, this_peer_id: PeerId, is_distributed: bool, ) -> CollectionResult<()> { - // decompress archive - let mut ar = open_snapshot_archive_with_validation(snapshot_path)?; - ar.unpack(target_dir)?; + match snapshot_data { + SnapshotData::Packed(snapshot_path) => { + tar_unpack_file(&snapshot_path, target_dir)?; + snapshot_path.close()?; + } + SnapshotData::Unpacked(snapshot_dir) => { + // already unpacked snapshot, validate files and move to target dir + let snapshot_dir_path = snapshot_dir.path(); + move_all(snapshot_dir_path, target_dir)?; + } + } let config = CollectionConfigInternal::load(target_dir)?; config.validate_and_warn(); @@ -299,7 +308,7 @@ impl Collection { pub async fn restore_shard_snapshot( &self, shard_id: ShardId, - snapshot_path: MaybeTempPath, + snapshot_data: SnapshotData, recovery_type: RecoveryType, this_peer_id: PeerId, is_distributed: bool, @@ -308,14 +317,7 @@ impl Collection { ) -> CollectionResult> + 'static> { // `ShardHolder::validate_shard_snapshot` is cancel safe, so we explicitly cancel it // when token is triggered - let shard_holder = cancel::future::cancel_on_token(cancel.clone(), async { - let shard_holder = self.shards_holder.clone().read_owned().await; - - shard_holder.validate_shard_snapshot(&snapshot_path).await?; - - CollectionResult::Ok(shard_holder) - }) - .await??; + let shard_holder = self.shards_holder.clone().read_owned().await; let collection_path = self.path.clone(); let collection_name = self.name().to_string(); @@ -327,7 +329,7 @@ impl Collection { let restore = self.update_runtime.spawn(async move { shard_holder .restore_shard_snapshot( - &snapshot_path, + snapshot_data, recovery_type, &collection_path, &collection_name, @@ -339,10 +341,6 @@ impl Collection { ) .await?; - if let Err(err) = snapshot_path.close() { - log::error!("Failed to remove downloaded snapshot archive after recovery: {err}"); - } - CollectionResult::Ok(()) }); diff --git a/lib/collection/src/shards/shard_holder/mod.rs b/lib/collection/src/shards/shard_holder/mod.rs index 89d7db4996..f845cc292c 100644 --- a/lib/collection/src/shards/shard_holder/mod.rs +++ b/lib/collection/src/shards/shard_holder/mod.rs @@ -13,16 +13,16 @@ use api::rest::ShardKeyWithFallback; use common::budget::ResourceBudget; use common::save_on_disk::SaveOnDisk; use common::tar_ext::BuilderExt; +use common::tar_unpack::tar_unpack_file; use fs_err as fs; use fs_err::{File, tokio as tokio_fs}; use futures::{Future, StreamExt, TryStreamExt as _, stream}; use io::safe_delete::sync_parent_dir_async; use itertools::Itertools; -use segment::common::validate_snapshot_archive::{ - open_snapshot_archive, validate_snapshot_archive, -}; use segment::json_path::JsonPath; use segment::types::{PayloadFieldSchema, ShardKey, SnapshotFormat}; +use segment::utils::fs::move_all; +use shard::snapshots::snapshot_data::SnapshotData; use shard::snapshots::snapshot_manifest::{RecoveryType, SnapshotManifest}; use shard_mapping::ShardKeyMapping; use tokio::runtime::Handle; @@ -1199,24 +1199,13 @@ impl ShardHolder { )) } - /// # Cancel safety - /// - /// This method is cancel safe. - pub async fn validate_shard_snapshot(&self, snapshot_path: &Path) -> CollectionResult<()> { - validate_snapshot_archive(snapshot_path)?; - - // TODO: Validate that shard/partial snapshot is compatible with collection config! - - Ok(()) - } - /// # Cancel safety /// /// This method is *not* cancel safe. #[allow(clippy::too_many_arguments)] pub async fn restore_shard_snapshot( &self, - snapshot_path: &Path, + snapshot_data: SnapshotData, recovery_type: RecoveryType, collection_path: &Path, collection_name: &str, @@ -1234,30 +1223,29 @@ impl ShardHolder { fs::create_dir_all(temp_dir)?; } - let snapshot_file_name = snapshot_path.file_name().unwrap().to_string_lossy(); - let snapshot_temp_dir = tempfile::Builder::new() - .prefix(&format!( - "{collection_name}-shard-{shard_id}-{snapshot_file_name}" - )) + .prefix(&format!("{collection_name}-shard-{shard_id}")) .tempdir_in(temp_dir)?; let extract = { - let snapshot_path = snapshot_path.to_path_buf(); let snapshot_temp_dir = snapshot_temp_dir.path().to_path_buf(); cancel::blocking::spawn_cancel_on_token( cancel.child_token(), move |cancel| -> CollectionResult<_> { - let mut ar = open_snapshot_archive(&snapshot_path)?; - - if cancel.is_cancelled() { - return Err(cancel::Error::Cancelled.into()); + match snapshot_data { + SnapshotData::Packed(snapshot_path) => { + if cancel.is_cancelled() { + return Err(cancel::Error::Cancelled.into()); + } + tar_unpack_file(&snapshot_path, &snapshot_temp_dir)?; + snapshot_path.close()?; + } + SnapshotData::Unpacked(snapshot_dir) => { + move_all(snapshot_dir.path(), &snapshot_temp_dir)?; + } } - ar.unpack(&snapshot_temp_dir)?; - drop(ar); - if cancel.is_cancelled() { return Err(cancel::Error::Cancelled.into()); } @@ -1289,9 +1277,7 @@ impl ShardHolder { .await?; if !recovered { - return Err(CollectionError::bad_request(format!( - "Invalid snapshot {snapshot_file_name}" - ))); + return Err(CollectionError::bad_request("Invalid snapshot")); } if recovery_type.is_partial() { diff --git a/lib/collection/src/tests/snapshot_test.rs b/lib/collection/src/tests/snapshot_test.rs index 60572d40f1..aaa639f5af 100644 --- a/lib/collection/src/tests/snapshot_test.rs +++ b/lib/collection/src/tests/snapshot_test.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use ahash::AHashMap; use common::budget::ResourceBudget; use segment::types::Distance; +use shard::snapshots::snapshot_data::SnapshotData; use tempfile::Builder; use crate::collection::{Collection, RequestShardTransfer}; @@ -109,15 +110,13 @@ async fn _test_snapshot_collection(node_type: NodeType) { .prefix("test_collection_rec") .tempdir() .unwrap(); + let snapshot_data = SnapshotData::new_packed_persistent( + snapshots_path.path().join(&snapshot_description.name), + ); + // Do not recover in local mode if some shards are remote assert!( - Collection::restore_snapshot( - &snapshots_path.path().join(&snapshot_description.name), - recover_dir.path(), - 0, - false, - ) - .is_err(), + Collection::restore_snapshot(snapshot_data, recover_dir.path(), 0, false,).is_err(), ); } @@ -125,13 +124,9 @@ async fn _test_snapshot_collection(node_type: NodeType) { .prefix("test_collection_rec") .tempdir() .unwrap(); - - if let Err(err) = Collection::restore_snapshot( - &snapshots_path.path().join(snapshot_description.name), - recover_dir.path(), - 0, - true, - ) { + let snapshot_data = + SnapshotData::new_packed_persistent(snapshots_path.path().join(&snapshot_description.name)); + if let Err(err) = Collection::restore_snapshot(snapshot_data, recover_dir.path(), 0, true) { panic!("Failed to restore snapshot: {err}") } diff --git a/lib/collection/tests/integration/snapshot_recovery_test.rs b/lib/collection/tests/integration/snapshot_recovery_test.rs index dedc481dba..a8817f645b 100644 --- a/lib/collection/tests/integration/snapshot_recovery_test.rs +++ b/lib/collection/tests/integration/snapshot_recovery_test.rs @@ -18,6 +18,7 @@ use collection::shards::replica_set::replica_set_state::ReplicaState; use common::budget::ResourceBudget; use common::counter::hardware_accumulator::HwMeasurementAcc; use segment::types::{Distance, WithPayloadInterface, WithVector}; +use shard::snapshots::snapshot_data::SnapshotData; use tempfile::Builder; use crate::common::{ @@ -128,12 +129,10 @@ async fn _test_snapshot_and_recover_collection(node_type: NodeType) { .await .unwrap(); - if let Err(err) = Collection::restore_snapshot( - &snapshots_path.path().join(snapshot_description.name), - recover_dir.path(), - 0, - false, - ) { + let snapshot_data = + SnapshotData::new_packed_persistent(snapshots_path.path().join(snapshot_description.name)); + + if let Err(err) = Collection::restore_snapshot(snapshot_data, recover_dir.path(), 0, false) { panic!("Failed to restore snapshot: {err}") } diff --git a/lib/common/common/src/lib.rs b/lib/common/common/src/lib.rs index 4fb2e2bd31..165ecc32e5 100644 --- a/lib/common/common/src/lib.rs +++ b/lib/common/common/src/lib.rs @@ -30,6 +30,7 @@ pub mod small_uint; pub mod sort_utils; pub mod stable_hash; pub mod tar_ext; +pub mod tar_unpack; pub mod tempfile_ext; pub mod top_k; pub mod toposort; diff --git a/lib/common/common/src/tar_unpack.rs b/lib/common/common/src/tar_unpack.rs new file mode 100644 index 0000000000..3e0e1b85b5 --- /dev/null +++ b/lib/common/common/src/tar_unpack.rs @@ -0,0 +1,41 @@ +//! Wrappers around [`tar::Archive::unpack()`] with extra safety checks. + +use std::io; +use std::path::Path; + +use fs_err as fs; +use tar::{Archive, EntryType}; + +pub fn tar_unpack_file(path: &Path, dst: &Path) -> Result<(), io::Error> { + let reader = io::BufReader::new(fs::File::open(path)?); + tar_unpack_reader(reader, dst)?; + Ok(()) +} + +/// Same as [`Archive::new()`] followed by [`Archive::unpack()`], but checks +/// that we don't unpack something beyond regular files and directories. +/// +/// Accepts a reader and returns the same reader. +pub fn tar_unpack_reader(reader: R, dst: &Path) -> Result { + let mut archive = Archive::new(reader); + archive.set_overwrite(false); + + fs::create_dir_all(dst)?; + let dst = &fs::canonicalize(dst).unwrap_or(dst.to_path_buf()); + + for entry in archive.entries()? { + let mut entry = entry?; + + match entry.header().entry_type() { + EntryType::Directory | EntryType::Regular | EntryType::GNUSparse => (), + entry_type => { + return Err(io::Error::other(format!( + "Invalid entry type in tar archive: {entry_type:?}" + ))); + } + } + entry.unpack_in(dst)?; + } + + Ok(archive.into_inner()) +} diff --git a/lib/segment/src/common/validate_snapshot_archive.rs b/lib/segment/src/common/validate_snapshot_archive.rs index 4303c204d0..8b13789179 100644 --- a/lib/segment/src/common/validate_snapshot_archive.rs +++ b/lib/segment/src/common/validate_snapshot_archive.rs @@ -1,64 +1 @@ -use std::io; -use std::path::Path; -use fs_err as fs; - -use crate::common::operation_error::{OperationError, OperationResult}; - -pub fn open_snapshot_archive_with_validation( - path: &Path, -) -> OperationResult> { - validate_snapshot_archive(path)?; - open_snapshot_archive(path) -} - -pub fn validate_snapshot_archive(path: &Path) -> OperationResult<()> { - let mut ar = open_snapshot_archive(path)?; - - let entries = ar.entries_with_seek().map_err(|err| { - OperationError::service_error(format!( - "failed to read snapshot archive {}: {err}", - path.display() - )) - })?; - - for entry in entries { - let entry = entry.map_err(|err| { - log::error!("Failed to read snapshot archive {}: {err}", path.display()); - - // Deliberately mask underlying error from API users, because it can expose arbitrary file contents - OperationError::service_error(format!( - "failed to read snapshot archive {}", - path.display(), - )) - })?; - - match entry.header().entry_type() { - tar::EntryType::Directory | tar::EntryType::Regular | tar::EntryType::GNUSparse => (), - entry_type => { - return Err(OperationError::validation_error(format!( - "malformed snapshot archive {}: archive contains {entry_type:?} entry", - path.display(), - ))); - } - } - } - - Ok(()) -} - -pub fn open_snapshot_archive( - path: &Path, -) -> OperationResult> { - let file = fs::File::open(path).map_err(|err| { - OperationError::service_error(format!( - "failed to open snapshot archive {}: {err}", - path.display() - )) - })?; - - let mut ar = tar::Archive::new(io::BufReader::new(file)); - ar.set_overwrite(false); - - Ok(ar) -} diff --git a/lib/segment/src/segment/segment_ops.rs b/lib/segment/src/segment/segment_ops.rs index 2397f3f4ca..098e63bc8c 100644 --- a/lib/segment/src/segment/segment_ops.rs +++ b/lib/segment/src/segment/segment_ops.rs @@ -4,6 +4,7 @@ use std::path::Path; use bitvec::prelude::BitVec; use common::counter::hardware_counter::HardwareCounterCell; +use common::tar_unpack::tar_unpack_file; use common::types::PointOffsetType; use fs_err as fs; use io::file_operations::{atomic_save_json, read_json}; @@ -12,7 +13,6 @@ use super::{SEGMENT_STATE_FILE, SNAPSHOT_FILES_PATH, SNAPSHOT_PATH, Segment}; use crate::common::operation_error::{ OperationError, OperationResult, SegmentFailedState, get_service_error, }; -use crate::common::validate_snapshot_archive::open_snapshot_archive_with_validation; use crate::common::{check_named_vectors, check_vector_name}; use crate::data_types::named_vectors::NamedVectors; use crate::data_types::vectors::VectorInternal; @@ -685,7 +685,7 @@ fn restore_snapshot_in_place(snapshot_path: &Path) -> OperationResult<()> { unpack_snapshot(snapshot_path)?; } else { let segment_path = segments_dir.join(segment_id); - open_snapshot_archive_with_validation(snapshot_path)?.unpack(&segment_path)?; + tar_unpack_file(snapshot_path, &segment_path)?; let inner_path = segment_path.join(SNAPSHOT_PATH); if inner_path.is_dir() { diff --git a/lib/segment/src/segment/tests.rs b/lib/segment/src/segment/tests.rs index 2c9625eccd..8192d7b346 100644 --- a/lib/segment/src/segment/tests.rs +++ b/lib/segment/src/segment/tests.rs @@ -2,6 +2,7 @@ use std::sync::atomic::AtomicBool; use common::counter::hardware_counter::HardwareCounterCell; use common::tar_ext; +use common::tar_unpack::tar_unpack_file; use fs_err as fs; use fs_err::File; use rstest::rstest; @@ -215,9 +216,7 @@ fn test_snapshot(#[case] format: SnapshotFormat) { tar.blocking_finish().unwrap(); let parent_snapshot_unpacked = Builder::new().prefix("parent_snapshot").tempdir().unwrap(); - tar::Archive::new(File::open(parent_snapshot_tar.path()).unwrap()) - .unpack(parent_snapshot_unpacked.path()) - .unwrap(); + tar_unpack_file(parent_snapshot_tar.path(), parent_snapshot_unpacked.path()).unwrap(); // Should be exactly one entry in the snapshot. let mut entries = fs::read_dir(parent_snapshot_unpacked.path()).unwrap(); diff --git a/lib/segment/tests/integration/segment_on_disk_snapshot.rs b/lib/segment/tests/integration/segment_on_disk_snapshot.rs index 7b4e10e727..890e31b3c6 100644 --- a/lib/segment/tests/integration/segment_on_disk_snapshot.rs +++ b/lib/segment/tests/integration/segment_on_disk_snapshot.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::sync::atomic::AtomicBool; use common::tar_ext; +use common::tar_unpack::tar_unpack_file; use fs_err as fs; use fs_err::File; use rstest::rstest; @@ -162,9 +163,8 @@ fn test_on_disk_segment_snapshot(#[case] format: SnapshotFormat) { tar.blocking_finish().unwrap(); let parent_snapshot_unpacked = Builder::new().prefix("parent_snapshot").tempdir().unwrap(); - tar::Archive::new(File::open(parent_snapshot_tar.path()).unwrap()) - .unpack(parent_snapshot_unpacked.path()) - .unwrap(); + + tar_unpack_file(parent_snapshot_tar.path(), parent_snapshot_unpacked.path()).unwrap(); // Should be exactly one entry in the snapshot. let mut entries = fs::read_dir(parent_snapshot_unpacked.path()).unwrap(); diff --git a/lib/shard/Cargo.toml b/lib/shard/Cargo.toml index ab43cd7334..e805cf064f 100644 --- a/lib/shard/Cargo.toml +++ b/lib/shard/Cargo.toml @@ -43,6 +43,7 @@ validator = { workspace = true } wal = { workspace = true } serde_json = { workspace = true } fs-err = { workspace = true } +tempfile = { workspace = true } [dev-dependencies] @@ -54,4 +55,3 @@ proptest = { workspace = true } rstest = { workspace = true } serde_cbor = { workspace = true } tar = { workspace = true } -tempfile = { workspace = true } diff --git a/lib/shard/src/snapshots/mod.rs b/lib/shard/src/snapshots/mod.rs index e90065070f..2ed68beaa6 100644 --- a/lib/shard/src/snapshots/mod.rs +++ b/lib/shard/src/snapshots/mod.rs @@ -1,2 +1,3 @@ +pub mod snapshot_data; pub mod snapshot_manifest; pub mod snapshot_utils; diff --git a/lib/shard/src/snapshots/snapshot_data.rs b/lib/shard/src/snapshots/snapshot_data.rs new file mode 100644 index 0000000000..292f6f80c1 --- /dev/null +++ b/lib/shard/src/snapshots/snapshot_data.rs @@ -0,0 +1,23 @@ +use common::tempfile_ext::MaybeTempPath; +use tempfile::TempDir; + +pub enum SnapshotData { + /// Tar file containing the snapshot, needs to be unpacked + Packed(MaybeTempPath), + /// Directory containing the unpacked snapshot + Unpacked(TempDir), +} + +impl SnapshotData { + pub fn new_packed_persistent>(path: P) -> Self { + SnapshotData::Packed(MaybeTempPath::Persistent(path.as_ref().to_path_buf())) + } + + /// Get path to the downloaded data + pub fn path(&self) -> &std::path::Path { + match self { + SnapshotData::Packed(maybe_path) => maybe_path, + SnapshotData::Unpacked(temp_dir) => temp_dir.path(), + } + } +} diff --git a/lib/shard/src/snapshots/snapshot_utils.rs b/lib/shard/src/snapshots/snapshot_utils.rs index 8b3f0a3031..37f5f4971b 100644 --- a/lib/shard/src/snapshots/snapshot_utils.rs +++ b/lib/shard/src/snapshots/snapshot_utils.rs @@ -1,8 +1,8 @@ use std::path::{Path, PathBuf}; +use common::tar_unpack::tar_unpack_file; use fs_err as fs; use segment::common::operation_error::OperationResult; -use segment::common::validate_snapshot_archive::open_snapshot_archive; use segment::segment::Segment; use crate::files::{ShardDataFiles, get_shard_data_files, segments_path}; @@ -20,8 +20,7 @@ impl SnapshotUtils { fs::create_dir_all(target_path)?; } - let mut ar = open_snapshot_archive(snapshot_path)?; - ar.unpack(target_path)?; + tar_unpack_file(snapshot_path, target_path)?; Self::restore_unpacked_snapshot(target_path)?; Ok(()) } diff --git a/lib/storage/Cargo.toml b/lib/storage/Cargo.toml index 0f5c41c94b..cbe6db5e97 100644 --- a/lib/storage/Cargo.toml +++ b/lib/storage/Cargo.toml @@ -18,6 +18,7 @@ staging = ["collection/staging"] [dev-dependencies] fs-err = { workspace = true, features = ["debug"] } proptest = { workspace = true } +mockito = { workspace = true } env_logger = "0.11" [dependencies] @@ -43,6 +44,7 @@ chrono = { workspace = true } validator = { workspace = true } dashmap = { workspace = true } semver = { workspace = true } +sha2 = { workspace = true } # Consensus related atomicwrites = { workspace = true } diff --git a/lib/storage/src/content_manager/snapshots/download.rs b/lib/storage/src/content_manager/snapshots/download.rs index 56013bc9c7..446990c5e1 100644 --- a/lib/storage/src/content_manager/snapshots/download.rs +++ b/lib/storage/src/content_manager/snapshots/download.rs @@ -1,17 +1,17 @@ use std::ffi::OsString; use std::path::Path; +use collection::common::sha_256::hash_file; use common::tempfile_ext::MaybeTempPath; -use fs_err::tokio as tokio_fs; -use futures::StreamExt; -use segment::common::BYTES_IN_MB; +use reqwest; +use shard::snapshots::snapshot_data::SnapshotData; use tap::Tap; -use tempfile::TempPath; -use tokio::io::AsyncWriteExt; +use tempfile::TempDir; use url::Url; -use {fs_err as fs, reqwest}; use crate::StorageError; +use crate::content_manager::snapshots::download_result::DownloadResult; +use crate::content_manager::snapshots::download_tar::download_and_unpack_tar; fn snapshot_prefix(url: &Url) -> OsString { Path::new(url.path()) @@ -20,72 +20,46 @@ fn snapshot_prefix(url: &Url) -> OsString { .unwrap_or_default() } -/// Download a remote file from `url` to `path` +/// Download and unpack a snapshot from `url` into a temporary directory. /// -/// Returns a `TempPath` that will delete the downloaded file once it is dropped. -/// To persist the file, use `download_file(...).keep()`. -#[must_use = "returns a TempPath, if dropped the downloaded file is deleted"] -async fn download_file( +/// Returns a `TempDir` that will delete the downloaded file once it is dropped. +/// To persist the file, use [`keep()`](TempDir::keep). +#[must_use = "returns a TempDir, if dropped the downloaded file is deleted"] +async fn _download_snapshot( client: &reqwest::Client, url: &Url, dir_path: &Path, -) -> Result { + compute_checksum: bool, +) -> Result<(TempDir, Option), StorageError> { let download_start_time = tokio::time::Instant::now(); - let (file, temp_path) = tempfile::Builder::new() - .prefix(&snapshot_prefix(url)) + + let snapshot_name = snapshot_prefix(url); + + let tempdir = tempfile::Builder::new() + .prefix(&snapshot_name) .suffix(".download") - .tempfile_in(dir_path)? - .into_parts(); - let file = fs::File::from_parts::<&Path>(file, temp_path.as_ref()); + .tempdir_in(dir_path)?; - log::debug!("Downloading snapshot from {url} to {temp_path:?}"); - - let mut file = tokio_fs::File::from_std(file); - - let response = client.get(url.clone()).send().await?; - - if !response.status().is_success() { - return Err(StorageError::bad_input(format!( - "Failed to download snapshot from {}: status - {}", - url, - response.status() - ))); - } - - let mut stream = response.bytes_stream(); - let mut total_bytes_downloaded = 0u64; - - while let Some(chunk_result) = stream.next().await { - let chunk = chunk_result?; - total_bytes_downloaded += chunk.len() as u64; - file.write_all(&chunk).await?; - } - - file.flush().await?; + let hash = download_and_unpack_tar(client, url, tempdir.path(), compute_checksum).await?; let download_duration = download_start_time.elapsed(); - let total_size_mb = total_bytes_downloaded as f64 / BYTES_IN_MB as f64; - let download_speed_mbps = total_size_mb / download_duration.as_secs_f64(); log::debug!( - "Snapshot download completed: path={}, size={:.2} MB, duration={:.2}s, speed={:.2} MB/s", - temp_path.display(), - total_size_mb, + "Snapshot download completed: path={tempdir:?}, duration={:.2}s", download_duration.as_secs_f64(), - download_speed_mbps ); - Ok(temp_path) + Ok((tempdir, hash)) } /// Download a snapshot from the given URI. /// -/// May returen a `TempPath` if a file was downloaded from a remote source. If it is dropped the -/// downloaded file is deleted automatically. To keep the file `keep()` may be used. +/// Returns a `DownloadResult` containing the snapshot data and optional checksum. pub async fn download_snapshot( client: &reqwest::Client, url: Url, snapshots_dir: &Path, -) -> Result { + compute_checksum: bool, +) -> Result { match url.scheme() { "file" => { let local_path = url.to_file_path().map_err(|_| { @@ -98,15 +72,28 @@ pub async fn download_snapshot( "Snapshot file {local_path:?} does not exist" ))); } - Ok(MaybeTempPath::Persistent(local_path)) + let hash = if compute_checksum { + Some(hash_file(&local_path).await?) + } else { + None + }; + + Ok(DownloadResult { + snapshot: SnapshotData::Packed(MaybeTempPath::Persistent(local_path)), + hash, + }) + } + "http" | "https" => { + let (snapshot_dir, hash) = + _download_snapshot(client, &url, snapshots_dir, compute_checksum).await?; + Ok(DownloadResult { + snapshot: SnapshotData::Unpacked(snapshot_dir), + hash, + }) } - "http" | "https" => Ok(MaybeTempPath::Temporary( - download_file(client, &url, snapshots_dir).await?, - )), _ => Err(StorageError::bad_request(format!( - "URL {} with schema {} is not supported", - url, - url.scheme() + "URL {url} with scheme {} is not supported", + url.scheme(), ))), } } diff --git a/lib/storage/src/content_manager/snapshots/download_result.rs b/lib/storage/src/content_manager/snapshots/download_result.rs new file mode 100644 index 0000000000..7dd2ec1556 --- /dev/null +++ b/lib/storage/src/content_manager/snapshots/download_result.rs @@ -0,0 +1,7 @@ +use shard::snapshots::snapshot_data::SnapshotData; + +pub struct DownloadResult { + pub snapshot: SnapshotData, + /// Sha256 hash of the downloaded snapshot file, if computed + pub hash: Option, +} diff --git a/lib/storage/src/content_manager/snapshots/download_tar.rs b/lib/storage/src/content_manager/snapshots/download_tar.rs new file mode 100644 index 0000000000..b525581324 --- /dev/null +++ b/lib/storage/src/content_manager/snapshots/download_tar.rs @@ -0,0 +1,194 @@ +use std::io::Read; +use std::path::Path; + +use cancel::CancellationToken; +use common::tar_unpack::tar_unpack_reader; +use futures::TryStreamExt; +use sha2::{Digest, Sha256}; +use tokio_util::io::StreamReader; +use url::Url; + +use crate::StorageError; + +/// A sync Read wrapper that checks a cancellation token before each read. +struct CancellableReader { + inner: R, + cancel: CancellationToken, +} + +impl CancellableReader { + fn new(inner: R, cancel: CancellationToken) -> Self { + Self { inner, cancel } + } +} + +impl Read for CancellableReader { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + if self.cancel.is_cancelled() { + return Err(std::io::Error::new( + std::io::ErrorKind::Interrupted, + "download cancelled", + )); + } + self.inner.read(buf) + } +} + +/// A sync Read wrapper that computes SHA-256 hash of the data as it's read. +struct HashingReader { + inner: R, + hasher: Option, +} + +impl HashingReader { + fn new(inner: R, compute_hash: bool) -> Self { + Self { + inner, + hasher: compute_hash.then(Sha256::new), + } + } + + /// Consume the reader and return the computed hash as a hex string. + /// Returns None if hashing was not enabled. + fn finalize(self) -> Option { + self.hasher.map(|h| { + let hash = h.finalize(); + format!("{hash:x}") + }) + } +} + +impl Read for HashingReader { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let n = self.inner.read(buf)?; + if let Some(ref mut hasher) = self.hasher { + hasher.update(&buf[..n]); + } + Ok(n) + } +} + +/// Download and unpack a tar file in streaming fashion without saving to disk first. +/// +/// This function streams the HTTP response directly into the tar extractor, +/// avoiding the need to store the entire tar file on disk before extraction. +/// +/// # Cancel safety +/// +/// This function is cancel safe. If cancelled, the cancellation token will be triggered +/// and the download will be interrupted at the next read operation. +/// +/// # Arguments +/// +/// * `client` - The reqwest HTTP client to use for the download +/// * `url` - The URL to download the tar file from +/// * `target_dir` - The directory to extract the tar contents into +/// * `compute_checksum` - If true, compute and return the SHA-256 hash of the downloaded data +/// +/// # Returns +/// +/// Returns `Ok(Some(hash))` if `compute_checksum` is true, `Ok(None)` otherwise. +/// Returns a `StorageError` if the download or extraction fails. +pub async fn download_and_unpack_tar( + client: &reqwest::Client, + url: &Url, + target_dir: &Path, + compute_checksum: bool, +) -> Result, StorageError> { + log::debug!( + "Streaming tar download from {url} to {}", + target_dir.display() + ); + + let response = client.get(url.clone()).send().await?; + + if !response.status().is_success() { + return Err(StorageError::bad_input(format!( + "Failed to download tar from {url}: status - {}", + response.status() + ))); + } + + // Convert the response body stream into an AsyncRead + let stream = response.bytes_stream().map_err(std::io::Error::other); + let async_reader = StreamReader::new(stream); + + let target_dir = target_dir.to_path_buf(); + let target_dir_for_log = target_dir.clone(); + + // Use spawn_cancel_on_drop to ensure the blocking task is cancelled when the future is dropped + let hash = cancel::blocking::spawn_cancel_on_drop(move |cancel| { + // SyncIoBridge converts an AsyncRead into a sync Read + // It must be used within a tokio runtime context (spawn_blocking provides this) + let sync_reader = tokio_util::io::SyncIoBridge::new(async_reader); + + // Wrap the reader with cancellation support + let cancellable_reader = CancellableReader::new(sync_reader, cancel); + + // Wrap the reader with optional hashing + let hashing_reader = HashingReader::new(cancellable_reader, compute_checksum); + + let mut reader = tar_unpack_reader(hashing_reader, &target_dir).map_err(|e| { + StorageError::service_error(format!("Failed to unpack tar archive: {e}")) + })?; + + // Drain any remaining bytes to ensure the full stream is hashed. + // Tar files have trailing padding that Archive doesn't read. + if reader.hasher.is_some() { + let mut buf = [0u8; 8192]; + while reader.read(&mut buf)? > 0 {} + } + + // Get the hash from the inner reader + let hash = reader.finalize(); + + Ok::, StorageError>(hash) + }) + .await + .map_err(|e| StorageError::service_error(format!("Download task failed: {e}")))??; + + log::debug!( + "Successfully unpacked tar from {url} to {}", + target_dir_for_log.display() + ); + + Ok(hash) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_download_and_unpack_tar() { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/test-shard.snapshot") + .with_body(include_bytes!("./test-shard.snapshot")) + .create(); + let url = Url::parse(&format!("{}/test-shard.snapshot", server.url())).unwrap(); + + let client = reqwest::Client::new(); + let temp_dir = tempfile::tempdir().unwrap(); + + let hash = download_and_unpack_tar(&client, &url, temp_dir.path(), true) + .await + .unwrap(); + + let hash = hash.expect("Hash should be computed"); + + // Verify the expected hash + assert_eq!( + hash, + "5d94eac5c1ede3994a28bc406120046c37370d5d45b489a0d2252531b4e3e1f2", + ); + + // Verify content was extracted + let entries: Vec<_> = fs_err::read_dir(temp_dir.path()) + .unwrap() + .map(|res| res.unwrap().file_name().to_string_lossy().to_string()) + .collect(); + + assert!(entries.contains(&"wal".to_string())); + } +} diff --git a/lib/storage/src/content_manager/snapshots/mod.rs b/lib/storage/src/content_manager/snapshots/mod.rs index b40ba6fc7d..6862eda22c 100644 --- a/lib/storage/src/content_manager/snapshots/mod.rs +++ b/lib/storage/src/content_manager/snapshots/mod.rs @@ -1,4 +1,6 @@ pub mod download; +pub mod download_result; +pub mod download_tar; pub mod recover; use std::collections::HashMap; diff --git a/lib/storage/src/content_manager/snapshots/recover.rs b/lib/storage/src/content_manager/snapshots/recover.rs index 772ee5431c..5c63d36f50 100644 --- a/lib/storage/src/content_manager/snapshots/recover.rs +++ b/lib/storage/src/content_manager/snapshots/recover.rs @@ -1,6 +1,6 @@ use collection::collection::Collection; use collection::collection::payload_index_schema::{PAYLOAD_INDEX_CONFIG_FILE, PayloadIndexSchema}; -use collection::common::sha_256::{hash_file, hashes_equal}; +use collection::common::sha_256::hashes_equal; use collection::config::CollectionConfigInternal; use collection::operations::snapshot_ops::{SnapshotPriority, SnapshotRecover}; use collection::operations::verification::new_unchecked_verification_pass; @@ -17,6 +17,7 @@ use crate::content_manager::collection_meta_ops::{ CollectionMetaOperations, CreateCollectionOperation, CreatePayloadIndex, }; use crate::content_manager::snapshots::download::download_snapshot; +use crate::content_manager::snapshots::download_result::DownloadResult; use crate::dispatcher::Dispatcher; use crate::rbac::{Access, AccessRequirements, CollectionPass}; use crate::{StorageError, TableOfContent}; @@ -115,15 +116,23 @@ async fn _do_recover_from_snapshot( let is_distributed = toc.is_distributed(); - let snapshot_path = download_snapshot( + let DownloadResult { + snapshot: snapshot_data, + hash: snapshot_hash, + } = download_snapshot( client, location, &toc.optional_temp_or_snapshot_temp_path()?, + checksum.is_some(), ) .await?; if let Some(checksum) = checksum { - let snapshot_checksum = hash_file(&snapshot_path).await?; + let Some(snapshot_checksum) = snapshot_hash else { + return Err(StorageError::service_error( + "Snapshot checksum was not computed during download", + )); + }; if !hashes_equal(&snapshot_checksum, &checksum) { return Err(StorageError::bad_input(format!( "Snapshot checksum mismatch: expected {checksum}, got {snapshot_checksum}" @@ -131,29 +140,17 @@ async fn _do_recover_from_snapshot( } } - log::debug!("Snapshot downloaded to {}", snapshot_path.display()); - let temp_storage_path = toc.optional_temp_or_storage_temp_path()?; let tmp_collection_dir = tempfile::Builder::new() .prefix(&format!("col-{collection_pass}-recovery-")) .tempdir_in(temp_storage_path)?; - log::debug!( - "Recovering collection {collection_pass} from snapshot {}", - snapshot_path.display(), - ); - - log::debug!( - "Unpacking snapshot to {}", - tmp_collection_dir.path().display(), - ); - let tmp_collection_dir_clone = tmp_collection_dir.path().to_path_buf(); - let snapshot_path_clone = snapshot_path.to_path_buf(); + let restoring = tokio::task::spawn_blocking(move || { Collection::restore_snapshot( - &snapshot_path_clone, + snapshot_data, &tmp_collection_dir_clone, this_peer_id, is_distributed, @@ -397,10 +394,5 @@ async fn _do_recover_from_snapshot( // Remove tmp collection dir tokio_fs::remove_dir_all(&tmp_collection_dir).await?; - // Remove snapshot after recovery if downloaded - if let Err(err) = snapshot_path.close() { - log::error!("Failed to remove downloaded collection snapshot after recovery: {err}"); - } - Ok(true) } diff --git a/lib/storage/src/content_manager/snapshots/test-shard.snapshot b/lib/storage/src/content_manager/snapshots/test-shard.snapshot new file mode 100644 index 0000000000..64c5d5860e Binary files /dev/null and b/lib/storage/src/content_manager/snapshots/test-shard.snapshot differ diff --git a/src/actix/api/snapshot_api.rs b/src/actix/api/snapshot_api.rs index f98cc20340..9f866889b9 100644 --- a/src/actix/api/snapshot_api.rs +++ b/src/actix/api/snapshot_api.rs @@ -21,6 +21,7 @@ use reqwest::Url; use schemars::JsonSchema; use segment::common::BYTES_IN_MB; use serde::{Deserialize, Serialize}; +use shard::snapshots::snapshot_data::SnapshotData; use shard::snapshots::snapshot_manifest::{RecoveryType, SnapshotManifest}; use storage::content_manager::errors::{StorageError, StorageResult}; use storage::content_manager::snapshots::recover::do_recover_from_snapshot; @@ -499,12 +500,15 @@ async fn upload_shard_snapshot( let collection = cancel::future::cancel_on_token(cancel.clone(), cancel_safe).await??; + let snapshot_data = + SnapshotData::Packed(MaybeTempPath::from(form.snapshot.file.into_temp_path())); + // `recover_shard_snapshot_impl` is *not* cancel safe common::snapshots::recover_shard_snapshot_impl( dispatcher.toc(&access, &pass), &collection, shard, - MaybeTempPath::from(form.snapshot.file.into_temp_path()), + snapshot_data, priority.unwrap_or_default(), RecoveryType::Full, cancel, @@ -657,12 +661,15 @@ async fn recover_partial_snapshot( let collection = cancel::future::cancel_on_token(cancel.clone(), cancel_safe).await??; + let snapshot_data = + SnapshotData::Packed(MaybeTempPath::from(form.snapshot.file.into_temp_path())); + // `recover_shard_snapshot_impl` is *not* cancel safe common::snapshots::recover_shard_snapshot_impl( dispatcher.toc(&access, &pass), &collection, shard, - MaybeTempPath::from(form.snapshot.file.into_temp_path()), + snapshot_data, priority.unwrap_or_default(), RecoveryType::Partial, cancel, @@ -818,11 +825,14 @@ async fn recover_partial_snapshot_from( shard_id ); + let snapshot_data = + SnapshotData::Packed(MaybeTempPath::from(partial_snapshot_temp_path)); + common::snapshots::recover_shard_snapshot_impl( dispatcher.toc(&access, &pass), &collection, shard_id, - MaybeTempPath::from(partial_snapshot_temp_path), + snapshot_data, SnapshotPriority::NoSync, RecoveryType::Partial, cancel, diff --git a/src/common/snapshots.rs b/src/common/snapshots.rs index e4c32fccbd..d3f92a2dc9 100644 --- a/src/common/snapshots.rs +++ b/src/common/snapshots.rs @@ -9,10 +9,11 @@ use collection::operations::snapshot_ops::{ use collection::operations::verification::VerificationPass; use collection::shards::replica_set::replica_set_state::ReplicaState; use collection::shards::shard::ShardId; -use common::tempfile_ext::MaybeTempPath; +use shard::snapshots::snapshot_data::SnapshotData; use shard::snapshots::snapshot_manifest::{RecoveryType, SnapshotManifest}; use storage::content_manager::errors::StorageError; use storage::content_manager::snapshots; +use storage::content_manager::snapshots::download_result::DownloadResult; use storage::content_manager::toc::TableOfContent; use storage::dispatcher::Dispatcher; use storage::rbac::{Access, AccessRequirements}; @@ -166,7 +167,10 @@ pub async fn recover_shard_snapshot( let download_dir = toc.optional_temp_or_snapshot_temp_path()?; - let snapshot_path = match snapshot_location { + let DownloadResult { + snapshot, + hash + } = match snapshot_location { ShardSnapshotLocation::Url(url) => { if !matches!(url.scheme(), "http" | "https") { let description = format!( @@ -179,7 +183,13 @@ pub async fn recover_shard_snapshot( let client = client.client(api_key.as_deref())?; - snapshots::download::download_snapshot(&client, url, &download_dir).await? + snapshots::download::download_snapshot( + &client, + url, + &download_dir, + checksum.is_some(), + ) + .await? } ShardSnapshotLocation::Path(snapshot_file_name) => { @@ -194,26 +204,42 @@ pub async fn recover_shard_snapshot( ) .await?; - collection + let snapshot_file = collection .get_snapshots_storage_manager()? .get_snapshot_file(&snapshot_path, &download_dir) - .await? + .await?; + + let hash = if checksum.is_some() { + Some(sha_256::hash_file(&snapshot_path).await?) + } else { + None + }; + + DownloadResult { + snapshot: SnapshotData::Packed(snapshot_file), + hash, + } } }; if let Some(checksum) = checksum { - let snapshot_checksum = sha_256::hash_file(&snapshot_path).await?; - if !sha_256::hashes_equal(&snapshot_checksum, &checksum) { - return Err(StorageError::bad_input(format!( - "Snapshot checksum mismatch: expected {checksum}, got {snapshot_checksum}" - ))); + if let Some(snapshot_checksum) = hash { + if !sha_256::hashes_equal(&snapshot_checksum, &checksum) { + return Err(StorageError::bad_input(format!( + "Snapshot checksum mismatch: expected {checksum}, got {snapshot_checksum}" + ))); + } + } else { + return Err(StorageError::service_error( + "Snapshot checksum could not be verified".to_string(), + )); } } - Ok((collection, snapshot_path)) + Ok((collection, snapshot)) }; - let (collection, snapshot_path) = + let (collection, snapshot_data) = cancel::future::cancel_on_token(cancel.clone(), cancel_safe).await??; // `recover_shard_snapshot_impl` is *not* cancel safe @@ -221,7 +247,7 @@ pub async fn recover_shard_snapshot( &toc, &collection, shard_id, - snapshot_path, + snapshot_data, snapshot_priority, RecoveryType::Full, cancel, @@ -240,7 +266,7 @@ pub async fn recover_shard_snapshot_impl( toc: &TableOfContent, collection: &Collection, shard: ShardId, - snapshot_path: MaybeTempPath, + snapshot_data: SnapshotData, priority: SnapshotPriority, recovery_type: RecoveryType, cancel: cancel::CancellationToken, @@ -259,7 +285,7 @@ pub async fn recover_shard_snapshot_impl( collection .restore_shard_snapshot( shard, - snapshot_path, + snapshot_data, recovery_type, toc.this_peer_id, toc.is_distributed(), diff --git a/src/snapshots.rs b/src/snapshots.rs index b1b9fd8c3c..3d5e6ea770 100644 --- a/src/snapshots.rs +++ b/src/snapshots.rs @@ -3,11 +3,12 @@ use std::path::{Path, PathBuf}; use collection::collection::Collection; use collection::shards::shard::PeerId; +use common::tar_unpack::tar_unpack_file; use fs_err as fs; use fs_err::File; use io::safe_delete::safe_delete_in_tmp; use log::info; -use segment::common::validate_snapshot_archive::open_snapshot_archive_with_validation; +use shard::snapshots::snapshot_data::SnapshotData; use storage::content_manager::alias_mapping::AliasPersistence; use storage::content_manager::snapshots::SnapshotConfig; use storage::content_manager::toc::{ALIASES_PATH, COLLECTIONS_DIR}; @@ -39,7 +40,8 @@ pub fn recover_snapshots( .next() .unwrap_or_else(|| panic!("Snapshot path is missing: {snapshot_params}")); - let snapshot_path = Path::new(path); + let snapshot_data = SnapshotData::new_packed_persistent(path); + let collection_name = split .next() .unwrap_or_else(|| panic!("Collection name is missing: {snapshot_params}")); @@ -65,7 +67,7 @@ pub fn recover_snapshots( let collection_temp_path = temp_dir.map_or_else(|| collection_path.with_extension("tmp"), PathBuf::from); if let Err(err) = Collection::restore_snapshot( - snapshot_path, + snapshot_data, &collection_temp_path, this_peer_id, is_distributed, @@ -98,8 +100,7 @@ pub fn recover_full_snapshot( fs::create_dir_all(&snapshot_temp_path).unwrap(); // Un-tar snapshot into temporary directory - let mut ar = open_snapshot_archive_with_validation(Path::new(snapshot_path)).unwrap(); - ar.unpack(&snapshot_temp_path).unwrap(); + tar_unpack_file(Path::new(snapshot_path), &snapshot_temp_path).unwrap(); // Read configuration file with snapshot-to-collection mapping let config_path = snapshot_temp_path.join("config.json");