path: root/tvix/glue/src/fetchers.rs

use futures::TryStreamExt;
use md5::Md5;
use nix_compat::{
    nixhash::{CAHash, HashAlgo, NixHash},
    store_path::{build_ca_path, BuildStorePathError, StorePathRef},
};
use sha1::Sha1;
use sha2::{digest::Output, Digest, Sha256, Sha512};
use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite};
use tokio_util::io::InspectReader;
use tracing::warn;
use tvix_castore::{
    blobservice::BlobService,
    directoryservice::DirectoryService,
    proto::{node::Node, FileNode},
};
use tvix_store::{pathinfoservice::PathInfoService, proto::PathInfo};

use crate::{builtins::FetcherError, decompression::DecompressedReader};

/// Representing options for doing a fetch.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Fetch {
    /// Fetch a literal file from the given URL, with an optional expected
    /// NixHash of it.
    /// TODO: check if this is *always* sha256, and if so, make it [u8; 32].
    URL(String, Option<NixHash>),

    /// Fetch a tarball from the given URL and unpack.
    /// The file must be a tape archive (.tar) compressed with gzip, bzip2 or xz.
    /// The top-level path component of the files in the tarball is removed,
    /// so it is best if the tarball contains a single directory at top level.
    /// Optionally, a sha256 digest can be provided to verify the unpacked
    /// contents against.
    Tarball(String, Option<[u8; 32]>),

    /// TODO
    Git(),
}

impl Fetch {
    /// If the [Fetch] contains an expected hash upfront, returns the resulting
    /// store path.
    /// This doesn't do any fetching.
    pub fn store_path<'a>(
        &self,
        name: &'a str,
    ) -> Result<Option<StorePathRef<'a>>, BuildStorePathError> {
        let ca_hash = match self {
            Fetch::URL(_, Some(nixhash)) => CAHash::Flat(nixhash.clone()),
            Fetch::Tarball(_, Some(nar_sha256)) => CAHash::Nar(NixHash::Sha256(*nar_sha256)),
            _ => return Ok(None),
        };

        // calculate the store path of this fetch
        build_ca_path(name, &ca_hash, Vec::<String>::new(), false).map(Some)
    }
}

/// Knows how to fetch a given [Fetch].
pub struct Fetcher<BS, DS, PS> {
    http_client: reqwest::Client,
    blob_service: BS,
    directory_service: DS,
    path_info_service: PS,
}

impl<BS, DS, PS> Fetcher<BS, DS, PS> {
    pub fn new(blob_service: BS, directory_service: DS, path_info_service: PS) -> Self {
        Self {
            http_client: reqwest::Client::new(),
            blob_service,
            directory_service,
            path_info_service,
        }
    }

    /// Constructs a HTTP request to the passed URL, and returns a AsyncReadBuf to it.
    async fn download<'a>(
        &self,
        url: &str,
    ) -> Result<impl AsyncBufRead + Unpin + 'a, reqwest::Error> {
        let resp = self.http_client.get(url).send().await?;
        Ok(tokio_util::io::StreamReader::new(
            resp.bytes_stream().map_err(|e| {
                let e = e.without_url();
                warn!(%e, "failed to get response body");
                std::io::Error::new(std::io::ErrorKind::BrokenPipe, e)
            }),
        ))
    }
}

/// Copies all data from the passed reader to the passed writer.
/// Afterwards, it also returns the resulting [Digest], as well as the number of
/// bytes copied.
/// The exact hash function used is left generic over all [Digest].
async fn hash<D: Digest + std::io::Write>(
    mut r: impl AsyncRead + Unpin,
    mut w: impl AsyncWrite + Unpin,
) -> std::io::Result<(Output<D>, u64)> {
    let mut hasher = D::new();
    let bytes_copied = tokio::io::copy(
        &mut InspectReader::new(&mut r, |d| hasher.write_all(d).unwrap()),
        &mut w,
    )
    .await?;
    Ok((hasher.finalize(), bytes_copied))
}

