Merge pull request #4294 from mhils/sans-io-adjustments

Add a switch for sans-io proxy core

mitmproxy/addons/dumper.py

@@ -1,18 +1,19 @@
 import itertools
+import shutil
 import sys
+from typing import Optional, TextIO
 
 import click
-import shutil
-
-import typing  # noqa
 
 from mitmproxy import contentviews
 from mitmproxy import ctx
+from mitmproxy import flow
 from mitmproxy import exceptions
 from mitmproxy import flowfilter
+from mitmproxy import http
+from mitmproxy.net import http as net_http
 from mitmproxy.utils import human
 from mitmproxy.utils import strutils
-from mitmproxy.net import http
 
 
 def indent(n: int, text: str) -> str:
@@ -29,9 +30,9 @@ def colorful(line, styles):
 
 class Dumper:
     def __init__(self, outfile=sys.stdout, errfile=sys.stderr):
-        self.filter: flowfilter.TFilter = None
+        self.filter: Optional[flowfilter.TFilter] = None
-        self.outfp: typing.io.TextIO = outfile
+        self.outfp: TextIO = outfile
-        self.errfp: typing.io.TextIO = errfile
+        self.errfp: TextIO = errfile
 
     def load(self, loader):
         loader.add_option(
@@ -47,10 +48,10 @@ class Dumper:
         loader.add_option(
             "dumper_default_contentview", str, "auto",
             "The default content view mode.",
-            choices = [i.name.lower() for i in contentviews.views]
+            choices=[i.name.lower() for i in contentviews.views]
         )
         loader.add_option(
-            "dumper_filter", typing.Optional[str], None,
+            "dumper_filter", Optional[str], None,
             "Limit which flows are dumped."
         )
 
@@ -65,19 +66,19 @@ class Dumper:
             else:
                 self.filter = None
 
-    def echo(self, text, ident=None, **style):
+    def echo(self, text: str, ident=None, **style):
         if ident:
             text = indent(ident, text)
         click.secho(text, file=self.outfp, **style)
         if self.outfp:
             self.outfp.flush()
 
-    def echo_error(self, text, **style):
+    def echo_error(self, text: str, **style):
         click.secho(text, file=self.errfp, **style)
         if self.errfp:
             self.errfp.flush()
 
-    def _echo_headers(self, headers):
+    def _echo_headers(self, headers: net_http.Headers):
         for k, v in headers.fields:
             k = strutils.bytes_to_escaped_str(k)
             v = strutils.bytes_to_escaped_str(v)
@@ -87,13 +88,13 @@ class Dumper:
             )
             self.echo(out, ident=4)
 
-    def _echo_trailers(self, trailers):
+    def _echo_trailers(self, trailers: Optional[net_http.Headers]):
-        if not trailers or not isinstance(trailers, http.Headers):
+        if not trailers:
             return
         self.echo(click.style("--- HTTP Trailers", fg="magenta"), ident=4)
         self._echo_headers(trailers)
 
