diff --git a/config/default.toml b/config/default.toml index 15f68640..59bcc20f 100644 --- a/config/default.toml +++ b/config/default.toml @@ -83,6 +83,30 @@ peer_discovery_refresh_interval_secs = 5 # Timeout for custom extension HTTP calls, in milliseconds. # timeout_ms = 5000 +[template_build] +# Build-context upload settings backing the E2B SDK's COPY support +# (GET /templates/{templateID}/files/{hash} plus the returned upload URL). +# Maximum accepted size for one uploaded build-context archive, in MiB. +# files_max_upload_mib = 1024 +# Maximum size one build-context archive may expand to once decompressed, in MiB. +# files_max_context_mib = 4096 +# Cap on the combined on-disk size of all build-context archives one build spec +# may reference, in MiB. +# files_max_build_context_mib = 4096 +# How long an issued upload URL stays valid, in seconds. +# files_url_ttl_secs = 3600 +# How long one build-context upload request may run before the server gives up +# and responds 408, in seconds. +# files_upload_timeout_secs = 300 +# Optional external base URL used when building upload URLs. Defaults to +# "http://{Host header}" of the upload-link request, which matches +# direct-node and bundled-gateway deployments. +# Any TLS-terminated, gateway-fronted, or multi-hop deployment MUST set this to +# the external origin clients reach: the fallback derives the URL from the +# request Host header with plain http, and that upload URL carries a bearer +# token in its query string. +# public_base_url = "https://agentenv.example.com" + [cluster] # Shared gRPC endpoint for cluster-level services such as scheduler heartbeat # reporting and P2P peer discovery (e.g. "http://127.0.0.1:9090"). diff --git a/src/api/build_files.rs b/src/api/build_files.rs new file mode 100644 index 00000000..d911fbe5 --- /dev/null +++ b/src/api/build_files.rs @@ -0,0 +1,287 @@ +//! Hand-written upload endpoint for template build-context archives. +//! +//! `GET /templates/{templateID}/files/{hash}` (generated API) hands the E2B +//! SDK a bearer URL pointing here; the SDK then `PUT`s a tar archive with no +//! authentication headers. The durable random token embedded in the URL is +//! therefore the credential, and this route stays outside the generated +//! router so the archive can stream to disk instead of buffering in memory. + +use std::time::Duration; + +use axum::extract::{Path, Request, State}; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use axum::routing::put; +use axum::{Json, Router}; +use futures::StreamExt; +use tokio::io::AsyncWriteExt; +use tracing::{debug, warn}; + +use agentenv_http_server::models; + +use super::ApiImpl; +use crate::cfg::ConfigManager; +use crate::snapshot::repository::build_files::is_valid_build_files_hash; + +pub(crate) fn router(api_impl: I) -> Router +where + I: AsRef + Clone + Send + Sync + 'static, +{ + Router::new() + .route( + "/templates/{template_id}/files/{hash}/content", + put(upload_build_archive::), + ) + .with_state(api_impl) +} + +struct UploadQuery { + expires: i64, + token: String, +} + +fn parse_upload_query(query: Option<&str>) -> Option { + let query = query?; + let mut expires: Option = None; + let mut token: Option = None; + for (key, value) in url::form_urlencoded::parse(query.as_bytes()) { + match key.as_ref() { + "expires" => expires = value.parse().ok(), + "token" => token = Some(value.into_owned()), + _ => {} + } + } + Some(UploadQuery { + expires: expires?, + token: token?, + }) +} + +fn error_response(code: StatusCode, message: impl Into) -> Response { + ( + code, + Json(models::Error::new(code.as_u16() as i32, message.into())), + ) + .into_response() +} + +async fn upload_build_archive( + State(api_impl): State, + Path((template_id, hash)): Path<(String, String)>, + request: Request, +) -> Response +where + I: AsRef + Clone + Send + Sync + 'static, +{ + let api: &ApiImpl = api_impl.as_ref(); + + if !is_valid_build_files_hash(&hash) { + return error_response( + StatusCode::BAD_REQUEST, + format!("invalid build files hash '{hash}'"), + ); + } + let Some(store) = api.snapshot_manager().template_build_files() else { + return error_response( + StatusCode::BAD_REQUEST, + "the configured snapshot backend does not support build-context uploads", + ); + }; + let Some(query) = parse_upload_query(request.uri().query()) else { + return error_response( + StatusCode::UNAUTHORIZED, + "upload URL is missing the expires/token query parameters", + ); + }; + + // Verification does not consume the grant: consumption happens only after + // the archive has been durably published, so an upload that fails while + // streaming, staging, or storing the body stays retryable with this URL. + let now_unix = chrono::Utc::now().timestamp(); + let authorized = match store + .verify_upload_grant(&query.token, &template_id, &hash, query.expires, now_unix) + .await + { + Ok(authorized) => authorized, + Err(error) => { + warn!(error = %error, "failed to verify build-file upload grant"); + return error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "failed to validate upload grant", + ); + } + }; + if !authorized { + return error_response( + StatusCode::UNAUTHORIZED, + "upload grant is invalid, expired, or already used; request a fresh upload link", + ); + } + + let max_bytes = ConfigManager::global_config() + .template_build + .files_max_upload_mib + .saturating_mul(1024 * 1024); + let upload_timeout = Duration::from_secs( + ConfigManager::global_config() + .template_build + .files_upload_timeout_secs, + ); + + // `staged` is the drop guard that removes the staging file on every early + // return below, so it must stay bound for the rest of the handler. + let staged = match tokio::task::spawn_blocking(tempfile::NamedTempFile::new).await { + Ok(Ok(staged)) => staged, + Ok(Err(error)) => { + warn!(error = %error, "failed to create staging file for build archive"); + return error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "failed to stage build archive", + ); + } + Err(error) => { + warn!(error = %error, "failed to join staging file creation for build archive"); + return error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "failed to stage build archive", + ); + } + }; + let staged_path = staged.path().to_path_buf(); + + let mut file = match tokio::fs::File::create(&staged_path).await { + Ok(file) => file, + Err(error) => { + warn!(error = %error, "failed to open staging file for build archive"); + return error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "failed to stage build archive", + ); + } + }; + + let consume_body = async { + let mut total: u64 = 0; + let mut stream = request.into_body().into_data_stream(); + while let Some(chunk) = stream.next().await { + let chunk = match chunk { + Ok(chunk) => chunk, + Err(error) => { + debug!(error = %error, "build archive upload stream aborted"); + return Err(error_response( + StatusCode::BAD_REQUEST, + "failed to read the uploaded archive body", + )); + } + }; + total += chunk.len() as u64; + if total > max_bytes { + return Err(error_response( + StatusCode::PAYLOAD_TOO_LARGE, + format!("build archive exceeds the configured limit of {max_bytes} bytes"), + )); + } + if let Err(error) = file.write_all(&chunk).await { + warn!(error = %error, "failed to write staged build archive"); + return Err(error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "failed to stage build archive", + )); + } + } + if let Err(error) = file.flush().await { + warn!(error = %error, "failed to flush staged build archive"); + return Err(error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "failed to stage build archive", + )); + } + Ok(total) + }; + + let total = match tokio::time::timeout(upload_timeout, consume_body).await { + Ok(Ok(total)) => total, + Ok(Err(response)) => return response, + Err(_) => { + debug!(template_id, hash, "build archive upload timed out"); + return error_response( + StatusCode::REQUEST_TIMEOUT, + format!( + "build archive upload did not complete within {} seconds", + upload_timeout.as_secs() + ), + ); + } + }; + drop(file); + + // Publishing before the grant is consumed keeps a failed store retryable + // with the same URL. An unclaimed replay reaching this point is harmless: + // the token authorizes exactly this template_id/hash and `import` is + // first-write-wins, so it can neither publish a different key nor change + // what is already stored. + // + // `hash` is the cache key the SDK computed for this build context, not a + // digest of the received bytes that the server verified. + if let Err(error) = store.import(&hash, &staged_path).await { + warn!(error = %error, hash, "failed to import build archive"); + return error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "failed to store build archive; the upload can be retried with the same link", + ); + } + + // The archive is published, so the claim only enforces single-use: the + // atomic remove/delete picks a single winner among concurrent replays, and + // a replay that loses the race is rejected even though the archive it + // uploaded is stored. `now_unix` is the timestamp taken before the body was + // read, so a slow but authorized upload is not rejected for aging past the + // TTL. + let claimed = match store + .claim_upload_grant(&query.token, &template_id, &hash, query.expires, now_unix) + .await + { + Ok(claimed) => claimed, + Err(error) => { + warn!(error = %error, "failed to claim build-file upload grant"); + return error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "failed to validate upload grant", + ); + } + }; + if !claimed { + return error_response( + StatusCode::UNAUTHORIZED, + "upload grant is invalid, expired, or already used; request a fresh upload link", + ); + } + + debug!( + template_id, + hash, + bytes = total, + "stored build-context archive" + ); + StatusCode::OK.into_response() +} + +#[cfg(test)] +mod tests { + use super::parse_upload_query; + + #[test] + fn upload_query_parses_bearer_token_and_expiry() { + let query = parse_upload_query(Some("expires=1234&token=upload-token")) + .expect("query should parse"); + assert_eq!(query.expires, 1234); + assert_eq!(query.token, "upload-token"); + } + + #[test] + fn upload_query_requires_both_fields() { + assert!(parse_upload_query(Some("expires=1234")).is_none()); + assert!(parse_upload_query(Some("token=upload-token")).is_none()); + assert!(parse_upload_query(None).is_none()); + } +} diff --git a/src/api/generated/src/apis/templates.rs b/src/api/generated/src/apis/templates.rs index a25ac144..fc727885 100644 --- a/src/api/generated/src/apis/templates.rs +++ b/src/api/generated/src/apis/templates.rs @@ -64,6 +64,22 @@ pub enum TemplatesTemplateIdDeleteResponse { Status500_ServerError(models::Error), } +#[derive(Debug, PartialEq, Serialize, Deserialize)] +#[must_use] +#[allow(clippy::large_enum_variant)] +pub enum TemplatesTemplateIdFilesHashGetResponse { + /// Successfully returned the upload link + Status201_SuccessfullyReturnedTheUploadLink(models::TemplateBuildFileUpload), + /// Bad request + Status400_BadRequest(models::Error), + /// Authentication error + Status401_AuthenticationError(models::Error), + /// Not found + Status404_NotFound(models::Error), + /// Server error + Status500_ServerError(models::Error), +} + #[derive(Debug, PartialEq, Serialize, Deserialize)] #[must_use] #[allow(clippy::large_enum_variant)] @@ -187,6 +203,19 @@ pub trait Templates: path_params: &models::TemplatesTemplateIdDeletePathParams, ) -> Result; + /// Template build file upload link. + /// + /// TemplatesTemplateIdFilesHashGet - GET /templates/{templateID}/files/{hash} + async fn templates_template_id_files_hash_get( + &self, + + method: &Method, + host: &Host, + cookies: &CookieJar, + claims: &Self::Claims, + path_params: &models::TemplatesTemplateIdFilesHashGetPathParams, + ) -> Result; + /// List template builds. /// /// TemplatesTemplateIdGet - GET /templates/{templateID} diff --git a/src/api/generated/src/models.rs b/src/api/generated/src/models.rs index 76c8704b..7457339a 100644 --- a/src/api/generated/src/models.rs +++ b/src/api/generated/src/models.rs @@ -249,6 +249,13 @@ pub struct TemplatesTemplateIdDeletePathParams { pub template_id: String, } +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, validator::Validate)] +#[cfg_attr(feature = "conversion", derive(frunk::LabelledGeneric))] +pub struct TemplatesTemplateIdFilesHashGetPathParams { + pub template_id: String, + pub hash: String, +} + #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, validator::Validate)] #[cfg_attr(feature = "conversion", derive(frunk::LabelledGeneric))] pub struct TemplatesTemplateIdGetPathParams { @@ -7665,6 +7672,157 @@ impl std::convert::TryFrom for header::IntoHeaderValue, +} + +impl TemplateBuildFileUpload { + #[allow(clippy::new_without_default, clippy::too_many_arguments)] + pub fn new(present: bool) -> TemplateBuildFileUpload { + TemplateBuildFileUpload { present, url: None } + } +} + +/// Converts the TemplateBuildFileUpload value to the Query Parameters representation (style=form, explode=false) +/// specified in https://swagger.io/docs/specification/serialization/ +/// Should be implemented in a serde serializer +impl std::fmt::Display for TemplateBuildFileUpload { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let params: Vec> = vec![ + Some("present".to_string()), + Some(self.present.to_string()), + self.url + .as_ref() + .map(|url| ["url".to_string(), url.to_string()].join(",")), + ]; + + write!( + f, + "{}", + params.into_iter().flatten().collect::>().join(",") + ) + } +} + +/// Converts Query Parameters representation (style=form, explode=false) to a TemplateBuildFileUpload value +/// as specified in https://swagger.io/docs/specification/serialization/ +/// Should be implemented in a serde deserializer +impl std::str::FromStr for TemplateBuildFileUpload { + type Err = String; + + fn from_str(s: &str) -> std::result::Result { + /// An intermediate representation of the struct to use for parsing. + #[derive(Default)] + #[allow(dead_code)] + struct IntermediateRep { + pub present: Vec, + pub url: Vec, + } + + let mut intermediate_rep = IntermediateRep::default(); + + // Parse into intermediate representation + let mut string_iter = s.split(','); + let mut key_result = string_iter.next(); + + while key_result.is_some() { + let val = match string_iter.next() { + Some(x) => x, + None => { + return std::result::Result::Err( + "Missing value while parsing TemplateBuildFileUpload".to_string(), + ); + } + }; + + if let Some(key) = key_result { + #[allow(clippy::match_single_binding)] + match key { + #[allow(clippy::redundant_clone)] + "present" => intermediate_rep.present.push( + ::from_str(val).map_err(|x| x.to_string())?, + ), + #[allow(clippy::redundant_clone)] + "url" => intermediate_rep.url.push( + ::from_str(val).map_err(|x| x.to_string())?, + ), + _ => { + return std::result::Result::Err( + "Unexpected key while parsing TemplateBuildFileUpload".to_string(), + ); + } + } + } + + // Get the next key + key_result = string_iter.next(); + } + + // Use the intermediate representation to return the struct + std::result::Result::Ok(TemplateBuildFileUpload { + present: intermediate_rep + .present + .into_iter() + .next() + .ok_or_else(|| "present missing in TemplateBuildFileUpload".to_string())?, + url: intermediate_rep.url.into_iter().next(), + }) + } +} + +// Methods for converting between header::IntoHeaderValue and HeaderValue + +#[cfg(feature = "server")] +impl std::convert::TryFrom> for HeaderValue { + type Error = String; + + fn try_from( + hdr_value: header::IntoHeaderValue, + ) -> std::result::Result { + let hdr_value = hdr_value.to_string(); + match HeaderValue::from_str(&hdr_value) { + std::result::Result::Ok(value) => std::result::Result::Ok(value), + std::result::Result::Err(e) => std::result::Result::Err(format!( + r#"Invalid header value for TemplateBuildFileUpload - value: {hdr_value} is invalid {e}"# + )), + } + } +} + +#[cfg(feature = "server")] +impl std::convert::TryFrom for header::IntoHeaderValue { + type Error = String; + + fn try_from(hdr_value: HeaderValue) -> std::result::Result { + match hdr_value.to_str() { + std::result::Result::Ok(value) => { + match ::from_str(value) { + std::result::Result::Ok(value) => { + std::result::Result::Ok(header::IntoHeaderValue(value)) + } + std::result::Result::Err(err) => std::result::Result::Err(format!( + r#"Unable to convert header value '{value}' into TemplateBuildFileUpload - {err}"# + )), + } + } + std::result::Result::Err(e) => std::result::Result::Err(format!( + r#"Unable to convert header: {hdr_value:?} to string: {e}"# + )), + } + } +} + #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, validator::Validate)] #[cfg_attr(feature = "conversion", derive(frunk::LabelledGeneric))] pub struct TemplateBuildInfo { diff --git a/src/api/generated/src/server/mod.rs b/src/api/generated/src/server/mod.rs index 766300a6..e5439fbf 100644 --- a/src/api/generated/src/server/mod.rs +++ b/src/api/generated/src/server/mod.rs @@ -111,6 +111,10 @@ where "/templates/{template_id}/builds/{build_id}/status", get(templates_template_id_builds_build_id_status_get::), ) + .route( + "/templates/{template_id}/files/{hash}", + get(templates_template_id_files_hash_get::), + ) .route("/v2/sandboxes", get(v2_sandboxes_get::)) .route("/v2/templates", get(v2_templates_get::)) .route( @@ -4160,6 +4164,174 @@ where }) } +#[tracing::instrument(skip_all)] +fn templates_template_id_files_hash_get_validation( + path_params: models::TemplatesTemplateIdFilesHashGetPathParams, +) -> std::result::Result<(models::TemplatesTemplateIdFilesHashGetPathParams,), ValidationErrors> { + path_params.validate()?; + + Ok((path_params,)) +} +/// TemplatesTemplateIdFilesHashGet - GET /templates/{templateID}/files/{hash} +#[tracing::instrument(skip_all)] +async fn templates_template_id_files_hash_get( + method: Method, + TypedHeader(host): TypedHeader, + cookies: CookieJar, + headers: HeaderMap, + Path(path_params): Path, + State(api_impl): State, +) -> Result +where + I: AsRef + Send + Sync, + A: apis::templates::Templates + + apis::ApiKeyAuthHeader + + apis::ApiAuthBasic + + Send + + Sync, + E: std::fmt::Debug + Send + Sync + 'static, +{ + // Authentication + let claims_in_header = api_impl + .as_ref() + .extract_claims_from_header(&headers, "X-Team-ID") + .await; + let claims_in_auth_header = api_impl + .as_ref() + .extract_claims_from_auth_header(apis::BasicAuthKind::Bearer, &headers, "authorization") + .await; + let claims = None.or(claims_in_header).or(claims_in_auth_header); + let Some(claims) = claims else { + return response_with_status_code_only(StatusCode::UNAUTHORIZED); + }; + + #[allow(clippy::redundant_closure)] + let validation = tokio::task::spawn_blocking(move || { + templates_template_id_files_hash_get_validation(path_params) + }) + .await + .unwrap(); + + let Ok((path_params,)) = validation else { + return Response::builder() + .status(StatusCode::BAD_REQUEST) + .body(Body::from(validation.unwrap_err().to_string())) + .map_err(|_| StatusCode::BAD_REQUEST); + }; + + let result = api_impl + .as_ref() + .templates_template_id_files_hash_get(&method, &host, &cookies, &claims, &path_params) + .await; + + let mut response = Response::builder(); + + let resp = match result { + Ok(rsp) => match rsp { + apis::templates::TemplatesTemplateIdFilesHashGetResponse::Status201_SuccessfullyReturnedTheUploadLink + (body) + => { + let mut response = response.status(201); + { + let mut response_headers = response.headers_mut().unwrap(); + response_headers.insert( + CONTENT_TYPE, + HeaderValue::from_static("application/json")); + } + + let body_content = tokio::task::spawn_blocking(move || + serde_json::to_vec(&body).map_err(|e| { + error!(error = ?e); + StatusCode::INTERNAL_SERVER_ERROR + })).await.unwrap()?; + response.body(Body::from(body_content)) + }, + apis::templates::TemplatesTemplateIdFilesHashGetResponse::Status400_BadRequest + (body) + => { + let mut response = response.status(400); + { + let mut response_headers = response.headers_mut().unwrap(); + response_headers.insert( + CONTENT_TYPE, + HeaderValue::from_static("application/json")); + } + + let body_content = tokio::task::spawn_blocking(move || + serde_json::to_vec(&body).map_err(|e| { + error!(error = ?e); + StatusCode::INTERNAL_SERVER_ERROR + })).await.unwrap()?; + response.body(Body::from(body_content)) + }, + apis::templates::TemplatesTemplateIdFilesHashGetResponse::Status401_AuthenticationError + (body) + => { + let mut response = response.status(401); + { + let mut response_headers = response.headers_mut().unwrap(); + response_headers.insert( + CONTENT_TYPE, + HeaderValue::from_static("application/json")); + } + + let body_content = tokio::task::spawn_blocking(move || + serde_json::to_vec(&body).map_err(|e| { + error!(error = ?e); + StatusCode::INTERNAL_SERVER_ERROR + })).await.unwrap()?; + response.body(Body::from(body_content)) + }, + apis::templates::TemplatesTemplateIdFilesHashGetResponse::Status404_NotFound + (body) + => { + let mut response = response.status(404); + { + let mut response_headers = response.headers_mut().unwrap(); + response_headers.insert( + CONTENT_TYPE, + HeaderValue::from_static("application/json")); + } + + let body_content = tokio::task::spawn_blocking(move || + serde_json::to_vec(&body).map_err(|e| { + error!(error = ?e); + StatusCode::INTERNAL_SERVER_ERROR + })).await.unwrap()?; + response.body(Body::from(body_content)) + }, + apis::templates::TemplatesTemplateIdFilesHashGetResponse::Status500_ServerError + (body) + => { + let mut response = response.status(500); + { + let mut response_headers = response.headers_mut().unwrap(); + response_headers.insert( + CONTENT_TYPE, + HeaderValue::from_static("application/json")); + } + + let body_content = tokio::task::spawn_blocking(move || + serde_json::to_vec(&body).map_err(|e| { + error!(error = ?e); + StatusCode::INTERNAL_SERVER_ERROR + })).await.unwrap()?; + response.body(Body::from(body_content)) + }, + }, + Err(why) => { + // Application code returned an error. This should not happen, as the implementation should + // return a valid response. + return api_impl.as_ref().handle_error(&method, &host, &cookies, why).await; + }, + }; + + resp.map_err(|e| { + error!(error = ?e); + StatusCode::INTERNAL_SERVER_ERROR + }) +} + #[tracing::instrument(skip_all)] fn templates_template_id_get_validation( path_params: models::TemplatesTemplateIdGetPathParams, diff --git a/src/api/impls/mod.rs b/src/api/impls/mod.rs index 1e1d50bc..454a9d29 100644 --- a/src/api/impls/mod.rs +++ b/src/api/impls/mod.rs @@ -59,6 +59,10 @@ impl ApiImpl { Arc::clone(&self.orchestrator) } + pub(crate) fn snapshot_manager(&self) -> &SnapshotManager { + &self.snapshot_manager + } + pub(crate) fn proxy_client(&self) -> &ProxyClient { &self.proxy_client } diff --git a/src/api/impls/template.rs b/src/api/impls/template.rs index bf88f6a1..3c8727b5 100644 --- a/src/api/impls/template.rs +++ b/src/api/impls/template.rs @@ -17,6 +17,7 @@ use super::template_helpers::{ }; use super::ApiImpl; use crate::image::ResolvedBlockImage; +use crate::snapshot::repository::build_files::is_valid_build_files_hash; use crate::snapshot::{ CommandContext, SnapshotAlias, SnapshotId, SnapshotListFilter, SnapshotRecord, SnapshotSource, TemplateBuildErrorReason, TemplateBuildStatus, @@ -319,6 +320,100 @@ impl Templates<()> for ApiImpl { } } + async fn templates_template_id_files_hash_get( + &self, + _method: &Method, + host: &Host, + _cookies: &CookieJar, + _claims: &Self::Claims, + path_params: &models::TemplatesTemplateIdFilesHashGetPathParams, + ) -> Result { + let template_id = &path_params.template_id; + let hash = &path_params.hash; + + if !is_valid_build_files_hash(hash) { + return Ok( + TemplatesTemplateIdFilesHashGetResponse::Status400_BadRequest(Self::error( + 400, + format!("invalid build files hash '{hash}'"), + )), + ); + } + let Some(store) = self.snapshot_manager.template_build_files() else { + return Ok( + TemplatesTemplateIdFilesHashGetResponse::Status400_BadRequest(Self::error( + 400, + "the configured snapshot backend does not support build-context uploads", + )), + ); + }; + + match self.snapshot_manager.get(template_id).await { + Ok(Some(_)) => {} + Ok(None) => { + return Ok(TemplatesTemplateIdFilesHashGetResponse::Status404_NotFound( + Self::error(404, format!("template {template_id} not found")), + )); + } + Err(err) => { + return Ok( + TemplatesTemplateIdFilesHashGetResponse::Status500_ServerError( + Self::snapshot_manager_error(&err), + ), + ); + } + } + + let present = match store.exists(hash).await { + Ok(present) => present, + Err(err) => { + return Ok( + TemplatesTemplateIdFilesHashGetResponse::Status500_ServerError(Self::error( + 500, + format!("failed to check build archive: {err}"), + )), + ); + } + }; + let config = &crate::cfg::ConfigManager::global_config().template_build; + let expires = chrono::Utc::now() + .timestamp() + .saturating_add(i64::try_from(config.files_url_ttl_secs).unwrap_or(i64::MAX)); + let token = match store.create_upload_grant(template_id, hash, expires).await { + Ok(token) => token, + Err(err) => { + return Ok( + TemplatesTemplateIdFilesHashGetResponse::Status500_ServerError(Self::error( + 500, + format!("failed to prepare upload link: {err}"), + )), + ); + } + }; + + // The SDK PUTs to this URL with a bare HTTP client (no auth headers), + // so the durable bearer token in the query string is the credential. + // Reusing the Host header keeps the URL valid across gateway and + // direct-node access. + let base = config + .public_base_url + .clone() + .unwrap_or_else(|| format!("http://{host}")); + let url = format!( + "{}/templates/{template_id}/files/{hash}/content?expires={expires}&token={token}", + base.trim_end_matches('/'), + ); + + Ok( + TemplatesTemplateIdFilesHashGetResponse::Status201_SuccessfullyReturnedTheUploadLink( + models::TemplateBuildFileUpload { + present, + url: Some(url), + }, + ), + ) + } + async fn templates_get( &self, _method: &Method, diff --git a/src/api/mod.rs b/src/api/mod.rs index 6cae7863..86bb400a 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -1,3 +1,4 @@ +mod build_files; mod impls; mod proxy; pub mod server; diff --git a/src/api/openapi.yml b/src/api/openapi.yml index 3b146348..8f60e802 100644 --- a/src/api/openapi.yml +++ b/src/api/openapi.yml @@ -904,6 +904,18 @@ components: type: boolean description: Whether the step should be forced to run regardless of the cache + TemplateBuildFileUpload: + description: Upload link for one build context archive, addressed by its files hash + required: + - present + properties: + present: + type: boolean + description: Whether the archive for this hash is already stored + url: + type: string + description: URL the client should PUT the tar archive to + TemplateBuildRequestV3: properties: name: @@ -2147,6 +2159,42 @@ paths: "500": $ref: "#/components/responses/500" + /templates/{templateID}/files/{hash}: + get: + summary: Template build file upload link + description: Get an upload link for a tar archive containing build context files for one COPY step + tags: [templates] + security: + - AccessTokenAuth: [] + - ApiKeyAuth: [] + - AuthProviderBearerAuth: [] + AuthProviderTeamAuth: [] + - AdminApiKeyAuth: [] + AdminTeamAuth: [] + parameters: + - $ref: "#/components/parameters/templateID" + - in: path + name: hash + required: true + schema: + type: string + description: Hash of the build context files + responses: + "201": + description: Successfully returned the upload link + content: + application/json: + schema: + $ref: "#/components/schemas/TemplateBuildFileUpload" + "400": + $ref: "#/components/responses/400" + "401": + $ref: "#/components/responses/401" + "404": + $ref: "#/components/responses/404" + "500": + $ref: "#/components/responses/500" + /templates/aliases/{alias}: get: summary: Check template alias diff --git a/src/api/server.rs b/src/api/server.rs index 2a738898..a1a39016 100644 --- a/src/api/server.rs +++ b/src/api/server.rs @@ -1,6 +1,6 @@ use axum::{middleware, routing::get, Router}; -use super::{proxy, ApiImpl}; +use super::{build_files, proxy, ApiImpl}; use crate::observability::prometheus; use agentenv_http_server::apis; use agentenv_observability::metrics_handler; @@ -23,9 +23,10 @@ where { // Keep the generated control-plane API as the primary router, then merge in // the hand-written `/proxy/*` entrypoints needed for the temporary reverse - // proxy contract. + // proxy contract and the streaming build-context upload endpoint. agentenv_http_server::server::new::(api_impl.clone()) .merge(proxy::router(api_impl.clone())) + .merge(build_files::router(api_impl.clone())) .route("/metrics", get(metrics_handler)) .layer(middleware::from_fn_with_state( api_impl, diff --git a/src/cfg.rs b/src/cfg.rs index 0ac185b4..1a07bd0e 100644 --- a/src/cfg.rs +++ b/src/cfg.rs @@ -132,6 +132,8 @@ pub struct AppConfig { pub network: NetworkConfig, #[config(nested)] pub custom_extension: CustomExtensionConfig, + #[config(nested)] + pub template_build: TemplateBuildConfig, } #[derive(Debug, Deserialize, Clone, Config)] @@ -544,6 +546,43 @@ pub struct CustomExtensionConfig { pub timeout_ms: u64, } +/// Settings for the template build-context upload path used by the E2B SDK's +/// `COPY` support (`GET /templates/{templateID}/files/{hash}` plus the upload +/// URL it returns). +#[derive(Debug, Config, Clone)] +pub struct TemplateBuildConfig { + /// Maximum accepted size for one uploaded build-context archive, in MiB. + #[config(default = 1024u64)] + pub files_max_upload_mib: u64, + /// Maximum size one build-context archive may expand to once + /// decompressed, in MiB. This bounds what a compressed upload can cost + /// the node that runs the build. + #[config(default = 4096u64)] + pub files_max_context_mib: u64, + /// Cap on the combined on-disk size of all build-context archives one + /// build spec may reference, in MiB. + #[config(default = 4096u64)] + pub files_max_build_context_mib: u64, + /// How long an issued upload URL stays valid, in seconds. + #[config(default = 3600u64)] + pub files_url_ttl_secs: u64, + /// How long one build-context upload request may run before the server + /// gives up and responds 408, in seconds. + #[config(default = 300u64)] + pub files_upload_timeout_secs: u64, + /// Optional external base URL (e.g. "https://agentenv.example.com") used + /// when building upload URLs. When unset, upload URLs reuse the Host + /// header of the upload-link request with plain http, which matches + /// direct-node and bundled-gateway deployments. + /// + /// Any TLS-terminated, gateway-fronted, or multi-hop deployment MUST set + /// this to the external origin clients reach: the fallback derives the URL + /// from the request Host header with plain http, and that upload URL + /// carries a bearer token in its query string. + #[config(env = "AENV_TEMPLATE_BUILD_PUBLIC_BASE_URL", parse_env = parse_trimmed_string)] + pub public_base_url: Option, +} + #[derive(Debug, Config, Clone)] pub struct P2pConfig { #[config(default = false)] @@ -858,6 +897,17 @@ impl AppConfig { self.cluster.normalize(); self.sandbox_proxy.normalize()?; + // An env var exported empty means unset, matching the custom-extension + // URL handling; validation and the upload-URL builder then agree on + // the exact value in use. + self.template_build.public_base_url = self + .template_build + .public_base_url + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_owned); + Ok(()) } @@ -888,6 +938,7 @@ impl AppConfig { self.validate_memory_snapshot_background_download()?; self.validate_overlaybd_global_config_paths()?; self.validate_disk_rate_limit()?; + self.validate_template_build()?; Ok(()) } @@ -955,6 +1006,69 @@ impl AppConfig { Ok(()) } + /// Reject template build-context settings that would only fail later, at + /// upload-link time: a base URL that cannot produce a usable upload URL, + /// or a TTL that makes the `now + ttl` expiry arithmetic overflow or the + /// grant effectively unexpirable. + fn validate_template_build(&self) -> Result<()> { + // 7 days. Upload grants are single-use credentials in a query string, + // so a longer window is always a misconfiguration. + const MAX_URL_TTL_SECS: u64 = 604_800; + let cfg = &self.template_build; + if cfg.files_url_ttl_secs == 0 { + bail!("template_build.files_url_ttl_secs must be > 0"); + } + if cfg.files_url_ttl_secs > MAX_URL_TTL_SECS { + bail!( + "template_build.files_url_ttl_secs must be <= {MAX_URL_TTL_SECS} (got {})", + cfg.files_url_ttl_secs + ); + } + if cfg.files_upload_timeout_secs == 0 { + bail!("template_build.files_upload_timeout_secs must be > 0"); + } + // An upload slower than the grant TTL would stage the whole body and + // then lose the grant to expiry-based pruning at claim time. + if cfg.files_upload_timeout_secs > cfg.files_url_ttl_secs { + bail!( + "template_build.files_upload_timeout_secs ({}) must be <= \ + files_url_ttl_secs ({})", + cfg.files_upload_timeout_secs, + cfg.files_url_ttl_secs + ); + } + if cfg.files_max_upload_mib == 0 { + bail!("template_build.files_max_upload_mib must be > 0"); + } + if cfg.files_max_context_mib == 0 { + bail!("template_build.files_max_context_mib must be > 0"); + } + if cfg.files_max_build_context_mib == 0 { + bail!("template_build.files_max_build_context_mib must be > 0"); + } + if let Some(base_url) = cfg.public_base_url.as_deref() { + let parsed = url::Url::parse(base_url).with_context(|| { + format!( + "invalid template_build.public_base_url {base_url:?}: must be an absolute \ + http/https URL" + ) + })?; + if !matches!(parsed.scheme(), "http" | "https") { + bail!( + "invalid template_build.public_base_url {base_url:?}: scheme must be http or \ + https" + ); + } + if parsed.query().is_some() { + bail!("invalid template_build.public_base_url {base_url:?}: must have no query"); + } + if parsed.fragment().is_some() { + bail!("invalid template_build.public_base_url {base_url:?}: must have no fragment"); + } + } + Ok(()) + } + /// Sanity-bound the memory-snapshot background download knobs so a legal /// config cannot allocate unbounded scratch or fan out unbounded requests. /// Peak scratch per active layer download is `block_size × concurrency` @@ -1568,6 +1682,115 @@ mod tests { assert!(config.validate().is_err()); } + #[test] + fn template_build_defaults_pass_validation() { + let config = AppConfig::default(); + assert_eq!(config.template_build.files_url_ttl_secs, 3600); + assert_eq!(config.template_build.files_upload_timeout_secs, 300); + assert_eq!(config.template_build.files_max_build_context_mib, 4096); + assert!(config.template_build.public_base_url.is_none()); + config.validate().expect("default config passes"); + } + + #[test] + fn validate_accepts_absolute_template_build_public_base_url() { + let mut config = AppConfig::default(); + config.template_build.public_base_url = Some("https://agentenv.example.com".to_string()); + + config.validate().expect("https base url passes"); + } + + #[test] + fn validate_rejects_template_build_public_base_url_with_query() { + let mut config = AppConfig::default(); + config.template_build.public_base_url = + Some("https://agentenv.example.com/?token=abc".to_string()); + + let err = config.validate().unwrap_err(); + let message = err.to_string(); + assert!(message.contains("public_base_url"), "{message}"); + assert!(message.contains("must have no query"), "{message}"); + } + + #[test] + fn validate_rejects_template_build_public_base_url_without_scheme() { + let mut config = AppConfig::default(); + config.template_build.public_base_url = Some("agentenv.example.com".to_string()); + + let err = config.validate().unwrap_err(); + assert!( + err.to_string().contains("public_base_url"), + "unexpected error: {err}" + ); + } + + #[test] + fn validate_bounds_template_build_files_url_ttl() { + let mut config = AppConfig::default(); + config.template_build.files_url_ttl_secs = 0; + let err = config.validate().unwrap_err(); + assert!( + err.to_string() + .contains("template_build.files_url_ttl_secs must be > 0"), + "unexpected error: {err}" + ); + + let mut config = AppConfig::default(); + config.template_build.files_url_ttl_secs = 604_801; + assert!(config.validate().is_err()); + + let mut config = AppConfig::default(); + config.template_build.files_url_ttl_secs = 604_800; + config.validate().expect("max ttl passes"); + } + + #[test] + fn validate_bounds_template_build_upload_timeout() { + let mut config = AppConfig::default(); + config.template_build.files_upload_timeout_secs = 0; + let err = config.validate().unwrap_err(); + assert!( + err.to_string() + .contains("template_build.files_upload_timeout_secs must be > 0"), + "unexpected error: {err}" + ); + + let mut config = AppConfig::default(); + config.template_build.files_upload_timeout_secs = 3601; + let err = config.validate().unwrap_err(); + assert!( + err.to_string().contains("must be <= files_url_ttl_secs"), + "unexpected error: {err}" + ); + + let mut config = AppConfig::default(); + config.template_build.files_max_upload_mib = 0; + let err = config.validate().unwrap_err(); + assert!( + err.to_string() + .contains("template_build.files_max_upload_mib must be > 0"), + "unexpected error: {err}" + ); + + let mut config = AppConfig::default(); + config.template_build.files_max_context_mib = 0; + let err = config.validate().unwrap_err(); + assert!( + err.to_string() + .contains("template_build.files_max_context_mib must be > 0"), + "unexpected error: {err}" + ); + + let mut config = AppConfig::default(); + config.template_build.files_max_build_context_mib = 0; + let err = config.validate().unwrap_err(); + assert!( + err.to_string() + .contains("template_build.files_max_build_context_mib must be > 0"), + "unexpected error: {err}" + ); + } + #[test] fn overlaybd_converter_cache_version_includes_tool_version() { let config = AppConfig::default(); diff --git a/src/snapshot/image_export/service.rs b/src/snapshot/image_export/service.rs index a94af4ff..473dca70 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 }) } @@ -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()); 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 0af7465f..c7047881 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 769c7440..b00a271d 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, } } @@ -238,6 +245,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, @@ -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), ) } 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};