def make_pipe(self) -> layer.CommandGenerator[None]:
self.state = self.passthrough
if self.buf:
already_received = self.buf.maybe_extract_at_most(len(self.buf))
yield from self.state(
events.DataReceived(self.conn, already_received))
self.buf.compress()
def receive_data(self, data: bytes) -> layer.CommandGenerator[None]:
if data:
self.tls.bio_write(data)
yield from self.tls_interact()
plaintext = bytearray()
close = False
while True:
try:
plaintext.extend(self.tls.recv(65535))
except SSL.WantReadError:
break
except SSL.ZeroReturnError:
close = True
break
if plaintext:
yield from self.event_to_child(
events.DataReceived(self.conn, bytes(plaintext))
)
if close:
self.conn.state &= ~context.ConnectionState.CAN_READ
if self.debug:
yield commands.Log(f"{self.debug}[tls] close_notify {self.conn}", level="debug")
yield from self.event_to_child(
events.ConnectionClosed(self.conn)
)
def passthrough(self, event: events.Event) -> layer.CommandGenerator[None]:
assert self.flow.response
assert self.child_layer
# HTTP events -> normal connection events
if isinstance(event, RequestData):
event = events.DataReceived(self.context.client, event.data)
elif isinstance(event, ResponseData):
event = events.DataReceived(self.context.server, event.data)
elif isinstance(event, RequestEndOfMessage):
event = events.ConnectionClosed(self.context.client)
elif isinstance(event, ResponseEndOfMessage):
event = events.ConnectionClosed(self.context.server)
for command in self.child_layer.handle_event(event):
# normal connection events -> HTTP events
if isinstance(command, commands.SendData):
if command.connection == self.context.client:
yield SendHttp(ResponseData(self.stream_id, command.data), self.context.client)
elif command.connection == self.context.server and self.flow.response.status_code == 101:
# there only is a HTTP server connection if we have switched protocols,
# not if a connection is established via CONNECT.
yield SendHttp(RequestData(self.stream_id, command.data), self.context.server)
else:
yield command
elif isinstance(command, commands.CloseConnection):
if command.connection == self.context.client:
yield SendHttp(ResponseProtocolError(self.stream_id, "EOF"), self.context.client)
elif command.connection == self.context.server and self.flow.response.status_code == 101:
yield SendHttp(RequestProtocolError(self.stream_id, "EOF"), self.context.server)
else:
# If we are running TCP over HTTP we want to be consistent with half-closes.
# The easiest approach for this is to just always full close for now.
# Alternatively, we could signal that we want a half close only through ResponseProtocolError,
# but that is more complex to implement.
command.half_close = False
yield command
else:
yield command
def _h2_response(chunks):
tctx = context.Context(
context.Client(("client", 1234), ("127.0.0.1", 8080), 1605699329),
opts)
playbook = Playbook(http.HttpLayer(tctx, HTTPMode.regular), hooks=False)
server = Placeholder(context.Server)
assert (playbook >> DataReceived(
tctx.client,
b"GET http://example.com/ HTTP/1.1\r\nHost: example.com\r\n\r\n") <<
OpenConnection(server) >> reply(
None, side_effect=make_h2) << SendData(server, Placeholder()))
for chunk in chunks:
for _ in playbook.layer.handle_event(
events.DataReceived(server(), chunk)):
pass
async def handle_connection(self, connection: Connection) -> None:
"""
Handle a connection for its entire lifetime.
This means we read until EOF,
but then possibly also keep on waiting for our side of the connection to be closed.
"""
cancelled = None
reader = self.transports[connection].reader
assert reader
while True:
try:
data = await reader.read(65535)
if not data:
raise OSError("Connection closed by peer.")
except OSError:
break
except asyncio.CancelledError as e:
cancelled = e
break
else:
self.server_event(events.DataReceived(connection, data))
if cancelled is None:
connection.state &= ~ConnectionState.CAN_READ
else:
connection.state = ConnectionState.CLOSED
self.server_event(events.ConnectionClosed(connection))
if cancelled is None and connection.state is ConnectionState.CAN_WRITE:
# we may still use this connection to *send* stuff,
# even though the remote has closed their side of the connection.
# to make this work we keep this task running and wait for cancellation.
await asyncio.Event().wait()
try:
writer = self.transports[connection].writer
assert writer
writer.close()
except OSError:
pass
self.transports.pop(connection)
if cancelled:
raise cancelled
def mark_done(self,
*,
request: bool = False,
response: bool = False) -> layer.CommandGenerator[None]:
if request:
self.request_done = True
if response:
self.response_done = True
if self.request_done and self.response_done:
assert self.request
assert self.response
if should_make_pipe(self.request, self.response):
yield from self.make_pipe()
return
connection_done = (
http1_sansio.expected_http_body_size(self.request,
self.response) == -1
or http1.connection_close(self.request.http_version,
self.request.headers)
or http1.connection_close(self.response.http_version,
self.response.headers)
# If we proxy HTTP/2 to HTTP/1, we only use upstream connections for one request.
# This simplifies our connection management quite a bit as we can rely on
# the proxyserver's max-connection-per-server throttling.
or (self.request.is_http2 and isinstance(self, Http1Client)))
if connection_done:
yield commands.CloseConnection(self.conn)
self.state = self.done
return
self.request_done = self.response_done = False
self.request = self.response = None
if isinstance(self, Http1Server):
self.stream_id += 2
else:
self.stream_id = None
self.state = self.read_headers
if self.buf:
yield from self.state(events.DataReceived(self.conn, b""))
def test_dataclasses(tconn):
assert repr(events.Start())
assert repr(events.DataReceived(tconn, b"foo"))
assert repr(events.ConnectionClosed(tconn))
def receive_data(self, data: bytes) -> layer.CommandGenerator[None]:
yield from self.event_to_child(events.DataReceived(self.conn, data))
def start_handshake(self) -> layer.CommandGenerator[None]:
yield from self._handle_event(
events.DataReceived(self.tunnel_connection, b""))