use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::sync::Mutex;
use arcbox_connect::sandbox_v1;
use arcbox_vm::template_catalog::{TemplateDefaultsSpec, TemplateEntry, compute_digest};
use buffa::Message;
use super::{SandboxService, convert, template};
use crate::error::SandboxError;
fn operation_key(name: &str) -> String {
format!("template:{name}")
}
fn dockerfile_tag(dockerfile: &[u8]) -> String {
use sha2::{Digest as _, Sha256};
let hex = format!("{:x}", Sha256::digest(dockerfile));
format!("arcbox-template-build:{}", &hex[..16])
}
struct BuildFlight<'a> {
name: String,
builds: &'a Mutex<HashSet<String>>,
}
impl Drop for BuildFlight<'_> {
fn drop(&mut self) {
self.builds.lock().unwrap().remove(&self.name);
}
}
impl SandboxService {
pub async fn build_template(
&self,
payload: &[u8],
) -> Result<sandbox_v1::Template, SandboxError> {
use sandbox_v1::build_template_request::Source;
let req = sandbox_v1::BuildTemplateRequest::decode_from_slice(payload)
.map_err(|e| SandboxError::Decode(e.to_string()))?;
arcbox_vm::template_catalog::validate_template_name(&req.name)
.map_err(SandboxError::from)?;
let defaults = convert::template_defaults_from_proto(&req.defaults)?;
let source = req.source.clone().ok_or_else(|| {
SandboxError::InvalidArgument(
"build source is required: docker_ref, dockerfile, or snapshot_id".into(),
)
})?;
if req.prewarm && !matches!(source, Source::SnapshotId(_)) {
return Err(SandboxError::Unsupported(
"prewarm is not implemented yet (CORE-107)".into(),
));
}
let _flight = self.begin_template_build(&req.name)?;
let labels: HashMap<String, String> = req.labels.clone().into_iter().collect();
match source {
Source::DockerRef(image) => {
self.build_template_from_docker(&req.name, &image, defaults, labels)
.await
}
Source::Dockerfile(dockerfile) => {
self.build_template_from_dockerfile(&req.name, &dockerfile, defaults, labels)
.await
}
Source::SnapshotId(_) => Err(SandboxError::Unsupported(
"snapshot promotion is not implemented yet (CORE-107)".into(),
)),
}
}
async fn build_template_from_dockerfile(
&self,
name: &str,
dockerfile: &str,
defaults: TemplateDefaultsSpec,
labels: HashMap<String, String>,
) -> Result<sandbox_v1::Template, SandboxError> {
if dockerfile.trim().is_empty() {
return Err(SandboxError::InvalidArgument(
"dockerfile content must not be empty".into(),
));
}
let tag = dockerfile_tag(dockerfile.as_bytes());
if template::inspect_image(&tag).await.is_ok() {
tracing::info!(template = name, %tag, "reusing previously built dockerfile image");
} else {
template::build_image(dockerfile.as_bytes(), &tag)
.await
.map_err(|e| SandboxError::Internal(format!("template build {name}: {e:#}")))?;
}
self.build_template_from_docker(name, &tag, defaults, labels)
.await
}
async fn build_template_from_docker(
&self,
name: &str,
image: &str,
defaults: TemplateDefaultsSpec,
labels: HashMap<String, String>,
) -> Result<sandbox_v1::Template, SandboxError> {
if image.trim().is_empty() {
return Err(SandboxError::InvalidArgument(
"docker_ref must name an image".into(),
));
}
let layout = template::export_docker_image(image)
.await
.map_err(|e| SandboxError::Internal(format!("template build {image}: {e:#}")))?;
let pinned = self
.manager
.pinned_rootfs_paths()
.map_err(SandboxError::from)?;
let rootfs = crate::rootfs_builder::convert_layer_to_rootfs(&layout, &pinned)
.await
.map_err(|e| SandboxError::Internal(format!("template build {image}: {e:#}")))?;
let source_identity = Path::new(&rootfs)
.file_stem()
.and_then(|stem| stem.to_str())
.ok_or_else(|| {
SandboxError::Internal(format!("unparsable rootfs cache path {rootfs}"))
})?
.to_owned();
let size_bytes = tokio::fs::metadata(&rootfs)
.await
.map_err(|e| SandboxError::Internal(format!("stat {rootfs}: {e}")))?
.len();
let digest = compute_digest(&source_identity, None, &defaults, &labels);
let entry = TemplateEntry {
version: String::new(),
digest,
rootfs_path: rootfs,
warm: None,
defaults,
labels,
created_at: chrono::Utc::now(),
size_bytes,
};
let _operation = self.operations.lock(&operation_key(name)).await;
let entry = self
.manager
.register_template_draft(name, entry)
.await
.map_err(SandboxError::from)?;
Ok(convert::template_to_proto(name, &entry))
}
fn begin_template_build(&self, name: &str) -> Result<BuildFlight<'_>, SandboxError> {
let mut builds = self.template_builds.lock().unwrap();
if !builds.insert(name.to_owned()) {
return Err(SandboxError::WrongState(format!(
"a build for template {name} is already running"
)));
}
Ok(BuildFlight {
name: name.to_owned(),
builds: &self.template_builds,
})
}
pub fn get_template(&self, payload: &[u8]) -> Result<sandbox_v1::Template, SandboxError> {
let req = sandbox_v1::GetTemplateRequest::decode_from_slice(payload)
.map_err(|e| SandboxError::Decode(e.to_string()))?;
let resolved = self
.manager
.get_template(&req.reference)
.map_err(SandboxError::from)?;
Ok(convert::template_to_proto(&resolved.name, &resolved.entry))
}
pub fn list_templates(
&self,
payload: &[u8],
) -> Result<sandbox_v1::ListTemplatesResponse, SandboxError> {
let req = sandbox_v1::ListTemplatesRequest::decode_from_slice(payload)
.map_err(|e| SandboxError::Decode(e.to_string()))?;
let mut rows = self.manager.list_templates().map_err(SandboxError::from)?;
if !req.labels.is_empty() {
rows.retain(|(_, entry)| {
req.labels
.iter()
.all(|(key, value)| entry.labels.get(key) == Some(value))
});
}
let rows: Vec<(String, sandbox_v1::Template)> = rows
.into_iter()
.map(|(name, entry)| {
let reference = format!("{name}:{}", entry.version);
(reference, convert::template_to_proto(&name, &entry))
})
.collect();
let (page, next_page_token) =
convert::paginate(rows, |row| row.0.as_str(), req.page_size, &req.page_token);
Ok(sandbox_v1::ListTemplatesResponse {
templates: page.into_iter().map(|row| row.1).collect(),
next_page_token,
..Default::default()
})
}
pub async fn publish_template(
&self,
payload: &[u8],
) -> Result<sandbox_v1::Template, SandboxError> {
let req = sandbox_v1::PublishTemplateRequest::decode_from_slice(payload)
.map_err(|e| SandboxError::Decode(e.to_string()))?;
let _operation = self.operations.lock(&operation_key(&req.name)).await;
let entry = self
.manager
.publish_template(&req.name, &req.version)
.await
.map_err(SandboxError::from)?;
Ok(convert::template_to_proto(&req.name, &entry))
}
pub async fn delete_template(&self, payload: &[u8]) -> Result<(), SandboxError> {
let req = sandbox_v1::DeleteTemplateRequest::decode_from_slice(payload)
.map_err(|e| SandboxError::Decode(e.to_string()))?;
let name = req
.reference
.split_once(':')
.map_or(&*req.reference, |(name, _)| name);
let _operation = self.operations.lock(&operation_key(name)).await;
self.manager
.delete_template(&req.reference)
.await
.map_err(SandboxError::from)
}
}
#[cfg(test)]
mod tests {
use super::dockerfile_tag;
#[test]
fn dockerfile_tag_is_content_addressed_and_valid() {
let a = dockerfile_tag(b"FROM alpine\n");
let b = dockerfile_tag(b"FROM alpine\n");
let c = dockerfile_tag(b"FROM debian\n");
assert_eq!(a, b, "identical content must reuse the tag");
assert_ne!(a, c, "different content must not share a tag");
let (repo, tag) = a.split_once(':').expect("repo:tag");
assert_eq!(repo, "arcbox-template-build");
assert_eq!(tag.len(), 16);
assert!(tag.bytes().all(|b| b.is_ascii_hexdigit()));
}
}