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 pathstream.py
More file actions
344 lines (288 loc) · 11.8 KB
/
Copy pathstream.py
File metadata and controls
344 lines (288 loc) · 11.8 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
342
343
344
# -*- coding: utf-8 -*-
"""
hyper/http20/stream
~~~~~~~~~~~~~~~~~~~
Objects that make up the stream-level abstraction of hyper's HTTP/2 support.
These objects are not expected to be part of the public HTTP/2 API: they're
intended purely for use inside hyper's HTTP/2 abstraction.
Conceptually, a single HTTP/2 connection is made up of many streams: each
stream is an independent, bi-directional sequence of HTTP headers and data.
Each stream is identified by a monotonically increasing integer, assigned to
the stream by the endpoint that initiated the stream.
"""
import h2.exceptions
from ..common.headers import HTTPHeaderMap
from .util import h2_safe_headers
import logging
log = logging.getLogger(__name__)
# Define the largest chunk of data we'll send in one go. Realistically, we
# should take the MSS into account but that's pretty dull, so let's just say
# 1kB and call it a day.
MAX_CHUNK = 1024
class Stream(object):
"""
A single HTTP/2 stream.
A stream is an independent, bi-directional sequence of HTTP headers and
data. Each stream is identified by a single integer. From a HTTP
perspective, a stream _approximately_ matches a single request-response
pair.
"""
def __init__(self,
stream_id,
window_manager,
connection,
send_outstanding_data,
recv_cb,
close_cb):
self.stream_id = stream_id
self.headers = HTTPHeaderMap()
# Set to a key-value set of the response headers once their
# HEADERS..CONTINUATION frame sequence finishes.
self.response_headers = None
# Set to a key-value set of the response trailers once their
# HEADERS..CONTINUATION frame sequence finishes.
self.response_trailers = None
# A dict mapping the promised stream ID of a pushed resource to a
# key-value set of its request headers. Entries are added once their
# PUSH_PROMISE..CONTINUATION frame sequence finishes.
self.promised_headers = {}
# Unconsumed response data chunks. Empties after every call to _read().
self.data = []
# Whether the remote side has completed the stream.
self.remote_closed = False
# Whether we have closed the stream.
self.local_closed = False
# There are two flow control windows: one for data we're sending,
# one for data being sent to us.
self._in_window_manager = window_manager
# Save off a reference to the state machine wrapped with lock.
self._conn = connection
# Save off a data callback.
self._send_outstanding_data = send_outstanding_data
self._recv_cb = recv_cb
self._close_cb = close_cb
def add_header(self, name, value, replace=False):
"""
Adds a single HTTP header to the headers to be sent on the request.
"""
if not replace:
self.headers[name] = value
else:
self.headers.replace(name, value)
def send_headers(self, end_stream=False):
"""
Sends the complete saved header block on the stream.
"""
headers = self.get_headers()
with self._conn as conn:
conn.send_headers(self.stream_id, headers, end_stream)
self._send_outstanding_data()
if end_stream:
self.local_closed = True
def send_data(self, data, final):
"""
Send some data on the stream. If this is the end of the data to be
sent, the ``final`` flag _must_ be set to True. If no data is to be
sent, set ``data`` to ``None``.
"""
# Define a utility iterator for file objects.
def file_iterator(fobj):
while True:
data = fobj.read(MAX_CHUNK)
yield data
if len(data) < MAX_CHUNK:
break
# Build the appropriate iterator for the data, in chunks of CHUNK_SIZE.
if hasattr(data, 'read'):
chunks = file_iterator(data)
else:
chunks = (data[i:i+MAX_CHUNK]
for i in range(0, len(data), MAX_CHUNK))
# since we need to know when we have a last package we need to know
# if there is another package in advance
cur_chunk = None
try:
cur_chunk = next(chunks)
while True:
next_chunk = next(chunks)
self._send_chunk(cur_chunk, False)
cur_chunk = next_chunk
except StopIteration:
if cur_chunk is not None: # cur_chunk none when no chunks to send
self._send_chunk(cur_chunk, final)
def _read(self, amt=None):
"""
Read data from the stream. Unlike a normal read behaviour, this
function returns _at least_ ``amt`` data, but may return more.
"""
def listlen(list):
return sum(map(len, list))
# Keep reading until the stream is closed or we get enough data.
while (not self.remote_closed and
(amt is None or listlen(self.data) < amt)):
self._recv_cb(stream_id=self.stream_id)
result = b''.join(self.data)
self.data = []
return result
def _read_one_frame(self):
"""
Reads a single data frame from the stream and returns it.
"""
# Keep reading until the stream is closed or we have a data frame.
while not self.remote_closed and not self.data:
self._recv_cb(stream_id=self.stream_id)
try:
return self.data.pop(0)
except IndexError:
return None
def receive_response(self, event):
"""
Receive response headers.
"""
# TODO: If this is called while we're still sending data, we may want
# to stop sending that data and check the response. Early responses to
# big uploads are almost always a problem.
self.response_headers = HTTPHeaderMap(event.headers)
def receive_trailers(self, event):
"""
Receive response trailers.
"""
self.response_trailers = HTTPHeaderMap(event.headers)
def receive_push(self, event):
"""
Receive the request headers for a pushed stream.
"""
self.promised_headers[event.pushed_stream_id] = event.headers
def receive_data(self, event):
"""
Receive a chunk of data.
"""
size = event.flow_controlled_length
increment = self._in_window_manager._handle_frame(size)
# Append the data to the buffer.
self.data.append(event.data)
if increment:
try:
with self._conn as conn:
conn.increment_flow_control_window(
increment, stream_id=self.stream_id
)
except h2.exceptions.StreamClosedError:
# We haven't got to it yet, but the stream is already
# closed. We don't need to increment the window in this
# case!
pass
else:
self._send_outstanding_data()
def receive_end_stream(self, event):
"""
All of the data is returned now.
"""
self.remote_closed = True
def receive_reset(self, event):
"""
Stream forcefully reset.
"""
self.remote_closed = True
self._close_cb(self.stream_id)
def get_headers(self):
"""
Provides the headers to the connection object.
"""
# Strip any headers invalid in H2.
return h2_safe_headers(self.headers)
def getheaders(self):
"""
Once all data has been sent on this connection, returns a key-value set
of the headers of the response to the original request.
"""
# Keep reading until all headers are received.
while self.response_headers is None:
self._recv_cb(stream_id=self.stream_id)
# Find the Content-Length header if present.
self._in_window_manager.document_size = (
int(self.response_headers.get(b'content-length', [0])[0])
)
return self.response_headers
def gettrailers(self):
"""
Once all data has been sent on this connection, returns a key-value set
of the trailers of the response to the original request.
.. warning:: Note that this method requires that the stream is
totally exhausted. This means that, if you have not
completely read from the stream, all stream data will be
read into memory.
:returns: The key-value set of the trailers, or ``None`` if no trailers
were sent.
"""
# Keep reading until the stream is done.
while not self.remote_closed:
self._recv_cb(stream_id=self.stream_id)
return self.response_trailers
def get_pushes(self, capture_all=False):
"""
Returns a generator that yields push promises from the server. Note
that this method is not idempotent; promises returned in one call will
not be returned in subsequent calls. Iterating through generators
returned by multiple calls to this method simultaneously results in
undefined behavior.
:param capture_all: If ``False``, the generator will yield all buffered
push promises without blocking. If ``True``, the generator will
first yield all buffered push promises, then yield additional ones
as they arrive, and terminate when the original stream closes.
"""
while True:
for pair in self.promised_headers.items():
yield pair
self.promised_headers = {}
if not capture_all or self.remote_closed:
break
self._recv_cb(stream_id=self.stream_id)
def close(self, error_code=None):
"""
Closes the stream. If the stream is currently open, attempts to close
it as gracefully as possible.
:param error_code: (optional) The error code to reset the stream with.
:returns: Nothing.
"""
# FIXME: I think this is overbroad, but for now it's probably ok.
if not (self.remote_closed and self.local_closed):
try:
with self._conn as conn:
conn.reset_stream(self.stream_id, error_code or 0)
except h2.exceptions.ProtocolError:
# If for any reason we can't reset the stream, just
# tolerate it.
pass
else:
self._send_outstanding_data(tolerate_peer_gone=True)
self.remote_closed = True
self.local_closed = True
self._close_cb(self.stream_id)
@property
def _out_flow_control_window(self):
"""
The size of our outbound flow control window.
"""
with self._conn as conn:
return conn.local_flow_control_window(self.stream_id)
def _send_chunk(self, data, final):
"""
Implements most of the sending logic.
Takes a single chunk of size at most MAX_CHUNK, wraps it in a frame and
sends it. Optionally sets the END_STREAM flag if this is the last chunk
(determined by being of size less than MAX_CHUNK) and no more data is
to be sent.
"""
# If we don't fit in the connection window, try popping frames off the
# connection in hope that one might be a window update frame.
while len(data) > self._out_flow_control_window:
self._recv_cb()
# Send the frame and decrement the flow control window.
with self._conn as conn:
conn.send_data(
stream_id=self.stream_id, data=data, end_stream=final
)
self._send_outstanding_data()
if final:
self.local_closed = True