start-os/core/startos/src/s9pk/v1/docker.rs

use std::collections::BTreeSet;
use std::io::SeekFrom;
use std::path::Path;

use color_eyre::eyre::eyre;
use futures::{FutureExt, TryStreamExt};
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncRead, AsyncSeek, AsyncSeekExt};
use tokio_tar::{Archive, Entry};

use crate::util::io::from_cbor_async_reader;
use crate::{Error, ErrorKind};

#[derive(Default, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DockerMultiArch {
    pub default: String,
    pub available: BTreeSet<String>,
}

#[pin_project::pin_project(project = DockerReaderProject)]
#[derive(Debug)]
pub enum DockerReader<R: AsyncRead + Unpin> {
    SingleArch(#[pin] R),
    MultiArch(#[pin] Entry<Archive<R>>),
}
impl<R: AsyncRead + AsyncSeek + Unpin + Send + Sync> DockerReader<R> {
    pub async fn list_arches(rdr: &mut R) -> Result<BTreeSet<String>, Error> {
        if let Some(multiarch) = tokio_tar::Archive::new(rdr)
            .entries()?
            .try_filter_map(|e| {
                async move {
                    Ok(if &*e.path()? == Path::new("multiarch.cbor") {
                        Some(e)
                    } else {
                        None
                    })
                }
                .boxed()
            })
            .try_next()
            .await?
        {
            let multiarch: DockerMultiArch = from_cbor_async_reader(multiarch).await?;
            Ok(multiarch.available)
        } else {
            Err(Error::new(
                eyre!("Single arch legacy s9pks not supported"),
                ErrorKind::ParseS9pk,
            ))
        }
    }
    pub async fn new(mut rdr: R, arch: &str) -> Result<Self, Error> {
        rdr.seek(SeekFrom::Start(0)).await?;
        if let Some(image) = tokio_tar::Archive::new(rdr)
            .entries()?
            .try_filter_map(|e| {
                async move {
                    Ok(if &*e.path()? == Path::new(&format!("{}.tar", arch)) {
                        Some(e)
                    } else {
                        None
                    })
                }
                .boxed()
            })
            .try_next()
            .await?
        {
            Ok(Self::MultiArch(image))
        } else {
            Err(Error::new(
                eyre!("Docker image section does not contain tarball for architecture"),
                ErrorKind::ParseS9pk,
            ))
        }
    }
}
impl<R: AsyncRead + Unpin + Send + Sync> AsyncRead for DockerReader<R> {
    fn poll_read(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        match self.project() {
            DockerReaderProject::SingleArch(r) => r.poll_read(cx, buf),
            DockerReaderProject::MultiArch(r) => r.poll_read(cx, buf),
        }
    }
}