Skip to content
Closed
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/snapshot/image_export/service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 })
}
Expand Down Expand Up @@ -447,6 +448,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());
Expand Down
8 changes: 8 additions & 0 deletions src/snapshot/manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Arc<dyn crate::snapshot::repository::TemplateBuildFileStore>> {
self.repository.template_build_files()
}

#[tracing::instrument(skip(self, metadata, manifest), fields(snapshot_id = %metadata.id))]
pub async fn publish(
&self,
Expand Down
169 changes: 169 additions & 0 deletions src/snapshot/repository/backends/oss/build_files.rs
Original file line number Diff line number Diff line change
@@ -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<OssClient>,
}

impl OssTemplateBuildFileStore {
pub(crate) fn new(client: Arc<OssClient>) -> Arc<Self> {
Arc::new(Self { client })
}

fn archive_key(hash: &str) -> RepositoryResult<String> {
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<String> {
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<Option<TemplateBuildUploadGrant>> {
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<bool> {
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))
Comment on lines +75 to +86

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[other · high]
This does not provide the trait's required first-write-wins immutability. OssClient::put_file performs an unconditional upload, so two concurrent imports can both observe the object as absent and the later completed PUT overwrites the first archive. That can change the bytes beneath an in-flight build, and it also makes concurrent replay of an OSS upload grant materially unsafe. Use an atomic create-only OSS operation; if multipart OSS uploads cannot support that, publish each upload to a unique temporary key and use a backend primitive that conditionally establishes the canonical key, or redesign the key/content verification so concurrent writers are guaranteed to upload identical bytes.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alibaba OSS does not support create-only semantics for multipart uploads. This protocol treats uploads sharing the SDK hash as equivalent input, so the contract now guarantees atomic complete-object publication rather than first-write-wins.

}

async fn materialize(
&self,
hash: &str,
scratch_dir: &Path,
) -> RepositoryResult<Option<PathBuf>> {
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<String> {
let token = generate_upload_token();
let key = Self::grant_key(&token).expect("generated token is valid");
Comment on lines +109 to +110

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[other · low]
This expect is on a production request path even though the method returns RepositoryResult. It is currently backed by the implementation invariant that generate_upload_token() always produces a token accepted by is_valid_upload_token, but a future change to either helper would turn a recoverable internal inconsistency into a process panic. Propagate the invariant failure as a contextual repository error instead.

Suggestion:

Suggested change
let token = generate_upload_token();
let key = Self::grant_key(&token).expect("generated token is valid");
let token = generate_upload_token();
let key = Self::grant_key(&token).ok_or_else(|| RepositoryError::backend(
"validate generated upload token",
std::io::Error::new(std::io::ErrorKind::InvalidData, "invalid generated token"),
))?;

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The token is generated internally by generate_upload_token, and generation and validation use the same URL-safe encoding and UPLOAD_TOKEN_LEN; a unit test covers that invariant. This is not request-derived data, so the expect remains an internal invariant rather than a recoverable backend condition.

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))?;
Comment on lines +117 to +120

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[performance · medium]
Expired or abandoned grants are never removed by this implementation. The repository contains no lifecycle policy/provisioning for this prefix, so installations without an externally configured bucket rule will accumulate one durable JSON object per issued upload URL indefinitely. Please provide an in-code cleanup mechanism or make a lifecycle rule for template-build-files/upload-grants/ an explicitly provisioned/validated requirement (and document its retention period).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cleanup remains an operator-managed bucket lifecycle policy because AgentENV does not provision the external OSS bucket. The backend documentation, sample configuration, and configuration reference now explicitly require a seven-day expiration rule for <prefix>/template-build-files/, covering cached archives and abandoned grants.

Ok(token)
}

async fn verify_upload_grant(
&self,
token: &str,
template_id: &str,
hash: &str,
expires_unix: i64,
now_unix: i64,
) -> RepositoryResult<bool> {
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<bool> {
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)
}
}
1 change: 1 addition & 0 deletions src/snapshot/repository/backends/oss/mod.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
mod build_files;
mod client;
mod config;
mod layout;
Expand Down
12 changes: 12 additions & 0 deletions src/snapshot/repository/backends/oss/repository.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ pub(crate) struct OssSnapshotRepository {
client: Arc<OssClient>,
snapshot_image_storage: SnapshotImageStoragePolicy,
acr_exporter: AcrDiskImageExporter,
build_files: Arc<super::build_files::OssTemplateBuildFileStore>,
}

const MAX_ALIAS_BIND_ATTEMPTS: usize = 5;
Expand All @@ -51,10 +52,12 @@ impl OssSnapshotRepository {
client: Arc<OssClient>,
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,
}
}

Expand Down Expand Up @@ -148,6 +151,15 @@ fn fallback_to_object_storage_would_mix_sources(

#[async_trait]
impl SnapshotRepository for OssSnapshotRepository {
fn template_build_files(
&self,
) -> Option<Arc<dyn crate::snapshot::repository::TemplateBuildFileStore>> {
Some(Arc::clone(&self.build_files)
as Arc<
dyn crate::snapshot::repository::TemplateBuildFileStore,
>)
}

async fn create(&self, record: SnapshotRecord) -> RepositoryResult<SnapshotRecord> {
if !matches!(record.source, SnapshotSource::Template { .. }) {
return Err(RepositoryError::InvalidRequest {
Expand Down
12 changes: 12 additions & 0 deletions src/snapshot/repository/backends/posixfs/backend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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<dyn SnapshotRepository> = Arc::new(PosixFsSnapshotRepository::new(
catalog_store,
artifact_store,
build_files,
));
let runtime_resolver: Arc<dyn SnapshotRuntimeResolver> = Arc::new(
PosixFsRuntimeResolver::new(root, runtime_cache_root, store, cache),
Expand Down Expand Up @@ -111,16 +115,19 @@ impl PosixFsBackend {
pub(crate) struct PosixFsSnapshotRepository {
catalog_store: Arc<PosixFsCatalogStore>,
artifact_store: Arc<PosixFsArtifactStore>,
build_files: Arc<PosixFsTemplateBuildFileStore>,
}

impl PosixFsSnapshotRepository {
pub(crate) fn new(
catalog_store: Arc<PosixFsCatalogStore>,
artifact_store: Arc<PosixFsArtifactStore>,
build_files: Arc<PosixFsTemplateBuildFileStore>,
) -> Self {
Self {
catalog_store,
artifact_store,
build_files,
}
}

Expand Down Expand Up @@ -238,6 +245,10 @@ impl SnapshotRepository for PosixFsSnapshotRepository {
.await
}

fn template_build_files(&self) -> Option<Arc<dyn TemplateBuildFileStore>> {
Some(Arc::clone(&self.build_files) as Arc<dyn TemplateBuildFileStore>)
}

async fn publish(
&self,
metadata: SnapshotPublishMetadata,
Expand Down Expand Up @@ -400,6 +411,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),
)
}

Expand Down
Loading
Loading