impl<BS, DS, PS> Fetcher<BS, DS, PS>
where
    BS: AsRef<(dyn BlobService + 'static)> + Send + Sync,
    DS: AsRef<(dyn DirectoryService + 'static)>,
    PS: PathInfoService,
{
    /// Ingest the data from a specified [Fetch].
    /// On success, return the root node, a content digest and length.
    /// Returns an error if there was a failure during fetching, or the contents
    /// didn't match the previously communicated hash contained inside the FetchArgs.
    pub async fn ingest(&self, fetch: Fetch) -> Result<(Node, CAHash, u64), FetcherError> {
        match fetch {
            Fetch::URL(url, exp_nixhash) => {
                // Construct a AsyncRead reading from the data as its downloaded.
                let mut r = self.download(&url).await?;

                // Construct a AsyncWrite to write into the BlobService.
                let mut blob_writer = self.blob_service.open_write().await;

                // Copy the contents from the download reader to the blob writer.
                // Calculate the digest of the file received, depending on the
                // communicated expected hash (or sha256 if none provided).
                let (actual_nixhash, blob_size) = match exp_nixhash
                    .as_ref()
                    .map(NixHash::algo)
                    .unwrap_or_else(|| HashAlgo::Sha256)
                {
                    HashAlgo::Sha256 => hash::<Sha256>(&mut r, &mut blob_writer).await.map(
                        |(digest, bytes_written)| (NixHash::Sha256(digest.into()), bytes_written),
                    )?,
                    HashAlgo::Md5 => hash::<Md5>(&mut r, &mut blob_writer).await.map(
                        |(digest, bytes_written)| (NixHash::Md5(digest.into()), bytes_written),
                    )?,
                    HashAlgo::Sha1 => hash::<Sha1>(&mut r, &mut blob_writer).await.map(
                        |(digest, bytes_written)| (NixHash::Sha1(digest.into()), bytes_written),
                    )?,
                    HashAlgo::Sha512 => hash::<Sha512>(&mut r, &mut blob_writer).await.map(
                        |(digest, bytes_written)| {
                            (NixHash::Sha512(Box::new(digest.into())), bytes_written)
                        },
                    )?,
                };

                if let Some(exp_nixhash) = exp_nixhash {
                    if exp_nixhash != actual_nixhash {
                        return Err(FetcherError::HashMismatch {
                            url,
                            wanted: exp_nixhash,
                            got: actual_nixhash,
                        });
                    }
                }

                // Construct and return the FileNode describing the downloaded contents.
                Ok((
                    Node::File(FileNode {
                        name: vec![].into(),
                        digest: blob_writer.close().await?.into(),
                        size: blob_size,
                        executable: false,
                    }),
                    CAHash::Flat(actual_nixhash),
                    blob_size,
                ))
            }
            Fetch::Tarball(url, exp_nar_sha256) => {
                // Construct a AsyncRead reading from the data as its downloaded.
                let r = self.download(&url).await?;

                // Pop compression.
                let r = DecompressedReader::new(r);
                // Open the archive.
                let archive = tokio_tar::Archive::new(r);

                // Ingest the archive, get the root node
                let node = tvix_castore::import::archive::ingest_archive(
                    &self.blob_service,
                    &self.directory_service,
                    archive,
                )
                .await?;

                // If an expected NAR sha256 was provided, compare with the one
                // calculated from our root node.
                // Even if no expected NAR sha256 has been provided, we need
                // the actual one later.
                let (nar_size, actual_nar_sha256) = self
                    .path_info_service
                    .calculate_nar(&node)
                    .await
                    .map_err(|e| {
                        // convert the generic Store error to an IO error.
                        FetcherError::Io(e.into())
                    })?;

                if let Some(exp_nar_sha256) = exp_nar_sha256 {
                    if exp_nar_sha256 != actual_nar_sha256 {
                        return Err(FetcherError::HashMismatch {
                            url,
                            wanted: NixHash::Sha256(exp_nar_sha256),
                            got: NixHash::Sha256(actual_nar_sha256),
                        });
                    }
                }

                Ok((
                    node,
                    CAHash::Nar(NixHash::Sha256(actual_nar_sha256)),
                    nar_size,
                ))
            }
            Fetch::Git() => todo!(),
        }
    }

    /// Ingests the data from a specified [Fetch], persists the returned node
    /// in the PathInfoService, and returns the calculated StorePath, as well as
    /// the root node pointing to the contents.
    /// The root node can be used to descend into the data without doing the
    /// lookup to the PathInfoService again.
    pub async fn ingest_and_persist<'a>(
        &self,
        name: &'a str,
        fetch: Fetch,
    ) -> Result<(StorePathRef<'a>, Node), FetcherError> {
        // Fetch file, return the (unnamed) (File)Node of its contents, ca hash and filesize.
        let (mut node, ca_hash, size) = self.ingest(fetch).await?;

        // Calculate the store path to return later, which is done with the ca_hash.
        let store_path = build_ca_path(name, &ca_hash, Vec::<String>::new(), false)?;

        // Rename the node name to match the Store Path.
        if let Node::File(file_node) = &mut node {
            file_node.name = store_path.to_string().into();
        } else {
            unreachable!("Tvix bug: do_fetch for URL returned non-FileNode");
        }

        // If the resulting hash is not a CAHash::Nar, we also need to invoke
        // `calculate_nar` to calculate this representation, as it's required in
        // the [PathInfo].
        let (nar_size, nar_sha256) = match &ca_hash {
            CAHash::Flat(_nix_hash) => self
                .path_info_service
                .calculate_nar(&node)
                .await
                .map_err(|e| FetcherError::Io(e.into()))?,
            CAHash::Nar(NixHash::Sha256(nar_sha256)) => (size, *nar_sha256),
            CAHash::Nar(_) => unreachable!("Tvix bug: fetch returned non-sha256 CAHash::Nar"),
            CAHash::Text(_) => unreachable!("Tvix bug: fetch returned CAHash::Text"),
        };

        // Construct the PathInfo and persist it.
        let path_info = PathInfo {
            node: Some(tvix_castore::proto::Node { node: Some(node) }),
            references: vec![],
            narinfo: Some(tvix_store::proto::NarInfo {
                nar_size,
                nar_sha256: nar_sha256.to_vec().into(),
                signatures: vec![],
                reference_names: vec![],
                deriver: None,
                ca: Some(ca_hash.into()),
            }),
        };

        let path_info = self
            .path_info_service
            .put(path_info)
            .await
            .map_err(|e| FetcherError::Io(e.into()))?;

        Ok((store_path, path_info.node.unwrap().node.unwrap()))
    }
}

/// Attempts to mimic `nix::libutil::baseNameOf`
pub(crate) fn url_basename(s: &str) -> &str {
    if s.is_empty() {
        return "";
    }

    let mut last = s.len() - 1;
    if s.chars().nth(last).unwrap() == '/' && last > 0 {
        last -= 1;
    }

    if last == 0 {
        return "";
    }

    let pos = match s[..=last].rfind('/') {
        Some(pos) => {
            if pos == last - 1 {
                0
            } else {
                pos
            }
        }
        None => 0,
    };

    &s[(pos + 1)..=last]
}

#[cfg(test)]
mod tests {
    mod fetch {
        use nix_compat::nixbase32;

        use crate::fetchers::Fetch;

        use super::super::*;

        #[test]
        fn fetchurl_store_path() {
            let url = "https://raw.githubusercontent.com/aaptel/notmuch-extract-patch/f732a53e12a7c91a06755ebfab2007adc9b3063b/notmuch-extract-patch";
            let exp_nixhash = NixHash::Sha256(
                nixbase32::decode_fixed("0nawkl04sj7psw6ikzay7kydj3dhd0fkwghcsf5rzaw4bmp4kbax")
                    .unwrap(),
            );

            let fetch = Fetch::URL(url.into(), Some(exp_nixhash));
            assert_eq!(
                "06qi00hylriyfm0nl827crgjvbax84mz-notmuch-extract-patch",
                &fetch
                    .store_path("notmuch-extract-patch")
                    .unwrap()
                    .unwrap()
                    .to_string(),
            )
        }

        #[test]
        fn fetch_tarball_store_path() {
            let url = "https://github.com/NixOS/nixpkgs/archive/91050ea1e57e50388fa87a3302ba12d188ef723a.tar.gz";
            let exp_nixbase32 =
                nixbase32::decode_fixed("1hf6cgaci1n186kkkjq106ryf8mmlq9vnwgfwh625wa8hfgdn4dm")
                    .unwrap();
            let fetch = Fetch::Tarball(url.into(), Some(exp_nixbase32));

            assert_eq!(
                "7adgvk5zdfq4pwrhsm3n9lzypb12gw0g-source",
                &fetch.store_path("source").unwrap().unwrap().to_string(),
            )
        }
    }

    mod url_basename {
        use super::super::*;

        #[test]
        fn empty_path() {
            assert_eq!(url_basename(""), "");
        }

        #[test]
        fn path_on_root() {
            assert_eq!(url_basename("/dir"), "dir");
        }

        #[test]
        fn relative_path() {
            assert_eq!(url_basename("dir/foo"), "foo");
        }

        #[test]
        fn root_with_trailing_slash() {
            assert_eq!(url_basename("/"), "");
        }

        #[test]
        fn trailing_slash() {
            assert_eq!(url_basename("/dir/"), "dir");
        }
    }
}