From 9160631d7fb957a319ce8985c97a858f8eab49a3 Mon Sep 17 00:00:00 2001
From: Josep Cugat <jcugat@gmail.com>
Date: Wed, 27 May 2020 21:18:51 +0200
Subject: [PATCH] Improve error when redirecting with custom schemes

Fixes #822
---
 httpx/_client.py               |  4 ++++
 tests/client/test_redirects.py | 20 ++++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/httpx/_client.py b/httpx/_client.py
index 7248f87ce9..2c56edc560 100644
--- a/httpx/_client.py
+++ b/httpx/_client.py
@@ -322,6 +322,10 @@ def redirect_url(self, request: Request, response: Response) -> URL:
 
         url = URL(location, allow_relative=True)
 
+        # Check that we can handle the scheme
+        if url.scheme and url.scheme not in ("http", "https"):
+            raise InvalidURL(f'Scheme "{url.scheme}" not supported.')
+
         # Handle malformed 'Location' headers that are "absolute" form, have no host.
         # See: https://github.com/encode/httpx/issues/771
         if url.scheme and not url.host:
diff --git a/tests/client/test_redirects.py b/tests/client/test_redirects.py
index e91bb7e66f..fa5ae4eb53 100644
--- a/tests/client/test_redirects.py
+++ b/tests/client/test_redirects.py
@@ -8,6 +8,7 @@
 from httpx import (
     URL,
     AsyncClient,
+    InvalidURL,
     NotRedirectResponse,
     RequestBodyUnavailable,
     TooManyRedirects,
@@ -140,6 +141,17 @@ async def body():
             else:
                 return b"HTTP/1.1", 200, b"OK", [], ByteStream(b"Hello, world!")
 
+        elif path == b"/redirect_custom_scheme":
+            status_code = codes.MOVED_PERMANENTLY
+            headers = [(b"location", b"market://details?id=42")]
+            return (
+                b"HTTP/1.1",
+                status_code,
+                b"Moved Permanently",
+                headers,
+                ByteStream(b""),
+            )
+
         return b"HTTP/1.1", 200, b"OK", [], ByteStream(b"Hello, world!")
 
 
@@ -431,3 +443,11 @@ async def test_redirect_cookie_behavior():
     response = await client.get("https://example.com/")
     assert response.url == "https://example.com/"
     assert response.text == "Not logged in"
+
+
+@pytest.mark.usefixtures("async_environment")
+async def test_redirect_custom_scheme():
+    client = AsyncClient(dispatch=MockDispatch())
+    with pytest.raises(InvalidURL) as e:
+        await client.post("https://example.org/redirect_custom_scheme")
+    assert str(e.value) == 'Scheme "market" not supported.'