From 9b9e27194a78b261eecc27cb7074f9110920e67e Mon Sep 17 00:00:00 2001 From: JoyboyBrian Date: Wed, 29 Jul 2026 12:27:42 -0700 Subject: [PATCH 1/3] fix(api): return the JSON error envelope for unmatched routes --- src/api/proxy.rs | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/src/api/proxy.rs b/src/api/proxy.rs index 207872b5..76eaefe6 100644 --- a/src/api/proxy.rs +++ b/src/api/proxy.rs @@ -224,7 +224,18 @@ where I: AsRef + Send + Sync, { if !has_routing_header(request.headers()) { - return StatusCode::NOT_FOUND.into_response(); + // Unmatched control-plane route: return the API error envelope so + // JSON clients surface "route not found" instead of failing to parse + // an empty 404 body. + let error = agentenv_http_server::models::Error::new( + 404, + format!( + "route not found: {} {}", + request.method(), + request.uri().path() + ), + ); + return (StatusCode::NOT_FOUND, axum::Json(error)).into_response(); } let forward_path = request.uri().path().to_owned(); with_route_source( @@ -2120,6 +2131,24 @@ mod tests { .unwrap(); assert_eq!(response.status(), StatusCode::NOT_FOUND); + let content_type = response + .headers() + .get(axum::http::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_owned(); + assert!( + content_type.starts_with("application/json"), + "unexpected content-type: {content_type}" + ); + let body = response.into_body().collect().await.unwrap().to_bytes(); + let payload: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(payload["code"], 404); + let message = payload["message"].as_str().unwrap(); + assert!( + message.contains("route not found: GET /nonexistent/path"), + "unexpected message: {message}" + ); } #[tokio::test] From f9e9c3eb63a8865753bfbbfa0b35f7626bbb8623 Mon Sep 17 00:00:00 2001 From: JoyboyBrian Date: Wed, 29 Jul 2026 12:29:03 -0700 Subject: [PATCH 2/3] feat(snapshot): E2B alias rebuild semantics on template publish --- .../repository/backends/oss/repository.rs | 201 +++++++++++++++--- .../repository/backends/posixfs/backend.rs | 145 ++++++++++++- .../repository/backends/posixfs/catalog.rs | 179 +++++++++++----- 3 files changed, 426 insertions(+), 99 deletions(-) diff --git a/src/snapshot/repository/backends/oss/repository.rs b/src/snapshot/repository/backends/oss/repository.rs index 06f747ef..8af954a8 100644 --- a/src/snapshot/repository/backends/oss/repository.rs +++ b/src/snapshot/repository/backends/oss/repository.rs @@ -164,25 +164,28 @@ impl SnapshotRepository for OssSnapshotRepository { reason: format!("snapshot '{}' already exists", record.id), }); } + // When the alias already points at a live snapshot, leave the binding + // untouched so the existing template keeps resolving while the new + // build runs; a successful publish moves the alias to the new snapshot + // (E2B rebuild semantics). + let mut bind_on_create = true; if let Some(alias) = record.alias.as_ref() { if let Some(existing) = self.load_alias_target(alias.as_ref()).await? { if existing != record.id && self.snapshot_exists(&existing).await? { - return Err(RepositoryError::AliasConflict { - alias: alias.to_string(), - existing, - new_id: record.id.clone(), - }); + bind_on_create = false; } } } self.write_record(&record).await?; - if let Some(alias) = record.alias.as_ref() { - if let Err(error) = self.bind_alias(alias.as_ref(), &record.id).await { - let _ = self - .client - .delete(&OssSnapshotArtifactLayout::record_key(&record.id)) - .await; - return Err(error); + if bind_on_create { + if let Some(alias) = record.alias.as_ref() { + if let Err(error) = self.bind_alias(alias.as_ref(), &record.id, false).await { + let _ = self + .client + .delete(&OssSnapshotArtifactLayout::record_key(&record.id)) + .await; + return Err(error); + } } } Ok(record) @@ -287,26 +290,63 @@ impl SnapshotRepository for OssSnapshotRepository { disk_publications: disk_publications.clone(), }; - // 5. Bind alias (if present) with conflict detection. + // 5. Commit the record before moving the alias. This prevents an + // alias from ever resolving to a snapshot whose catalog record + // has not been published yet. The tradeoff is a crash window: + // dying after the record write but before `bind_alias` leaves a + // committed record whose `alias` field names an alias that still + // resolves to the previous snapshot, so readers of `record.alias` + // (listings) may observe the stale claim until the next rebind. + // Nothing reconciles that state automatically. + let previous_record = self.read_record(id).await?; + let previous_alias_target = if let Some(alias) = metadata.alias.as_ref() { + match self.load_alias_target(alias.as_ref()).await? { + Some(existing) if self.snapshot_exists(&existing).await? => Some(existing), + _ => None, + } + } else { + None + }; + let record = self + .write_committed_record( + metadata.id.clone(), + metadata.alias.clone(), + metadata.resources, + committed, + metadata.source.clone(), + ) + .await?; + + // 6. Move the alias only after the new record is readable. If the + // bind fails, restore the pending record and old alias. if let Some(ref alias) = metadata.alias { - if let Err(e) = self.bind_alias(alias.as_ref(), id).await { - // Best-effort rollback. Content-addressed managed layers are intentionally left - // in place; they are shared across snapshots and require separate GC. - if let Err(error) = self.client.delete_prefix(&layout.artifact_prefix()).await { - warn!(snapshot_id = %id, error = %error, "failed to roll back snapshot artifacts after alias bind failure"); + if let Err(error) = self.bind_alias(alias.as_ref(), id, true).await { + self.restore_alias_after_failed_bind( + alias.as_ref(), + id, + previous_alias_target.as_ref(), + ) + .await; + self.restore_record_after_failed_publish(id, previous_record.as_ref()) + .await; + return Err(error); + } + if let Some(previous_id) = previous_alias_target + .as_ref() + .filter(|previous_id| *previous_id != id) + { + if let Err(error) = self.clear_record_alias(previous_id, alias.as_ref()).await { + warn!( + alias = %alias, + previous_snapshot_id = %previous_id, + error = %error, + "failed to clear previous snapshot alias metadata" + ); } - return Err(e); } } - self.write_committed_record( - metadata.id.clone(), - metadata.alias.clone(), - metadata.resources, - committed, - metadata.source.clone(), - ) - .await + Ok(record) } .await; @@ -584,7 +624,8 @@ impl OssSnapshotRepository { /// Instead the algorithm is: /// 1. Read the current alias target. /// 2. If it already points to `id`, return success (idempotent). - /// 3. If it points to a live snapshot, return `AliasConflict`. + /// 3. If it points to a live snapshot: with `rebind` move the alias to + /// `id` (E2B rebuild semantics), otherwise return `AliasConflict`. /// 4. If it points to a deleted snapshot, remove the stale alias. /// 5. Write our binding unconditionally. /// 6. Read back and verify we won the race. If someone else wrote a @@ -595,7 +636,7 @@ impl OssSnapshotRepository { /// interval between our write and the subsequent read. This is weaker /// than a true CAS but sufficient for the current deployment model /// where concurrent publishes for the *same alias* are rare. - async fn bind_alias(&self, alias: &str, id: &SnapshotId) -> RepositoryResult<()> { + async fn bind_alias(&self, alias: &str, id: &SnapshotId, rebind: bool) -> RepositoryResult<()> { let key = validated_alias_key(alias)?; let payload = serde_json::to_vec(id) .map_err(|e| RepositoryError::backend("serialize alias binding", e))?; @@ -608,18 +649,19 @@ impl OssSnapshotRepository { } let still_exists = self.snapshot_exists(&existing_id).await?; - if still_exists { + if still_exists && !rebind { return Err(RepositoryError::AliasConflict { alias: alias.to_string(), existing: existing_id, new_id: id.clone(), }); } - - self.client - .delete(&key) - .await - .map_err(|e| RepositoryError::backend("delete stale alias", e))?; + if !still_exists { + self.client + .delete(&key) + .await + .map_err(|e| RepositoryError::backend("delete stale alias", e))?; + } } // Step 5: write our binding (unconditional — OSS does not @@ -669,6 +711,95 @@ impl OssSnapshotRepository { }) } + async fn restore_record_after_failed_publish( + &self, + id: &SnapshotId, + previous_record: Option<&SnapshotRecord>, + ) { + let result = match previous_record { + Some(record) => self.write_record(record).await, + None => self + .client + .delete(&OssSnapshotArtifactLayout::record_key(id)) + .await + .map_err(|error| RepositoryError::backend("remove failed snapshot record", error)), + }; + if let Err(error) = result { + warn!(snapshot_id = %id, error = %error, "failed to restore snapshot record after publish failure"); + } + } + + async fn restore_alias_after_failed_bind( + &self, + alias: &str, + id: &SnapshotId, + previous_id: Option<&SnapshotId>, + ) { + let current = match self.load_alias_target(alias).await { + Ok(current) => current, + Err(error) => { + warn!(alias, snapshot_id = %id, error = %error, "failed to inspect alias during publish rollback"); + return; + } + }; + // Skip when a concurrent publisher already moved the alias elsewhere. + // This only narrows the lost-update window: like `bind_alias`, the + // rollback cannot be atomic on a store without conditional writes, so a + // publisher that rebinds between this read and the write below is still + // clobbered. + if current.as_ref() != Some(id) { + return; + } + + let key = match validated_alias_key(alias) { + Ok(key) => key, + Err(error) => { + warn!(alias, snapshot_id = %id, error = %error, "failed to validate alias during publish rollback"); + return; + } + }; + let result = + match previous_id { + Some(previous_id) => match serde_json::to_vec(previous_id) { + Ok(payload) => { + self.client.put_bytes(&key, payload).await.map_err(|error| { + RepositoryError::backend("restore alias binding", error) + }) + } + Err(error) => Err(RepositoryError::backend( + "serialize restored alias binding", + error, + )), + }, + None => self.client.delete(&key).await.map_err(|error| { + RepositoryError::backend("remove failed alias binding", error) + }), + }; + if let Err(error) = result { + warn!(alias, snapshot_id = %id, error = %error, "failed to restore alias after publish failure"); + } + } + + /// Clears the alias field on the record that previously owned a rebound + /// alias so template listings do not report the moved name twice. + /// + /// Only `moved_alias` is cleared; a previous owner that already claims a + /// different name keeps it. + async fn clear_record_alias(&self, id: &SnapshotId, moved_alias: &str) -> RepositoryResult<()> { + if let Some(mut previous) = self.read_record(id).await? { + let claims_moved_alias = previous + .alias + .as_ref() + .is_some_and(|alias| alias.as_ref() == moved_alias); + if claims_moved_alias { + previous.alias = None; + previous.updated_at_unix_ms = now_unix_ms(); + self.write_record(&previous).await?; + } + } + Ok(()) + } + async fn export_managed_disk_image( &self, image_config_path: &Path, diff --git a/src/snapshot/repository/backends/posixfs/backend.rs b/src/snapshot/repository/backends/posixfs/backend.rs index 4dc7f7ba..459c1547 100644 --- a/src/snapshot/repository/backends/posixfs/backend.rs +++ b/src/snapshot/repository/backends/posixfs/backend.rs @@ -474,37 +474,162 @@ mod tests { } #[tokio::test] - async fn failed_commit_cleans_uncommitted_snapshot_directory() { + async fn publish_rebinds_existing_alias_to_new_snapshot() { let tempdir = TempDir::new().expect("tempdir should exist"); - let repository_root = tempdir.path().to_path_buf(); let repository = test_backend(tempdir.path()).repository(); let first_id = SnapshotId::generate(); let local_artifacts = seed_built_snapshot(tempdir.path()); - let first_metadata = sample_metadata(first_id.clone(), Some("conflict")); repository - .publish(first_metadata, local_artifacts) + .publish( + sample_metadata(first_id.clone(), Some("rebind")), + local_artifacts, + ) .await .expect("first publish should work"); let second_id = SnapshotId::generate(); let local_artifacts = seed_built_snapshot(tempdir.path()); - let err = repository + repository + .publish( + sample_metadata(second_id.clone(), Some("rebind")), + local_artifacts, + ) + .await + .expect("second publish should rebind the alias"); + + let resolved = repository + .resolve_alias("rebind") + .await + .expect("resolve should work") + .expect("alias should resolve"); + assert_eq!(resolved, second_id, "alias should move to the new snapshot"); + + let previous = repository + .get(first_id.to_string().as_str()) + .await + .expect("get should work") + .expect("previous snapshot should stay addressable by id"); + assert_eq!( + previous.alias, None, + "previous snapshot should lose the rebound alias" + ); + } + + #[tokio::test] + async fn failed_publish_keeps_previous_alias_and_removes_snapshot_dir() { + let tempdir = TempDir::new().expect("tempdir should exist"); + let repository = test_backend(tempdir.path()).repository(); + + let first_id = SnapshotId::generate(); + let local_artifacts = seed_built_snapshot(tempdir.path()); + repository .publish( - sample_metadata(second_id.clone(), Some("conflict")), + sample_metadata(first_id.clone(), Some("rebind")), local_artifacts, ) .await - .expect_err("second publish should fail"); + .expect("first publish should work"); + + let second_id = SnapshotId::generate(); + let broken_artifacts = seed_built_snapshot(tempdir.path()); + // `import_built_artifacts` copies `vm_state.bin` first, so removing it + // fails the publish before any catalog state is committed. + fs::remove_file(&broken_artifacts.vm_state.path).expect("remove seeded vm state"); + repository + .publish( + sample_metadata(second_id.clone(), Some("rebind")), + broken_artifacts, + ) + .await + .expect_err("publish should fail when the vm state artifact is missing"); - assert!(matches!(err, RepositoryError::AliasConflict { .. })); assert!( - !repository_root + !tempdir + .path() .join("snapshots") .join(second_id.to_string()) .exists(), - "failed publish should not leave a committed revision directory" + "failed publish should not leave a snapshot directory behind" + ); + + let resolved = repository + .resolve_alias("rebind") + .await + .expect("resolve should work") + .expect("alias should still resolve"); + assert_eq!( + resolved, first_id, + "alias should stay bound to the previously committed snapshot" ); + + let previous = repository + .get(first_id.to_string().as_str()) + .await + .expect("get should work") + .expect("previous snapshot should stay addressable by id"); + assert_eq!( + previous.alias.as_ref().map(ToString::to_string), + Some("rebind".to_string()), + "previous snapshot should keep the alias after a failed rebind" + ); + } + + #[tokio::test] + async fn create_keeps_existing_alias_until_new_build_commits() { + let tempdir = TempDir::new().expect("tempdir should exist"); + let repository = test_backend(tempdir.path()).repository(); + + let committed_id = SnapshotId::generate(); + let local_artifacts = seed_built_snapshot(tempdir.path()); + repository + .publish( + sample_metadata(committed_id.clone(), Some("stable")), + local_artifacts, + ) + .await + .expect("publish should work"); + + let waiting = SnapshotRecord::template_waiting( + SnapshotId::generate(), + Some(SnapshotAlias::parse("stable").expect("alias should parse")), + crate::types::SandboxResources { + cpu_count: 1, + memory_mib: 256, + disk_size_mib: 0, + }, + ); + let waiting_id = waiting.id.clone(); + repository + .create(waiting) + .await + .expect("create with an existing alias should be allowed"); + + let resolved = repository + .resolve_alias("stable") + .await + .expect("resolve should work") + .expect("alias should resolve"); + assert_eq!( + resolved, committed_id, + "alias should keep pointing at the committed snapshot while the rebuild is pending" + ); + + let local_artifacts = seed_built_snapshot(tempdir.path()); + repository + .publish( + sample_metadata(waiting_id.clone(), Some("stable")), + local_artifacts, + ) + .await + .expect("publishing the rebuild should rebind the alias"); + + let resolved = repository + .resolve_alias("stable") + .await + .expect("resolve should work") + .expect("alias should resolve"); + assert_eq!(resolved, waiting_id, "alias should move after commit"); } #[tokio::test] diff --git a/src/snapshot/repository/backends/posixfs/catalog.rs b/src/snapshot/repository/backends/posixfs/catalog.rs index 6c36d6bd..76b7486c 100644 --- a/src/snapshot/repository/backends/posixfs/catalog.rs +++ b/src/snapshot/repository/backends/posixfs/catalog.rs @@ -6,6 +6,7 @@ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use serde::de::DeserializeOwned; use serde::Serialize; +use tracing::warn; use super::layout::PosixFsSnapshotArtifactLayout; use crate::snapshot::repository::SnapshotListFilter; @@ -79,9 +80,9 @@ impl PosixFsCatalogStore { /// /// Flow: /// 1. acquire the alias lock when an alias is present - /// 2. bind the alias - /// 3. write the commit marker - /// 4. write the committed snapshot record + /// 2. write the commit marker + /// 3. write the committed snapshot record + /// 4. atomically bind the alias as the final visible operation pub(crate) fn commit_publish( &self, session: &PublishSession, @@ -90,25 +91,30 @@ impl PosixFsCatalogStore { ) -> RepositoryResult { let now = now_unix_ms(); let snapshot_id = metadata.id.clone(); + let previous_record = self.load_record_by_id_unlocked(&snapshot_id)?; let write_result = if let Some(alias) = metadata.alias.as_ref() { self.with_alias_lock(alias, |store| { let record = store.committed_record_unlocked(&metadata, committed.clone(), now)?; let alias_path = PosixFsSnapshotArtifactLayout::alias_path(&store.root, alias); - if let Some(existing) = store.load_alias_target(alias)? { - if existing != snapshot_id { - if store.load_record_by_id_unlocked(&existing)?.is_some() { - return Err(RepositoryError::AliasConflict { - alias: alias.to_string(), - existing, - new_id: snapshot_id.clone(), - }); - } - store.remove_file_if_exists(&alias_path)?; - } - } - store.write_json(&alias_path, &snapshot_id)?; + let existing = store.load_alias_target(alias)?; store.write_commit_marker(&session.snapshot_id)?; store.write_committed_record_unlocked(&record)?; + // `write_json` uses an atomic rename. Keeping this as the final + // fallible operation means a failed rebuild leaves the old + // alias binding untouched. The tradeoff is a crash window: dying + // after the record write but before the alias write leaves a + // committed record whose `alias` field names an alias that still + // resolves to the previous snapshot, so readers of `record.alias` + // (listings) may observe the stale claim until the next rebind. + store.write_json(&alias_path, &snapshot_id)?; + + if let Some(existing) = existing.filter(|existing| existing != &snapshot_id) { + // The previous snapshot stays addressable by id, so running + // sandboxes and explicit id references keep working. Alias + // metadata cleanup is best effort because the binding has + // already moved successfully. + store.clear_moved_alias_on_previous_record(&existing, alias.as_ref(), now); + } Ok(record) }) } else { @@ -123,17 +129,7 @@ impl PosixFsCatalogStore { match write_result { Ok(record) => Ok(record), Err(error) => { - if let Some(alias) = metadata.alias.as_ref() { - let _ = self.with_alias_lock(alias, |store| { - let alias_path = - PosixFsSnapshotArtifactLayout::alias_path(&store.root, alias); - if store.load_alias_target(alias)?.as_ref() == Some(&snapshot_id) { - store.remove_file_if_exists(&alias_path)?; - } - Ok(()) - }); - } - let _ = self.cleanup_uncommitted_snapshot_dir(&session.snapshot_id); + self.rollback_failed_publish(&session.snapshot_id, previous_record.as_ref()); Err(error) } } @@ -164,12 +160,35 @@ impl PosixFsCatalogStore { if let Some(alias) = record.alias.as_ref() { self.with_alias_lock(alias, |store| { - store.ensure_alias_available(alias, &record.id)?; store.write_record_unlocked(&record)?; - store.write_json( - &PosixFsSnapshotArtifactLayout::alias_path(&store.root, alias), - &record.id, - ) + let alias_path = PosixFsSnapshotArtifactLayout::alias_path(&store.root, alias); + let bind = (|| -> RepositoryResult<()> { + match store.load_alias_target(alias)? { + // The alias currently points at a live snapshot. Leave the + // binding untouched so the existing template keeps resolving + // while the new build runs; a successful commit moves the + // alias to the new snapshot (E2B rebuild semantics). + Some(existing) + if existing != record.id + && store.load_record_by_id_unlocked(&existing)?.is_some() => + { + Ok(()) + } + Some(existing) if existing != record.id => { + store.remove_file_if_exists(&alias_path)?; + store.write_json(&alias_path, &record.id) + } + _ => store.write_json(&alias_path, &record.id), + } + })(); + if let Err(error) = bind { + // Keep creation all-or-nothing under the alias lock: a record + // that survives a failed binding claims the alias in listings + // with nothing left to reconcile it. + let _ = store.remove_file_if_exists(&store.record_path(&record.id)); + return Err(error); + } + Ok(()) })?; } else { self.write_record_unlocked(&record)?; @@ -376,6 +395,67 @@ impl PosixFsCatalogStore { self.write_record_unlocked(&record) } + /// Clears `moved_alias` from the record that owned it before a rebind. + /// + /// Best effort: the alias binding has already moved, so a failure here only + /// leaves stale alias metadata on the previous owner's record. + fn clear_moved_alias_on_previous_record( + &self, + previous_id: &SnapshotId, + moved_alias: &str, + now: i64, + ) { + // Lock order is alias lock first, then record lock; nothing takes them + // in the reverse order today. + let _guard = match self.acquire_record_lock(previous_id) { + Ok(guard) => guard, + Err(error) => { + warn!( + alias = %moved_alias, + previous_snapshot_id = %previous_id, + error = %error, + "failed to lock previous snapshot record for alias cleanup" + ); + return; + } + }; + + let mut previous = match self.load_record_by_id_unlocked(previous_id) { + Ok(Some(previous)) => previous, + Ok(None) => return, + Err(error) => { + warn!( + alias = %moved_alias, + previous_snapshot_id = %previous_id, + error = %error, + "failed to load previous snapshot alias metadata" + ); + return; + } + }; + + // Only clear the alias this publish actually moved; the previous owner + // may already claim a different name. + let claims_moved_alias = previous + .alias + .as_ref() + .is_some_and(|alias| alias.as_ref() == moved_alias); + if !claims_moved_alias { + return; + } + + previous.alias = None; + previous.updated_at_unix_ms = now; + if let Err(error) = self.write_record_unlocked(&previous) { + warn!( + alias = %moved_alias, + previous_snapshot_id = %previous_id, + error = %error, + "failed to clear previous snapshot alias metadata" + ); + } + } + fn read_json(&self, path: &Path) -> RepositoryResult where T: DeserializeOwned, @@ -498,6 +578,19 @@ impl PosixFsCatalogStore { self.remove_dir_if_exists(&snapshot_layout.snapshot_dir()) } + fn rollback_failed_publish(&self, id: &SnapshotId, previous_record: Option<&SnapshotRecord>) { + if let Err(error) = self.remove_dir_if_exists(&self.layout(id).snapshot_dir()) { + warn!(snapshot_id = %id, error = %error, "failed to remove snapshot artifacts after publish failure"); + } + let restore_result = match previous_record { + Some(record) => self.write_record_unlocked(record), + None => self.remove_file_if_exists(&self.record_path(id)), + }; + if let Err(error) = restore_result { + warn!(snapshot_id = %id, error = %error, "failed to restore snapshot record after publish failure"); + } + } + fn load_record_by_id_unlocked( &self, id: &SnapshotId, @@ -626,28 +719,6 @@ impl PosixFsCatalogStore { action(self) } - fn ensure_alias_available( - &self, - alias: &SnapshotAlias, - new_id: &SnapshotId, - ) -> RepositoryResult<()> { - let alias_path = PosixFsSnapshotArtifactLayout::alias_path(&self.root, alias); - if let Some(existing) = self.load_alias_target(alias)? { - if &existing == new_id { - return Ok(()); - } - if self.load_record_by_id_unlocked(&existing)?.is_some() { - return Err(RepositoryError::AliasConflict { - alias: alias.to_string(), - existing, - new_id: new_id.clone(), - }); - } - self.remove_file_if_exists(&alias_path)?; - } - Ok(()) - } - fn write_record_unlocked(&self, record: &SnapshotRecord) -> RepositoryResult<()> { self.write_json(&self.record_path(&record.id), record) } From b33dc13b54c9461752e9a870cbda88ad8a86f062 Mon Sep 17 00:00:00 2001 From: JoyboyBrian Date: Wed, 29 Jul 2026 12:30:05 -0700 Subject: [PATCH 3/3] feat(snapshot): template build-context archive store (posixfs + oss) --- src/snapshot/image_export/service.rs | 2 + src/snapshot/manager.rs | 8 + .../repository/backends/oss/build_files.rs | 169 ++++ src/snapshot/repository/backends/oss/mod.rs | 1 + .../repository/backends/oss/repository.rs | 12 + .../repository/backends/posixfs/backend.rs | 12 + .../backends/posixfs/build_files.rs | 726 ++++++++++++++++++ .../repository/backends/posixfs/mod.rs | 2 + src/snapshot/repository/build_files.rs | 197 +++++ src/snapshot/repository/interfaces.rs | 10 + src/snapshot/repository/mod.rs | 2 + 11 files changed, 1141 insertions(+) create mode 100644 src/snapshot/repository/backends/oss/build_files.rs create mode 100644 src/snapshot/repository/backends/posixfs/build_files.rs create mode 100644 src/snapshot/repository/build_files.rs diff --git a/src/snapshot/image_export/service.rs b/src/snapshot/image_export/service.rs index 01a088ee..40caa497 100644 --- a/src/snapshot/image_export/service.rs +++ b/src/snapshot/image_export/service.rs @@ -67,6 +67,7 @@ impl SnapshotImageService { let repository = Arc::new(posixfs::PosixFsSnapshotRepository::new( Arc::new(posixfs::PosixFsCatalogStore::new(root.clone())), Arc::new(posixfs::PosixFsArtifactStore::new(root.clone())), + posixfs::PosixFsTemplateBuildFileStore::new(&root), )); (repository, ManagedLayerLocator::PosixFs { root }) } @@ -431,6 +432,7 @@ mod tests { let repository = posixfs::PosixFsSnapshotRepository::new( Arc::new(posixfs::PosixFsCatalogStore::new(root.clone())), Arc::new(posixfs::PosixFsArtifactStore::new(root.clone())), + posixfs::PosixFsTemplateBuildFileStore::new(&root), ); let uncommitted = SnapshotRecord::template_waiting(SnapshotId::generate(), None, Default::default()); diff --git a/src/snapshot/manager.rs b/src/snapshot/manager.rs index 634bf25d..5167f354 100644 --- a/src/snapshot/manager.rs +++ b/src/snapshot/manager.rs @@ -84,6 +84,14 @@ impl SnapshotManager { self.repository.create(record).await } + /// Returns the shared build-context archive store, when the configured + /// repository backend provides one. + pub fn template_build_files( + &self, + ) -> Option> { + self.repository.template_build_files() + } + #[tracing::instrument(skip(self, metadata, manifest), fields(snapshot_id = %metadata.id))] pub async fn publish( &self, diff --git a/src/snapshot/repository/backends/oss/build_files.rs b/src/snapshot/repository/backends/oss/build_files.rs new file mode 100644 index 00000000..2ac034dd --- /dev/null +++ b/src/snapshot/repository/backends/oss/build_files.rs @@ -0,0 +1,169 @@ +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use async_trait::async_trait; + +use super::client::OssClient; +use crate::snapshot::repository::build_files::{ + generate_upload_token, is_valid_build_files_hash, is_valid_upload_token, + TemplateBuildFileStore, TemplateBuildUploadGrant, +}; +use crate::snapshot::repository::{RepositoryError, RepositoryResult}; + +const BUILD_FILES_PREFIX: &str = "template-build-files"; + +/// Build-context archive store backed by the OSS repository bucket. +/// +/// Layout: `template-build-files/{hash}.tar` plus durable bearer grants under +/// `template-build-files/upload-grants/`. Retention is delegated to bucket +/// lifecycle rules; archives are cache entries the SDK re-uploads when absent. +pub(crate) struct OssTemplateBuildFileStore { + client: Arc, +} + +impl OssTemplateBuildFileStore { + pub(crate) fn new(client: Arc) -> Arc { + Arc::new(Self { client }) + } + + fn archive_key(hash: &str) -> RepositoryResult { + if !is_valid_build_files_hash(hash) { + return Err(RepositoryError::InvalidRequest { + reason: format!("invalid build files hash '{hash}'"), + }); + } + Ok(format!("{BUILD_FILES_PREFIX}/{hash}.tar")) + } + + fn grant_key(token: &str) -> Option { + is_valid_upload_token(token) + .then(|| format!("{BUILD_FILES_PREFIX}/upload-grants/{token}.json")) + } + + /// Reads a grant record, mapping an absent object to `None`. + async fn read_grant(&self, key: &str) -> RepositoryResult> { + let bytes = match self.client.get_bytes(key).await { + Ok(bytes) => bytes, + Err(error) if OssClient::is_not_found_error(&error) => return Ok(None), + Err(error) => return Err(RepositoryError::backend("read upload grant", error)), + }; + serde_json::from_slice(&bytes) + .map(Some) + .map_err(|error| RepositoryError::backend("parse upload grant", error)) + } +} + +#[async_trait] +impl TemplateBuildFileStore for OssTemplateBuildFileStore { + async fn exists(&self, hash: &str) -> RepositoryResult { + let key = Self::archive_key(hash)?; + self.client + .exists(&key) + .await + .map_err(|error| RepositoryError::backend("check build archive", error)) + } + + async fn import(&self, hash: &str, staged: &Path) -> RepositoryResult<()> { + let key = Self::archive_key(hash)?; + // Archives are immutable so a repeat upload cannot change what an + // in-flight build reads. This fast path is not atomic against a + // concurrent import: the loser's bytes are dropped, and since the hash + // is a caller-supplied cache key rather than a verified digest, which + // racing upload wins is undefined — first-write-wins stability, not + // content authenticity. + if self + .client + .exists(&key) + .await + .map_err(|error| RepositoryError::backend("check build archive", error))? + { + return Ok(()); + } + self.client + .put_file(&key, staged) + .await + .map_err(|error| RepositoryError::backend("upload build archive", error)) + } + + async fn materialize( + &self, + hash: &str, + scratch_dir: &Path, + ) -> RepositoryResult> { + let key = Self::archive_key(hash)?; + let dest = scratch_dir.join(format!("{hash}.tar")); + match self.client.get_to_file(&key, &dest).await { + Ok(_) => Ok(Some(dest)), + Err(error) if OssClient::is_not_found_error(&error) => Ok(None), + Err(error) => Err(RepositoryError::backend("download build archive", error)), + } + } + + async fn create_upload_grant( + &self, + template_id: &str, + hash: &str, + expires_unix: i64, + ) -> RepositoryResult { + let token = generate_upload_token(); + let key = Self::grant_key(&token).expect("generated token is valid"); + let grant = serde_json::to_vec(&TemplateBuildUploadGrant::new( + template_id, + hash, + expires_unix, + )) + .map_err(|error| RepositoryError::backend("serialize upload grant", error))?; + self.client + .put_bytes(&key, grant) + .await + .map_err(|error| RepositoryError::backend("write upload grant", error))?; + Ok(token) + } + + async fn verify_upload_grant( + &self, + token: &str, + template_id: &str, + hash: &str, + expires_unix: i64, + now_unix: i64, + ) -> RepositoryResult { + let Some(key) = Self::grant_key(token) else { + return Ok(false); + }; + // Deliberately does not delete the object: verification must leave the + // upload URL usable for a retry. + let Some(grant) = self.read_grant(&key).await? else { + return Ok(false); + }; + Ok(grant.authorizes(template_id, hash, expires_unix, now_unix)) + } + + async fn claim_upload_grant( + &self, + token: &str, + template_id: &str, + hash: &str, + expires_unix: i64, + now_unix: i64, + ) -> RepositoryResult { + let Some(key) = Self::grant_key(token) else { + return Ok(false); + }; + let Some(grant) = self.read_grant(&key).await? else { + return Ok(false); + }; + if !grant.authorizes(template_id, hash, expires_unix, now_unix) { + return Ok(false); + } + // Consume the grant so the upload URL cannot be replayed. S3-compatible + // stores offer no conditional delete, so simultaneous replays of one + // token can both observe the grant; archives are immutable, which is + // what keeps that from mattering. + self.client + .delete(&key) + .await + .map_err(|error| RepositoryError::backend("consume upload grant", error))?; + Ok(true) + } +} diff --git a/src/snapshot/repository/backends/oss/mod.rs b/src/snapshot/repository/backends/oss/mod.rs index 837c2540..435342e7 100644 --- a/src/snapshot/repository/backends/oss/mod.rs +++ b/src/snapshot/repository/backends/oss/mod.rs @@ -1,3 +1,4 @@ +mod build_files; mod client; mod config; mod layout; diff --git a/src/snapshot/repository/backends/oss/repository.rs b/src/snapshot/repository/backends/oss/repository.rs index 8af954a8..b01d187a 100644 --- a/src/snapshot/repository/backends/oss/repository.rs +++ b/src/snapshot/repository/backends/oss/repository.rs @@ -42,6 +42,7 @@ pub(crate) struct OssSnapshotRepository { client: Arc, snapshot_image_storage: SnapshotImageStoragePolicy, acr_exporter: AcrDiskImageExporter, + build_files: Arc, } const MAX_ALIAS_BIND_ATTEMPTS: usize = 5; @@ -51,10 +52,12 @@ impl OssSnapshotRepository { client: Arc, snapshot_image_storage: SnapshotImageStoragePolicy, ) -> Self { + let build_files = super::build_files::OssTemplateBuildFileStore::new(Arc::clone(&client)); Self { client, snapshot_image_storage, acr_exporter: AcrDiskImageExporter::new(), + build_files, } } @@ -148,6 +151,15 @@ fn fallback_to_object_storage_would_mix_sources( #[async_trait] impl SnapshotRepository for OssSnapshotRepository { + fn template_build_files( + &self, + ) -> Option> { + Some(Arc::clone(&self.build_files) + as Arc< + dyn crate::snapshot::repository::TemplateBuildFileStore, + >) + } + async fn create(&self, record: SnapshotRecord) -> RepositoryResult { if !matches!(record.source, SnapshotSource::Template { .. }) { return Err(RepositoryError::InvalidRequest { diff --git a/src/snapshot/repository/backends/posixfs/backend.rs b/src/snapshot/repository/backends/posixfs/backend.rs index 459c1547..713a4961 100644 --- a/src/snapshot/repository/backends/posixfs/backend.rs +++ b/src/snapshot/repository/backends/posixfs/backend.rs @@ -7,11 +7,13 @@ use tokio::task; use super::super::shared_runtime_cache_root; use super::artifacts::{CollectedBuiltArtifacts, PosixFsArtifactStore}; +use super::build_files::PosixFsTemplateBuildFileStore; use super::catalog::PosixFsCatalogStore; use super::runtime::PosixFsRuntimeResolver; use crate::image::cache::{local_image_services_from_global_config, OverlaybdLayerStore}; use crate::sandbox::FirecrackerSnapshotManifest; use crate::snapshot::artifact_cache::LocalArtifactCache; +use crate::snapshot::repository::build_files::TemplateBuildFileStore; use crate::snapshot::repository::interfaces::{SnapshotRepository, SnapshotRuntimeResolver}; use crate::snapshot::repository::{RepositoryError, RepositoryResult, SnapshotListFilter}; use crate::snapshot::types::{ @@ -72,9 +74,11 @@ impl PosixFsBackend { let runtime_cache_root = runtime_cache_root.unwrap_or_else(|| cache_root.join("runtime")); let catalog_store = Arc::new(PosixFsCatalogStore::new(root.clone())); let artifact_store = Arc::new(PosixFsArtifactStore::new(root.clone())); + let build_files = PosixFsTemplateBuildFileStore::new(&root); let repository: Arc = Arc::new(PosixFsSnapshotRepository::new( catalog_store, artifact_store, + build_files, )); let runtime_resolver: Arc = Arc::new( PosixFsRuntimeResolver::new(root, runtime_cache_root, store, cache), @@ -111,16 +115,19 @@ impl PosixFsBackend { pub(crate) struct PosixFsSnapshotRepository { catalog_store: Arc, artifact_store: Arc, + build_files: Arc, } impl PosixFsSnapshotRepository { pub(crate) fn new( catalog_store: Arc, artifact_store: Arc, + build_files: Arc, ) -> Self { Self { catalog_store, artifact_store, + build_files, } } @@ -237,6 +244,10 @@ impl SnapshotRepository for PosixFsSnapshotRepository { .await } + fn template_build_files(&self) -> Option> { + Some(Arc::clone(&self.build_files) as Arc) + } + async fn publish( &self, metadata: SnapshotPublishMetadata, @@ -399,6 +410,7 @@ mod tests { PosixFsSnapshotRepository::new( Arc::new(PosixFsCatalogStore::new(root.to_path_buf())), Arc::new(PosixFsArtifactStore::new(root.to_path_buf())), + super::super::build_files::PosixFsTemplateBuildFileStore::new(root), ) } diff --git a/src/snapshot/repository/backends/posixfs/build_files.rs b/src/snapshot/repository/backends/posixfs/build_files.rs new file mode 100644 index 00000000..c486ee8b --- /dev/null +++ b/src/snapshot/repository/backends/posixfs/build_files.rs @@ -0,0 +1,726 @@ +use std::fs; +use std::io::Write; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::{Duration, SystemTime}; + +use async_trait::async_trait; +use tokio::task; +use tracing::{debug, warn}; + +use crate::snapshot::repository::build_files::{ + generate_upload_token, is_valid_build_files_hash, is_valid_upload_token, + TemplateBuildFileStore, TemplateBuildUploadGrant, +}; +use crate::snapshot::repository::{RepositoryError, RepositoryResult}; + +/// How long imported build-context archives and upload grants are retained. +/// Archives are cache entries keyed by content hash; the SDK re-uploads any +/// archive that has been pruned, so expiry only costs one extra upload. +/// Grants expire after `template_build.files_url_ttl_secs` anyway, so this +/// only bounds how long the spent grant files linger on disk. +const BUILD_FILE_RETENTION: Duration = Duration::from_secs(7 * 24 * 60 * 60); + +const GRANTS_DIR_NAME: &str = "upload-grants"; + +/// Build-context archive store rooted on the shared POSIX repository. +/// +/// Layout: `{repository_root}/template-build-files/{hash}.tar` plus durable +/// upload grants under `upload-grants/`. Both live on the shared filesystem, +/// so every node observes the same archives and verifies the same upload URLs. +pub(crate) struct PosixFsTemplateBuildFileStore { + root: PathBuf, +} + +impl PosixFsTemplateBuildFileStore { + pub(crate) fn new(repository_root: &Path) -> Arc { + Arc::new(Self { + root: repository_root.join("template-build-files"), + }) + } + + fn archive_path(&self, hash: &str) -> RepositoryResult { + if !is_valid_build_files_hash(hash) { + return Err(RepositoryError::InvalidRequest { + reason: format!("invalid build files hash '{hash}'"), + }); + } + Ok(self.root.join(format!("{hash}.tar"))) + } + + fn ensure_root(root: &Path) -> RepositoryResult<()> { + fs::create_dir_all(root).map_err(|error| { + RepositoryError::backend( + format!("create template build files dir '{}'", root.display()), + error, + ) + }) + } + + /// Removes archives whose modification time is older than the retention + /// window. Runs opportunistically on import and scans a bounded number of + /// entries per call; failures only log. + fn prune_expired(root: &Path) { + let cutoff = SystemTime::now() - BUILD_FILE_RETENTION; + Self::prune_dir_older_than(root, "tar", cutoff); + } + + /// Removes upload grants that have passed their own `expires_unix`. Runs + /// opportunistically whenever a new grant is written, so the grants + /// directory stays bounded by upload-link traffic; the scan is bounded per + /// call and drains the backlog over successive requests, and failures only + /// log. + /// + /// Pruning by the record rather than by mtime keeps grants alive for + /// exactly their TTL even when `template_build.files_url_ttl_secs` is + /// configured beyond the retention window. + fn prune_expired_grants(root: &Path) { + let cutoff = SystemTime::now() - BUILD_FILE_RETENTION; + let now_unix = chrono::Utc::now().timestamp(); + Self::prune_dir(&Self::grants_dir(root), "json", |path, modified| { + match fs::read(path) + .ok() + .and_then(|bytes| serde_json::from_slice::(&bytes).ok()) + { + Some(grant) => grant.expires_unix < now_unix, + // Unparseable leftovers fall back to the mtime rule. + None => modified.is_some_and(|modified| modified < cutoff), + } + }); + } + + fn prune_dir_older_than(dir: &Path, extension: &str, cutoff: SystemTime) { + Self::prune_dir(dir, extension, |_, modified| { + modified.is_some_and(|modified| modified < cutoff) + }); + } + + /// Pruning is opportunistic and bounded: at most `MAX_PRUNE_SCAN` matching + /// entries are inspected per call, so the cost a request pays stays + /// constant no matter how many records the directory holds. Anything left + /// over is reclaimed by later calls. + fn prune_dir( + dir: &Path, + extension: &str, + is_expired: impl Fn(&Path, Option) -> bool, + ) { + const MAX_PRUNE_SCAN: usize = 256; + + let Ok(entries) = fs::read_dir(dir) else { + return; + }; + let mut scanned: usize = 0; + for entry in entries.flatten() { + let path = entry.path(); + if path.extension().is_none_or(|ext| ext != extension) { + continue; + } + if scanned >= MAX_PRUNE_SCAN { + break; + } + scanned += 1; + let modified = entry + .metadata() + .and_then(|metadata| metadata.modified()) + .ok(); + if is_expired(&path, modified) { + if let Err(error) = fs::remove_file(&path) { + warn!( + path = %path.display(), + error = %error, + "failed to prune expired template build file" + ); + } else { + debug!(path = %path.display(), "pruned expired template build file"); + } + } + } + } + + fn grants_dir(root: &Path) -> PathBuf { + root.join(GRANTS_DIR_NAME) + } + + fn grant_path(root: &Path, token: &str) -> Option { + is_valid_upload_token(token).then(|| Self::grants_dir(root).join(format!("{token}.json"))) + } + + /// Reads a grant record, mapping an absent file to `None`. + fn read_grant(path: &Path) -> RepositoryResult> { + let bytes = match fs::read(path) { + Ok(bytes) => bytes, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(RepositoryError::backend("read upload grant", error)), + }; + serde_json::from_slice(&bytes) + .map(Some) + .map_err(|error| RepositoryError::backend("parse upload grant", error)) + } + + /// Best-effort mtime refresh, so retention means "unused for the window" + /// and an archive a build is still reading stays outside the prune + /// horizon. Read-only repository mounts must keep working, so failures + /// only log. + fn touch(path: &Path) { + let refreshed = fs::File::options() + .write(true) + .open(path) + .and_then(|file| file.set_times(fs::FileTimes::new().set_modified(SystemTime::now()))); + if let Err(error) = refreshed { + debug!( + path = %path.display(), + error = %error, + "failed to refresh build archive mtime" + ); + } + } + + fn write_grant( + root: &Path, + template_id: &str, + hash: &str, + expires_unix: i64, + ) -> RepositoryResult { + let grants_dir = Self::grants_dir(root); + fs::create_dir_all(&grants_dir).map_err(|error| { + RepositoryError::backend( + format!("create upload grants dir '{}'", grants_dir.display()), + error, + ) + })?; + Self::prune_expired_grants(root); + let bytes = serde_json::to_vec(&TemplateBuildUploadGrant::new( + template_id, + hash, + expires_unix, + )) + .map_err(|error| RepositoryError::backend("serialize upload grant", error))?; + + for _ in 0..3 { + let token = generate_upload_token(); + let path = Self::grant_path(root, &token).expect("generated token is valid"); + match fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&path) + { + Ok(mut file) => { + file.write_all(&bytes) + .and_then(|()| file.sync_all()) + .map_err(|error| { + let _ = fs::remove_file(&path); + RepositoryError::backend("write upload grant", error) + })?; + return Ok(token); + } + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => continue, + Err(error) => return Err(RepositoryError::backend("create upload grant", error)), + } + } + Err(RepositoryError::Backend { + message: "failed to allocate a unique upload grant token".to_string(), + source: None, + }) + } +} + +#[async_trait] +impl TemplateBuildFileStore for PosixFsTemplateBuildFileStore { + async fn exists(&self, hash: &str) -> RepositoryResult { + let path = self.archive_path(hash)?; + task::spawn_blocking(move || -> RepositoryResult { + match fs::metadata(&path) { + Ok(_) => Ok(true), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(error) => Err(RepositoryError::backend( + format!("stat build archive '{}'", path.display()), + error, + )), + } + }) + .await + .map_err(|error| RepositoryError::backend("join build file exists task", error))? + } + + async fn import(&self, hash: &str, staged: &Path) -> RepositoryResult<()> { + let final_path = self.archive_path(hash)?; + let root = self.root.clone(); + let staged = staged.to_path_buf(); + task::spawn_blocking(move || -> RepositoryResult<()> { + // Archives are immutable: the hash addresses the content, so a + // repeat upload cannot change what an in-flight build reads. + if final_path.exists() { + return Ok(()); + } + Self::ensure_root(&root)?; + Self::prune_expired(&root); + // Copy into the store filesystem first (the staged file usually + // lives on node-local tmp), then link it into place within the + // store directory so readers only ever observe complete archives. + let store_staged = root.join(format!(".import-{}.tmp", uuid::Uuid::new_v4())); + fs::copy(&staged, &store_staged).map_err(|error| { + let _ = fs::remove_file(&store_staged); + RepositoryError::backend("copy build archive into store", error) + })?; + // The archive is only ever published once, so its data must reach + // stable storage before the name does: a directory entry that + // outlives the bytes would pin a truncated archive forever behind + // the `exists` fast path. + fs::File::open(&store_staged) + .and_then(|file| file.sync_all()) + .map_err(|error| { + let _ = fs::remove_file(&store_staged); + RepositoryError::backend("sync build archive", error) + })?; + // Link rather than rename so a concurrent import cannot replace an + // archive a running build is already reading: the first writer + // wins and everyone else observes `AlreadyExists`. + let published = match fs::hard_link(&store_staged, &final_path) { + Ok(()) => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => Ok(()), + Err(error) => Err(RepositoryError::backend("publish build archive", error)), + }; + if published.is_ok() { + // Best effort: filesystems that reject a directory fsync must + // keep working, and a lost entry only costs one re-upload. + if let Err(error) = fs::File::open(&root).and_then(|dir| dir.sync_all()) { + debug!( + path = %root.display(), + error = %error, + "failed to sync build archive store directory" + ); + } + } + let _ = fs::remove_file(&store_staged); + published + }) + .await + .map_err(|error| RepositoryError::backend("join build file import task", error))? + } + + async fn materialize( + &self, + hash: &str, + _scratch_dir: &Path, + ) -> RepositoryResult> { + let path = self.archive_path(hash)?; + task::spawn_blocking(move || -> RepositoryResult> { + match fs::metadata(&path) { + Ok(_) => { + Self::touch(&path); + Ok(Some(path)) + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(error) => Err(RepositoryError::backend( + format!("stat build archive '{}'", path.display()), + error, + )), + } + }) + .await + .map_err(|error| RepositoryError::backend("join build file materialize task", error))? + } + + async fn create_upload_grant( + &self, + template_id: &str, + hash: &str, + expires_unix: i64, + ) -> RepositoryResult { + let root = self.root.clone(); + let template_id = template_id.to_string(); + let hash = hash.to_string(); + task::spawn_blocking(move || Self::write_grant(&root, &template_id, &hash, expires_unix)) + .await + .map_err(|error| RepositoryError::backend("join create upload grant task", error))? + } + + async fn verify_upload_grant( + &self, + token: &str, + template_id: &str, + hash: &str, + expires_unix: i64, + now_unix: i64, + ) -> RepositoryResult { + let Some(path) = Self::grant_path(&self.root, token) else { + return Ok(false); + }; + let template_id = template_id.to_string(); + let hash = hash.to_string(); + task::spawn_blocking(move || -> RepositoryResult { + // Reads only: the grant file must survive so an upload that fails + // before the archive is stored can be retried with the same URL. + let Some(grant) = Self::read_grant(&path)? else { + return Ok(false); + }; + Ok(grant.authorizes(&template_id, &hash, expires_unix, now_unix)) + }) + .await + .map_err(|error| RepositoryError::backend("join verify upload grant task", error))? + } + + async fn claim_upload_grant( + &self, + token: &str, + template_id: &str, + hash: &str, + expires_unix: i64, + now_unix: i64, + ) -> RepositoryResult { + let Some(path) = Self::grant_path(&self.root, token) else { + return Ok(false); + }; + let template_id = template_id.to_string(); + let hash = hash.to_string(); + task::spawn_blocking(move || -> RepositoryResult { + let Some(grant) = Self::read_grant(&path)? else { + return Ok(false); + }; + if !grant.authorizes(&template_id, &hash, expires_unix, now_unix) { + return Ok(false); + } + // Consume the grant. `remove_file` succeeds for exactly one + // caller, so it is the claim: concurrent replays of the same + // token lose the race and are rejected. + match fs::remove_file(&path) { + Ok(()) => Ok(true), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(error) => Err(RepositoryError::backend("consume upload grant", error)), + } + }) + .await + .map_err(|error| RepositoryError::backend("join claim upload grant task", error))? + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::TempDir; + + const HASH: &str = "9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08"; + + fn staged_file(dir: &Path, contents: &[u8]) -> PathBuf { + let path = dir.join("staged.tar"); + fs::write(&path, contents).expect("write staged file"); + path + } + + #[tokio::test] + async fn import_then_exists_and_materialize() { + let tempdir = TempDir::new().expect("tempdir"); + let store = PosixFsTemplateBuildFileStore::new(tempdir.path()); + + assert!(!store.exists(HASH).await.expect("exists should work")); + assert_eq!( + store + .materialize(HASH, tempdir.path()) + .await + .expect("materialize should work"), + None + ); + + let staged = staged_file(tempdir.path(), b"tar-bytes"); + store + .import(HASH, &staged) + .await + .expect("import should work"); + + assert!(store.exists(HASH).await.expect("exists should work")); + let materialized = store + .materialize(HASH, tempdir.path()) + .await + .expect("materialize should work") + .expect("archive should exist"); + assert_eq!( + fs::read(materialized).expect("read materialized"), + b"tar-bytes" + ); + } + + #[tokio::test] + async fn import_rejects_invalid_hash() { + let tempdir = TempDir::new().expect("tempdir"); + let store = PosixFsTemplateBuildFileStore::new(tempdir.path()); + let staged = staged_file(tempdir.path(), b"tar-bytes"); + + let err = store + .import("../escape", &staged) + .await + .expect_err("invalid hash should fail"); + assert!(matches!(err, RepositoryError::InvalidRequest { .. })); + } + + #[tokio::test] + async fn writing_a_grant_prunes_expired_grant_files() { + let tempdir = TempDir::new().expect("tempdir"); + let store = PosixFsTemplateBuildFileStore::new(tempdir.path()); + + let fresh_token = store + .create_upload_grant("template", HASH, i64::MAX) + .await + .expect("fresh grant should be created"); + + // Plant a grant file that predates the retention window. + let grants_dir = tempdir + .path() + .join("template-build-files") + .join("upload-grants"); + let stale_path = grants_dir.join(format!("{}.json", generate_upload_token())); + fs::write(&stale_path, b"{}").expect("write stale grant"); + let stale_mtime = SystemTime::now() - BUILD_FILE_RETENTION - Duration::from_secs(60); + let stale_file = fs::File::options() + .write(true) + .open(&stale_path) + .expect("open stale grant"); + stale_file + .set_times(fs::FileTimes::new().set_modified(stale_mtime)) + .expect("set stale mtime"); + drop(stale_file); + + store + .create_upload_grant("template", HASH, i64::MAX) + .await + .expect("new grant should be created"); + + assert!(!stale_path.exists(), "expired grant file should be pruned"); + assert!( + store + .claim_upload_grant(&fresh_token, "template", HASH, i64::MAX, 0) + .await + .expect("validation should work"), + "unexpired grants must survive pruning" + ); + } + + #[tokio::test] + async fn upload_grant_is_shared_across_instances() { + let tempdir = TempDir::new().expect("tempdir"); + let first = PosixFsTemplateBuildFileStore::new(tempdir.path()); + let second = PosixFsTemplateBuildFileStore::new(tempdir.path()); + + // A mismatched or expired claim leaves the grant usable. + let token = first + .create_upload_grant("template", HASH, 1000) + .await + .expect("grant should be created"); + assert!(!second + .claim_upload_grant(&token, "other", HASH, 1000, 999) + .await + .expect("mismatched grant should be rejected")); + assert!(!second + .claim_upload_grant(&token, "template", HASH, 1000, 1001) + .await + .expect("expired grant should be rejected")); + assert!(second + .claim_upload_grant(&token, "template", HASH, 1000, 999) + .await + .expect("grant issued by another instance should claim")); + } + + #[tokio::test] + async fn upload_grant_is_single_use() { + let tempdir = TempDir::new().expect("tempdir"); + let store = PosixFsTemplateBuildFileStore::new(tempdir.path()); + let token = store + .create_upload_grant("template", HASH, 1000) + .await + .expect("grant should be created"); + + assert!(store + .claim_upload_grant(&token, "template", HASH, 1000, 999) + .await + .expect("first claim should succeed")); + assert!( + !store + .claim_upload_grant(&token, "template", HASH, 1000, 999) + .await + .expect("replay should be rejected"), + "an upload URL must not be replayable" + ); + } + + #[tokio::test] + async fn archives_are_immutable_once_stored() { + let tempdir = TempDir::new().expect("tempdir"); + let store = PosixFsTemplateBuildFileStore::new(tempdir.path()); + + let first = staged_file(tempdir.path(), b"original"); + store.import(HASH, &first).await.expect("first import"); + + let replacement = tempdir.path().join("replacement.tar"); + fs::write(&replacement, b"replaced").expect("write replacement"); + store + .import(HASH, &replacement) + .await + .expect("repeat import should be accepted"); + + // Two imports racing for a hash neither has stored yet must both + // succeed; the loser's hard link hits AlreadyExists and is dropped. + // A fresh hash keeps both calls off the exists() fast path. + const FRESH_HASH: &str = "f00ff00ff00ff00ff00ff00ff00ff00f"; + let concurrent = tempdir.path().join("concurrent.tar"); + fs::write(&concurrent, b"concurrent").expect("write concurrent"); + let (left, right) = tokio::join!( + store.import(FRESH_HASH, &replacement), + store.import(FRESH_HASH, &concurrent) + ); + left.expect("concurrent import should be accepted"); + right.expect("concurrent import should be accepted"); + let winner = store + .materialize(FRESH_HASH, tempdir.path()) + .await + .expect("materialize should work") + .expect("archive should exist"); + let winner_bytes = fs::read(winner).expect("read winner"); + assert!( + winner_bytes == b"replaced" || winner_bytes == b"concurrent", + "stored bytes must come from one of the racing imports" + ); + + let materialized = store + .materialize(HASH, tempdir.path()) + .await + .expect("materialize should work") + .expect("archive should exist"); + assert_eq!( + fs::read(materialized).expect("read materialized"), + b"original", + "a stored archive must never be replaced underneath a build" + ); + + let leftovers = fs::read_dir(tempdir.path().join("template-build-files")) + .expect("read store dir") + .flatten() + .filter(|entry| entry.file_name().to_string_lossy().starts_with(".import-")) + .count(); + assert_eq!(leftovers, 0, "import must not leak staging files"); + } + + #[tokio::test] + async fn materialize_refreshes_the_archive_mtime() { + let tempdir = TempDir::new().expect("tempdir"); + let store = PosixFsTemplateBuildFileStore::new(tempdir.path()); + + let staged = staged_file(tempdir.path(), b"tar-bytes"); + store.import(HASH, &staged).await.expect("import"); + + let archive = tempdir + .path() + .join("template-build-files") + .join(format!("{HASH}.tar")); + let stale = SystemTime::now() - BUILD_FILE_RETENTION - Duration::from_secs(60); + let file = fs::File::options() + .write(true) + .open(&archive) + .expect("open archive"); + file.set_times(fs::FileTimes::new().set_modified(stale)) + .expect("set stale mtime"); + drop(file); + + store + .materialize(HASH, tempdir.path()) + .await + .expect("materialize should work") + .expect("archive should exist"); + + let modified = fs::metadata(&archive) + .and_then(|metadata| metadata.modified()) + .expect("read archive mtime"); + assert!( + modified > stale, + "materializing an archive must keep it outside the prune horizon" + ); + } + + #[tokio::test] + async fn verifying_a_grant_does_not_consume_it() { + let tempdir = TempDir::new().expect("tempdir"); + let store = PosixFsTemplateBuildFileStore::new(tempdir.path()); + let token = store + .create_upload_grant("template", HASH, 1000) + .await + .expect("grant should be created"); + + for _ in 0..2 { + assert!( + store + .verify_upload_grant(&token, "template", HASH, 1000, 999) + .await + .expect("verification should work"), + "verification must not consume the grant" + ); + } + assert!(!store + .verify_upload_grant(&token, "other", HASH, 1000, 999) + .await + .expect("mismatched grant should be rejected")); + + // An upload that failed after verification can still be retried. + assert!(store + .claim_upload_grant(&token, "template", HASH, 1000, 999) + .await + .expect("claim should succeed")); + assert!( + !store + .verify_upload_grant(&token, "template", HASH, 1000, 999) + .await + .expect("verification should work"), + "a consumed grant must no longer verify" + ); + } + + #[tokio::test] + async fn concurrent_claims_pick_a_single_winner() { + let tempdir = TempDir::new().expect("tempdir"); + let store = PosixFsTemplateBuildFileStore::new(tempdir.path()); + let token = store + .create_upload_grant("template", HASH, 1000) + .await + .expect("grant should be created"); + + let (left, right) = tokio::join!( + store.claim_upload_grant(&token, "template", HASH, 1000, 999), + store.claim_upload_grant(&token, "template", HASH, 1000, 999) + ); + let claims = [ + left.expect("claim should work"), + right.expect("claim should work"), + ]; + assert_eq!( + claims.iter().filter(|claimed| **claimed).count(), + 1, + "exactly one concurrent claim may win" + ); + } + + #[tokio::test] + async fn grants_are_pruned_once_their_own_expiry_passes() { + let tempdir = TempDir::new().expect("tempdir"); + let store = PosixFsTemplateBuildFileStore::new(tempdir.path()); + + // Expired long ago in grant terms, but freshly written on disk, so the + // mtime rule alone would keep it for the whole retention window. + let expired_token = store + .create_upload_grant("template", HASH, 1000) + .await + .expect("grant should be created"); + let expired_path = tempdir + .path() + .join("template-build-files") + .join("upload-grants") + .join(format!("{expired_token}.json")); + assert!(expired_path.exists()); + + store + .create_upload_grant("template", HASH, i64::MAX) + .await + .expect("new grant should be created"); + + assert!( + !expired_path.exists(), + "a grant past its own expiry should be pruned" + ); + } +} diff --git a/src/snapshot/repository/backends/posixfs/mod.rs b/src/snapshot/repository/backends/posixfs/mod.rs index 1afa7803..5964e4f6 100644 --- a/src/snapshot/repository/backends/posixfs/mod.rs +++ b/src/snapshot/repository/backends/posixfs/mod.rs @@ -1,5 +1,6 @@ mod artifacts; mod backend; +mod build_files; mod catalog; mod layout; mod runtime; @@ -7,5 +8,6 @@ mod runtime; pub(crate) use artifacts::PosixFsArtifactStore; pub(crate) use backend::PosixFsSnapshotRepository; pub use backend::{PosixFsBackend, PosixFsBackendConfig}; +pub(crate) use build_files::PosixFsTemplateBuildFileStore; pub(crate) use catalog::PosixFsCatalogStore; pub(crate) use layout::PosixFsSnapshotArtifactLayout; diff --git a/src/snapshot/repository/build_files.rs b/src/snapshot/repository/build_files.rs new file mode 100644 index 00000000..f693af92 --- /dev/null +++ b/src/snapshot/repository/build_files.rs @@ -0,0 +1,197 @@ +use std::path::{Path, PathBuf}; + +use async_trait::async_trait; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use base64::Engine; +use serde::{Deserialize, Serialize}; + +use super::errors::RepositoryResult; + +/// Number of random bytes in an upload bearer token. +pub const UPLOAD_TOKEN_LEN: usize = 32; + +/// Durable authorization record for one build-context upload URL. +/// +/// Grants live in the same shared repository as build archives. That makes a +/// URL issued by one node verifiable by any other node without coordinating a +/// deployment-wide in-memory signing secret. +#[derive(Debug, Deserialize, Serialize)] +pub struct TemplateBuildUploadGrant { + pub template_id: String, + pub hash: String, + pub expires_unix: i64, +} + +impl TemplateBuildUploadGrant { + pub fn new(template_id: &str, hash: &str, expires_unix: i64) -> Self { + Self { + template_id: template_id.to_string(), + hash: hash.to_string(), + expires_unix, + } + } + + pub fn authorizes( + &self, + template_id: &str, + hash: &str, + expires_unix: i64, + now_unix: i64, + ) -> bool { + now_unix <= expires_unix + && self.expires_unix == expires_unix + && self.template_id == template_id + && self.hash == hash + } +} + +/// Durable store for template build-context archives. +/// +/// The E2B SDK resolves every `COPY` step through +/// `GET /templates/{templateID}/files/{hash}` and then `PUT`s a tar archive of +/// the matching context files to the returned URL. This store owns those +/// archives, addressed by the SDK-computed content hash, so that: +/// +/// - any node can answer the upload-link request (`exists`), +/// - any node can accept the upload (`import`), and +/// - the node that runs the build can read the archive back (`materialize`). +/// +/// Implementations must place the archives in storage shared by all nodes of +/// the deployment, mirroring the visibility rules of committed snapshots. +#[async_trait] +pub trait TemplateBuildFileStore: Send + Sync { + /// Returns whether an archive for `hash` is already stored. + async fn exists(&self, hash: &str) -> RepositoryResult; + + /// Imports a fully written local file as the archive for `hash`. + /// + /// Implementations must publish atomically: concurrent readers never + /// observe a partially imported archive. `hash` is the cache key supplied + /// by the authenticated caller, not a digest the store verifies, so + /// immutability here means first-write-wins stability rather than content + /// authenticity: importing a hash that is already stored keeps the stored + /// archive, so an in-flight build can never observe its build context + /// change underneath it. + async fn import(&self, hash: &str, staged: &Path) -> RepositoryResult<()>; + + /// Materializes the archive for `hash` as a node-local file. + /// + /// `scratch_dir` is a caller-owned directory the implementation may use + /// for downloads; implementations backed by a shared filesystem may return + /// the shared path directly. Callers must treat the returned file as + /// read-only. Returns `None` when no archive is stored for `hash`. + async fn materialize( + &self, + hash: &str, + scratch_dir: &Path, + ) -> RepositoryResult>; + + /// Creates a durable bearer grant for one upload URL and returns its + /// URL-safe token. + async fn create_upload_grant( + &self, + template_id: &str, + hash: &str, + expires_unix: i64, + ) -> RepositoryResult; + + /// Verifies a durable bearer grant without consuming it, returning + /// whether it authorizes this upload. + /// + /// Verification never removes the grant, so a request that fails before + /// the archive is stored can be retried with the same upload URL. Callers + /// must `claim_upload_grant` after publishing the archive, so a failed + /// publication leaves the URL retryable. + async fn verify_upload_grant( + &self, + token: &str, + template_id: &str, + hash: &str, + expires_unix: i64, + now_unix: i64, + ) -> RepositoryResult; + + /// Claims a durable bearer grant, returning whether it authorized this + /// upload. + /// + /// Grants are single-use: a successful claim consumes the grant, so an + /// upload URL cannot be replayed within its TTL. Implementations must + /// make the claim itself the atomic step wherever the backend offers an + /// atomic primitive (a POSIX filesystem does, via rename/unlink), so + /// concurrent requests carrying the same token cannot both succeed. + /// S3-compatible backends have no conditional delete and therefore + /// degrade to best-effort single-use within the grant TTL; archive + /// immutability is what keeps a lost race from mattering: both uploads are + /// bound to the same (template_id, hash), and `import` is first-write-wins, + /// so neither can change an archive that is already stored — which upload + /// wins a first store is undefined. + async fn claim_upload_grant( + &self, + token: &str, + template_id: &str, + hash: &str, + expires_unix: i64, + now_unix: i64, + ) -> RepositoryResult; +} + +/// Returns whether `hash` is acceptable as a build-file content hash. +/// +/// The E2B SDK sends a lowercase hex SHA-256, but the value is treated as an +/// opaque cache key; this only enforces a path- and URL-safe shape. +pub fn is_valid_build_files_hash(hash: &str) -> bool { + (16..=128).contains(&hash.len()) && hash.bytes().all(|b| b.is_ascii_hexdigit()) +} + +/// Generates a cryptographically random URL-safe upload bearer token. +pub fn generate_upload_token() -> String { + let mut token = [0u8; UPLOAD_TOKEN_LEN]; + rand::fill(&mut token); + URL_SAFE_NO_PAD.encode(token) +} + +/// Returns whether `token` has the exact shape generated for upload grants. +pub fn is_valid_upload_token(token: &str) -> bool { + URL_SAFE_NO_PAD + .decode(token) + .is_ok_and(|decoded| decoded.len() == UPLOAD_TOKEN_LEN) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hash_validation_accepts_sha256_hex() { + assert!(is_valid_build_files_hash( + "9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08" + )); + assert!(is_valid_build_files_hash("ABCDEF0123456789")); + } + + #[test] + fn hash_validation_rejects_path_unsafe_values() { + assert!(!is_valid_build_files_hash("")); + assert!(!is_valid_build_files_hash("short")); + assert!(!is_valid_build_files_hash("../../../../etc/passwd")); + assert!(!is_valid_build_files_hash("deadbeef/deadbeef")); + assert!(!is_valid_build_files_hash(&"a".repeat(129))); + } + + #[test] + fn upload_token_has_expected_shape() { + let token = generate_upload_token(); + assert!(is_valid_upload_token(&token)); + assert!(!is_valid_upload_token("not-a-valid-token")); + } + + #[test] + fn upload_grant_is_bound_to_request_and_expiry() { + let grant = TemplateBuildUploadGrant::new("tmpl", "aabbccddeeff0011", 1000); + assert!(grant.authorizes("tmpl", "aabbccddeeff0011", 1000, 999)); + assert!(!grant.authorizes("tmpl", "aabbccddeeff0011", 1000, 1001)); + assert!(!grant.authorizes("other", "aabbccddeeff0011", 1000, 999)); + assert!(!grant.authorizes("tmpl", "aabbccddeeff0012", 1000, 999)); + assert!(!grant.authorizes("tmpl", "aabbccddeeff0011", 2000, 999)); + } +} diff --git a/src/snapshot/repository/interfaces.rs b/src/snapshot/repository/interfaces.rs index 16c387f4..604b07a6 100644 --- a/src/snapshot/repository/interfaces.rs +++ b/src/snapshot/repository/interfaces.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use async_trait::async_trait; +use super::build_files::TemplateBuildFileStore; use super::errors::RepositoryResult; use crate::sandbox::FirecrackerSnapshotManifest; use crate::snapshot::types::{ @@ -173,6 +174,15 @@ pub trait SnapshotRepository: Send + Sync { id: &SnapshotId, reason: TemplateBuildErrorReason, ) -> RepositoryResult<()>; + + /// Returns the shared store for template build-context archives. + /// + /// Returns `None` when this backend does not support build-context + /// uploads; the template files API then reports the capability as + /// unavailable instead of failing at build time. + fn template_build_files(&self) -> Option> { + None + } } #[async_trait] diff --git a/src/snapshot/repository/mod.rs b/src/snapshot/repository/mod.rs index 788c94e8..aa45e140 100644 --- a/src/snapshot/repository/mod.rs +++ b/src/snapshot/repository/mod.rs @@ -1,6 +1,8 @@ pub mod backends; +pub mod build_files; pub mod errors; pub mod interfaces; +pub use build_files::TemplateBuildFileStore; pub use errors::{RepositoryError, RepositoryResult}; pub use interfaces::{SnapshotListFilter, SnapshotRepository, SnapshotRuntimeResolver};