def wait(self, event: events.Event) -> layer.CommandGenerator[None]:
"""
We wait for the current flow to be finished before parsing the next message,
as we may want to upgrade to WebSocket or plain TCP before that.
"""
assert self.stream_id
if isinstance(event, events.DataReceived):
return
elif isinstance(event, events.ConnectionClosed):
# for practical purposes, we assume that a peer which sent at least a FIN
# is not interested in any more data from us, see
# see https://github.com/httpwg/http-core/issues/22
if event.connection.state is not ConnectionState.CLOSED:
yield commands.CloseConnection(event.connection)
yield ReceiveHttp(
self.ReceiveProtocolError(
self.stream_id,
f"Client disconnected.",
code=status_codes.CLIENT_CLOSED_REQUEST))
else: # pragma: no cover
raise AssertionError(f"Unexpected event: {event}")
def send(self, event: HttpEvent) -> layer.CommandGenerator[None]:
assert event.stream_id == self.stream_id
if isinstance(event, ResponseHeaders):
self.response = response = event.response
if response.is_http2:
response = response.copy()
# Convert to an HTTP/1 response.
response.http_version = "HTTP/1.1"
# not everyone supports empty reason phrases, so we better make up one.
response.reason = status_codes.RESPONSES.get(
response.status_code, "")
# Shall we set a Content-Length header here if there is none?
# For now, let's try to modify as little as possible.
raw = http1.assemble_response_head(response)
yield commands.SendData(self.conn, raw)
elif isinstance(event, ResponseData):
assert self.response
if "chunked" in self.response.headers.get("transfer-encoding",
"").lower():
raw = b"%x\r\n%s\r\n" % (len(event.data), event.data)
else:
raw = event.data
if raw:
yield commands.SendData(self.conn, raw)
elif isinstance(event, ResponseEndOfMessage):
assert self.response
if "chunked" in self.response.headers.get("transfer-encoding",
"").lower():
yield commands.SendData(self.conn, b"0\r\n\r\n")
yield from self.mark_done(response=True)
elif isinstance(event, ResponseProtocolError):
if not self.response:
resp = http.make_error_response(event.code, event.message)
raw = http1.assemble_response(resp)
yield commands.SendData(self.conn, raw)
yield commands.CloseConnection(self.conn)
else:
raise AssertionError(f"Unexpected event: {event}")
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(commands.SendData(tconn, b"foo"))
assert repr(commands.OpenConnection(tconn))
assert repr(commands.CloseConnection(tconn))
assert repr(commands.GetSocket(tconn))
assert repr(commands.Log("hello", "info"))
def relay_messages(
self,
event: events.ConnectionEvent) -> layer.CommandGenerator[None]:
from_client = event.connection == self.context.client
from_str = 'client' if from_client else 'server'
if from_client:
src_ws = self.client_ws
dst_ws = self.server_ws
else:
src_ws = self.server_ws
dst_ws = self.client_ws
if isinstance(event, events.DataReceived):
src_ws.receive_data(event.data)
elif isinstance(event, events.ConnectionClosed):
src_ws.receive_data(None)
else: # pragma: no cover
raise AssertionError(f"Unexpected event: {event}")
for ws_event in src_ws.events():
if isinstance(ws_event, wsproto.events.Message):
src_ws.frame_buf.append(ws_event.data)
if ws_event.message_finished:
if isinstance(ws_event, wsproto.events.TextMessage):
frame_type = Opcode.TEXT
content = "".join(src_ws.frame_buf) # type: ignore
else:
frame_type = Opcode.BINARY
content = b"".join(src_ws.frame_buf) # type: ignore
fragmentizer = Fragmentizer(src_ws.frame_buf)
src_ws.frame_buf.clear()
message = websocket.WebSocketMessage(
frame_type, from_client, content)
self.flow.messages.append(message)
yield WebsocketMessageHook(self.flow)
assert not message.killed # this is deprecated, instead we should have .content set to emptystr.
for msg in fragmentizer(message.content):
yield dst_ws.send2(msg)
elif isinstance(ws_event,
(wsproto.events.Ping, wsproto.events.Pong)):
yield commands.Log(
f"Received WebSocket {ws_event.__class__.__name__.lower()} from {from_str} "
f"(payload: {bytes(ws_event.payload)!r})")
yield dst_ws.send2(ws_event)
elif isinstance(ws_event, wsproto.events.CloseConnection):
self.flow.close_sender = from_str
self.flow.close_code = ws_event.code
self.flow.close_reason = ws_event.reason
for ws in [self.server_ws, self.client_ws]:
if ws.state in {
ConnectionState.OPEN,
ConnectionState.REMOTE_CLOSING
}:
# response == original event, so no need to differentiate here.
yield ws.send2(ws_event)
yield commands.CloseConnection(ws.conn)
if ws_event.code in {1000, 1001, 1005}:
yield WebsocketEndHook(self.flow)
else:
self.flow.error = flow.Error(
f"WebSocket Error: {format_close_event(ws_event)}")
yield WebsocketErrorHook(self.flow)
self._handle_event = self.done
else: # pragma: no cover
raise AssertionError(f"Unexpected WebSocket event: {ws_event}")
def send_close(self, half_close: bool) -> layer.CommandGenerator[None]:
yield commands.CloseConnection(self.tunnel_connection,
half_close=half_close)
def on_handshake_error(self, err: str) -> layer.CommandGenerator[None]:
"""Called if either receive_handshake_data returns an error or we receive a close during handshake."""
yield commands.CloseConnection(self.tunnel_connection)
def _handle_event(self,
event: events.Event) -> layer.CommandGenerator[None]:
if isinstance(event, events.DataReceived):
yield commands.SendData(event.connection, event.data.lower())
if isinstance(event, events.ConnectionClosed):
yield commands.CloseConnection(event.connection)