Options unification: sizes

Start dealing with corner cases:

- Sizes are always stored in options as strings
- Add a new core addon that's responsible for verifying settings that don't
belong to an addon
- Add a _processed scratch space on the Options object for processed core
values to be stored in. This is pretty dirty, but less dirty than re-parsing
values every time. We'll come up with something better down the track.

mitmproxy/addons/__init__.py

@@ -3,6 +3,7 @@ from mitmproxy.addons import anticomp
 from mitmproxy.addons import check_alpn
 from mitmproxy.addons import check_ca
 from mitmproxy.addons import clientplayback
+from mitmproxy.addons import core
 from mitmproxy.addons import disable_h2c_upgrade
 from mitmproxy.addons import onboarding
 from mitmproxy.addons import proxyauth
@@ -19,6 +20,7 @@ from mitmproxy.addons import upstream_auth
 
 def default_addons():
     return [
+        core.Core(),
         anticache.AntiCache(),
         anticomp.AntiComp(),
         check_alpn.CheckALPN(),

mitmproxy/addons/core.py

@@ -0,0 +1,20 @@
+"""
+    The core addon is responsible for verifying core settings that are not
+    checked by other addons.
+"""
+from mitmproxy import exceptions
+from mitmproxy.utils import human
+
+
+class Core:
+    def configure(self, options, updated):
+        if "body_size_limit" in updated and options.body_size_limit:
+            try:
+                options._processed["body_size_limit"] = human.parse_size(
+                    options.body_size_limit
+                )
+            except ValueError as e:
+                raise exceptions.OptionsError(
+                    "Invalid body size limit specification: %s" %
+                    options.body_size_limit
+                )

mitmproxy/addons/streambodies.py

@@ -1,6 +1,7 @@
 from mitmproxy.net.http import http1
 from mitmproxy import exceptions
 from mitmproxy import ctx
+from mitmproxy.utils import human
 
 
 class StreamBodies:
@@ -8,7 +9,11 @@ class StreamBodies:
         self.max_size = None
 
     def configure(self, options, updated):
-        self.max_size = options.stream_large_bodies
+        if "stream_large_bodies" in updated and options.stream_large_bodies:
+            try:
+                self.max_size = human.parse_size(options.stream_large_bodies)
+            except ValueError as e:
+                raise exceptions.OptionsError(e)
 
     def run(self, f, is_request):
         if self.max_size:

mitmproxy/options.py

@@ -77,7 +77,14 @@ class Options(optmanager.OptManager):
         self.add_option("server_replay", [], Sequence[str])
         self.add_option("stickycookie", None, Optional[str])
         self.add_option("stickyauth", None, Optional[str])
-        self.add_option("stream_large_bodies", None, Optional[int])
+        self.add_option(
+            "stream_large_bodies", None, Optional[str],
+            """
+                Stream data to the client if response body exceeds the given
+                threshold. If streamed, the body will not be stored in any way.
+                Understands k/m/g suffixes, i.e. 3m for 3 megabytes.
+            """
+        )
         self.add_option(
             "verbosity", 2, int,
             "Log verbosity."
@@ -109,7 +116,11 @@ class Options(optmanager.OptManager):
             "Add all certificates of the upstream server to the certificate chain "
             "that will be served to the proxy client, as extras."
         )
-        self.add_option("body_size_limit", None, Optional[int])
+        self.add_option(
+            "body_size_limit", None, Optional[str],
+            "Byte size limit of HTTP request and response bodies."
+            " Understands k/m/g suffixes, i.e. 3m for 3 megabytes."
+        )
         self.add_option("cadir", CA_DIR, str)
         self.add_option("certs", [], Sequence[Tuple[str, str]])
         self.add_option("ciphers_client", DEFAULT_CLIENT_CIPHERS, str)

mitmproxy/optmanager.py

@@ -12,7 +12,6 @@ import ruamel.yaml
 from mitmproxy import exceptions
 from mitmproxy.utils import typecheck
 
-
 """
     The base implementation for Options.
 """
@@ -92,6 +91,7 @@ class OptManager:
         self.__dict__["_options"] = {}
         self.__dict__["changed"] = blinker.Signal()
         self.__dict__["errored"] = blinker.Signal()
+        self.__dict__["_processed"] = {}
 
     def add_option(
         self,
@@ -330,7 +330,7 @@ class OptManager:
                 help=o.help
             )
             parser.set_defaults(**{option: o.default})
-        elif o.typespec == int:
+        elif o.typespec in (int, typing.Optional[int]):
             parser.add_argument(
                 "--%s" % f,
                 action="store",
@@ -339,5 +339,14 @@ class OptManager:
                 help=o.help,
                 metavar=metavar
             )
+        elif o.typespec in (str, typing.Optional[str]):
+            parser.add_argument(
+                "--%s" % f,
+                action="store",
+                type=str,
+                dest=option,
+                help=o.help,
+                metavar=metavar
+            )
         else:
             raise ValueError("Unsupported option type: %s", o.typespec)

mitmproxy/proxy/protocol/http1.py

@@ -19,7 +19,7 @@ class Http1Layer(httpbase._HttpTransmissionLayer):
         return http1.read_body(
             self.client_conn.rfile,
             expected_size,
-            self.config.options.body_size_limit
+            self.config.options._processed.get("body_size_limit")
         )
 
     def send_request(self, request):
@@ -35,7 +35,7 @@ class Http1Layer(httpbase._HttpTransmissionLayer):
         return http1.read_body(
             self.server_conn.rfile,
             expected_size,
-            self.config.options.body_size_limit
+            self.config.options._processed.get("body_size_limit")
         )
 
     def send_response_headers(self, response):

mitmproxy/proxy/protocol/http2.py

@@ -183,7 +183,7 @@ class Http2Layer(base.Layer):
         return True
 
     def _handle_data_received(self, eid, event, source_conn):
-        bsl = self.config.options.body_size_limit
+        bsl = self.config.options._processed.get("body_size_limit")
         if bsl and self.streams[eid].queued_data_length > bsl:
             self.streams[eid].kill()
             self.connections[source_conn].safe_reset_stream(

mitmproxy/proxy/protocol/http_replay.py

@@ -31,6 +31,7 @@ class RequestReplayThread(basethread.BaseThread):
 
     def run(self):
         r = self.f.request
+        bsl = self.config.options._processed.get("body_size_limit")
         first_line_format_backup = r.first_line_format
         server = None
         try:
@@ -55,7 +56,7 @@ class RequestReplayThread(basethread.BaseThread):
                         resp = http1.read_response(
                             server.rfile,
                             connect_request,
-                            body_size_limit=self.config.options.body_size_limit
+                            body_size_limit=bsl
                         )
                         if resp.status_code != 200:
                             raise exceptions.ReplayException("Upstream server refuses CONNECT request")
@@ -87,7 +88,7 @@ class RequestReplayThread(basethread.BaseThread):
                     http1.read_response(
                         server.rfile,
                         r,
-                        body_size_limit=self.config.options.body_size_limit
+                        body_size_limit=bsl
                     )
                 )
             if self.channel:

mitmproxy/tools/cmdline.py

@@ -4,7 +4,6 @@ import os
 from mitmproxy import exceptions
 from mitmproxy import options
 from mitmproxy import platform
-from mitmproxy.utils import human
 from mitmproxy.net import tcp
 from mitmproxy import version
 from mitmproxy.addons import view
@@ -25,10 +24,6 @@ def get_common_options(args):
     if args.stickyauth_filt:
         stickyauth = args.stickyauth_filt
 
-    stream_large_bodies = args.stream_large_bodies
-    if stream_large_bodies:
-        stream_large_bodies = human.parse_size(stream_large_bodies)
-
     if args.streamfile and args.streamfile[0] == args.rfile:
         if args.streamfile[1] == "wb":
             raise exceptions.OptionsError(
@@ -49,15 +44,6 @@ def get_common_options(args):
             parts = ["*", parts[0]]
         certs.append(parts)
 
-    body_size_limit = args.body_size_limit
-    if body_size_limit:
-        try:
-            body_size_limit = human.parse_size(body_size_limit)
-        except ValueError as e:
-            raise exceptions.OptionsError(
-                "Invalid body size limit specification: %s" % body_size_limit
-            )
-
     # Establish proxy mode
     c = 0
     mode, upstream_server = "regular", None
@@ -117,7 +103,7 @@ def get_common_options(args):
         scripts=args.scripts,
         stickycookie=stickycookie,
         stickyauth=stickyauth,
-        stream_large_bodies=stream_large_bodies,
+        stream_large_bodies=args.stream_large_bodies,
         showhost=args.showhost,
         streamfile=args.streamfile[0] if args.streamfile else None,
         streamfile_append=True if args.streamfile and args.streamfile[1] == "a" else False,
@@ -132,7 +118,7 @@ def get_common_options(args):
         auth_singleuser = args.auth_singleuser,
         auth_htpasswd = args.auth_htpasswd,
         add_upstream_certs_to_client_chain = args.add_upstream_certs_to_client_chain,
-        body_size_limit = body_size_limit,
+        body_size_limit = args.body_size_limit,
         cadir = args.cadir,
         certs = certs,
         ciphers_client = args.ciphers_client,
@@ -229,23 +215,8 @@ def basic_options(parser, opts):
         help="Append flows to file."
     )
     opts.make_parser(parser, "anticomp")
-    parser.add_argument(
-        "-Z", "--body-size-limit",
-        action="store", dest="body_size_limit",
-        metavar="SIZE",
-        help="Byte size limit of HTTP request and response bodies."
-             " Understands k/m/g suffixes, i.e. 3m for 3 megabytes."
-    )
-    parser.add_argument(
-        "--stream",
-        action="store", dest="stream_large_bodies",
-        metavar="SIZE",
-        help="""
-            Stream data to the client if response body exceeds the given
-            threshold. If streamed, the body will not be stored in any way.
-            Understands k/m/g suffixes, i.e. 3m for 3 megabytes.
-         """
-    )
+    opts.make_parser(parser, "body_size_limit", metavar="SIZE")
+    opts.make_parser(parser, "stream_large_bodies")
 
 
 def proxy_modes(parser, opts):

test/mitmproxy/addons/test_core.py

@@ -0,0 +1,13 @@
+from mitmproxy import exceptions
+from mitmproxy.addons import core
+from mitmproxy.test import taddons
+import pytest
+
+
+def test_simple():
+    sa = core.Core()
+    with taddons.context() as tctx:
+        with pytest.raises(exceptions.OptionsError):
+            tctx.configure(sa, body_size_limit = "invalid")
+        tctx.configure(sa, body_size_limit = "1m")
+        assert tctx.options._processed["body_size_limit"]

test/mitmproxy/addons/test_streambodies.py

@@ -1,13 +1,16 @@
+from mitmproxy import exceptions
 from mitmproxy.test import tflow
 from mitmproxy.test import taddons
-
 from mitmproxy.addons import streambodies
+import pytest
 
 
 def test_simple():
     sa = streambodies.StreamBodies()
     with taddons.context() as tctx:
-        tctx.configure(sa, stream_large_bodies = 10)
+        with pytest.raises(exceptions.OptionsError):
+            tctx.configure(sa, stream_large_bodies = "invalid")
+        tctx.configure(sa, stream_large_bodies = "10")
 
         f = tflow.tflow()
         f.request.content = b""

test/mitmproxy/proxy/protocol/test_http2.py

@@ -499,7 +499,8 @@ class TestBodySizeLimit(_Http2Test):
         return True
 
     def test_body_size_limit(self):
-        self.config.options.body_size_limit = 20
+        self.config.options.body_size_limit = "20"
+        self.config.options._processed["body_size_limit"] = 20
 
         client, h2_conn = self._setup_connection()