def relay_messages(self, event: events.ConnectionEvent) -> layer.CommandGenerator[None]:
from_client = event.connection == self.context.client
send_to: Connection
if from_client:
send_to = self.context.server
else:
send_to = self.context.client
if isinstance(event, events.DataReceived):
if self.flow:
tcp_message = tcp.TCPMessage(from_client, event.data)
self.flow.messages.append(tcp_message)
yield TcpMessageHook(self.flow)
yield commands.SendData(send_to, tcp_message.content)
else:
yield commands.SendData(send_to, event.data)
elif isinstance(event, events.ConnectionClosed):
all_done = not (
(self.context.client.state & ConnectionState.CAN_READ)
or
(self.context.server.state & ConnectionState.CAN_READ)
)
if all_done:
if self.context.server.state is not ConnectionState.CLOSED:
yield commands.CloseConnection(self.context.server)
if self.context.client.state is not ConnectionState.CLOSED:
yield commands.CloseConnection(self.context.client)
self._handle_event = self.done
if self.flow:
yield TcpEndHook(self.flow)
else:
yield commands.CloseConnection(send_to, half_close=True)
def tls_interact(self) -> layer.CommandGenerator[None]:
while True:
try:
data = self.tls.bio_read(65535)
except SSL.WantReadError:
return # Okay, nothing more waiting to be sent.
else:
yield commands.SendData(self.conn, data)
def send(self, event: HttpEvent) -> layer.CommandGenerator[None]:
if not self.stream_id:
assert isinstance(event, RequestHeaders)
self.stream_id = event.stream_id
self.request = event.request
assert self.stream_id == event.stream_id
if isinstance(event, RequestHeaders):
request = event.request
if request.is_http2:
# Convert to an HTTP/1 request.
request = request.copy(
) # (we could probably be a bit more efficient here.)
request.http_version = "HTTP/1.1"
if "Host" not in request.headers and request.authority:
request.headers.insert(0, "Host", request.authority)
request.authority = ""
raw = http1.assemble_request_head(request)
yield commands.SendData(self.conn, raw)
elif isinstance(event, RequestData):
assert self.request
if "chunked" in self.request.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, RequestEndOfMessage):
assert self.request
if "chunked" in self.request.headers.get("transfer-encoding",
"").lower():
yield commands.SendData(self.conn, b"0\r\n\r\n")
elif http1_sansio.expected_http_body_size(self.request,
self.response) == -1:
yield commands.CloseConnection(self.conn, half_close=True)
yield from self.mark_done(request=True)
elif isinstance(event, RequestProtocolError):
yield commands.CloseConnection(self.conn)
return
else:
raise AssertionError(f"Unexpected event: {event}")
def start_handshake(self) -> layer.CommandGenerator[None]:
if self.tunnel_connection.tls:
# "Secure Web Proxy": We may have negotiated an ALPN when connecting to the upstream proxy.
# The semantics are not really clear here, but we make sure that if we negotiated h2,
# we act as an h2 client.
self.conn.alpn = self.tunnel_connection.alpn
if not self.send_connect:
return (yield from super().start_handshake())
assert self.conn.address
req = http.make_connect_request(self.conn.address)
raw = http1.assemble_request(req)
yield commands.SendData(self.tunnel_connection, raw)
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 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 send2(self, event: wsproto.events.Event) -> commands.SendData:
data = self.send(event)
return commands.SendData(self.conn, data)
def send_data(self, data: bytes) -> layer.CommandGenerator[None]:
yield commands.SendData(self.tunnel_connection, data)
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)