urlutils: Handle data URLs in filename_from_url

2 files changed, 14 insertions, 0 deletions

diff --git a/qutebrowser/utils/urlutils.py b/qutebrowser/utils/urlutils.py
index 24285404a..fa7867d4d 100644
--- a/qutebrowser/utils/urlutils.py
+++ b/qutebrowser/utils/urlutils.py
@@ -25,6 +25,7 @@ import os.path
 import ipaddress
 import posixpath
 import urllib.parse
+import mimetypes
 from typing import Optional, Tuple, Union
 
 from PyQt5.QtCore import QUrl
@@ -452,6 +453,14 @@ def filename_from_url(url: QUrl, fallback: str = None) -> Optional[str]:
     if not url.isValid():
         return fallback
 
+    if url.scheme().lower() == 'data':
+        media_type, _params, _data = parse_data_url(url)
+        if not media_type:
+            return fallback
+
+        ext = mimetypes.guess_extension(media_type, strict=False) or ''
+        return 'download' + ext
+
     pathname = posixpath.basename(url.path())
     if pathname:
         return pathname
diff --git a/tests/unit/utils/test_urlutils.py b/tests/unit/utils/test_urlutils.py
index 54d195d1a..3372c67a1 100644
--- a/tests/unit/utils/test_urlutils.py
+++ b/tests/unit/utils/test_urlutils.py
@@ -651,6 +651,11 @@ def test_parse_data_url_hypothesis(s):
     (QUrl('http://user:password@qutebrowser.org/foo?bar=baz#fish'), 'foo'),
     (QUrl('http://qutebrowser.org/'), 'qutebrowser.org.html'),
     (QUrl('qute://'), None),
+    # data URL support
+    (QUrl('data:text/plain,'), 'download.txt'),
+    (QUrl('data:application/pdf,'), 'download.pdf'),
+    (QUrl('data:foo/bar,'), 'download'),  # unknown extension
+    (QUrl('data:text/xul,'), 'download.xul'),  # strict=False
 ])
 def test_filename_from_url(qurl, output):
     assert urlutils.filename_from_url(qurl) == output