This repository was archived by the owner on Jan 13, 2021. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 197
Expand file tree
/
Copy pathresponse.py
More file actions
341 lines (279 loc) · 12.5 KB
/
Copy pathresponse.py
File metadata and controls
341 lines (279 loc) · 12.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
# -*- coding: utf-8 -*-
"""
hyper/http11/response
~~~~~~~~~~~~~~~~~~~~~
Contains the HTTP/1.1 equivalent of the HTTPResponse object defined in
httplib/http.client.
"""
import logging
import weakref
import zlib
import brotli
from ..common.decoder import DeflateDecoder
from ..common.exceptions import ChunkedDecodeError, InvalidResponseError
from ..common.exceptions import ConnectionResetError
from ..common.util import HTTPVersion
log = logging.getLogger(__name__)
class HTTP11Response(object):
"""
An ``HTTP11Response`` wraps the HTTP/1.1 response from the server. It
provides access to the response headers and the entity body. The response
is an iterable object and can be used in a with statement.
"""
version = HTTPVersion.http11
def __init__(self, code, reason, headers, sock, connection=None,
request_method=None):
#: The reason phrase returned by the server.
self.reason = reason
#: The status code returned by the server.
self.status = code
#: The response headers. These are determined upon creation, assigned
#: once, and never assigned again.
self.headers = headers
#: The response trailers. These are always intially ``None``.
self.trailers = None
# The socket this response is being sent over.
self._sock = sock
# Whether we expect the connection to be closed. If we do, we don't
# bother checking for content-length, we just keep reading until
# we no longer can.
self._expect_close = False
if b'close' in self.headers.get(b'connection', []):
self._expect_close = True
# The expected length of the body.
if request_method != b'HEAD':
try:
self._length = int(self.headers[b'content-length'][0])
except KeyError:
self._length = None
else:
self._length = 0
# Whether we expect a chunked response.
self._chunked = (
b'chunked' in self.headers.get(b'transfer-encoding', [])
)
# When content-length is absent and response is not chunked,
# body length is determined by connection closure.
# https://tools.ietf.org/html/rfc7230#section-3.3.3
if self._length is None and not self._chunked:
# 200 response to a CONNECT request means that proxy has connected
# to the target host and it will start forwarding everything sent
# from the either side. Thus we must not try to read body of this
# response. Socket of current connection will be taken over by
# the code that has sent a CONNECT request.
if not (request_method is not None and
b'CONNECT' == request_method.upper() and
code == 200):
self._expect_close = True
# This object is used for decompressing gzipped request bodies. Right
# now we only support gzip because that's all the RFC mandates of us.
# Later we'll add support for more encodings.
# This 16 + MAX_WBITS nonsense is to force gzip. See this
# Stack Overflow answer for more:
# http://stackoverflow.com/a/2695466/1401686
if b'gzip' in self.headers.get(b'content-encoding', []):
self._decompressobj = zlib.decompressobj(16 + zlib.MAX_WBITS)
elif b'br' in self.headers.get(b'content-encoding', []):
self._decompressobj = brotli.Decompressor()
elif b'deflate' in self.headers.get(b'content-encoding', []):
self._decompressobj = DeflateDecoder()
else:
self._decompressobj = None
# This is a reference that allows for the Response class to tell the
# parent connection object to throw away its socket object. This is to
# be used when the connection is genuinely closed, so that the user
# can keep using the Connection object.
# Strictly, we take a weakreference to this so that we don't set up a
# reference cycle.
if connection is not None:
self._parent = weakref.ref(connection)
else:
self._parent = None
self._buffered_data = b''
self._chunker = None
def read(self, amt=None, decode_content=True):
"""
Reads the response body, or up to the next ``amt`` bytes.
:param amt: (optional) The amount of data to read. If not provided, all
the data will be read from the response.
:param decode_content: (optional) If ``True``, will transparently
decode the response data.
:returns: The read data. Note that if ``decode_content`` is set to
``True``, the actual amount of data returned may be different to
the amount requested.
"""
# Return early if we've lost our connection.
if self._sock is None:
return b''
if self._chunked:
return self._normal_read_chunked(amt, decode_content)
# If we're asked to do a read without a length, we need to read
# everything. That means either the entire content length, or until the
# socket is closed, depending.
if amt is None:
if self._length is not None:
amt = self._length
elif self._expect_close:
return self._read_expect_closed(decode_content)
else: # pragma: no cover
raise InvalidResponseError(
"Response must either have length or Connection: close"
)
# Otherwise, we've been asked to do a bounded read. We should read no
# more than the remaining length, obviously.
# FIXME: Handle cases without _length
if self._length is not None:
amt = min(amt, self._length)
# Now, issue reads until we read that length. This is to account for
# the fact that it's possible that we'll be asked to read more than
# 65kB in one shot.
to_read = amt
chunks = []
# Ideally I'd like this to read 'while to_read', but I want to be
# defensive against the admittedly unlikely case that the socket
# returns *more* data than I want.
while to_read > 0:
chunk = self._sock.recv(amt).tobytes()
# If we got an empty read, but were expecting more, the remote end
# has hung up. Raise an exception if we were expecting more data,
# but if we were expecting the remote end to close then it's ok.
if not chunk:
if self._length is not None or not self._expect_close:
self.close(socket_close=True)
raise ConnectionResetError("Remote end hung up!")
break
to_read -= len(chunk)
chunks.append(chunk)
data = b''.join(chunks)
if self._length is not None:
self._length -= len(data)
# If we're at the end of the request, we have some cleaning up to do.
# Close the stream, and if necessary flush the buffer. Checking that
# we're at the end is actually obscenely complex: either we've read the
# full content-length or, if we were expecting a closed connection,
# we've had a read shorter than the requested amount. We also have to
# do this before we try to decompress the body.
end_of_request = (self._length == 0 or
(self._expect_close and len(data) < amt))
# We may need to decode the body.
if decode_content and self._decompressobj and data:
data = self._decompressobj.decompress(data)
if decode_content and self._decompressobj and end_of_request:
data += self._decompressobj.flush()
# We're at the end. Close the connection. Explicit check for zero here
# because self._length might be None.
if end_of_request:
self.close(socket_close=self._expect_close)
return data
def read_chunked(self, decode_content=True):
"""
Reads chunked transfer encoded bodies. This method returns a generator:
each iteration of which yields one chunk *unless* the chunks are
compressed, in which case it yields whatever the decompressor provides
for each chunk.
.. warning:: This may yield the empty string, without that being the
end of the body!
"""
if not self._chunked:
raise ChunkedDecodeError(
"Attempted chunked read of non-chunked body."
)
# Return early if possible.
if self._sock is None:
return
while True:
# Read to the newline to get the chunk length. This is a
# hexadecimal integer.
chunk_length = int(self._sock.readline().tobytes().strip(), 16)
data = b''
# If the chunk length is zero, consume the newline and then we're
# done. If we were decompressing data, return the remaining data.
if not chunk_length:
self._sock.readline()
if decode_content and self._decompressobj:
yield self._decompressobj.flush()
self.close(socket_close=self._expect_close)
break
# Then read that many bytes.
while chunk_length > 0:
chunk = self._sock.recv(chunk_length).tobytes()
data += chunk
chunk_length -= len(chunk)
assert chunk_length == 0
# Now, consume the newline.
self._sock.readline()
# We may need to decode the body.
if decode_content and self._decompressobj and data:
data = self._decompressobj.decompress(data)
yield data
return
def close(self, socket_close=False):
"""
Close the response. This causes the Response to lose access to the
backing socket. In some cases, it can also cause the backing connection
to be torn down.
:param socket_close: Whether to close the backing socket.
:returns: Nothing.
"""
if socket_close and self._parent is not None:
# The double call is necessary because we need to dereference the
# weakref. If the weakref is no longer valid, that's fine, there's
# no connection object to tell.
parent = self._parent()
if parent is not None:
parent.close()
self._sock = None
def _read_expect_closed(self, decode_content):
"""
Implements the logic for an unbounded read on a socket that we expect
to be closed by the remote end.
"""
# In this case, just read until we cannot read anymore. Then, close the
# socket, becuase we know we have to.
chunks = []
while True:
try:
chunk = self._sock.recv(65535).tobytes()
if not chunk:
break
except ConnectionResetError:
break
else:
chunks.append(chunk)
self.close(socket_close=True)
# We may need to decompress the data.
data = b''.join(chunks)
if decode_content and self._decompressobj:
data = self._decompressobj.decompress(data)
data += self._decompressobj.flush()
return data
def _normal_read_chunked(self, amt, decode_content):
"""
Implements the logic for calling ``read()`` on a chunked response.
"""
# If we're doing a full read, read it as chunked and then just join
# the chunks together!
if amt is None:
return self._buffered_data + b''.join(self.read_chunked())
if self._chunker is None:
self._chunker = self.read_chunked()
# Otherwise, we have a certain amount of data we want to read.
current_amount = len(self._buffered_data)
extra_data = [self._buffered_data]
while current_amount < amt:
try:
chunk = next(self._chunker)
except StopIteration:
self.close(socket_close=self._expect_close)
break
current_amount += len(chunk)
extra_data.append(chunk)
data = b''.join(extra_data)
self._buffered_data = data[amt:]
return data[:amt]
# The following methods implement the context manager protocol.
def __enter__(self):
return self
def __exit__(self, *args):
self.close()
return False # Never swallow exceptions.