fix: parse Content-Disposition properly, closes #414

Author: felixmosh

runpod/serverless/utils/rp_download.py

@@ -11,8 +11,7 @@
 import uuid
 import zipfile
 from concurrent.futures import ThreadPoolExecutor
-from email import message_from_string
-from typing import List, Union
+from typing import List, Union, Dict
 from urllib.parse import urlparse
 
 import backoff
@@ -34,6 +33,18 @@ def calculate_chunk_size(file_size: int) -> int:
 
     return 1024 * 1024 * 10  # 10 MB
 
+def extract_disposition_params(content_disposition: str)-> Dict[str, str]:
+    parts = (p.strip() for p in content_disposition.split(";"))
+
+    params = {
+        key.strip().lower(): value.strip().strip('"')
+        for part in parts
+        if "=" in part
+        for key, value in [part.split("=", 1)]
+    }
+
+    return params
+
 
 def download_files_from_urls(job_id: str, urls: Union[str, List[str]]) -> List[str]:
     """
@@ -55,8 +66,7 @@ def download_file(url: str, path_to_save: str) -> str:
             content_disposition = response.headers.get("Content-Disposition")
             file_extension = ""
             if content_disposition:
-                msg = message_from_string(f"Content-Disposition: {content_disposition}")
-                params = dict(msg.items())
+                params = extract_disposition_params(content_disposition)
                 file_extension = os.path.splitext(params.get("filename", ""))[1]
 
             # If no extension could be determined from 'Content-Disposition', get it from the URL
@@ -113,15 +123,15 @@ def file(file_url: str) -> dict:
 
     download_response = SyncClientSession().get(file_url, headers=HEADERS, timeout=30)
 
-    original_file_name = []
-    if "Content-Disposition" in download_response.headers.keys():
-        original_file_name = re.findall(
-            "filename=(.+)", download_response.headers["Content-Disposition"]
-        )
+    content_disposition = download_response.headers.get("Content-Disposition")
 
-    if len(original_file_name) > 0:
-        original_file_name = original_file_name[0]
-    else:
+    original_file_name = ""
+    if content_disposition:
+        params = extract_disposition_params(content_disposition)
+
+        original_file_name = params.get("filename", "")
+    
+    if not original_file_name:
         download_path = urlparse(file_url).path
         original_file_name = os.path.basename(download_path)
 

tests/test_serverless/test_utils/test_download.py

@@ -17,6 +17,7 @@
 URL_LIST = [
     "https://example.com/picture.jpg",
     "https://example.com/picture.jpg?X-Amz-Signature=123",
+    "https://example.com/file_without_extension"
 ]
 
 JOB_ID = "job_123"
@@ -86,22 +87,27 @@ def test_download_files_from_urls(self, mock_open_file, mock_get, mock_makedirs)
         """
         Tests download_files_from_urls
         """
+        urls = [
+            "https://example.com/picture.jpg",
+            "https://example.com/file_without_extension"
+        ]
         downloaded_files = download_files_from_urls(
             JOB_ID,
-            [
-                "https://example.com/picture.jpg",
-            ],
+            urls,
         )
 
-        self.assertEqual(len(downloaded_files), 1)
+        self.assertEqual(len(downloaded_files), len(urls))
 
-        # Check that the url was called with SyncClientSession.get
-        self.assertIn("https://example.com/picture.jpg", mock_get.call_args_list[0][0])
+        for index, url in enumerate(urls):
+            # Check that the url was called with SyncClientSession.get
+            self.assertIn(url, mock_get.call_args_list[index][0])
 
-        # Check that the file has the correct extension
-        self.assertTrue(downloaded_files[0].endswith(".jpg"))
+            # Check that the file has the correct extension
+            self.assertTrue(downloaded_files[index].endswith(".jpg"))
+
+            mock_open_file.assert_any_call(downloaded_files[index], "wb")
+        
 
-        mock_open_file.assert_called_once_with(downloaded_files[0], "wb")
         mock_makedirs.assert_called_once_with(
             os.path.abspath(f"jobs/{JOB_ID}/downloaded_files"), exist_ok=True
         )