From b29a8c50de7c6e1bccde8518b5a7675cd3f36195 Mon Sep 17 00:00:00 2001
From: Florian Bruhin <me@the-compiler.org>
Date: Sat, 16 Jan 2021 11:34:12 +0100
Subject: downloads: Handle data: URLs with explicit filename

This is a follow-up to 9ae08c0f159e8c23e5eb24dc499dc806cc255ae6.

See #1099
Fixes #6041
---
 .../browser/webengine/webenginedownloads.py        | 29 +++++++++++++++-------
 .../browser/webengine/test_webenginedownloads.py   | 25 ++++++++++++++++---
 2 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/qutebrowser/browser/webengine/webenginedownloads.py b/qutebrowser/browser/webengine/webenginedownloads.py
index 206f7b9ad..8814a34b9 100644
--- a/qutebrowser/browser/webengine/webenginedownloads.py
+++ b/qutebrowser/browser/webengine/webenginedownloads.py
@@ -201,7 +201,7 @@ class DownloadItem(downloads.AbstractDownloadItem):
         return None
 
 
-def _get_suggested_filename(path):
+def _strip_suffix(filename):
     """Convert a path we got from chromium to a suggested filename.
 
     Chromium thinks we want to download stuff to ~/Download, so even if we
@@ -211,8 +211,6 @@ def _get_suggested_filename(path):
 
     See https://bugreports.qt.io/browse/QTBUG-56978
     """
-    filename = os.path.basename(path)
-
     suffix_re = re.compile(r"""
       \ ?  # Optional space between filename and suffix
       (
@@ -249,15 +247,28 @@ class DownloadManager(downloads.AbstractDownloadManager):
     @pyqtSlot(QWebEngineDownloadItem)
     def handle_download(self, qt_item):
         """Start a download coming from a QWebEngineProfile."""
-        if qt_item.url().scheme().lower() == 'data':
-            # WORKAROUND for an unknown QtWebEngine bug (?) which gives us base64 data
-            # as filename.
+        qt_filename = os.path.basename(qt_item.path())   # FIXME use 5.14 API
+        mime_type = qt_item.mimeType()
+        url = qt_item.url()
+
+        # WORKAROUND for https://bugreports.qt.io/browse/QTBUG-90355
+        if url.scheme().lower() == 'data':
+            if '/' in url.path().split(',')[-1]:  # e.g. a slash in base64
+                wrong_filename = url.path().split('/')[-1]
+            else:
+                wrong_filename = mime_type.split('/')[1]
+
+            needs_workaround = qt_filename == wrong_filename
+        else:
+            needs_workaround = False
+
+        if needs_workaround:
             suggested_filename = urlutils.filename_from_url(
-                qt_item.url(), fallback='qutebrowser-download')
+                url, fallback='qutebrowser-download')
         else:
-            suggested_filename = _get_suggested_filename(qt_item.path())
+            suggested_filename = _strip_suffix(qt_filename)
 
-        use_pdfjs = pdfjs.should_use_pdfjs(qt_item.mimeType(), qt_item.url())
+        use_pdfjs = pdfjs.should_use_pdfjs(mime_type, url)
 
         download = DownloadItem(qt_item, manager=self)
         self._init_item(download, auto_remove=use_pdfjs,
diff --git a/tests/unit/browser/webengine/test_webenginedownloads.py b/tests/unit/browser/webengine/test_webenginedownloads.py
index b1b18003c..80aada5e5 100644
--- a/tests/unit/browser/webengine/test_webenginedownloads.py
+++ b/tests/unit/browser/webengine/test_webenginedownloads.py
@@ -24,12 +24,11 @@ import pytest
 pytest.importorskip('PyQt5.QtWebEngineWidgets')
 from PyQt5.QtWebEngineWidgets import QWebEngineProfile
 
-from qutebrowser.utils import urlutils
+from qutebrowser.utils import urlutils, usertypes
 from qutebrowser.browser.webengine import webenginedownloads
 
 
 @pytest.mark.parametrize('path, expected', [
-    (os.path.join('subfolder', 'foo'), 'foo'),
     ('foo(1)', 'foo'),
     ('foo (1)', 'foo'),
     ('foo - 1970-01-01T00:00:00.000Z', 'foo'),
@@ -38,8 +37,8 @@ from qutebrowser.browser.webengine import webenginedownloads
     ('foo%20bar', 'foo%20bar'),
     ('foo%2Fbar', 'foo%2Fbar'),
 ])
-def test_get_suggested_filename(path, expected):
-    assert webenginedownloads._get_suggested_filename(path) == expected
+def test_strip_suffix(path, expected):
+    assert webenginedownloads._strip_suffix(path) == expected
 
 
 class TestDataUrlWorkaround:
@@ -95,6 +94,24 @@ class TestDataUrlWorkaround:
         question = message_mock.get_question()
         assert question.default == 'download.pdf'
 
+    def test_explicit_filename(self, webengine_tab, message_mock, qtbot,
+                               pdf_url, download_manager):
+        """If a website sets an explicit filename, we should respect that."""
+        pdf_url_str = pdf_url.toDisplayString()
+        html = f'<a href="{pdf_url_str}" download="filename.pdf" id="link">'
+
+        with qtbot.waitSignal(webengine_tab.load_finished):
+            webengine_tab.set_html(html)
+
+        with qtbot.waitSignal(message_mock.got_question):
+            webengine_tab.elements.find_id(
+                "link",
+                lambda elem: elem.click(usertypes.ClickTarget.normal),
+            )
+
+        question = message_mock.get_question()
+        assert question.default == 'filename.pdf'
+
     @pytest.fixture
     def expected_wrong_filename(self, pdf_bytes):
         with_slash = b'% ?' in pdf_bytes
-- 
cgit v1.2.3-54-g00ecf