use chrono::{DateTime, Utc};
use flate2::{write::GzEncoder, Compression};
use oci_spec::image::*;
use std::{collections::HashMap, convert::TryFrom, fs, io, path::Path};
use crate::{
digest::{Digest, DigestBuf},
error::*,
image::annotations::flat::Annotations,
ImageName,
};
pub struct Builder<W: io::Write> {
builder: Option<tar::Builder<W>>,
name: Option<ImageName>,
created: Option<DateTime<Utc>>,
author: Option<String>,
annotations: Option<Annotations>,
platform: Option<Platform>,
diff_ids: Vec<Digest>,
layers: Vec<Descriptor>,
}
impl<W: io::Write> Builder<W> {
pub fn new(writer: W) -> Self {
Builder {
builder: Some(tar::Builder::new(writer)),
name: None,
created: None,
author: None,
platform: None,
annotations: None,
diff_ids: Vec::new(),
layers: Vec::new(),
}
}
pub fn set_name(&mut self, name: &ImageName) {
self.name = Some(name.clone());
}
pub fn set_created(&mut self, created: DateTime<Utc>) {
self.created = Some(created);
}
pub fn set_annotations(&mut self, annotations: Annotations) {
self.annotations = Some(annotations);
}
pub fn set_author(&mut self, author: &str) {
self.author = Some(author.to_string());
}
pub fn set_platform(&mut self, platform: &Platform) {
self.platform = Some(platform.clone());
}
pub fn append_files(&mut self, ps: &[impl AsRef<Path>]) -> Result<()> {
let mut ar = tar::Builder::new(DigestBuf::new(GzEncoder::new(
Vec::new(),
Compression::default(),
)));
for path in ps {
let path = path.as_ref();
if !path.is_file() {
return Err(Error::NotAFile(path.to_owned()));
}
let name = path
.file_name()
.expect("This never fails since checked above")
.to_str()
.expect("Non-UTF8 file name");
let mut f = fs::File::open(path)?;
ar.append_file(name, &mut f)?;
}
let (gz, digest) = ar
.into_inner()
.expect("This never fails since tar arhive is creating on memory")
.finish();
self.diff_ids.push(digest);
let buf = gz
.finish()
.expect("This never fails since zip is creating on memory");
let layer_desc = self.save_blob(MediaType::ImageLayerGzip, &buf)?;
self.layers.push(layer_desc);
Ok(())
}
pub fn append_dir_all(&mut self, path: &Path) -> Result<()> {
if !path.is_dir() {
return Err(Error::NotADirectory(path.to_owned()));
}
let mut ar = tar::Builder::new(DigestBuf::new(GzEncoder::new(
Vec::new(),
Compression::default(),
)));
ar.append_dir_all("", path)?;
let (gz, digest) = ar
.into_inner()
.expect("This never fails since tar arhive is creating on memory")
.finish();
self.diff_ids.push(digest);
let buf = gz
.finish()
.expect("This never fails since zip is creating on memory");
let layer_desc = self.save_blob(MediaType::ImageLayerGzip, &buf)?;
self.layers.push(layer_desc);
Ok(())
}
pub fn into_inner(mut self) -> Result<W> {
self.finish()?;
Ok(self.builder.take().unwrap().into_inner()?)
}
fn create_config(&self) -> ImageConfiguration {
let mut builder = ImageConfigurationBuilder::default();
let created = self.created.unwrap_or_else(Utc::now);
builder = builder.created(created.to_rfc3339());
if let Some(ref author) = self.author {
builder = builder.author(author);
}
if let Some(ref platform) = self.platform {
builder = builder.os(platform.os().clone());
builder = builder.architecture(platform.architecture().clone());
}
let rootfs = RootFsBuilder::default()
.typ("layers".to_string())
.diff_ids(
self.diff_ids
.iter()
.map(|digest| digest.to_string())
.collect::<Vec<_>>(),
)
.build()
.unwrap();
builder = builder.rootfs(rootfs);
let config = ConfigBuilder::default()
.labels(self.create_annotations_as_map())
.build()
.unwrap();
builder = builder.config(config);
builder.build().unwrap()
}
fn create_annotations_as_map(&self) -> HashMap<String, String> {
if let Some(mut a) = self.annotations.clone() {
if self.created.is_some() && a.created.is_none() {
a.created = self.created.as_ref().map(|date| date.to_string());
}
if self.author.is_some() && a.authors.is_none() {
a.authors = self.author.clone();
}
a.to_map()
} else {
HashMap::new()
}
}
fn finish(&mut self) -> Result<()> {
let cfg = self.create_config();
let mut buf = Vec::new();
cfg.to_writer(&mut buf)?;
let cfg_desc = self.save_blob(MediaType::ImageConfig, &buf)?;
let image_manifest = ImageManifestBuilder::default()
.schema_version(SCHEMA_VERSION)
.config(cfg_desc)
.layers(std::mem::take(&mut self.layers))
.build()
.unwrap();
let mut buf = Vec::new();
image_manifest.to_writer(&mut buf)?;
let mut image_manifest_desc = self.save_blob(MediaType::ImageManifest, &buf)?;
image_manifest_desc.set_annotations(Some(HashMap::from([(
"org.opencontainers.image.ref.name".to_string(),
self.name.clone().unwrap_or_default().to_string(),
)])));
let index = ImageIndexBuilder::default()
.schema_version(SCHEMA_VERSION)
.manifests(vec![image_manifest_desc])
.build()?;
let mut index_json = Vec::new();
index.to_writer(&mut index_json)?;
let index_json = String::from_utf8(index_json).expect("ImageIndex must returns valid JSON");
self.save_file(Path::new("index.json"), &index_json)?;
let version = r#"{"imageLayoutVersion":"1.0.0"}"#;
self.save_file(Path::new("oci-layout"), version)?;
self.builder
.as_mut()
.expect("builder never becomes None except on Drop")
.finish()?;
Ok(())
}
fn save_blob(&mut self, media_type: MediaType, buf: &[u8]) -> Result<Descriptor> {
let digest = Digest::from_buf_sha256(buf);
let mut header = create_header(buf.len());
self.builder
.as_mut()
.expect("builder never becomes None except on Drop")
.append_data(&mut header, digest.as_path(), buf)?;
Ok(DescriptorBuilder::default()
.media_type(media_type)
.size(buf.len() as i64)
.digest(format!("{}", digest))
.build()
.expect("Requirement for descriptor is mediaType, digest, and size."))
}
fn save_file(&mut self, dest: &Path, input: &str) -> Result<()> {
let mut header = create_header(input.len());
self.builder
.as_mut()
.expect("builder never becomes None except on Drop")
.append_data(&mut header, dest, input.as_bytes())?;
Ok(())
}
}
impl<W: io::Write> Drop for Builder<W> {
fn drop(&mut self) {
if self.builder.is_some() {
let _ = self.finish();
}
}
}
fn create_header(size: usize) -> tar::Header {
let mut header = tar::Header::new_gnu();
header.set_size(u64::try_from(size).unwrap());
header.set_cksum();
header.set_mode(0b110100100); header.set_mtime(Utc::now().timestamp() as u64);
header
}