diff options
| -rw-r--r-- | requests/compat.py | 2 | ||||
| -rw-r--r-- | requests/exceptions.py | 2 | ||||
| -rw-r--r-- | requests/models.py | 14 | ||||
| -rw-r--r-- | requests/sessions.py | 14 | ||||
| -rw-r--r-- | requests/utils.py | 22 | ||||
| -rwxr-xr-x | tests/test_requests.py | 2 |
6 files changed, 51 insertions, 5 deletions
diff --git a/requests/compat.py b/requests/compat.py index eb6530d6..f88e600d 100644 --- a/requests/compat.py +++ b/requests/compat.py @@ -50,6 +50,7 @@ if is_py2: str = unicode basestring = basestring numeric_types = (int, long, float) + integer_types = (int, long) elif is_py3: from urllib.parse import urlparse, urlunparse, urljoin, urlsplit, urlencode, quote, unquote, quote_plus, unquote_plus, urldefrag @@ -64,3 +65,4 @@ elif is_py3: bytes = bytes basestring = (str, bytes) numeric_types = (int, float) + integer_types = (int,) diff --git a/requests/exceptions.py b/requests/exceptions.py index b89e0cc6..0658e7ec 100644 --- a/requests/exceptions.py +++ b/requests/exceptions.py @@ -100,6 +100,8 @@ class StreamConsumedError(RequestException, TypeError): class RetryError(RequestException): """Custom retries logic failed""" +class UnrewindableBodyError(RequestException): + """Requests encountered an error when trying to rewind a body""" # Warnings diff --git a/requests/models.py b/requests/models.py index ba791721..eac7c7d5 100644 --- a/requests/models.py +++ b/requests/models.py @@ -291,6 +291,8 @@ class PreparedRequest(RequestEncodingMixin, RequestHooksMixin): self.body = None #: dictionary of callback hooks, for internal usage. self.hooks = default_hooks() + #: integer denoting starting position of a readable file-like body. + self._body_position = None def prepare(self, method=None, url=None, headers=None, files=None, data=None, params=None, auth=None, cookies=None, hooks=None, json=None): @@ -320,6 +322,7 @@ class PreparedRequest(RequestEncodingMixin, RequestHooksMixin): p._cookies = _copy_cookie_jar(self._cookies) p.body = self.body p.hooks = self.hooks + p._body_position = self._body_position return p def prepare_method(self, method): @@ -447,6 +450,17 @@ class PreparedRequest(RequestEncodingMixin, RequestHooksMixin): if is_stream: body = data + if getattr(body, 'tell', None) is not None: + # Record the current file position before reading. + # This will allow us to rewind a file in the event + # of a redirect. + try: + self._body_position = body.tell() + except (IOError, OSError): + # This differentiates from None, allowing us to catch + # a failed `tell()` later when trying to rewind the body + self._body_position = object() + if files: raise NotImplementedError('Streamed bodies and files are mutually exclusive.') diff --git a/requests/sessions.py b/requests/sessions.py index 971575b2..7983282a 100644 --- a/requests/sessions.py +++ b/requests/sessions.py @@ -28,7 +28,7 @@ from .adapters import HTTPAdapter from .utils import ( requote_uri, get_environ_proxies, get_netrc_auth, should_bypass_proxies, - get_auth_from_url + get_auth_from_url, rewind_body ) from .status_codes import codes @@ -164,6 +164,18 @@ class SessionRedirectMixin(object): proxies = self.rebuild_proxies(prepared_request, proxies) self.rebuild_auth(prepared_request, resp) + # A failed tell() sets `_body_position` to `object()`. This non-None + # value ensures `rewindable` will be True, allowing us to raise an + # UnrewindableBodyError, instead of hanging the connection. + rewindable = ( + prepared_request._body_position is not None and + ('Content-Length' in headers or 'Transfer-Encoding' in headers) + ) + + # Attempt to rewind consumed file-like object. + if rewindable: + rewind_body(prepared_request) + # Override the original request. req = prepared_request diff --git a/requests/utils.py b/requests/utils.py index 227f7c4b..5a2ec34d 100644 --- a/requests/utils.py +++ b/requests/utils.py @@ -23,11 +23,13 @@ from . import certs # to_native_string is unused here, but imported here for backwards compatibility from ._internal_utils import to_native_string from .compat import parse_http_list as _parse_list_header -from .compat import (quote, urlparse, bytes, str, OrderedDict, unquote, - getproxies, proxy_bypass, urlunparse, basestring) +from .compat import ( + quote, urlparse, bytes, str, OrderedDict, unquote, getproxies, + proxy_bypass, urlunparse, basestring, integer_types) from .cookies import RequestsCookieJar, cookiejar_from_dict from .structures import CaseInsensitiveDict -from .exceptions import InvalidURL, InvalidHeader, FileModeWarning +from .exceptions import ( + InvalidURL, InvalidHeader, FileModeWarning, UnrewindableBodyError) _hush_pyflakes = (RequestsCookieJar,) @@ -811,3 +813,17 @@ def urldefragauth(url): netloc = netloc.rsplit('@', 1)[-1] return urlunparse((scheme, netloc, path, params, query, '')) + +def rewind_body(prepared_request): + """Move file pointer back to its recorded starting position + so it can be read again on redirect. + """ + body_seek = getattr(prepared_request.body, 'seek', None) + if body_seek is not None and isinstance(prepared_request._body_position, integer_types): + try: + body_seek(prepared_request._body_position) + except (IOError, OSError): + raise UnrewindableBodyError("An error occured when rewinding request " + "body for redirect.") + else: + raise UnrewindableBodyError("Unable to rewind request body for redirect.") diff --git a/tests/test_requests.py b/tests/test_requests.py index 5a2608e2..b089e4b7 100755 --- a/tests/test_requests.py +++ b/tests/test_requests.py @@ -24,7 +24,7 @@ from requests.cookies import ( from requests.exceptions import ( ConnectionError, ConnectTimeout, InvalidSchema, InvalidURL, MissingSchema, ReadTimeout, Timeout, RetryError, TooManyRedirects, - ProxyError, InvalidHeader) + ProxyError, InvalidHeader, UnrewindableBodyError) from requests.models import PreparedRequest from requests.structures import CaseInsensitiveDict from requests.sessions import SessionRedirectMixin |