-    def _echo_message(self, message, flow):
+    def _echo_message(self, message, flow: flow.Flow):
         _, lines, error = contentviews.get_message_content_view(
             ctx.options.dumper_default_contentview,
             message,
@@ -127,7 +128,7 @@ class Dumper:
         if ctx.options.flow_detail >= 2:
             self.echo("")
 
-    def _echo_request_line(self, flow):
+    def _echo_request_line(self, flow: http.HTTPFlow) -> None:
         if flow.client_conn:
             client = click.style(
                 strutils.escape_control_characters(
@@ -154,46 +155,54 @@ class Dumper:
             url = flow.request.pretty_url
         else:
             url = flow.request.url
-        terminalWidthLimit = max(shutil.get_terminal_size()[0] - 25, 50)
+
-        if ctx.options.flow_detail < 1 and len(url) > terminalWidthLimit:
+        if ctx.options.flow_detail <= 1:
-            url = url[:terminalWidthLimit] + "…"
+            # We need to truncate before applying styles, so we just focus on the URL.
+            terminal_width_limit = max(shutil.get_terminal_size()[0] - 25, 50)
+            if len(url) > terminal_width_limit:
+                url = url[:terminal_width_limit] + "…"
         url = click.style(strutils.escape_control_characters(url), bold=True)
 
         http_version = ""
-        if flow.request.http_version not in ("HTTP/1.1", "HTTP/1.0"):
+        if (
-            # We hide "normal" HTTP 1.
+                flow.request.http_version not in ("HTTP/1.1", "HTTP/1.0")
+                or flow.request.http_version != getattr(flow.response, "http_version", "HTTP/1.1")
+        ):
+            # Hide version for h1 <-> h1 connections.
             http_version = " " + flow.request.http_version
 
-        line = "{client}: {method} {url}{http_version}".format(
+        self.echo(f"{client}: {method} {url}{http_version}")
-            client=client,
-            method=method,
-            url=url,
-            http_version=http_version
-        )
-        self.echo(line)
 
-    def _echo_response_line(self, flow):
+    def _echo_response_line(self, flow: http.HTTPFlow) -> None:
         if flow.is_replay == "response":
-            replay = click.style("[replay] ", fg="yellow", bold=True)
+            replay_str = "[replay]"
+            replay = click.style(replay_str, fg="yellow", bold=True)
         else:
+            replay_str = ""
             replay = ""
 
-        code = flow.response.status_code
+        assert flow.response
+        code_int = flow.response.status_code
         code_color = None
-        if 200 <= code < 300:
+        if 200 <= code_int < 300:
             code_color = "green"
-        elif 300 <= code < 400:
+        elif 300 <= code_int < 400:
             code_color = "magenta"
-        elif 400 <= code < 600:
+        elif 400 <= code_int < 600:
             code_color = "red"
         code = click.style(
-            str(code),
+            str(code_int),
             fg=code_color,
             bold=True,
-            blink=(code == 418)
+            blink=(code_int == 418),
         )
+
+        if not flow.response.is_http2:
+            reason = flow.response.reason
+        else:
+            reason = net_http.status_codes.RESPONSES.get(flow.response.status_code, "")
         reason = click.style(
-            strutils.escape_control_characters(flow.response.reason),
+            strutils.escape_control_characters(reason),
             fg=code_color,
             bold=True
         )
@@ -204,23 +213,25 @@ class Dumper:
             size = human.pretty_size(len(flow.response.raw_content))
         size = click.style(size, bold=True)
 
+        http_version = ""
+        if (
+                flow.response.http_version not in ("HTTP/1.1", "HTTP/1.0")
+                or flow.request.http_version != flow.response.http_version
+        ):
+            # Hide version for h1 <-> h1 connections.
+            http_version = f"{flow.response.http_version} "
+
         arrows = click.style(" <<", bold=True)
         if ctx.options.flow_detail == 1:
             # This aligns the HTTP response code with the HTTP request method:
             # 127.0.0.1:59519: GET http://example.com/
             #               << 304 Not Modified 0b
-            arrows = " " * (len(human.format_address(flow.client_conn.address)) - 2) + arrows
+            pad = max(0, len(human.format_address(flow.client_conn.address)) - (2 + len(http_version) + len(replay_str)))
+            arrows = " " * pad + arrows
 
-        line = "{replay}{arrows} {code} {reason} {size}".format(
+        self.echo(f"{replay}{arrows} {http_version}{code} {reason} {size}")
-            replay=replay,
-            arrows=arrows,
-            code=code,
-            reason=reason,
-            size=size
-        )
-        self.echo(line)
 
-    def echo_flow(self, f):
+    def echo_flow(self, f: http.HTTPFlow) -> None:
         if f.request:
             self._echo_request_line(f)
             if ctx.options.flow_detail >= 2:
@@ -286,12 +297,13 @@ class Dumper:
                 f.close_reason))
 
     def tcp_error(self, f):
-        self.echo_error(
+        if self.match(f):
-            "Error in TCP connection to {}: {}".format(
+            self.echo_error(
-                human.format_address(f.server_conn.address), f.error
+                "Error in TCP connection to {}: {}".format(
-            ),
+                    human.format_address(f.server_conn.address), f.error
-            fg="red"
+                ),
-        )
+                fg="red"
+            )
 
     def tcp_message(self, f):
         if self.match(f):

mitmproxy/addons/termstatus.py

@@ -9,7 +9,7 @@ from mitmproxy.utils import human
 
 class TermStatus:
     def running(self):
-        if ctx.options.server:
+        if ctx.master.server.bound:
             ctx.log.info(
                 "Proxy server listening at http://{}".format(
                     human.format_address(ctx.master.server.address)

mitmproxy/flow.py

@@ -2,10 +2,11 @@ import time
 import typing  # noqa
 import uuid
 
-from mitmproxy import connections
+from mitmproxy import controller
-from mitmproxy import controller, exceptions  # noqa
+from mitmproxy import exceptions
 from mitmproxy import stateobject
 from mitmproxy import version
+from mitmproxy.utils import compat
 
 
 class Error(stateobject.StateObject):
@@ -63,8 +64,8 @@ class Flow(stateobject.StateObject):
     def __init__(
             self,
             type: str,
-            client_conn: connections.ClientConnection,
+            client_conn: compat.Client,
-            server_conn: connections.ServerConnection,
+            server_conn: compat.Server,
             live: bool=None
     ) -> None:
         self.type = type
@@ -84,8 +85,8 @@ class Flow(stateobject.StateObject):
     _stateobject_attributes = dict(
         id=str,
         error=Error,
-        client_conn=connections.ClientConnection,
+        client_conn=compat.Client,
-        server_conn=connections.ServerConnection,
+        server_conn=compat.Server,
         type=str,
         intercepted=bool,
         is_replay=str,

mitmproxy/http.py

@@ -1,10 +1,10 @@
 import html
 import time
 from typing import Optional, Tuple
-from mitmproxy import connections
 from mitmproxy import flow
 from mitmproxy import version
 from mitmproxy.net import http
+from mitmproxy.utils import compat
 
 HTTPRequest = http.Request
 HTTPResponse = http.Response
@@ -15,16 +15,16 @@ class HTTPFlow(flow.Flow):
     An HTTPFlow is a collection of objects representing a single HTTP
     transaction.
     """
-    request: HTTPRequest
+    request: http.Request
-    response: Optional[HTTPResponse] = None
+    response: Optional[http.Response] = None
     error: Optional[flow.Error] = None
     """
     Note that it's possible for a Flow to have both a response and an error
     object. This might happen, for instance, when a response was received
     from the server, but there was an error sending it back to the client.
     """
-    server_conn: connections.ServerConnection
+    server_conn: compat.Server
-    client_conn: connections.ClientConnection
+    client_conn: compat.Client
     intercepted: bool = False
     """ Is this flow currently being intercepted? """
     mode: str
@@ -37,8 +37,8 @@ class HTTPFlow(flow.Flow):
     _stateobject_attributes = flow.Flow._stateobject_attributes.copy()
     # mypy doesn't support update with kwargs
     _stateobject_attributes.update(dict(
-        request=HTTPRequest,
+        request=http.Request,
-        response=HTTPResponse,
+        response=http.Response,
         mode=str
     ))
 
@@ -67,7 +67,7 @@ def make_error_response(
         status_code: int,
         message: str = "",
         headers: Optional[http.Headers] = None,
-) -> HTTPResponse:
+) -> http.Response:
     body: bytes = """
         <html>
             <head>
@@ -92,11 +92,11 @@ def make_error_response(
             Content_Type="text/html"
         )
 
-    return HTTPResponse.make(status_code, body, headers)
+    return http.Response.make(status_code, body, headers)
 
 
-def make_connect_request(address: Tuple[str, int]) -> HTTPRequest:
+def make_connect_request(address: Tuple[str, int]) -> http.Request:
-    return HTTPRequest(
+    return http.Request(
         host=address[0],
         port=address[1],
         method=b"CONNECT",
@@ -115,7 +115,7 @@ def make_connect_request(address: Tuple[str, int]) -> HTTPRequest:
 def make_connect_response(http_version):
     # Do not send any response headers as it breaks proxying non-80 ports on
     # Android emulators using the -http-proxy option.
-    return HTTPResponse(
+    return http.Response(
         http_version,
         200,
         b"Connection established",
@@ -128,4 +128,4 @@ def make_connect_response(http_version):
 
 
 def make_expect_continue_response():
-    return HTTPResponse.make(100)
+    return http.Response.make(100)

mitmproxy/tools/_main.py

@@ -15,7 +15,7 @@ from mitmproxy import exceptions, master
 from mitmproxy import options
 from mitmproxy import optmanager
 from mitmproxy import proxy
-from mitmproxy.utils import debug, arg_check
+from mitmproxy.utils import compat, debug, arg_check
 
 
 def assert_utf8_env():
@@ -92,7 +92,7 @@ def run(
         )
         pconf = process_options(parser, opts, args)
         server: typing.Any = None
-        if pconf.options.server:
+        if pconf.options.server and not compat.new_proxy_core:  # new core initializes itself as an addon
             try:
                 server = proxy.server.ProxyServer(pconf)
             except exceptions.ServerException as v:

mitmproxy/utils/compat.py

@@ -0,0 +1,13 @@
+new_proxy_core = False
+"""If true, use mitmproxy's new sans-io proxy core."""
+
+if new_proxy_core:  # pragma: no cover
+    from mitmproxy.proxy2 import context
+
+    Client = context.Client
+    Server = context.Server
+else:  # pragma: no cover
+    from mitmproxy import connections
+
+    Client = connections.ClientConnection
+    Server = connections.ServerConnection

test/mitmproxy/addons/test_dumper.py

@@ -236,3 +236,14 @@ def test_websocket():
         f = tflow.twebsocketflow(client_conn=True, err=True)
         d.websocket_error(f)
         assert "Error in WebSocket" in sio_err.getvalue()
+
+
+def test_http2():
+    sio = io.StringIO()
+    sio_err = io.StringIO()
+    d = dumper.Dumper(sio, sio_err)
+    with taddons.context(d):
+        f = tflow.tflow(resp=True)
+        f.response.http_version = b"HTTP/2.0"
+        d.response(f)
+        assert "HTTP/2.0 200 OK" in sio.getvalue()

test/mitmproxy/addons/test_termstatus.py

@@ -10,6 +10,7 @@ async def test_configure():
     ts = termstatus.TermStatus()
     with taddons.context() as ctx:
         ctx.master.server = proxy.DummyServer()
+        ctx.master.server.bound = True
         ctx.configure(ts, server=False)
         ts.running()
         ctx.configure(ts, server=True)

test/mitmproxy/utils/test_compat.py

@@ -0,0 +1,6 @@
+from mitmproxy.utils import compat
+
+
+def test_simple():
+    assert compat.Server
+    assert compat.Client