codekingpro/portable-devtools
114k
1"""2This module contains the proxy layers for raw QUIC proxying.3This is used if we want to speak QUIC, but we do not want to do HTTP.4"""5 6from __future__ import annotations7 8import time9 10from aioquic.quic.connection import QuicErrorCode11from aioquic.quic.connection import stream_is_client_initiated12from aioquic.quic.connection import stream_is_unidirectional13 14from ._commands import CloseQuicConnection15from ._commands import ResetQuicStream16from ._commands import SendQuicStreamData17from ._commands import StopSendingQuicStream18from ._events import QuicConnectionClosed19from ._events import QuicStreamDataReceived20from ._events import QuicStreamEvent21from ._events import QuicStreamReset22from mitmproxy import connection23from mitmproxy.connection import Connection24from mitmproxy.proxy import commands25from mitmproxy.proxy import context26from mitmproxy.proxy import events27from mitmproxy.proxy import layer28from mitmproxy.proxy import tunnel29from mitmproxy.proxy.layers.tcp import TCPLayer30from mitmproxy.proxy.layers.udp import UDPLayer31 32 33class QuicStreamNextLayer(layer.NextLayer):34 """`NextLayer` variant that callbacks `QuicStreamLayer` after layer decision."""35 36 def __init__(37 self,38 context: context.Context,39 stream: QuicStreamLayer,40 ask_on_start: bool = False,41 ) -> None:42 super().__init__(context, ask_on_start)43 self._stream = stream44 self._layer: layer.Layer | None = None45 46 @property # type: ignore47 def layer(self) -> layer.Layer | None: # type: ignore48 return self._layer49 50 @layer.setter51 def layer(self, value: layer.Layer | None) -> None:52 self._layer = value53 if self._layer:54 self._stream.refresh_metadata()55 56 57class QuicStreamLayer(layer.Layer):58 """59 Layer for QUIC streams.60 Serves as a marker for NextLayer and keeps track of the connection states.61 """62 63 client: connection.Client64 """Virtual client connection for this stream. Use this in QuicRawLayer instead of `context.client`."""65 server: connection.Server66 """Virtual server connection for this stream. Use this in QuicRawLayer instead of `context.server`."""67 child_layer: layer.Layer68 """The stream's child layer."""69 70 def __init__(71 self, context: context.Context, force_raw: bool, stream_id: int72 ) -> None:73 # we mustn't reuse the client from the QUIC connection, as the state and protocol differs74 self.client = context.client = context.client.copy()75 self.client.transport_protocol = "tcp"76 self.client.state = connection.ConnectionState.OPEN77 78 # unidirectional client streams are not fully open, set the appropriate state79 if stream_is_unidirectional(stream_id):80 self.client.state = (81 connection.ConnectionState.CAN_READ82 if stream_is_client_initiated(stream_id)83 else connection.ConnectionState.CAN_WRITE84 )85 self._client_stream_id = stream_id86 87 # start with a closed server88 self.server = context.server = connection.Server(89 address=context.server.address,90 transport_protocol="tcp",91 )92 self._server_stream_id: int | None = None93 94 super().__init__(context)95 self.child_layer = (96 TCPLayer(context) if force_raw else QuicStreamNextLayer(context, self)97 )98 self.refresh_metadata()99 100 # we don't handle any events, pass everything to the child layer101 self.handle_event = self.child_layer.handle_event # type: ignore102 self._handle_event = self.child_layer._handle_event # type: ignore103 104 def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]:105 raise AssertionError # pragma: no cover106 107 def open_server_stream(self, server_stream_id) -> None:108 assert self._server_stream_id is None109 self._server_stream_id = server_stream_id110 self.server.timestamp_start = time.time()111 self.server.state = (112 (113 connection.ConnectionState.CAN_WRITE114 if stream_is_client_initiated(server_stream_id)115 else connection.ConnectionState.CAN_READ116 )117 if stream_is_unidirectional(server_stream_id)118 else connection.ConnectionState.OPEN119 )120 self.refresh_metadata()121 122 def refresh_metadata(self) -> None:123 # find the first transport layer124 child_layer: layer.Layer | None = self.child_layer125 while True:126 if isinstance(child_layer, layer.NextLayer):127 child_layer = child_layer.layer128 elif isinstance(child_layer, tunnel.TunnelLayer):129 child_layer = child_layer.child_layer130 else:131 break # pragma: no cover132 if isinstance(child_layer, (UDPLayer, TCPLayer)) and child_layer.flow:133 child_layer.flow.metadata["quic_is_unidirectional"] = (134 stream_is_unidirectional(self._client_stream_id)135 )136 child_layer.flow.metadata["quic_initiator"] = (137 "client"138 if stream_is_client_initiated(self._client_stream_id)139 else "server"140 )141 child_layer.flow.metadata["quic_stream_id_client"] = self._client_stream_id142 child_layer.flow.metadata["quic_stream_id_server"] = self._server_stream_id143 144 def stream_id(self, client: bool) -> int | None:145 return self._client_stream_id if client else self._server_stream_id146 147 148class RawQuicLayer(layer.Layer):149 """150 This layer is responsible for de-multiplexing QUIC streams into an individual layer stack per stream.151 """152 153 force_raw: bool154 """Indicates whether traffic should be treated as raw TCP/UDP without further protocol detection."""155 datagram_layer: layer.Layer156 """157 The layer that is handling datagrams over QUIC. It's like a child_layer, but with a forked context.158 Instead of having a datagram-equivalent for all `QuicStream*` classes, we use `SendData` and `DataReceived` instead.159 There is also no need for another `NextLayer` marker, as a missing `QuicStreamLayer` implies UDP,160 and the connection state is the same as the one of the underlying QUIC connection.161 """162 client_stream_ids: dict[int, QuicStreamLayer]163 """Maps stream IDs from the client connection to stream layers."""164 server_stream_ids: dict[int, QuicStreamLayer]165 """Maps stream IDs from the server connection to stream layers."""166 connections: dict[connection.Connection, layer.Layer]167 """Maps connections to layers."""168 command_sources: dict[commands.Command, layer.Layer]169 """Keeps track of blocking commands and wakeup requests."""170 next_stream_id: list[int]171 """List containing the next stream ID for all four is_unidirectional/is_client combinations."""172 173 def __init__(self, context: context.Context, force_raw: bool = False) -> None:174 super().__init__(context)175 self.force_raw = force_raw176 self.datagram_layer = (177 UDPLayer(self.context.fork())178 if force_raw179 else layer.NextLayer(self.context.fork())180 )181 self.client_stream_ids = {}182 self.server_stream_ids = {}183 self.connections = {184 context.client: self.datagram_layer,185 context.server: self.datagram_layer,186 }187 self.command_sources = {}188 self.next_stream_id = [0, 1, 2, 3]189 190 def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]:191 # we treat the datagram layer as child layer, so forward Start192 if isinstance(event, events.Start):193 if self.context.server.timestamp_start is None:194 err = yield commands.OpenConnection(self.context.server)195 if err:196 yield commands.CloseConnection(self.context.client)197 self._handle_event = self.done # type: ignore198 return199 yield from self.event_to_child(self.datagram_layer, event)200 201 # properly forward completion events based on their command202 elif isinstance(event, events.CommandCompleted):203 yield from self.event_to_child(204 self.command_sources.pop(event.command), event205 )206 207 # route injected messages based on their connections (prefer client, fallback to server)208 elif isinstance(event, events.MessageInjected):209 if event.flow.client_conn in self.connections:210 yield from self.event_to_child(211 self.connections[event.flow.client_conn], event212 )213 elif event.flow.server_conn in self.connections:214 yield from self.event_to_child(215 self.connections[event.flow.server_conn], event216 )217 else:218 raise AssertionError(f"Flow not associated: {event.flow!r}")219 220 # handle stream events targeting this context221 elif isinstance(event, QuicStreamEvent) and (222 event.connection is self.context.client223 or event.connection is self.context.server224 ):225 from_client = event.connection is self.context.client226 227 # fetch or create the layer228 stream_ids = (229 self.client_stream_ids if from_client else self.server_stream_ids230 )231 if event.stream_id in stream_ids:232 stream_layer = stream_ids[event.stream_id]233 else:234 # ensure we haven't just forgotten to register the ID235 assert stream_is_client_initiated(event.stream_id) == from_client236 237 # for server-initiated streams we need to open the client as well238 if from_client:239 client_stream_id = event.stream_id240 server_stream_id = None241 else:242 client_stream_id = self.get_next_available_stream_id(243 is_client=False,244 is_unidirectional=stream_is_unidirectional(event.stream_id),245 )246 server_stream_id = event.stream_id247 248 # create, register and start the layer249 stream_layer = QuicStreamLayer(250 self.context.fork(),251 force_raw=self.force_raw,252 stream_id=client_stream_id,253 )254 self.client_stream_ids[client_stream_id] = stream_layer255 if server_stream_id is not None:256 stream_layer.open_server_stream(server_stream_id)257 self.server_stream_ids[server_stream_id] = stream_layer258 self.connections[stream_layer.client] = stream_layer259 self.connections[stream_layer.server] = stream_layer260 yield from self.event_to_child(stream_layer, events.Start())261 262 # forward data and close events263 conn: Connection = (264 stream_layer.client if from_client else stream_layer.server265 )266 if isinstance(event, QuicStreamDataReceived):267 if event.data:268 yield from self.event_to_child(269 stream_layer, events.DataReceived(conn, event.data)270 )271 if event.end_stream:272 yield from self.close_stream_layer(stream_layer, from_client)273 elif isinstance(event, QuicStreamReset):274 # preserve stream resets275 for command in self.close_stream_layer(stream_layer, from_client):276 if (277 isinstance(command, SendQuicStreamData)278 and command.stream_id == stream_layer.stream_id(not from_client)279 and command.end_stream280 and not command.data281 ):282 yield ResetQuicStream(283 command.connection, command.stream_id, event.error_code284 )285 else:286 yield command287 else:288 raise AssertionError(f"Unexpected stream event: {event!r}")289 290 # handle close events that target this context291 elif isinstance(event, QuicConnectionClosed) and (292 event.connection is self.context.client293 or event.connection is self.context.server294 ):295 from_client = event.connection is self.context.client296 other_conn = self.context.server if from_client else self.context.client297 298 # be done if both connections are closed299 if other_conn.connected:300 yield CloseQuicConnection(301 other_conn, event.error_code, event.frame_type, event.reason_phrase302 )303 else:304 self._handle_event = self.done # type: ignore305 306 # always forward to the datagram layer and swallow `CloseConnection` commands307 for command in self.event_to_child(self.datagram_layer, event):308 if (309 not isinstance(command, commands.CloseConnection)310 or command.connection is not other_conn311 ):312 yield command313 314 # forward to either the client or server connection of stream layers and swallow empty stream end315 for conn, child_layer in self.connections.items():316 if isinstance(child_layer, QuicStreamLayer) and (317 (conn is child_layer.client)318 if from_client319 else (conn is child_layer.server)320 ):321 conn.state &= ~connection.ConnectionState.CAN_WRITE322 for command in self.close_stream_layer(child_layer, from_client):323 if not isinstance(command, SendQuicStreamData) or command.data:324 yield command325 326 # all other connection events are routed to their corresponding layer327 elif isinstance(event, events.ConnectionEvent):328 yield from self.event_to_child(self.connections[event.connection], event)329 330 else:331 raise AssertionError(f"Unexpected event: {event!r}")332 333 def close_stream_layer(334 self, stream_layer: QuicStreamLayer, client: bool335 ) -> layer.CommandGenerator[None]:336 """Closes the incoming part of a connection."""337 338 conn = stream_layer.client if client else stream_layer.server339 conn.state &= ~connection.ConnectionState.CAN_READ340 assert conn.timestamp_start is not None341 if conn.timestamp_end is None:342 conn.timestamp_end = time.time()343 yield from self.event_to_child(stream_layer, events.ConnectionClosed(conn))344 345 def event_to_child(346 self, child_layer: layer.Layer, event: events.Event347 ) -> layer.CommandGenerator[None]:348 """Forwards events to child layers and translates commands."""349 350 for command in child_layer.handle_event(event):351 # intercept commands for streams connections352 if (353 isinstance(child_layer, QuicStreamLayer)354 and isinstance(command, commands.ConnectionCommand)355 and (356 command.connection is child_layer.client357 or command.connection is child_layer.server358 )359 ):360 # get the target connection and stream ID361 to_client = command.connection is child_layer.client362 quic_conn = self.context.client if to_client else self.context.server363 stream_id = child_layer.stream_id(to_client)364 365 # write data and check CloseConnection wasn't called before366 if isinstance(command, commands.SendData):367 assert stream_id is not None368 if command.connection.state & connection.ConnectionState.CAN_WRITE:369 yield SendQuicStreamData(quic_conn, stream_id, command.data)370 371 # send a FIN and optionally also a STOP frame372 elif isinstance(command, commands.CloseConnection):373 assert stream_id is not None374 if command.connection.state & connection.ConnectionState.CAN_WRITE:375 command.connection.state &= (376 ~connection.ConnectionState.CAN_WRITE377 )378 yield SendQuicStreamData(379 quic_conn, stream_id, b"", end_stream=True380 )381 # XXX: Use `command.connection.state & connection.ConnectionState.CAN_READ` instead?382 only_close_our_half = (383 isinstance(command, commands.CloseTcpConnection)384 and command.half_close385 )386 if not only_close_our_half:387 if stream_is_client_initiated(388 stream_id389 ) == to_client or not stream_is_unidirectional(stream_id):390 yield StopSendingQuicStream(391 quic_conn, stream_id, QuicErrorCode.NO_ERROR392 )393 yield from self.close_stream_layer(child_layer, to_client)394 395 # open server connections by reserving the next stream ID396 elif isinstance(command, commands.OpenConnection):397 assert not to_client398 assert stream_id is None399 client_stream_id = child_layer.stream_id(client=True)400 assert client_stream_id is not None401 stream_id = self.get_next_available_stream_id(402 is_client=True,403 is_unidirectional=stream_is_unidirectional(client_stream_id),404 )405 child_layer.open_server_stream(stream_id)406 self.server_stream_ids[stream_id] = child_layer407 yield from self.event_to_child(408 child_layer, events.OpenConnectionCompleted(command, None)409 )410 411 else:412 raise AssertionError(413 f"Unexpected stream connection command: {command!r}"414 )415 416 # remember blocking and wakeup commands417 else:418 if command.blocking or isinstance(command, commands.RequestWakeup):419 self.command_sources[command] = child_layer420 if isinstance(command, commands.OpenConnection):421 self.connections[command.connection] = child_layer422 yield command423 424 def get_next_available_stream_id(425 self, is_client: bool, is_unidirectional: bool = False426 ) -> int:427 index = (int(is_unidirectional) << 1) | int(not is_client)428 stream_id = self.next_stream_id[index]429 self.next_stream_id[index] = stream_id + 4430 return stream_id431 432 def done(self, _) -> layer.CommandGenerator[None]: # pragma: no cover433 yield from ()434