add support for download OCI packages

Author: QaidVoid

Cargo.lock

@@ -21,6 +21,10 @@ pub struct Args {
     #[arg(required = false, long)]
     pub gitlab: Vec<String>,
 
+    /// OCI reference
+    #[arg(required = false, long)]
+    pub ghcr: Vec<String>,
+
     /// Links to files
     #[arg(required = false)]
     pub links: Vec<String>,

src/bin/soar-dl/cli.rs

@@ -31,6 +31,7 @@ impl DownloadManager {
 
     pub async fn execute(&self) {
         let _ = self.handle_github_downloads().await;
+        let _ = self.handle_oci_downloads().await;
         let _ = self.handle_gitlab_downloads().await;
         let _ = self.handle_direct_downloads().await;
     }
@@ -120,6 +121,28 @@ impl DownloadManager {
         Ok(())
     }
 
+    async fn handle_oci_downloads(&self) -> Result<(), PlatformError> {
+        if self.args.ghcr.is_empty() {
+            return Ok(());
+        }
+
+        let downloader = Downloader::default();
+        for reference in &self.args.ghcr {
+            println!("Downloading using OCI reference: {}", reference);
+
+            let options = DownloadOptions {
+                url: reference.clone(),
+                output_path: self.args.output.clone(),
+                progress_callback: Some(self.progress_callback.clone()),
+            };
+            let _ = downloader
+                .download_oci(options)
+                .await
+                .map_err(|e| eprintln!("{}", e));
+        }
+        Ok(())
+    }
+
     async fn handle_direct_downloads(&self) -> Result<(), DownloadError> {
         let downloader = Downloader::default();
         for link in &self.args.links {
@@ -161,6 +184,19 @@ impl DownloadManager {
                         eprintln!("{}", e);
                     }
                 }
+                Ok(PlatformUrl::Oci(url)) => {
+                    println!("Downloading using OCI reference: {}", url);
+
+                    let options = DownloadOptions {
+                        url: link.clone(),
+                        output_path: self.args.output.clone(),
+                        progress_callback: Some(self.progress_callback.clone()),
+                    };
+                    let _ = downloader
+                        .download_oci(options)
+                        .await
+                        .map_err(|e| eprintln!("{}", e));
+                }
                 Err(err) => eprintln!("Error parsing URL '{}' : {}", link, err),
             };
         }

src/bin/soar-dl/download_manager.rs

@@ -1,15 +1,22 @@
-use std::{fs::Permissions, os::unix::fs::PermissionsExt, path::Path, sync::Arc};
+use std::{
+    fs::Permissions,
+    os::unix::fs::PermissionsExt,
+    path::Path,
+    sync::{Arc, Mutex},
+};
 
-use futures::StreamExt;
+use futures::{future::join_all, StreamExt};
 use reqwest::header::USER_AGENT;
 use tokio::{
     fs::{self, OpenOptions},
     io::AsyncWriteExt,
+    task,
 };
 use url::Url;
 
 use crate::{
     error::DownloadError,
+    oci::{OciClient, Reference},
     utils::{extract_filename, is_elf},
 };
 
@@ -120,4 +127,63 @@ impl Downloader {
 
         Ok(filename)
     }
+
+    pub async fn download_oci(&self, options: DownloadOptions) -> Result<(), DownloadError> {
+        let url = options.url.clone();
+        let reference: Reference = url.into();
+        let oci_client = OciClient::new(reference);
+
+        let manifest = oci_client.manifest().await.unwrap();
+
+        let mut tasks = Vec::new();
+        let total_bytes: u64 = manifest.layers.iter().map(|layer| layer.size).sum();
+
+        if let Some(ref callback) = options.progress_callback {
+            callback(DownloadState::Progress(DownloadProgress {
+                bytes_downloaded: 0,
+                total_bytes: Some(total_bytes),
+                url: options.url.clone(),
+                file_path: String::new(),
+            }));
+        }
+
+        let downloaded_bytes = Arc::new(Mutex::new(0u64));
+        let outdir = options.output_path;
+
+        for layer in manifest.layers {
+            let client_clone = oci_client.clone();
+            let cb_clone = options.progress_callback.clone();
+            let downloaded_bytes = downloaded_bytes.clone();
+            let url = options.url.clone();
+            let outdir = outdir.clone();
+
+            let task = task::spawn(async move {
+                let chunk_size = client_clone
+                    .pull_layer(&layer, outdir, move |bytes| {
+                        if let Some(ref callback) = cb_clone {
+                            let mut current = downloaded_bytes.lock().unwrap();
+                            *current = bytes;
+                            callback(DownloadState::Progress(DownloadProgress {
+                                bytes_downloaded: *current,
+                                total_bytes: Some(total_bytes),
+                                url: url.clone(),
+                                file_path: String::new(),
+                            }));
+                        }
+                    })
+                    .await?;
+
+                Ok::<u64, DownloadError>(chunk_size)
+            });
+            tasks.push(task);
+        }
+
+        let _ = join_all(tasks).await;
+
+        if let Some(ref callback) = options.progress_callback {
+            callback(DownloadState::Complete);
+        }
+
+        Ok(())
+    }
 }

