532 lines · 18.4 KB · Rust
RawDownload
use std::{    path::{Path, PathBuf},    pin::Pin,    task::Poll,}; use anyhow::{Context, Result};use async_compression::futures::bufread::{BzDecoder, GzipDecoder};use futures::{AsyncRead, AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt, io::BufReader};use sha2::{Digest, Sha256}; use crate::{HttpClient, github::AssetKind}; fn sha256_matches(actual: &str, expected: &str) -> bool {    actual.eq_ignore_ascii_case(expected)} #[derive(serde::Deserialize, serde::Serialize, Debug)]pub struct GithubBinaryMetadata {    pub metadata_version: u64,    pub digest: Option<String>,} impl GithubBinaryMetadata {    pub async fn read_from_file(metadata_path: &Path) -> Result<GithubBinaryMetadata> {        let metadata_content = async_fs::read_to_string(metadata_path)            .await            .with_context(|| format!("reading metadata file at {metadata_path:?}"))?;        serde_json::from_str(&metadata_content)            .with_context(|| format!("parsing metadata file at {metadata_path:?}"))    }     pub async fn write_to_file(&self, metadata_path: &Path) -> Result<()> {        let metadata_content = serde_json::to_string(self)            .with_context(|| format!("serializing metadata for {metadata_path:?}"))?;        async_fs::write(metadata_path, metadata_content.as_bytes())            .await            .with_context(|| format!("writing metadata file at {metadata_path:?}"))?;        Ok(())    }} pub async fn download_server_binary(    http_client: &dyn HttpClient,    url: &str,    digest: Option<&str>,    destination_path: &Path,    asset_kind: AssetKind,) -> Result<(), anyhow::Error> {    log::info!("downloading github artifact from {url}");    let Some(destination_parent) = destination_path.parent() else {        anyhow::bail!("destination path has no parent: {destination_path:?}");    };     let staging_path = staging_path(destination_parent, asset_kind)?;    let mut response = http_client        .get(url, Default::default(), true)        .await        .with_context(|| format!("downloading release from {url}"))?;    let body = response.body_mut();     if let Err(err) = extract_to_staging(body, digest, url, &staging_path, asset_kind).await {        cleanup_staging_path(&staging_path, asset_kind).await;        return Err(err);    }     if let Err(err) = finalize_download(&staging_path, destination_path).await {        cleanup_staging_path(&staging_path, asset_kind).await;        return Err(err);    }     Ok(())} pub async fn download_server_raw_binary(    http_client: &dyn HttpClient,    url: &str,    digest: Option<&str>,    destination_path: &Path,    binary_file_name: &str,) -> Result<(), anyhow::Error> {    log::info!("downloading raw binary from {url}");    let Some(destination_parent) = destination_path.parent() else {        anyhow::bail!("destination path has no parent: {destination_path:?}");    };     let staging_path = staging_dir_path(destination_parent)?;    let result = async {        let mut response = http_client            .get(url, Default::default(), true)            .await            .with_context(|| format!("downloading release from {url}"))?;         let binary_path = staging_path.join(binary_file_name);        let mut writer = HashingWriter {            writer: async_fs::File::create(&binary_path)                .await                .with_context(|| format!("creating a file {binary_path:?} for {url}"))?,            hasher: Sha256::new(),        };        futures::io::copy(&mut BufReader::new(response.body_mut()), &mut writer)            .await            .with_context(|| format!("saving binary contents from {url}"))?;        let asset_sha_256 = writer            .finish()            .await            .with_context(|| format!("flushing binary contents for {url}"))?;         if let Some(expected_sha_256) = digest {            anyhow::ensure!(                sha256_matches(&asset_sha_256, expected_sha_256),                "{url} asset got SHA-256 mismatch. Expected: {expected_sha_256}, Got: {asset_sha_256}",            );        }         util::fs::make_file_executable(&binary_path)            .await            .with_context(|| format!("marking {binary_path:?} as executable"))?;        finalize_download(&staging_path, destination_path).await    }    .await;     if let Err(err) = result {        if let Err(err) = async_fs::remove_dir_all(&staging_path).await {            log::warn!("failed to remove staging directory {staging_path:?}: {err:?}");        }        return Err(err);    }     Ok(())} async fn extract_to_staging(    body: impl AsyncRead + Unpin,    digest: Option<&str>,    url: &str,    staging_path: &Path,    asset_kind: AssetKind,) -> Result<()> {    match digest {        Some(expected_sha_256) => {            let temp_asset_file = tempfile::NamedTempFile::new()                .with_context(|| format!("creating a temporary file for {url}"))?;            let (temp_asset_file, _temp_guard) = temp_asset_file.into_parts();            let mut writer = HashingWriter {                writer: async_fs::File::from(temp_asset_file),                hasher: Sha256::new(),            };            futures::io::copy(&mut BufReader::new(body), &mut writer)                .await                .with_context(|| {                    format!("saving archive contents into the temporary file for {url}")                })?;            let asset_sha_256 = format!("{:x}", writer.hasher.finalize());             anyhow::ensure!(                sha256_matches(&asset_sha_256, expected_sha_256),                "{url} asset got SHA-256 mismatch. Expected: {expected_sha_256}, Got: {asset_sha_256}",            );            writer                .writer                .seek(std::io::SeekFrom::Start(0))                .await                .with_context(|| format!("seeking temporary file for {url}"))?;            stream_file_archive(&mut writer.writer, url, staging_path, asset_kind)                .await                .with_context(|| {                    format!("extracting downloaded asset for {url} into {staging_path:?}")                })?;        }        None => {            stream_response_archive(body, url, staging_path, asset_kind)                .await                .with_context(|| {                    format!("extracting response for asset {url} into {staging_path:?}")                })?;        }    }    Ok(())} fn staging_dir_path(parent: &Path) -> Result<PathBuf> {    let dir = tempfile::Builder::new()        .prefix(".tmp-github-download-")        .tempdir_in(parent)        .with_context(|| format!("creating staging directory in {parent:?}"))?;    Ok(dir.keep())} fn staging_path(parent: &Path, asset_kind: AssetKind) -> Result<PathBuf> {    match asset_kind {        AssetKind::TarGz | AssetKind::TarBz2 | AssetKind::Zip => staging_dir_path(parent),        AssetKind::Gz => {            let path = tempfile::Builder::new()                .prefix(".tmp-github-download-")                .tempfile_in(parent)                .with_context(|| format!("creating staging file in {parent:?}"))?                .into_temp_path()                .keep()                .with_context(|| format!("persisting staging file in {parent:?}"))?;            Ok(path)        }    }} async fn cleanup_staging_path(staging_path: &Path, asset_kind: AssetKind) {    match asset_kind {        AssetKind::TarGz | AssetKind::TarBz2 | AssetKind::Zip => {            if let Err(err) = async_fs::remove_dir_all(staging_path).await {                log::warn!("failed to remove staging directory {staging_path:?}: {err:?}");            }        }        AssetKind::Gz => {            if let Err(err) = async_fs::remove_file(staging_path).await {                log::warn!("failed to remove staging file {staging_path:?}: {err:?}");            }        }    }} async fn finalize_download(staging_path: &Path, destination_path: &Path) -> Result<()> {    _ = async_fs::remove_dir_all(destination_path).await;    async_fs::rename(staging_path, destination_path)        .await        .with_context(|| format!("renaming {staging_path:?} to {destination_path:?}"))?;    Ok(())} async fn stream_response_archive(    response: impl AsyncRead + Unpin,    url: &str,    destination_path: &Path,    asset_kind: AssetKind,) -> Result<()> {    match asset_kind {        AssetKind::TarGz => extract_tar_gz(destination_path, url, response).await?,        AssetKind::TarBz2 => extract_tar_bz2(destination_path, url, response).await?,        AssetKind::Gz => extract_gz(destination_path, url, response).await?,        AssetKind::Zip => {            util::archive::extract_zip(destination_path, response).await?;        }    };    Ok(())} async fn stream_file_archive(    file_archive: impl AsyncRead + AsyncSeek + Unpin,    url: &str,    destination_path: &Path,    asset_kind: AssetKind,) -> Result<()> {    match asset_kind {        AssetKind::TarGz => extract_tar_gz(destination_path, url, file_archive).await?,        AssetKind::TarBz2 => extract_tar_bz2(destination_path, url, file_archive).await?,        AssetKind::Gz => extract_gz(destination_path, url, file_archive).await?,        #[cfg(not(windows))]        AssetKind::Zip => {            util::archive::extract_seekable_zip(destination_path, file_archive).await?;        }        #[cfg(windows)]        AssetKind::Zip => {            util::archive::extract_zip(destination_path, file_archive).await?;        }    };    Ok(())} async fn extract_tar_gz(    destination_path: &Path,    url: &str,    from: impl AsyncRead + Unpin,) -> Result<(), anyhow::Error> {    let decompressed_bytes = GzipDecoder::new(BufReader::new(from));    unpack_tar_archive(destination_path, url, decompressed_bytes).await?;    Ok(())} async fn extract_tar_bz2(    destination_path: &Path,    url: &str,    from: impl AsyncRead + Unpin,) -> Result<(), anyhow::Error> {    let decompressed_bytes = BzDecoder::new(BufReader::new(from));    unpack_tar_archive(destination_path, url, decompressed_bytes).await?;    Ok(())} async fn unpack_tar_archive(    destination_path: &Path,    url: &str,    archive_bytes: impl AsyncRead + Unpin,) -> Result<(), anyhow::Error> {    // We don't need to set the modified time. It's irrelevant to downloaded    // archive verification, and some filesystems return errors when asked to    // apply it after extraction.    let archive = async_tar::ArchiveBuilder::new(archive_bytes)        .set_preserve_mtime(false)        .build();    archive        .unpack(&destination_path)        .await        .with_context(|| format!("extracting {url} to {destination_path:?}"))?;    Ok(())} async fn extract_gz(    destination_path: &Path,    url: &str,    from: impl AsyncRead + Unpin,) -> Result<(), anyhow::Error> {    let mut decompressed_bytes = GzipDecoder::new(BufReader::new(from));    let mut file = async_fs::File::create(&destination_path)        .await        .with_context(|| {            format!("creating a file {destination_path:?} for a download from {url}")        })?;    futures::io::copy(&mut decompressed_bytes, &mut file)        .await        .with_context(|| format!("extracting {url} to {destination_path:?}"))?;    Ok(())} struct HashingWriter<W: AsyncWrite + Unpin> {    writer: W,    hasher: Sha256,} impl<W: AsyncWrite + Unpin> HashingWriter<W> {    /// Closes and drops the inner writer, returning the hex SHA-256 digest of    /// everything written.    ///    /// Taking `self` by value guarantees the writer is dropped before this    /// returns. For file writers this releases the OS handle, which Windows    /// requires before an ancestor directory can be renamed or deleted; note    /// that closing alone is not enough, as `async_fs::File` holds its handle    /// until dropped.    async fn finish(mut self) -> std::io::Result<String> {        self.writer.close().await?;        drop(self.writer);        Ok(format!("{:x}", self.hasher.finalize()))    }} impl<W: AsyncWrite + Unpin> AsyncWrite for HashingWriter<W> {    fn poll_write(        mut self: Pin<&mut Self>,        cx: &mut std::task::Context<'_>,        buf: &[u8],    ) -> Poll<std::result::Result<usize, std::io::Error>> {        match Pin::new(&mut self.writer).poll_write(cx, buf) {            Poll::Ready(Ok(n)) => {                self.hasher.update(&buf[..n]);                Poll::Ready(Ok(n))            }            other => other,        }    }     fn poll_flush(        mut self: Pin<&mut Self>,        cx: &mut std::task::Context<'_>,    ) -> Poll<Result<(), std::io::Error>> {        Pin::new(&mut self.writer).poll_flush(cx)    }     fn poll_close(        mut self: Pin<&mut Self>,        cx: &mut std::task::Context<'_>,    ) -> Poll<std::result::Result<(), std::io::Error>> {        Pin::new(&mut self.writer).poll_close(cx)    }} #[cfg(test)]mod tests {    use super::*;    use crate::{AsyncBody, Response};    use futures::future::BoxFuture;    use http::HeaderValue;    use url::Url;     struct StaticResponseClient {        body: Vec<u8>,    }     impl HttpClient for StaticResponseClient {        fn send(            &self,            _req: http::Request<AsyncBody>,        ) -> BoxFuture<'static, anyhow::Result<Response<AsyncBody>>> {            let body = self.body.clone();            Box::pin(async move {                Ok(Response::builder()                    .status(200)                    .body(AsyncBody::from(body))                    .unwrap())            })        }         fn user_agent(&self) -> Option<&HeaderValue> {            None        }         fn proxy(&self) -> Option<&Url> {            None        }    }     #[test]    fn downloads_raw_binary_with_uppercase_digest_into_destination_dir() {        futures::executor::block_on(async {            let temp_dir = tempfile::tempdir().unwrap();            let destination_path = temp_dir.path().join("v_1");            let contents = b"#!/bin/sh\necho hello\n".to_vec();            let expected_sha_256 = format!("{:X}", Sha256::digest(&contents));            let client = StaticResponseClient { body: contents };             download_server_raw_binary(                &client,                "https://example.com/agent-binary",                Some(&expected_sha_256),                &destination_path,                "agent-binary",            )            .await            .unwrap();             let binary_path = destination_path.join("agent-binary");            assert_eq!(                std::fs::read(&binary_path).unwrap(),                b"#!/bin/sh\necho hello\n"            );            #[cfg(unix)]            {                use std::os::unix::fs::PermissionsExt;                let mode = std::fs::metadata(&binary_path)                    .unwrap()                    .permissions()                    .mode();                assert_eq!(mode & 0o111, 0o111, "binary should be executable");            }        });    }     #[test]    fn raw_binary_digest_mismatch_cleans_up_staging() {        futures::executor::block_on(async {            let temp_dir = tempfile::tempdir().unwrap();            let destination_path = temp_dir.path().join("v_1");            let client = StaticResponseClient {                body: b"some binary".to_vec(),            };             let error = download_server_raw_binary(                &client,                "https://example.com/agent-binary",                Some("0000000000000000000000000000000000000000000000000000000000000000"),                &destination_path,                "agent-binary",            )            .await            .unwrap_err();             assert!(error.to_string().contains("SHA-256 mismatch"));            assert!(!destination_path.exists());            let leftover_entries = std::fs::read_dir(temp_dir.path()).unwrap().count();            assert_eq!(leftover_entries, 0, "staging directory should be removed");        });    }     #[test]    fn downloads_archive_with_uppercase_digest_and_extracts_contents() {        futures::executor::block_on(async {            let archive = vec![                0x50, 0x4b, 0x03, 0x04, 0x14, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x21, 0x00,                0x86, 0xa6, 0x10, 0x36, 0x05, 0x00, 0x00, 0x00, 0x05, 0x00, 0x00, 0x00, 0x05, 0x00,                0x00, 0x00, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x68, 0x65, 0x6c, 0x6c, 0x6f, 0x50, 0x4b,                0x01, 0x02, 0x14, 0x03, 0x14, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x21, 0x00,                0x86, 0xa6, 0x10, 0x36, 0x05, 0x00, 0x00, 0x00, 0x05, 0x00, 0x00, 0x00, 0x05, 0x00,                0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x01, 0x00, 0x00,                0x00, 0x00, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x50, 0x4b, 0x05, 0x06, 0x00, 0x00, 0x00,                0x00, 0x01, 0x00, 0x01, 0x00, 0x33, 0x00, 0x00, 0x00, 0x28, 0x00, 0x00, 0x00, 0x00,                0x00,            ];            let expected_sha_256 = format!("{:X}", Sha256::digest(&archive));            let client = StaticResponseClient { body: archive };            let temp_dir = tempfile::tempdir().unwrap();            let destination_path = temp_dir.path().join("v_1");             download_server_binary(                &client,                "https://example.com/agent.zip",                Some(&expected_sha_256),                &destination_path,                AssetKind::Zip,            )            .await            .unwrap();             assert_eq!(

Lines 1–500 of 532.