533 lines
18 KiB
Rust
533 lines
18 KiB
Rust
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!(
|
|
std::fs::read(destination_path.join("agent")).unwrap(),
|
|
b"hello"
|
|
);
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn archive_digest_mismatch_prevents_extraction_and_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"not an archive".to_vec(),
|
|
};
|
|
|
|
let error = download_server_binary(
|
|
&client,
|
|
"https://example.com/agent.zip",
|
|
Some("0000000000000000000000000000000000000000000000000000000000000000"),
|
|
&destination_path,
|
|
AssetKind::Zip,
|
|
)
|
|
.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");
|
|
});
|
|
}
|
|
}
|