prioritize filename from response header if not provided

Author: QaidVoid

src/downloader.rs

@@ -8,7 +8,7 @@ use std::{
 
 use futures::{future::join_all, StreamExt};
 use regex::Regex;
-use reqwest::header::USER_AGENT;
+use reqwest::header::{CONTENT_DISPOSITION, USER_AGENT};
 use tokio::{
     fs::{self, OpenOptions},
     io::AsyncWriteExt,
@@ -20,7 +20,7 @@ use url::Url;
 use crate::{
     error::DownloadError,
     oci::{OciClient, OciLayer, OciManifest, Reference},
-    utils::{extract_filename, is_elf, matches_pattern},
+    utils::{extract_filename, extract_filename_from_url, is_elf, matches_pattern},
 };
 
 #[derive(Debug, Clone)]
@@ -83,13 +83,25 @@ impl Downloader {
 
         let filename = options
             .output_path
-            .unwrap_or_else(|| extract_filename(&options.url));
+            .or_else(|| {
+                response
+                    .headers()
+                    .get(CONTENT_DISPOSITION)
+                    .and_then(|header| header.to_str().ok())
+                    .and_then(|header| extract_filename(header))
+            })
+            .or_else(|| extract_filename_from_url(&options.url))
+            .ok_or(DownloadError::FileNameNotFound)?;
+
         let filename = if filename.ends_with('/') {
-            format!(
-                "{}/{}",
-                filename.trim_end_matches('/'),
-                extract_filename(&options.url)
-            )
+            let extracted_name = response
+                .headers()
+                .get(CONTENT_DISPOSITION)
+                .and_then(|header| header.to_str().ok())
+                .and_then(|header| extract_filename(header))
+                .or_else(|| extract_filename_from_url(&options.url))
+                .ok_or(DownloadError::FileNameNotFound)?;
+            format!("{}/{}", filename.trim_end_matches('/'), extracted_name)
         } else {
             filename
         };

src/error.rs

@@ -17,6 +17,7 @@ pub enum DownloadError {
     InvalidResponse,
     LayersNotFound,
     ChunkError,
+    FileNameNotFound,
 }
 
 impl Display for DownloadError {
@@ -31,6 +32,12 @@ impl Display for DownloadError {
             DownloadError::LayersNotFound => write!(f, "No downloadable layers found"),
             DownloadError::InvalidResponse => write!(f, "Failed to parse response"),
             DownloadError::ChunkError => write!(f, "Failed to read chunk"),
+            DownloadError::FileNameNotFound => {
+                write!(
+                    f,
+                    "Couldn't find filename. Please provide filename explicitly."
+                )
+            }
         }
     }
 }
@@ -41,10 +48,7 @@ impl Error for DownloadError {
             DownloadError::IoError(err) => Some(err),
             DownloadError::InvalidUrl { source, .. } => Some(source),
             DownloadError::NetworkError { source } => Some(source),
-            DownloadError::ResourceError { .. } => None,
-            DownloadError::LayersNotFound => None,
-            DownloadError::InvalidResponse => None,
-            DownloadError::ChunkError => None,
+            _ => None,
         }
     }
 }

src/utils.rs

@@ -1,28 +1,37 @@
-use std::{
-    path::Path,
-    time::{SystemTime, UNIX_EPOCH},
-};
+use std::path::Path;
 
 use regex::Regex;
 use reqwest::StatusCode;
 use tokio::{
     fs::File,
     io::{AsyncReadExt, BufReader},
 };
+use url::Url;
 
 pub const ELF_MAGIC_BYTES: [u8; 4] = [0x7f, 0x45, 0x4c, 0x46];
 
-pub fn extract_filename(url: &str) -> String {
-    Path::new(url)
+pub fn extract_filename_from_url(url: &str) -> Option<String> {
+    let url = Url::parse(url)
+        .map(|u| u.path().to_string())
+        .unwrap_or_else(|_| url.to_string());
+    Path::new(&url)
         .file_name()
         .map(|name| name.to_string_lossy().to_string())
-        .unwrap_or_else(|| {
-            let dt = SystemTime::now()
-                .duration_since(UNIX_EPOCH)
-                .unwrap()
-                .as_millis();
-            dt.to_string()
-        })
+}
+
+pub fn extract_filename(header_value: &str) -> Option<String> {
+    header_value.split(';').find_map(|part| {
+        let part = part.trim();
+        if part.starts_with("filename=") {
+            Some(
+                part.trim_start_matches("filename=")
+                    .trim_matches('"')
+                    .to_string(),
+            )
+        } else {
+            None
+        }
+    })
 }
 
 pub async fn is_elf<P: AsRef<Path>>(file_path: P) -> bool {