Merge pull request #3105 from cortesi/opts

Add deferred options, tweak benchmarks, document done event

examples/addons/events.py

@@ -142,8 +142,12 @@ class Events:
 
     def done(self):
         """
-            Called when the addon shuts down, either by being removed from the
-            mitmproxy instance, or when mitmproxy itself shuts down.
+            Called when the addon shuts down, either by being removed from
+            the mitmproxy instance, or when mitmproxy itself shuts down. On
+            shutdown, this event is called after the event loop is
+            terminated, guaranteeing that it will be the final event an addon
+            sees. Note that log handlers are shut down at this point, so
+            calls to log functions will produce no output.
         """
 
     def load(self, entry: mitmproxy.addonmanager.Loader):

mitmproxy/addonmanager.py

@@ -173,6 +173,7 @@ class AddonManager:
             self.lookup[name] = a
         for a in traverse([addon]):
             self.master.commands.collect_commands(a)
+        self.master.options.process_deferred()
         return addon
 
     def add(self, *addons):

mitmproxy/optmanager.py

@@ -91,9 +91,12 @@ class OptManager:
         mutation doesn't change the option state inadvertently.
     """
     def __init__(self):
-        self.__dict__["_options"] = {}
-        self.__dict__["changed"] = blinker.Signal()
-        self.__dict__["errored"] = blinker.Signal()
+        self._deferred: typing.Dict[str, str] = {}
+        self.changed = blinker.Signal()
+        self.errored = blinker.Signal()
+        # Options must be the last attribute here - after that, we raise an
+        # error for attribute assigment to unknown options.
+        self._options: typing.Dict[str, typing.Any] = {}
 
     def add_option(
         self,
@@ -168,7 +171,14 @@ class OptManager:
             raise AttributeError("No such option: %s" % attr)
 
     def __setattr__(self, attr, value):
-        self.update(**{attr: value})
+        # This is slightly tricky. We allow attributes to be set on the instance
+        # until we have an _options attribute. After that, assignment is sent to
+        # the update function, and will raise an error for unknown options.
+        opts = self.__dict__.get("_options")
+        if not opts:
+            super().__setattr__(attr, value)
+        else:
+            self.update(**{attr: value})
 
     def keys(self):
         return set(self._options.keys())
@@ -272,20 +282,48 @@ class OptManager:
             options=options
         )
 
-    def set(self, *spec):
+    def set(self, *spec, defer=False):
+        """
+            Takes a list of set specification in standard form (option=value).
+            Options that are known are updated immediately. If defer is true,
+            options that are not known are deferred, and will be set once they
+            are added.
+        """
         vals = {}
+        unknown = {}
         for i in spec:
-            vals.update(self._setspec(i))
+            parts = i.split("=", maxsplit=1)
+            if len(parts) == 1:
+                optname, optval = parts[0], None
+            else:
+                optname, optval = parts[0], parts[1]
+            if optname in self._options:
+                vals[optname] = self.parse_setval(self._options[optname], optval)
+            else:
+                unknown[optname] = optval
+        if defer:
+            self._deferred.update(unknown)
+        elif unknown:
+            raise exceptions.OptionsError("Unknown options: %s" % ", ".join(unknown.keys()))
         self.update(**vals)
 
-    def parse_setval(self, optname: str, optstr: typing.Optional[str]) -> typing.Any:
+    def process_deferred(self):
+        """
+            Processes options that were deferred in previous calls to set, and
+            have since been added.
+        """
+        update = {}
+        for optname, optval in self._deferred.items():
+            if optname in self._options:
+                update[optname] = self.parse_setval(self._options[optname], optval)
+        self.update(**update)
+        for k in update.keys():
+            del self._deferred[k]
+
+    def parse_setval(self, o: _Option, optstr: typing.Optional[str]) -> typing.Any:
         """
             Convert a string to a value appropriate for the option type.
         """
-        if optname not in self._options:
-            raise exceptions.OptionsError("No such option %s" % optname)
-        o = self._options[optname]
-
         if o.typespec in (str, typing.Optional[str]):
             return optstr
         elif o.typespec in (int, typing.Optional[int]):
@@ -295,7 +333,7 @@ class OptManager:
                 except ValueError:
                     raise exceptions.OptionsError("Not an integer: %s" % optstr)
             elif o.typespec == int:
-                raise exceptions.OptionsError("Option is required: %s" % optname)
+                raise exceptions.OptionsError("Option is required: %s" % o.name)
             else:
                 return None
         elif o.typespec == bool:
@@ -313,19 +351,9 @@ class OptManager:
             if not optstr:
                 return []
             else:
-                return getattr(self, optname) + [optstr]
+                return getattr(self, o.name) + [optstr]
         raise NotImplementedError("Unsupported option type: %s", o.typespec)
 
-    def _setspec(self, spec):
-        d = {}
-        parts = spec.split("=", maxsplit=1)
-        if len(parts) == 1:
-            optname, optval = parts[0], None
-        else:
-            optname, optval = parts[0], parts[1]
-        d[optname] = self.parse_setval(optname, optval)
-        return d
-
     def make_parser(self, parser, optname, metavar=None, short=None):
         """
             Auto-Create a command-line parser entry for a named option. If the

mitmproxy/tools/main.py

@@ -110,7 +110,7 @@ def run(
         if args.commands:
             master.commands.dump()
             sys.exit(0)
-        opts.set(*args.setoptions)
+        opts.set(*args.setoptions, defer=True)
         if extra:
             opts.update(**extra(args))
 

test/bench/benchmark.py

@@ -11,6 +11,15 @@ class Benchmark:
         self.pr = cProfile.Profile()
         self.started = False
 
+        self.resps = 0
+        self.reqs = 0
+
+    def request(self, f):
+        self.reqs += 1
+
+    def response(self, f):
+        self.resps += 1
+
     async def procs(self):
         ctx.log.error("starting benchmark")
         backend = await asyncio.create_subprocess_exec("devd", "-q", "-p", "10001", ".")
@@ -23,6 +32,7 @@ class Benchmark:
         )
         stdout, _ = await traf.communicate()
         open(ctx.options.benchmark_save_path + ".bench", mode="wb").write(stdout)
+        ctx.log.error("Proxy saw %s requests, %s responses" % (self.reqs, self.resps))
         ctx.log.error(stdout.decode("ascii"))
         backend.kill()
         ctx.master.shutdown()

test/bench/run-mitmproxy

@@ -0,0 +1,4 @@
+#!/bin/sh
+
+mkdir -p results
+mitmproxy -p0 -q --set benchmark_save_path=./results/mitmproxy -s ./benchmark.py

test/mitmproxy/test_optmanager.py

@@ -70,7 +70,7 @@ def test_defaults():
 def test_required_int():
     o = TO()
     with pytest.raises(exceptions.OptionsError):
-        o.parse_setval("required_int", None)
+        o.parse_setval(o._options["required_int"], None)
 
 
 def test_deepcopy():
@@ -426,4 +426,13 @@ def test_set():
     assert opts.seqstr == []
 
     with pytest.raises(exceptions.OptionsError):
-        opts.set("nonexistent=wobble")
+        opts.set("deferred=wobble")
+
+    opts.set("deferred=wobble", defer=True)
+    assert "deferred" in opts._deferred
+    opts.process_deferred()
+    assert "deferred" in opts._deferred
+    opts.add_option("deferred", str, "default", "help")
+    opts.process_deferred()
+    assert "deferred" not in opts._deferred
+    assert opts.deferred == "wobble"