diff --git a/dstack/verifier/src/verification.rs b/dstack/verifier/src/verification.rs index f70a2946f..081730e0c 100644 --- a/dstack/verifier/src/verification.rs +++ b/dstack/verifier/src/verification.rs @@ -3,6 +3,7 @@ // SPDX-License-Identifier: Apache-2.0 use std::{ + collections::HashMap, ffi::OsStr, path::{Component, Path, PathBuf}, sync::Arc, @@ -189,6 +190,7 @@ fn collect_rtmr_mismatch( // `os_image_hash` -- part of the `VmConfig` this key hashes -- has never been // cached. const MEASUREMENT_CACHE_VERSION: u32 = 3; +const IMAGE_DOWNLOAD_ATTEMPTS: u32 = 3; #[derive(Clone, Serialize, Deserialize)] struct CachedMeasurement { @@ -212,6 +214,7 @@ pub struct CvmVerifier { pub download_url: String, pub download_timeout: Duration, pub attestation_verifier: Arc, + download_locks: tokio::sync::Mutex>>>, } impl CvmVerifier { @@ -226,9 +229,18 @@ impl CvmVerifier { download_url, download_timeout, attestation_verifier, + download_locks: tokio::sync::Mutex::new(HashMap::new()), } } + async fn lock_for_image(&self, hex_os_image_hash: &str) -> Arc> { + let mut locks = self.download_locks.lock().await; + locks + .entry(hex_os_image_hash.to_string()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone() + } + fn measurement_cache_dir(&self) -> PathBuf { Path::new(&self.image_cache_dir).join("measurements") } @@ -630,14 +642,18 @@ impl CvmVerifier { let metadata_path = image_dir.join("metadata.json"); if !metadata_path.exists() { - info!("Image {hex_os_image_hash} not found, downloading"); - tokio::time::timeout( - self.download_timeout, - self.download_image(&hex_os_image_hash, &image_dir), - ) - .await - .context("Download image timeout")? - .with_context(|| format!("Failed to download image {hex_os_image_hash}"))?; + let lock = self.lock_for_image(&hex_os_image_hash).await; + let _guard = lock.lock().await; + if !metadata_path.exists() { + info!("Image {hex_os_image_hash} not found, downloading"); + tokio::time::timeout( + self.download_timeout, + self.download_image(&hex_os_image_hash, &image_dir), + ) + .await + .context("Download image timeout")? + .with_context(|| format!("Failed to download image {hex_os_image_hash}"))?; + } } let image_info = @@ -1277,24 +1293,26 @@ impl CvmVerifier { Ok(()) } - pub async fn download_image(&self, hex_os_image_hash: &str, dst_dir: &Path) -> Result<()> { - let url = self - .download_url - .replace("{OS_IMAGE_HASH}", hex_os_image_hash); - - // Create a temporary directory for extraction within the cache directory - let cache_dir = Path::new(&self.image_cache_dir).join("images").join("tmp"); - fs_err::create_dir_all(&cache_dir).context("Failed to create cache directory")?; - let auto_delete_temp_dir = tempfile::Builder::new() - .prefix("tmp-download-") - .tempdir_in(&cache_dir) - .context("Failed to create temporary directory")?; - let tmp_dir = auto_delete_temp_dir.path(); + fn is_truncated_download_error(err: &anyhow::Error) -> bool { + for cause in err.chain() { + if let Some(reqwest_err) = cause.downcast_ref::() { + if reqwest_err.is_body() || reqwest_err.is_decode() { + return true; + } + } + } + let text = format!("{err:#}"); + text.contains("truncated image download") + || text.contains("end of file before message length reached") + } - info!("Downloading image from {}", url); - let client = reqwest::Client::new(); + async fn download_image_tarball( + client: &reqwest::Client, + url: &str, + tarball_path: &Path, + ) -> Result<()> { let response = client - .get(&url) + .get(url) .send() .await .context("Failed to download image")?; @@ -1306,25 +1324,73 @@ impl CvmVerifier { ); } - // Save the tarball to a temporary file using streaming - let tarball_path = tmp_dir.join("image.tar.gz"); - let mut file = tokio::fs::File::create(&tarball_path) + let expected_len = response.content_length(); + let mut file = tokio::fs::File::create(tarball_path) .await .context("Failed to create tarball file")?; + let mut written = 0u64; let mut response = response; - while let Some(chunk) = response.chunk().await? { + while let Some(chunk) = response + .chunk() + .await + .context("failed to read image download body")? + { + written += chunk.len() as u64; file.write_all(&chunk) .await .context("Failed to write chunk to file")?; } - - let extracted_dir = tmp_dir.join("extracted"); - fs_err::create_dir_all(&extracted_dir).context("Failed to create extraction directory")?; - file.flush() .await .context("Failed to flush image archive")?; drop(file); + + if let Some(expected) = expected_len { + if written != expected { + bail!( + "truncated image download: got {written} bytes, content-length {expected}, url: {url}" + ); + } + } + Ok(()) + } + + pub async fn download_image(&self, hex_os_image_hash: &str, dst_dir: &Path) -> Result<()> { + let url = self + .download_url + .replace("{OS_IMAGE_HASH}", hex_os_image_hash); + + // Create a temporary directory for extraction within the cache directory + let cache_dir = Path::new(&self.image_cache_dir).join("images").join("tmp"); + fs_err::create_dir_all(&cache_dir).context("Failed to create cache directory")?; + let auto_delete_temp_dir = tempfile::Builder::new() + .prefix("tmp-download-") + .tempdir_in(&cache_dir) + .context("Failed to create temporary directory")?; + let tmp_dir = auto_delete_temp_dir.path(); + + info!("Downloading image from {}", url); + let client = reqwest::Client::new(); + let tarball_path = tmp_dir.join("image.tar.gz"); + let mut attempt = 1; + loop { + match Self::download_image_tarball(&client, &url, &tarball_path).await { + Ok(()) => break, + Err(err) + if Self::is_truncated_download_error(&err) + && attempt < IMAGE_DOWNLOAD_ATTEMPTS => + { + warn!( + "truncated image download from {url}, retrying ({attempt}/{IMAGE_DOWNLOAD_ATTEMPTS}): {err:#}" + ); + attempt += 1; + } + Err(err) => return Err(err), + } + } + + let extracted_dir = tmp_dir.join("extracted"); + fs_err::create_dir_all(&extracted_dir).context("Failed to create extraction directory")?; Self::extract_image_archive(&tarball_path, &extracted_dir)?; let sha256sum_path = extracted_dir.join("sha256sum.txt"); @@ -1421,7 +1487,11 @@ impl Mrs { #[cfg(test)] mod tests { - use std::collections::BTreeMap; + use std::{ + collections::BTreeMap, + io::Write, + sync::atomic::{AtomicUsize, Ordering}, + }; use super::*; @@ -2177,6 +2247,143 @@ mod tests { ); } + fn sample_os_image_archive() -> (Vec, VmConfig) { + let metadata = br#"{"cmdline":"console=ttyS0","kernel":"bzImage","initrd":"initrd.img","bios":"OVMF.fd"}"#; + let files: [(&str, &[u8]); 4] = [ + ("OVMF.fd", b"bios"), + ("bzImage", b"kernel"), + ("initrd.img", b"initrd"), + ("metadata.json", metadata), + ]; + let mut files_doc = String::new(); + for (name, payload) in files { + files_doc.push_str(&format!( + "{} {name}\n", + hex::encode(Sha256::digest(payload)) + )); + } + let os_image_hash = Sha256::digest(files_doc.as_bytes()); + let mut tarball = tar::Builder::new(Vec::new()); + for (name, payload) in files { + let mut header = tar::Header::new_gnu(); + header.set_size(payload.len() as u64); + header.set_mode(0o644); + header.set_cksum(); + tarball.append_data(&mut header, name, payload).unwrap(); + } + let payload = files_doc.as_bytes(); + let mut header = tar::Header::new_gnu(); + header.set_size(payload.len() as u64); + header.set_mode(0o644); + header.set_cksum(); + tarball + .append_data(&mut header, "sha256sum.txt", payload) + .unwrap(); + let tar_bytes = tarball.into_inner().unwrap(); + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder.write_all(&tar_bytes).unwrap(); + let archive = encoder.finish().unwrap(); + let vm_config = serde_json::from_value(serde_json::json!({ + "os_image_hash": hex::encode(os_image_hash), + })) + .unwrap(); + (archive, vm_config) + } + + async fn spawn_scripted_http_server( + responses: Vec<(u64, Vec)>, + ) -> (String, Arc, tokio::task::JoinHandle<()>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let hits = Arc::new(AtomicUsize::new(0)); + let hits_clone = hits.clone(); + let handle = tokio::spawn(async move { + loop { + let Ok((mut socket, _)) = listener.accept().await else { + break; + }; + let index = hits_clone.fetch_add(1, Ordering::SeqCst); + let (content_length, body) = responses + .get(index) + .or_else(|| responses.last()) + .cloned() + .expect("scripted HTTP response"); + let mut buf = vec![0u8; 1024]; + let mut collected = Vec::new(); + loop { + match tokio::io::AsyncReadExt::read(&mut socket, &mut buf).await { + Ok(0) => break, + Ok(n) => { + collected.extend_from_slice(&buf[..n]); + if collected.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + Err(_) => break, + } + } + let header = format!( + "HTTP/1.1 200 OK\r\nContent-Length: {content_length}\r\nConnection: close\r\n\r\n" + ); + let _ = tokio::io::AsyncWriteExt::write_all(&mut socket, header.as_bytes()).await; + let _ = tokio::io::AsyncWriteExt::write_all(&mut socket, &body).await; + let _ = tokio::io::AsyncWriteExt::shutdown(&mut socket).await; + } + }); + ( + format!("http://{addr}/{{OS_IMAGE_HASH}}.tar.gz"), + hits, + handle, + ) + } + + #[tokio::test] + async fn concurrent_same_hash_downloads_single_flight() { + let (archive, vm_config) = sample_os_image_archive(); + let (url, hits, server) = + spawn_scripted_http_server(vec![(archive.len() as u64, archive)]).await; + let cache = tempfile::tempdir().unwrap(); + let mut verifier = test_verifier(); + verifier.image_cache_dir = cache.path().display().to_string(); + verifier.download_url = url; + verifier.download_timeout = Duration::from_secs(30); + let verifier = Arc::new(verifier); + + let mut handles = Vec::new(); + for _ in 0..8 { + let verifier = verifier.clone(); + let vm_config = vm_config.clone(); + handles.push(tokio::spawn(async move { + verifier.ensure_image_downloaded(&vm_config).await + })); + } + for handle in handles { + handle.await.unwrap().unwrap(); + } + assert_eq!(hits.load(Ordering::SeqCst), 1); + server.abort(); + } + + #[tokio::test] + async fn truncated_image_download_retries_after_content_length_mismatch() { + let (archive, vm_config) = sample_os_image_archive(); + let truncated = archive[..archive.len() / 2].to_vec(); + let (url, hits, server) = spawn_scripted_http_server(vec![ + (archive.len() as u64, truncated), + (archive.len() as u64, archive), + ]) + .await; + let cache = tempfile::tempdir().unwrap(); + let mut verifier = test_verifier(); + verifier.image_cache_dir = cache.path().display().to_string(); + verifier.download_url = url; + verifier.download_timeout = Duration::from_secs(30); + + verifier.ensure_image_downloaded(&vm_config).await.unwrap(); + assert!(hits.load(Ordering::SeqCst) >= 2); + server.abort(); + } + #[tokio::test] async fn verifies_sev_snp_attestation_fixture_without_image_download() { let request: VerificationRequest =