src/downloader.rs

@@ -14,6 +14,7 @@ pub enum DownloadError {
         url: String,
         status: reqwest::StatusCode,
     },
+    InvalidResponse,
 }
 
 impl Display for DownloadError {
@@ -25,6 +26,7 @@ impl Display for DownloadError {
             DownloadError::ResourceError { url, status } => {
                 write!(f, "Failed to fetch resource from {} [{}]", url, status)
             }
+            DownloadError::InvalidResponse => write!(f, "Failed to parse response"),
         }
     }
 }
@@ -36,6 +38,7 @@ impl Error for DownloadError {
             DownloadError::InvalidUrl { source, .. } => Some(source),
             DownloadError::NetworkError { source } => Some(source),
             DownloadError::ResourceError { .. } => None,
+            DownloadError::InvalidResponse => None,
         }
     }
 }

src/error.rs

@@ -2,5 +2,6 @@ pub mod downloader;
 pub mod error;
 pub mod github;
 pub mod gitlab;
+pub mod oci;
 pub mod platform;
 pub mod utils;

src/lib.rs

@@ -0,0 +1,188 @@
+use std::{
+    collections::HashMap,
+    path::{Path, PathBuf},
+};
+
+use futures::StreamExt;
+use reqwest::header::{self, HeaderMap};
+use serde::Deserialize;
+use tokio::{
+    fs::{self, OpenOptions},
+    io::AsyncWriteExt,
+};
+
+use crate::error::DownloadError;
+
+#[derive(Deserialize)]
+pub struct OciLayer {
+    #[serde(rename = "mediaType")]
+    pub media_type: String,
+    pub digest: String,
+    pub size: u64,
+    pub annotations: HashMap<String, String>,
+}
+
+#[derive(Deserialize)]
+pub struct OciConfig {
+    #[serde(rename = "mediaType")]
+    pub media_type: String,
+    pub digest: String,
+    pub size: u64,
+}
+
+#[derive(Deserialize)]
+pub struct OciManifest {
+    #[serde(rename = "mediaType")]
+    pub media_type: String,
+    pub config: OciConfig,
+    pub layers: Vec<OciLayer>,
+}
+
+#[derive(Clone)]
+pub struct OciClient {
+    client: reqwest::Client,
+    reference: Reference,
+}
+
+#[derive(Clone)]
+pub struct Reference {
+    package: String,
+    tag: String,
+}
+
+impl From<&str> for Reference {
+    fn from(value: &str) -> Self {
+        let paths = value.trim_start_matches("ghcr.io/");
+        let (package, tag) = paths.split_once(':').unwrap_or((paths, "latest"));
+
+        Self {
+            package: package.to_string(),
+            tag: tag.to_string(),
+        }
+    }
+}
+
+impl From<String> for Reference {
+    fn from(value: String) -> Self {
+        let paths = value.trim_start_matches("ghcr.io/");
+        let (package, tag) = paths.split_once(':').unwrap_or((paths, "latest"));
+
+        Self {
+            package: package.to_string(),
+            tag: tag.to_string(),
+        }
+    }
+}
+
+impl OciClient {
+    pub fn new(reference: Reference) -> Self {
+        let client = reqwest::Client::new();
+        Self { client, reference }
+    }
+
+    pub fn headers(&self) -> HeaderMap {
+        let mut header_map = HeaderMap::new();
+        header_map.insert(header::ACCEPT, "application/vnd.docker.distribution.manifest.v2+json, application/vnd.docker.distribution.manifest.list.v2+json, application/vnd.oci.image.manifest.v1+json, application/vnd.oci.image.index.v1+json, application/vnd.oci.artifact.manifest.v1+json".parse().unwrap());
+        header_map.insert(header::AUTHORIZATION, "Bearer QQ==".parse().unwrap());
+        header_map
+    }
+
+    pub async fn manifest(&self) -> Result<OciManifest, DownloadError> {
+        let manifest_url = format!(
+            "https://ghcr.io/v2/{}/manifests/{}",
+            self.reference.package, self.reference.tag
+        );
+        let resp = self
+            .client
+            .get(&manifest_url)
+            .headers(self.headers())
+            .send()
+            .await
+            .map_err(|err| DownloadError::NetworkError { source: err })?;
+
+        if !resp.status().is_success() {
+            return Err(DownloadError::ResourceError {
+                status: resp.status(),
+                url: manifest_url,
+            });
+        }
+
+        let manifest: OciManifest = resp
+            .json()
+            .await
+            .map_err(|_| DownloadError::InvalidResponse)?;
+        Ok(manifest)
+    }
+
+    pub async fn pull_layer<F, P: AsRef<Path>>(
+        &self,
+        layer: &OciLayer,
+        output_dir: Option<P>,
+        progress_callback: F,
+    ) -> Result<u64, DownloadError>
+    where
+        F: Fn(u64) + Send + 'static,
+    {
+        let blob_url = format!(
+            "https://ghcr.io/v2/{}/blobs/{}",
+            self.reference.package, layer.digest
+        );
+        let resp = self
+            .client
+            .get(&blob_url)
+            .headers(self.headers())
+            .send()
+            .await
+            .map_err(|err| DownloadError::NetworkError { source: err })?;
+
+        if !resp.status().is_success() {
+            return Err(DownloadError::ResourceError {
+                status: resp.status(),
+                url: blob_url,
+            });
+        }
+
+        let filename = layer.get_title().unwrap();
+        let (temp_path, final_path) = if let Some(output_dir) = output_dir {
+            let output_dir = output_dir.as_ref();
+            fs::create_dir_all(output_dir).await?;
+            let final_path = output_dir.join(format!("{filename}"));
+            let temp_path = output_dir.join(format!("{filename}.part"));
+            (temp_path, final_path)
+        } else {
+            let final_path = PathBuf::from(&filename);
+            let temp_path = PathBuf::from(format!("{filename}.part"));
+            (temp_path, final_path)
+        };
+
+        let mut file = OpenOptions::new()
+            .create(true)
+            .append(true)
+            .open(&temp_path)
+            .await?;
+
+        let mut stream = resp.bytes_stream();
+        let mut total_bytes_downloaded = 0;
+
+        while let Some(chunk) = stream.next().await {
+            let chunk = chunk.unwrap();
+            let chunk_size = chunk.len() as u64;
+            file.write_all(&chunk).await.unwrap();
+
+            progress_callback(chunk_size);
+            total_bytes_downloaded += chunk_size;
+        }
+
+        fs::rename(&temp_path, &final_path).await?;
+
+        Ok(total_bytes_downloaded)
+    }
+}
+
+impl OciLayer {
+    pub fn get_title(&self) -> Option<String> {
+        self.annotations
+            .get("org.opencontainers.image.title")
+            .cloned()
+    }
+}