tenant.openagents/omega

1use std::{
2    path::{Path, PathBuf},
3    pin::Pin,
4    task::Poll,
5};
6
7use anyhow::{Context, Result};
8use async_compression::futures::bufread::{BzDecoder, GzipDecoder};
9use futures::{AsyncRead, AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt, io::BufReader};
10use sha2::{Digest, Sha256};
11
12use crate::{HttpClient, github::AssetKind};
13
14fn sha256_matches(actual: &str, expected: &str) -> bool {
15    actual.eq_ignore_ascii_case(expected)
16}
17
18#[derive(serde::Deserialize, serde::Serialize, Debug)]
19pub struct GithubBinaryMetadata {
20    pub metadata_version: u64,
21    pub digest: Option<String>,
22}
23
24impl GithubBinaryMetadata {
25    pub async fn read_from_file(metadata_path: &Path) -> Result<GithubBinaryMetadata> {
26        let metadata_content = async_fs::read_to_string(metadata_path)
27            .await
28            .with_context(|| format!("reading metadata file at {metadata_path:?}"))?;
29        serde_json::from_str(&metadata_content)
30            .with_context(|| format!("parsing metadata file at {metadata_path:?}"))
31    }
32
33    pub async fn write_to_file(&self, metadata_path: &Path) -> Result<()> {
34        let metadata_content = serde_json::to_string(self)
35            .with_context(|| format!("serializing metadata for {metadata_path:?}"))?;
36        async_fs::write(metadata_path, metadata_content.as_bytes())
37            .await
38            .with_context(|| format!("writing metadata file at {metadata_path:?}"))?;
39        Ok(())
40    }
41}
42
43pub async fn download_server_binary(
44    http_client: &dyn HttpClient,
45    url: &str,
46    digest: Option<&str>,
47    destination_path: &Path,
48    asset_kind: AssetKind,
49) -> Result<(), anyhow::Error> {
50    log::info!("downloading github artifact from {url}");
51    let Some(destination_parent) = destination_path.parent() else {
52        anyhow::bail!("destination path has no parent: {destination_path:?}");
53    };
54
55    let staging_path = staging_path(destination_parent, asset_kind)?;
56    let mut response = http_client
57        .get(url, Default::default(), true)
58        .await
59        .with_context(|| format!("downloading release from {url}"))?;
60    let body = response.body_mut();
61
62    if let Err(err) = extract_to_staging(body, digest, url, &staging_path, asset_kind).await {
63        cleanup_staging_path(&staging_path, asset_kind).await;
64        return Err(err);
65    }
66
67    if let Err(err) = finalize_download(&staging_path, destination_path).await {
68        cleanup_staging_path(&staging_path, asset_kind).await;
69        return Err(err);
70    }
71
72    Ok(())
73}
74
75pub async fn download_server_raw_binary(
76    http_client: &dyn HttpClient,
77    url: &str,
78    digest: Option<&str>,
79    destination_path: &Path,
80    binary_file_name: &str,
81) -> Result<(), anyhow::Error> {
82    log::info!("downloading raw binary from {url}");
83    let Some(destination_parent) = destination_path.parent() else {
84        anyhow::bail!("destination path has no parent: {destination_path:?}");
85    };
86
87    let staging_path = staging_dir_path(destination_parent)?;
88    let result = async {
89        let mut response = http_client
90            .get(url, Default::default(), true)
91            .await
92            .with_context(|| format!("downloading release from {url}"))?;
93
94        let binary_path = staging_path.join(binary_file_name);
95        let mut writer = HashingWriter {
96            writer: async_fs::File::create(&binary_path)
97                .await
98                .with_context(|| format!("creating a file {binary_path:?} for {url}"))?,
99            hasher: Sha256::new(),
100        };
101        futures::io::copy(&mut BufReader::new(response.body_mut()), &mut writer)
102            .await
103            .with_context(|| format!("saving binary contents from {url}"))?;
104        let asset_sha_256 = writer
105            .finish()
106            .await
107            .with_context(|| format!("flushing binary contents for {url}"))?;
108
109        if let Some(expected_sha_256) = digest {
110            anyhow::ensure!(
111                sha256_matches(&asset_sha_256, expected_sha_256),
112                "{url} asset got SHA-256 mismatch. Expected: {expected_sha_256}, Got: {asset_sha_256}",
113            );
114        }
115
116        util::fs::make_file_executable(&binary_path)
117            .await
118            .with_context(|| format!("marking {binary_path:?} as executable"))?;
119        finalize_download(&staging_path, destination_path).await
120    }
121    .await;
122
123    if let Err(err) = result {
124        if let Err(err) = async_fs::remove_dir_all(&staging_path).await {
125            log::warn!("failed to remove staging directory {staging_path:?}: {err:?}");
126        }
127        return Err(err);
128    }
129
130    Ok(())
131}
132
133async fn extract_to_staging(
134    body: impl AsyncRead + Unpin,
135    digest: Option<&str>,
136    url: &str,
137    staging_path: &Path,
138    asset_kind: AssetKind,
139) -> Result<()> {
140    match digest {
141        Some(expected_sha_256) => {
142            let temp_asset_file = tempfile::NamedTempFile::new()
143                .with_context(|| format!("creating a temporary file for {url}"))?;
144            let (temp_asset_file, _temp_guard) = temp_asset_file.into_parts();
145            let mut writer = HashingWriter {
146                writer: async_fs::File::from(temp_asset_file),
147                hasher: Sha256::new(),
148            };
149            futures::io::copy(&mut BufReader::new(body), &mut writer)
150                .await
151                .with_context(|| {
152                    format!("saving archive contents into the temporary file for {url}")
153                })?;
154            let asset_sha_256 = format!("{:x}", writer.hasher.finalize());
155
156            anyhow::ensure!(
157                sha256_matches(&asset_sha_256, expected_sha_256),
158                "{url} asset got SHA-256 mismatch. Expected: {expected_sha_256}, Got: {asset_sha_256}",
159            );
160            writer
161                .writer
162                .seek(std::io::SeekFrom::Start(0))
163                .await
164                .with_context(|| format!("seeking temporary file for {url}"))?;
165            stream_file_archive(&mut writer.writer, url, staging_path, asset_kind)
166                .await
167                .with_context(|| {
168                    format!("extracting downloaded asset for {url} into {staging_path:?}")
169                })?;
170        }
171        None => {
172            stream_response_archive(body, url, staging_path, asset_kind)
173                .await
174                .with_context(|| {
175                    format!("extracting response for asset {url} into {staging_path:?}")
176                })?;
177        }
178    }
179    Ok(())
180}
181
182fn staging_dir_path(parent: &Path) -> Result<PathBuf> {
183    let dir = tempfile::Builder::new()
184        .prefix(".tmp-github-download-")
185        .tempdir_in(parent)
186        .with_context(|| format!("creating staging directory in {parent:?}"))?;
187    Ok(dir.keep())
188}
189
190fn staging_path(parent: &Path, asset_kind: AssetKind) -> Result<PathBuf> {
191    match asset_kind {
192        AssetKind::TarGz | AssetKind::TarBz2 | AssetKind::Zip => staging_dir_path(parent),
193        AssetKind::Gz => {
194            let path = tempfile::Builder::new()
195                .prefix(".tmp-github-download-")
196                .tempfile_in(parent)
197                .with_context(|| format!("creating staging file in {parent:?}"))?
198                .into_temp_path()
199                .keep()
200                .with_context(|| format!("persisting staging file in {parent:?}"))?;
201            Ok(path)
202        }
203    }
204}
205
206async fn cleanup_staging_path(staging_path: &Path, asset_kind: AssetKind) {
207    match asset_kind {
208        AssetKind::TarGz | AssetKind::TarBz2 | AssetKind::Zip => {
209            if let Err(err) = async_fs::remove_dir_all(staging_path).await {
210                log::warn!("failed to remove staging directory {staging_path:?}: {err:?}");
211            }
212        }
213        AssetKind::Gz => {
214            if let Err(err) = async_fs::remove_file(staging_path).await {
215                log::warn!("failed to remove staging file {staging_path:?}: {err:?}");
216            }
217        }
218    }
219}
220
221async fn finalize_download(staging_path: &Path, destination_path: &Path) -> Result<()> {
222    _ = async_fs::remove_dir_all(destination_path).await;
223    async_fs::rename(staging_path, destination_path)
224        .await
225        .with_context(|| format!("renaming {staging_path:?} to {destination_path:?}"))?;
226    Ok(())
227}
228
229async fn stream_response_archive(
230    response: impl AsyncRead + Unpin,
231    url: &str,
232    destination_path: &Path,
233    asset_kind: AssetKind,
234) -> Result<()> {
235    match asset_kind {
236        AssetKind::TarGz => extract_tar_gz(destination_path, url, response).await?,
237        AssetKind::TarBz2 => extract_tar_bz2(destination_path, url, response).await?,
238        AssetKind::Gz => extract_gz(destination_path, url, response).await?,
239        AssetKind::Zip => {
240            util::archive::extract_zip(destination_path, response).await?;
241        }
242    };
243    Ok(())
244}
245
246async fn stream_file_archive(
247    file_archive: impl AsyncRead + AsyncSeek + Unpin,
248    url: &str,
249    destination_path: &Path,
250    asset_kind: AssetKind,
251) -> Result<()> {
252    match asset_kind {
253        AssetKind::TarGz => extract_tar_gz(destination_path, url, file_archive).await?,
254        AssetKind::TarBz2 => extract_tar_bz2(destination_path, url, file_archive).await?,
255        AssetKind::Gz => extract_gz(destination_path, url, file_archive).await?,
256        #[cfg(not(windows))]
257        AssetKind::Zip => {
258            util::archive::extract_seekable_zip(destination_path, file_archive).await?;
259        }
260        #[cfg(windows)]
261        AssetKind::Zip => {
262            util::archive::extract_zip(destination_path, file_archive).await?;
263        }
264    };
265    Ok(())
266}
267
268async fn extract_tar_gz(
269    destination_path: &Path,
270    url: &str,
271    from: impl AsyncRead + Unpin,
272) -> Result<(), anyhow::Error> {
273    let decompressed_bytes = GzipDecoder::new(BufReader::new(from));
274    unpack_tar_archive(destination_path, url, decompressed_bytes).await?;
275    Ok(())
276}
277
278async fn extract_tar_bz2(
279    destination_path: &Path,
280    url: &str,
281    from: impl AsyncRead + Unpin,
282) -> Result<(), anyhow::Error> {
283    let decompressed_bytes = BzDecoder::new(BufReader::new(from));
284    unpack_tar_archive(destination_path, url, decompressed_bytes).await?;
285    Ok(())
286}
287
288async fn unpack_tar_archive(
289    destination_path: &Path,
290    url: &str,
291    archive_bytes: impl AsyncRead + Unpin,
292) -> Result<(), anyhow::Error> {
293    // We don't need to set the modified time. It's irrelevant to downloaded
294    // archive verification, and some filesystems return errors when asked to
295    // apply it after extraction.
296    let archive = async_tar::ArchiveBuilder::new(archive_bytes)
297        .set_preserve_mtime(false)
298        .build();
299    archive
300        .unpack(&destination_path)
301        .await
302        .with_context(|| format!("extracting {url} to {destination_path:?}"))?;
303    Ok(())
304}
305
306async fn extract_gz(
307    destination_path: &Path,
308    url: &str,
309    from: impl AsyncRead + Unpin,
310) -> Result<(), anyhow::Error> {
311    let mut decompressed_bytes = GzipDecoder::new(BufReader::new(from));
312    let mut file = async_fs::File::create(&destination_path)
313        .await
314        .with_context(|| {
315            format!("creating a file {destination_path:?} for a download from {url}")
316        })?;
317    futures::io::copy(&mut decompressed_bytes, &mut file)
318        .await
319        .with_context(|| format!("extracting {url} to {destination_path:?}"))?;
320    Ok(())
321}
322
323struct HashingWriter<W: AsyncWrite + Unpin> {
324    writer: W,
325    hasher: Sha256,
326}
327
328impl<W: AsyncWrite + Unpin> HashingWriter<W> {
329    /// Closes and drops the inner writer, returning the hex SHA-256 digest of
330    /// everything written.
331    ///
332    /// Taking `self` by value guarantees the writer is dropped before this
333    /// returns. For file writers this releases the OS handle, which Windows
334    /// requires before an ancestor directory can be renamed or deleted; note
335    /// that closing alone is not enough, as `async_fs::File` holds its handle
336    /// until dropped.
337    async fn finish(mut self) -> std::io::Result<String> {
338        self.writer.close().await?;
339        drop(self.writer);
340        Ok(format!("{:x}", self.hasher.finalize()))
341    }
342}
343
344impl<W: AsyncWrite + Unpin> AsyncWrite for HashingWriter<W> {
345    fn poll_write(
346        mut self: Pin<&mut Self>,
347        cx: &mut std::task::Context<'_>,
348        buf: &[u8],
349    ) -> Poll<std::result::Result<usize, std::io::Error>> {
350        match Pin::new(&mut self.writer).poll_write(cx, buf) {
351            Poll::Ready(Ok(n)) => {
352                self.hasher.update(&buf[..n]);
353                Poll::Ready(Ok(n))
354            }
355            other => other,
356        }
357    }
358
359    fn poll_flush(
360        mut self: Pin<&mut Self>,
361        cx: &mut std::task::Context<'_>,
362    ) -> Poll<Result<(), std::io::Error>> {
363        Pin::new(&mut self.writer).poll_flush(cx)
364    }
365
366    fn poll_close(
367        mut self: Pin<&mut Self>,
368        cx: &mut std::task::Context<'_>,
369    ) -> Poll<std::result::Result<(), std::io::Error>> {
370        Pin::new(&mut self.writer).poll_close(cx)
371    }
372}
373
374#[cfg(test)]
375mod tests {
376    use super::*;
377    use crate::{AsyncBody, Response};
378    use futures::future::BoxFuture;
379    use http::HeaderValue;
380    use url::Url;
381
382    struct StaticResponseClient {
383        body: Vec<u8>,
384    }
385
386    impl HttpClient for StaticResponseClient {
387        fn send(
388            &self,
389            _req: http::Request<AsyncBody>,
390        ) -> BoxFuture<'static, anyhow::Result<Response<AsyncBody>>> {
391            let body = self.body.clone();
392            Box::pin(async move {
393                Ok(Response::builder()
394                    .status(200)
395                    .body(AsyncBody::from(body))
396                    .unwrap())
397            })
398        }
399
400        fn user_agent(&self) -> Option<&HeaderValue> {
401            None
402        }
403
404        fn proxy(&self) -> Option<&Url> {
405            None
406        }
407    }
408
409    #[test]
410    fn downloads_raw_binary_with_uppercase_digest_into_destination_dir() {
411        futures::executor::block_on(async {
412            let temp_dir = tempfile::tempdir().unwrap();
413            let destination_path = temp_dir.path().join("v_1");
414            let contents = b"#!/bin/sh\necho hello\n".to_vec();
415            let expected_sha_256 = format!("{:X}", Sha256::digest(&contents));
416            let client = StaticResponseClient { body: contents };
417
418            download_server_raw_binary(
419                &client,
420                "https://example.com/agent-binary",
421                Some(&expected_sha_256),
422                &destination_path,
423                "agent-binary",
424            )
425            .await
426            .unwrap();
427
428            let binary_path = destination_path.join("agent-binary");
429            assert_eq!(
430                std::fs::read(&binary_path).unwrap(),
431                b"#!/bin/sh\necho hello\n"
432            );
433            #[cfg(unix)]
434            {
435                use std::os::unix::fs::PermissionsExt;
436                let mode = std::fs::metadata(&binary_path)
437                    .unwrap()
438                    .permissions()
439                    .mode();
440                assert_eq!(mode & 0o111, 0o111, "binary should be executable");
441            }
442        });
443    }
444
445    #[test]
446    fn raw_binary_digest_mismatch_cleans_up_staging() {
447        futures::executor::block_on(async {
448            let temp_dir = tempfile::tempdir().unwrap();
449            let destination_path = temp_dir.path().join("v_1");
450            let client = StaticResponseClient {
451                body: b"some binary".to_vec(),
452            };
453
454            let error = download_server_raw_binary(
455                &client,
456                "https://example.com/agent-binary",
457                Some("0000000000000000000000000000000000000000000000000000000000000000"),
458                &destination_path,
459                "agent-binary",
460            )
461            .await
462            .unwrap_err();
463
464            assert!(error.to_string().contains("SHA-256 mismatch"));
465            assert!(!destination_path.exists());
466            let leftover_entries = std::fs::read_dir(temp_dir.path()).unwrap().count();
467            assert_eq!(leftover_entries, 0, "staging directory should be removed");
468        });
469    }
470
471    #[test]
472    fn downloads_archive_with_uppercase_digest_and_extracts_contents() {
473        futures::executor::block_on(async {
474            let archive = vec![
475                0x50, 0x4b, 0x03, 0x04, 0x14, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x21, 0x00,
476                0x86, 0xa6, 0x10, 0x36, 0x05, 0x00, 0x00, 0x00, 0x05, 0x00, 0x00, 0x00, 0x05, 0x00,
477                0x00, 0x00, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x68, 0x65, 0x6c, 0x6c, 0x6f, 0x50, 0x4b,
478                0x01, 0x02, 0x14, 0x03, 0x14, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x21, 0x00,
479                0x86, 0xa6, 0x10, 0x36, 0x05, 0x00, 0x00, 0x00, 0x05, 0x00, 0x00, 0x00, 0x05, 0x00,
480                0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x01, 0x00, 0x00,
481                0x00, 0x00, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x50, 0x4b, 0x05, 0x06, 0x00, 0x00, 0x00,
482                0x00, 0x01, 0x00, 0x01, 0x00, 0x33, 0x00, 0x00, 0x00, 0x28, 0x00, 0x00, 0x00, 0x00,
483                0x00,
484            ];
485            let expected_sha_256 = format!("{:X}", Sha256::digest(&archive));
486            let client = StaticResponseClient { body: archive };
487            let temp_dir = tempfile::tempdir().unwrap();
488            let destination_path = temp_dir.path().join("v_1");
489
490            download_server_binary(
491                &client,
492                "https://example.com/agent.zip",
493                Some(&expected_sha_256),
494                &destination_path,
495                AssetKind::Zip,
496            )
497            .await
498            .unwrap();
499
500            assert_eq!(
501                std::fs::read(destination_path.join("agent")).unwrap(),
502                b"hello"
503            );
504        });
505    }
506
507    #[test]
508    fn archive_digest_mismatch_prevents_extraction_and_cleans_up_staging() {
509        futures::executor::block_on(async {
510            let temp_dir = tempfile::tempdir().unwrap();
511            let destination_path = temp_dir.path().join("v_1");
512            let client = StaticResponseClient {
513                body: b"not an archive".to_vec(),
514            };
515
516            let error = download_server_binary(
517                &client,
518                "https://example.com/agent.zip",
519                Some("0000000000000000000000000000000000000000000000000000000000000000"),
520                &destination_path,
521                AssetKind::Zip,
522            )
523            .await
524            .unwrap_err();
525
526            assert!(error.to_string().contains("SHA-256 mismatch"));
527            assert!(!destination_path.exists());
528            let leftover_entries = std::fs::read_dir(temp_dir.path()).unwrap().count();
529            assert_eq!(leftover_entries, 0, "staging directory should be removed");
530        });
531    }
532}
533