Handle body in redirects

This commit is contained in:
Tom Christie 2019-04-29 15:25:44 +01:00
parent 4f23ab4f0d
commit 83d322f736
5 changed files with 54 additions and 3 deletions

View File

@ -10,6 +10,7 @@ from .exceptions import (
PoolTimeout,
ProtocolError,
ReadTimeout,
RedirectBodyUnavailable,
RedirectLoop,
ResponseClosed,
StreamConsumed,

View File

@ -2,7 +2,7 @@ import typing
from urllib.parse import urljoin, urlparse
from ..config import DEFAULT_MAX_REDIRECTS
from ..exceptions import RedirectLoop, TooManyRedirects
from ..exceptions import RedirectBodyUnavailable, RedirectLoop, TooManyRedirects
from ..interfaces import Adapter
from ..models import URL, Headers, Request, Response
from ..status_codes import codes
@ -44,7 +44,8 @@ class RedirectAdapter(Adapter):
method = self.redirect_method(request, response)
url = self.redirect_url(request, response)
headers = self.redirect_headers(request, url)
return Request(method=method, url=url, headers=headers)
body = self.redirect_body(request, method)
return Request(method=method, url=url, headers=headers, body=body)
def redirect_method(self, request: Request, response: Response) -> str:
"""
@ -97,3 +98,10 @@ class RedirectAdapter(Adapter):
if url.origin != request.url.origin:
del headers["Authorization"]
return headers
def redirect_body(self, request: Request, method: str) -> bytes:
if method != request.method and method == "GET":
return b""
if request.is_streaming:
raise RedirectBodyUnavailable()
return request.body

View File

@ -40,6 +40,13 @@ class TooManyRedirects(RedirectError):
"""
class RedirectBodyUnavailable(RedirectError):
"""
Got a redirect response, but the request body was streaming, and is
no longer available.
"""
class RedirectLoop(RedirectError):
"""
Infinite redirect loop.

View File

@ -5,6 +5,7 @@ codes = enum.IntEnum(
[
("continue", 100),
("switching_protocols", 101),
("processing", 102),
("ok", 200),
("created", 201),
("accepted", 202),
@ -13,14 +14,16 @@ codes = enum.IntEnum(
("reset_content", 205),
("partial_content", 206),
("multi_status", 207),
("already_reported", 208),
("im_used", 226),
("multiple_choices", 300),
("moved_permanently", 301),
("found", 302),
("see_other", 303),
("not_modified", 304),
("use_proxy", 305),
("reserved", 306),
("temporary_redirect", 307),
("permanent_redirect", 308),
("bad_request", 400),
("unauthorized", 401),
("payment_required", 402),

View File

@ -7,6 +7,7 @@ from httpcore import (
URL,
Adapter,
RedirectAdapter,
RedirectBodyUnavailable,
RedirectLoop,
Request,
Response,
@ -74,6 +75,14 @@ class MockDispatch(Adapter):
body = json.dumps({"headers": headers}).encode()
return Response(codes.ok, body=body, request=request)
elif request.url.path == "/redirect_body":
headers = [(b"location", b"/redirect_body_target")]
return Response(codes.permanent_redirect, headers=headers, request=request)
elif request.url.path == "/redirect_body_target":
body = json.dumps({"body": request.body.decode()}).encode()
return Response(codes.ok, body=body, request=request)
return Response(codes.ok, body=b"Hello, world!", request=request)
@ -178,3 +187,26 @@ async def test_same_domain_redirect():
data = json.loads(response.body.decode())
assert response.url == URL("https://example.org/cross_domain_target")
assert data == {"headers": {"authorization": "abc"}}
@pytest.mark.asyncio
async def test_body_redirect():
client = RedirectAdapter(MockDispatch())
body = b"Example request body"
url = "https://example.org/redirect_body"
response = await client.request("POST", url, body=body)
data = json.loads(response.body.decode())
assert response.url == URL("https://example.org/redirect_body_target")
assert data == {"body": "Example request body"}
@pytest.mark.asyncio
async def test_cannot_redirect_streaming_body():
client = RedirectAdapter(MockDispatch())
url = "https://example.org/redirect_body"
async def body():
yield b"Example request body"
with pytest.raises(RedirectBodyUnavailable):
await client.request("POST", url, body=body())