From fd4332916fd980f522ef13790a8ae68e55b8dde9 Mon Sep 17 00:00:00 2001 From: Nate Prewitt Date: Fri, 15 Jul 2016 01:31:14 -0400 Subject: [PATCH] raise InvalidHeader on multiple Location values --- requests/exceptions.py | 4 ++++ requests/sessions.py | 7 ++++++- tests/test_requests.py | 16 +++++++++++++++- 3 files changed, 25 insertions(+), 2 deletions(-) diff --git a/requests/exceptions.py b/requests/exceptions.py index 4a6a41c4..72fdfd05 100644 --- a/requests/exceptions.py +++ b/requests/exceptions.py @@ -83,6 +83,10 @@ class InvalidURL(RequestException, ValueError): """ The URL provided was somehow invalid. """ +class InvalidHeader(RequestException, ValueError): + """The header value provided was somehow invalid.""" + + class ChunkedEncodingError(RequestException): """The server declared chunked encoding but sent an invalid chunk.""" diff --git a/requests/sessions.py b/requests/sessions.py index 795c914f..d51897fa 100644 --- a/requests/sessions.py +++ b/requests/sessions.py @@ -20,7 +20,8 @@ from .models import Request, PreparedRequest, DEFAULT_REDIRECT_LIMIT from .hooks import default_hooks, dispatch_hook from .utils import to_key_val_list, default_headers, to_native_string from .exceptions import ( - TooManyRedirects, InvalidScheme, ChunkedEncodingError, ContentDecodingError) + TooManyRedirects, InvalidScheme, ChunkedEncodingError, + ContentDecodingError, InvalidHeader) from .packages.urllib3._collections import RecentlyUsedContainer from .structures import CaseInsensitiveDict @@ -98,6 +99,10 @@ class SessionRedirectMixin(object): request = response.request while response.is_redirect: + if len(response.raw.headers.getlist('location')) > 1: + raise InvalidHeader('Response contains multiple Location headers. ' + 'Unable to perform redirect.') + prepared_request = request.copy() if redirect_count > 0: diff --git a/tests/test_requests.py b/tests/test_requests.py index 3effcf94..f8536f32 100755 --- a/tests/test_requests.py +++ b/tests/test_requests.py @@ -22,7 +22,7 @@ from requests.cookies import cookiejar_from_dict, morsel_to_cookie from requests.exceptions import ( ConnectionError, ConnectTimeout, InvalidScheme, InvalidURL, MissingScheme, ReadTimeout, Timeout, RetryError, TooManyRedirects, - ProxyError) + ProxyError, InvalidHeader) from requests.models import PreparedRequest from requests.structures import CaseInsensitiveDict from requests.sessions import SessionRedirectMixin @@ -222,6 +222,20 @@ class TestRequests: assert r.history[0].status_code == 303 assert r.history[0].is_redirect + def test_multiple_location_headers(self, httpbin): + headers = [('Location', 'http://example.com'), + ('Location', 'https://example.com/1')] + params = '&'.join(['%s=%s' % (k, v) for k, v in headers]) + ses = requests.Session() + req = requests.Request('GET', httpbin('response-headers?%s' % params)) + prep = ses.prepare_request(req) + resp = ses.send(prep) + # change response to redirect + resp.status_code = 302 + with pytest.raises(InvalidHeader): + # next triggers yield on generator + next(ses.resolve_redirects(resp, prep)) + # def test_HTTP_302_ALLOW_REDIRECT_POST(self): # r = requests.post(httpbin('status', '302'), data={'some': 'data'}) # self.assertEqual(r.status_code, 200)