fixed handling of Transfer-Encoding header during streaming; wrote tests for streaming support

This commit is contained in:
Brad Peabody 2014-07-19 19:10:14 -07:00
parent c47ddaa3a0
commit 560e23af09
3 changed files with 84 additions and 16 deletions

View File

@ -632,24 +632,21 @@ class HTTPResponse(HTTPMessage):
return 'HTTP/%s.%s %s %s' % \
(self.httpversion[0], self.httpversion[1], self.code, self.msg)
def _assemble_headers(self):
def _assemble_headers(self, preserve_transfer_encoding=False):
headers = self.headers.copy()
utils.del_all(
headers,
[
'Proxy-Connection',
'Transfer-Encoding'
]
)
utils.del_all(headers,['Proxy-Connection'])
if not preserve_transfer_encoding:
utils.del_all(headers,['Transfer-Encoding'])
if self.content:
headers["Content-Length"] = [str(len(self.content))]
elif 'Transfer-Encoding' in self.headers: # add content-length for chuncked transfer-encoding with no content
elif not preserve_transfer_encoding and 'Transfer-Encoding' in self.headers: # add content-length for chuncked transfer-encoding with no content
headers["Content-Length"] = ["0"]
return str(headers)
def _assemble_head(self):
return '%s\r\n%s\r\n' % (self._assemble_first_line(), self._assemble_headers())
def _assemble_head(self, preserve_transfer_encoding=False):
return '%s\r\n%s\r\n' % (self._assemble_first_line(), self._assemble_headers(preserve_transfer_encoding=preserve_transfer_encoding))
def _assemble(self):
"""
@ -928,7 +925,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
# flow.stream_expecting_body now contains the state info of whether or not
# body still remains to be read
# call the appropriate script hook - this is an opportunity for
# call the appropriate script hook - this is an opportunity for an inline script to set flow.stream = True
responseheaders_reply = self.c.channel.ask("responseheaders", flow.response)
# hm - do we need to do something with responseheaders_reply??
@ -944,6 +941,8 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
if response_reply is None or response_reply == KILL:
return False
disconnected_while_streaming = False
if not flow.stream or not flow.stream_expecting_body:
# if not streaming or there is no body to be read, we'll already have the body, just send it
self.c.client_conn.send(flow.response._assemble())
@ -952,15 +951,17 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
# if streaming, we still need to read the body and stream its bits back to the client
# start with head
h = flow.response._assemble_head()
h = flow.response._assemble_head(preserve_transfer_encoding=True)
self.c.client_conn.send(h)
# if chunked then we send back each chunk
if http.has_chunked_encoding(flow.response.headers):
while 1:
content = http.read_next_chunk(self.c.server_conn.rfile, flow.response.headers, False)
if not http.write_chunk(self.c.client_conn.rfile, content):
if not http.write_chunk(self.c.client_conn.wfile, content):
break
self.c.client_conn.wfile.flush()
self.c.client_conn.wfile.flush()
else: # not chunked, we send back 4k at a time
clen = http.expected_http_body_size(flow.response.headers, False)
@ -969,17 +970,21 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
blocksize = 4096
while 1:
bytes_to_read = min(blocksize, clen - rcount)
if bytes_to_read == 0:
break
content = self.c.server_conn.rfile.read(bytes_to_read)
if content == "": # check for EOF
disconnected_while_streaming = True
break
rcount += len(content)
self.c.client_conn.rfile.write(content)
self.c.client_conn.wfile.write(content)
self.c.client_conn.wfile.flush()
if rcount >= clen: # check for having read up to clen
break
flow.timestamp_end = utils.timestamp()
if (http.connection_close(flow.request.httpversion, flow.request.headers) or
if (disconnected_while_streaming or http.connection_close(flow.request.httpversion, flow.request.headers) or
http.connection_close(flow.response.httpversion, flow.response.headers)):
return False
@ -989,6 +994,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
# If the user has changed the target server on this connection,
# restore the original target server
self.restore_server()
return True
except (HttpAuthenticationError, http.HttpError, proxy.ProxyError, tcp.NetLibError), e:
self.handle_error(e, flow)

View File

@ -15,6 +15,10 @@ def response(ctx, r):
ctx.log("XRESPONSE")
log.append("response")
def responseheaders(ctx, r):
ctx.log("XRESPONSEHEADERS")
log.append("responseheaders")
def clientdisconnect(ctx, cc):
ctx.log("XCLIENTDISCONNECT")
log.append("clientdisconnect")

View File

@ -383,6 +383,64 @@ class TestRedirectRequest(tservers.HTTPProxTest):
assert r3.content == r2.content == r1.content
# Make sure that we actually use the same connection in this test case
class MasterStreamRequest(tservers.TestMaster):
"""
Enables the stream flag on the flow for all requests
"""
def handle_responseheaders(self, r):
f = self.state.add_response(r)
f.stream = True
r.reply()
return f
class TestStreamRequest(tservers.HTTPProxTest):
masterclass = MasterStreamRequest
def test_stream_simple(self):
p = self.pathoc()
# a request with 100k of data but without content-length
self.server.clear_log()
r1 = p.request("get:'%s/p/200:r:b@100k:d102400'"%self.server.urlbase)
assert r1.status_code == 200
assert len(r1.content) > 100000
assert self.server.last_log()
def test_stream_multiple(self):
p = self.pathoc()
# simple request with streaming turned on
self.server.clear_log()
r1 = p.request("get:'%s/p/200'"%self.server.urlbase)
assert r1.status_code == 200
assert self.server.last_log()
# now send back 100k of data, streamed but not chunked
self.server.clear_log()
r1 = p.request("get:'%s/p/200:b@100k'"%self.server.urlbase)
assert r1.status_code == 200
assert self.server.last_log()
def test_stream_chunked(self):
connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
connection.connect(("127.0.0.1", self.proxy.port))
fconn = connection.makefile()
spec = '200:h"Transfer-Encoding"="chunked":r:b"4\\r\\nthis\\r\\n7\\r\\nisatest\\r\\n0\\r\\n\\r\\n"'
connection.send("GET %s/p/%s HTTP/1.1\r\n"%(self.server.urlbase, spec))
connection.send("\r\n");
httpversion, code, msg, headers, content = http.read_response(fconn, "GET", 100000, include_body=False)
assert headers["Transfer-Encoding"][0] == 'chunked'
assert code == 200
assert http.read_next_chunk(fconn, headers, False) == "this"
assert http.read_next_chunk(fconn, headers, False) == "isatest"
assert http.read_next_chunk(fconn, headers, False) == None
connection.close()
class MasterFakeResponse(tservers.TestMaster):
def handle_request(self, m):