preserve invalid itms-services url scheme (#2695)

pallets · May 6, 2023 · 5a149fa · 5a149fa
2 parents 3072610 + 6822773
commit 5a149fa
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 45 deletions.
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -10,6 +10,8 @@ Unreleased
 -   Remove usage of ``warnings.catch_warnings``. :issue:`2690`
 -   Remove ``max_form_parts`` restriction from standard form data parsing and only use
     if for multipart content. :pr:`2694`
+-   ``Response`` will avoid converting the ``Location`` header in some cases to preserve
+    invalid URL schemes like ``itms-services``. :issue:`2691`
 
 
 Version 2.3.3

diff --git a/src/werkzeug/urls.py b/src/werkzeug/urls.py
@@ -1053,6 +1053,27 @@ def iri_to_uri(
     return urlunsplit((parts.scheme, netloc, path, query, fragment))
 
 
+def _invalid_iri_to_uri(iri: str) -> str:
+    """The URL scheme ``itms-services://`` must contain the ``//`` even though it does
+    not have a host component. There may be other invalid schemes as well. Currently,
+    responses will always call ``iri_to_uri`` on the redirect ``Location`` header, which
+    removes the ``//``. For now, if the IRI only contains ASCII and does not contain
+    spaces, pass it on as-is. In Werkzeug 2.4, this should become a
+    ``response.process_location`` flag.
+
+    :meta private:
+    """
+    try:
+        iri.encode("ascii")
+    except UnicodeError:
+        pass
+    else:
+        if len(iri.split(None, 1)) == 1:
+            return iri
+
+    return iri_to_uri(iri)
+
+
 def url_decode(
     s: t.AnyStr,
     charset: str = "utf-8",

diff --git a/src/werkzeug/utils.py b/src/werkzeug/utils.py
@@ -260,21 +260,17 @@ def redirect(
         response. The default is :class:`werkzeug.wrappers.Response` if
         unspecified.
     """
-    from .urls import iri_to_uri
-
     if Response is None:
         from .wrappers import Response
 
-    display_location = escape(location)
-    location = iri_to_uri(location)
+    html_location = escape(location)
     response = Response(  # type: ignore[misc]
         "<!doctype html>\n"
         "<html lang=en>\n"
         "<title>Redirecting...</title>\n"
         "<h1>Redirecting...</h1>\n"
         "<p>You should be redirected automatically to the target URL: "
-        f'<a href="{escape(location)}">{display_location}</a>. If'
-        " not, click the link.\n",
+        f'<a href="{html_location}">{html_location}</a>. If not, click the link.\n',
         code,
         mimetype="text/html",
     )

diff --git a/src/werkzeug/wrappers/response.py b/src/werkzeug/wrappers/response.py
@@ -8,6 +8,7 @@
 from ..datastructures import Headers
 from ..http import remove_entity_headers
 from ..sansio.response import Response as _SansIOResponse
+from ..urls import _invalid_iri_to_uri
 from ..urls import iri_to_uri
 from ..utils import cached_property
 from ..wsgi import ClosingIterator
@@ -478,11 +479,11 @@ def get_wsgi_headers(self, environ: WSGIEnvironment) -> Headers:
             elif ikey == "content-length":
                 content_length = value
 
-        # make sure the location header is an absolute URL
         if location is not None:
-            location = iri_to_uri(location)
+            location = _invalid_iri_to_uri(location)
 
             if self.autocorrect_location_header:
+                # Make the location header an absolute URL.
                 current_url = get_current_url(environ, strip_querystring=True)
                 current_url = iri_to_uri(current_url)
                 location = urljoin(current_url, location)

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import inspect
 from datetime import datetime
 
@@ -9,48 +11,32 @@
 from werkzeug.http import http_date
 from werkzeug.http import parse_date
 from werkzeug.test import Client
+from werkzeug.test import EnvironBuilder
 from werkzeug.wrappers import Response
 
 
-def test_redirect():
-    resp = utils.redirect("/füübär")
-    assert resp.headers["Location"] == "/f%C3%BC%C3%BCb%C3%A4r"
-    assert resp.status_code == 302
-    assert resp.get_data() == (
-        b"<!doctype html>\n"
-        b"<html lang=en>\n"
-        b"<title>Redirecting...</title>\n"
-        b"<h1>Redirecting...</h1>\n"
-        b"<p>You should be redirected automatically to the target URL: "
-        b'<a href="/f%C3%BC%C3%BCb%C3%A4r">/f\xc3\xbc\xc3\xbcb\xc3\xa4r</a>. '
-        b"If not, click the link.\n"
-    )
+@pytest.mark.parametrize(
+    ("url", "code", "expect"),
+    [
+        ("http://example.com", None, "http://example.com"),
+        ("/füübär", 305, "/f%C3%BC%C3%BCb%C3%A4r"),
+        ("http://☃.example.com/", 307, "http://xn--n3h.example.com/"),
+        ("itms-services://?url=abc", None, "itms-services://?url=abc"),
+    ],
+)
+def test_redirect(url: str, code: int | None, expect: str) -> None:
+    environ = EnvironBuilder().get_environ()
 
-    resp = utils.redirect("http://☃.net/", 307)
-    assert resp.headers["Location"] == "http://xn--n3h.net/"
-    assert resp.status_code == 307
-    assert resp.get_data() == (
-        b"<!doctype html>\n"
-        b"<html lang=en>\n"
-        b"<title>Redirecting...</title>\n"
-        b"<h1>Redirecting...</h1>\n"
-        b"<p>You should be redirected automatically to the target URL: "
-        b'<a href="http://xn--n3h.net/">http://\xe2\x98\x83.net/</a>. '
-        b"If not, click the link.\n"
-    )
+    if code is None:
+        resp = utils.redirect(url)
+        assert resp.status_code == 302
+    else:
+        resp = utils.redirect(url, code)
+        assert resp.status_code == code
 
-    resp = utils.redirect("http://example.com/", 305)
-    assert resp.headers["Location"] == "http://example.com/"
-    assert resp.status_code == 305
-    assert resp.get_data() == (
-        b"<!doctype html>\n"
-        b"<html lang=en>\n"
-        b"<title>Redirecting...</title>\n"
-        b"<h1>Redirecting...</h1>\n"
-        b"<p>You should be redirected automatically to the target URL: "
-        b'<a href="http://example.com/">http://example.com/</a>. '
-        b"If not, click the link.\n"
-    )
+    assert resp.headers["Location"] == url
+    assert resp.get_wsgi_headers(environ)["Location"] == expect
+    assert resp.get_data(as_text=True).count(url) == 2
 
 
 def test_redirect_xss():