Merge pull request #1694 from cortesi/cachekey

console: add caching sort keys

8 files changed, 171 insertions, 93 deletions

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index b4ba2315..c09ff454 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -19,39 +19,83 @@ import mitmproxy.flow
 from mitmproxy import flowfilter
 from mitmproxy import exceptions
 
+# The underlying sorted list implementation expects the sort key to be stable
+# for the lifetime of the object. However, if we sort by size, for instance,
+# the sort order changes as the flow progresses through its lifecycle. We
+# address this through two means:
+#
+# - Let order keys cache the sort value by flow ID.
+#
+# - Add a facility to refresh items in the list by removing and re-adding them
+# when they are updated.
+
+
+class _OrderKey:
+    def __init__(self, view):
+        self.view = view
 
-def key_request_start(f: mitmproxy.flow.Flow) -> datetime.datetime:
-    return f.request.timestamp_start or 0
+    def generate(self, f: mitmproxy.flow.Flow) -> typing.Any:
+        pass
 
+    def refresh(self, f):
+        k = self._key()
+        old = self.view.settings[f][k]
+        new = self.generate(f)
+        if old != new:
+            self.view._view.remove(f)
+            self.view.settings[f][k] = new
+            self.view._view.add(f)
+            self.view.sig_refresh.send(self.view)
 
-def key_request_method(f: mitmproxy.flow.Flow) -> str:
-    return f.request.method
+    def _key(self):
+        return "_order_%s" % id(self)
 
+    def __call__(self, f):
+        k = self._key()
+        s = self.view.settings[f]
+        if k in s:
+            return s[k]
+        val = self.generate(f)
+        s[k] = val
+        return val
 
-def key_request_url(f: mitmproxy.flow.Flow) -> str:
-    return f.request.url
 
+class OrderRequestStart(_OrderKey):
+    def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime:
+        return f.request.timestamp_start or 0
 
-def key_size(f: mitmproxy.flow.Flow) -> int:
-    s = 0
-    if f.request.raw_content:
-        s += len(f.request.raw_content)
-    if f.response and f.response.raw_content:
-        s += len(f.response.raw_content)
-    return s
 
+class OrderRequestMethod(_OrderKey):
+    def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime:
+        return f.request.method
 
-orders = [
-    ("t", "time", key_request_start),
-    ("m", "method", key_request_method),
-    ("u", "url", key_request_url),
-    ("z", "size", key_size),
-]
+
+class OrderRequestURL(_OrderKey):
+    def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime:
+        return f.request.url
+
+
+class OrderKeySize(_OrderKey):
+    def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime:
+        s = 0
+        if f.request.raw_content:
+            s += len(f.request.raw_content)
+        if f.response and f.response.raw_content:
+            s += len(f.response.raw_content)
+        return s
 
 
 matchall = flowfilter.parse(".")
 
 
+orders = [
+    ("t", "time"),
+    ("m", "method"),
+    ("u", "url"),
+    ("z", "size"),
+]
+
+
 class View(collections.Sequence):
     def __init__(self):
         super().__init__()
@@ -59,8 +103,18 @@ class View(collections.Sequence):
         self.filter = matchall
         # Should we show only marked flows?
         self.show_marked = False
-        self.order_key = key_request_start
+
+        self.default_order = OrderRequestStart(self)
+        self.orders = dict(
+            time = self.default_order,
+            method = OrderRequestMethod(self),
+            url = OrderRequestURL(self),
+            size = OrderKeySize(self),
+        )
+        self.order_key = self.default_order
         self.order_reversed = False
+        self.focus_follow = False
+
         self._view = sortedcontainers.SortedListWithKey(key = self.order_key)
 
         # These signals broadcast events that affect the view. That is, an
@@ -115,13 +169,23 @@ class View(collections.Sequence):
     def index(self, f: mitmproxy.flow.Flow) -> int:
         return self._rev(self._view.index(f))
 
+    def __contains__(self, f: mitmproxy.flow.Flow) -> bool:
+        return self._view.__contains__(f)
+
+    def _order_key_name(self):
+        return "_order_%s" % id(self.order_key)
+
+    def _base_add(self, f):
+        self.settings[f][self._order_key_name()] = self.order_key(f)
+        self._view.add(f)
+
     def _refilter(self):
         self._view.clear()
         for i in self._store.values():
             if self.show_marked and not i.marked:
                 continue
             if self.filter(i):
-                self._view.add(i)
+                self._base_add(i)
         self.sig_refresh.send(self)
 
     # API
@@ -155,9 +219,10 @@ class View(collections.Sequence):
         """
         self._store.clear()
         self._view.clear()
+        self.settings.clear()
         self.sig_refresh.send(self)
 
-    def add(self, f: mitmproxy.flow.Flow):
+    def add(self, f: mitmproxy.flow.Flow) -> bool:
         """
             Adds a flow to the state. If the flow already exists, it is
             ignored.
@@ -165,7 +230,9 @@ class View(collections.Sequence):
         if f.id not in self._store:
             self._store[f.id] = f
             if self.filter(f):
-                self._view.add(f)
+                self._base_add(f)
+                if self.focus_follow:
+                    self.focus.flow = f
                 self.sig_add.send(self, flow=f)
 
     def remove(self, f: mitmproxy.flow.Flow):
@@ -173,10 +240,10 @@ class View(collections.Sequence):
             Removes the flow from the underlying store and the view.
         """
         if f.id in self._store:
-            del self._store[f.id]
             if f in self._view:
                 self._view.remove(f)
                 self.sig_remove.send(self, flow=f)
+            del self._store[f.id]
 
     def update(self, f: mitmproxy.flow.Flow):
         """
@@ -185,9 +252,16 @@ class View(collections.Sequence):
         if f.id in self._store:
             if self.filter(f):
                 if f not in self._view:
-                    self._view.add(f)
+                    self._base_add(f)
+                    if self.focus_follow:
+                        self.focus.flow = f
                     self.sig_add.send(self, flow=f)
                 else:
+                    # This is a tad complicated. The sortedcontainers
+                    # implementation assumes that the order key is stable. If
+                    # it changes mid-way Very Bad Things happen. We detect when
+                    # this happens, and re-fresh the item.
+                    self.order_key.refresh(f)
                     self.sig_update.send(self, flow=f)
             else:
                 try:
@@ -210,32 +284,31 @@ class View(collections.Sequence):
             self.set_filter(filt)
         if "order" in updated:
             if opts.order is None:
-                self.set_order(key_request_start)
+                self.set_order(self.default_order)
             else:
-                for _, name, func in orders:
-                    if name == opts.order:
-                        self.set_order(func)
-                        break
-                else:
+                if opts.order not in self.orders:
                     raise exceptions.OptionsError(
                         "Unknown flow order: %s" % opts.order
                     )
+                self.set_order(self.orders[opts.order])
         if "order_reversed" in updated:
             self.set_reversed(opts.order_reversed)
+        if "focus_follow" in updated:
+            self.focus_follow = opts.focus_follow
 
     def request(self, f):
         self.add(f)
 
-    def intercept(self, f):
+    def error(self, f):
         self.update(f)
 
-    def resume(self, f):
+    def response(self, f):
         self.update(f)
 
-    def error(self, f):
+    def intercept(self, f):
         self.update(f)
 
-    def response(self, f):
+    def resume(self, f):
         self.update(f)
 
 
@@ -258,7 +331,7 @@ class Focus:
         return self._flow
 
     @flow.setter
-    def flow(self, f: mitmproxy.flow.Flow):
+    def flow(self, f: typing.Optional[mitmproxy.flow.Flow]):
         if f is not None and f not in self.view:
             raise ValueError("Attempt to set focus to flow not in view")
         self._flow = f
@@ -305,6 +378,9 @@ class Settings(collections.Mapping):
         view.sig_remove.connect(self._sig_remove)
         view.sig_refresh.connect(self._sig_refresh)
 
+    def clear(self):
+        self.values.clear()
+
     def __iter__(self) -> typing.Iterable:
         return iter(self.values)
 
@@ -321,6 +397,6 @@ class Settings(collections.Mapping):
             del self.values[flow.id]
 
     def _sig_refresh(self, view):
-        for fid in self.values.keys():
+        for fid in list(self.values.keys()):
             if fid not in view._store:
                 del self.values[fid]
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 31ce17a3..ffbfb0cb 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -92,10 +92,14 @@ class Master:
         try:
             mtype, obj = self.event_queue.get(timeout=timeout)
             if mtype not in events.Events:
-                raise exceptions.ControlException("Unknown event %s" % repr(mtype))
+                raise exceptions.ControlException(
+                    "Unknown event %s" % repr(mtype)
+                )
             handle_func = getattr(self, mtype)
             if not callable(handle_func):
-                raise exceptions.ControlException("Handler %s not callable" % mtype)
+                raise exceptions.ControlException(
+                    "Handler %s not callable" % mtype
+                )
             if not handle_func.__dict__.get("__handler"):
                 raise exceptions.ControlException(
                     "Handler function %s is not decorated with controller.handler" % (
diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py
index cf6e1d35..e4b29d0f 100644
--- a/mitmproxy/tools/cmdline.py
+++ b/mitmproxy/tools/cmdline.py
@@ -784,8 +784,8 @@ def mitmproxy():
     )
     parser.add_argument(
         "--follow",
-        action="store_true", dest="follow",
-        help="Follow flow list."
+        action="store_true", dest="focus_follow",
+        help="Focus follows new flows."
     )
     parser.add_argument(
         "--order",
diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py
index a0a3dc94..76545893 100644
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -175,18 +175,6 @@ class FlowItem(urwid.WidgetWrap):
             signals.flowlist_change.send(self)
         elif key == "M":
             self.master.view.toggle_marked()
-        elif key == "o":
-            orders = [(i[1], i[0]) for i in view.orders]
-            lookup = dict([(i[0], i[1]) for i in view.orders])
-
-            def change_order(k):
-                self.master.options.order = lookup[k]
-
-            signals.status_prompt_onekey.send(
-                prompt = "Order",
-                keys = orders,
-                callback = change_order
-            )
         elif key == "r":
             try:
                 self.master.replay_request(self.flow)
@@ -220,9 +208,6 @@ class FlowItem(urwid.WidgetWrap):
             for f in self.master.view:
                 f.marked = False
             signals.flowlist_change.send(self)
-        elif key == "v":
-            val = not self.master.options.order_reversed
-            self.master.options.order_reversed = val
         elif key == "V":
             if not self.flow.modified():
                 signals.status_message.send(message="Flow not modified.")
@@ -389,8 +374,24 @@ class FlowListBox(urwid.ListBox):
                 keys = common.METHOD_OPTIONS,
                 callback = self.get_method
             )
+        elif key == "o":
+            orders = [(i[1], i[0]) for i in view.orders]
+            lookup = dict([(i[0], i[1]) for i in view.orders])
+
+            def change_order(k):
+                self.master.options.order = lookup[k]
+
+            signals.status_prompt_onekey.send(
+                prompt = "Order",
+                keys = orders,
+                callback = change_order
+            )
         elif key == "F":
-            self.master.toggle_follow_flows()
+            o = self.master.options
+            o.focus_follow = not o.focus_follow
+        elif key == "v":
+            val = not self.master.options.order_reversed
+            self.master.options.order_reversed = val
         elif key == "W":
             if self.master.options.outfile:
                 self.master.options.outfile = None
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index 932dc151..4c5e8c8c 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -44,25 +44,23 @@ class Options(mitmproxy.options.Options):
             self,
             *,  # all args are keyword-only.
             eventlog: bool = False,
-            follow: bool = False,
+            focus_follow: bool = False,
             intercept: Optional[str] = None,
             filter: Optional[str] = None,
             palette: Optional[str] = None,
             palette_transparent: bool = False,
             no_mouse: bool = False,
-            follow_focus: bool = False,
             order: Optional[str] = None,
             order_reversed: bool = False,
             **kwargs
     ):
         self.eventlog = eventlog
-        self.follow = follow
+        self.focus_follow = focus_follow
         self.intercept = intercept
         self.filter = filter
         self.palette = palette
         self.palette_transparent = palette_transparent
         self.no_mouse = no_mouse
-        self.follow_focus = follow_focus
         self.order = order
         self.order_reversed = order_reversed
         super().__init__(**kwargs)
@@ -83,7 +81,6 @@ class ConsoleMaster(master.Master):
         self.palette_transparent = options.palette_transparent
 
         self.logbuffer = urwid.SimpleListWalker([])
-        self.follow = options.follow
 
         self.view_stack = []
 
diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py
index b358f711..e292cbd7 100644
--- a/mitmproxy/tools/console/statusbar.py
+++ b/mitmproxy/tools/console/statusbar.py
@@ -199,7 +199,7 @@ class StatusBar(urwid.WidgetWrap):
             opts.append("killextra")
         if self.master.options.no_upstream_cert:
             opts.append("no-upstream-cert")
-        if self.master.options.follow_focus:
+        if self.master.options.focus_follow:
             opts.append("following")
         if self.master.options.stream_large_bodies:
             opts.append(
diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py
index 772841f5..08e5f2ca 100644
--- a/mitmproxy/tools/main.py
+++ b/mitmproxy/tools/main.py
@@ -69,7 +69,7 @@ def mitmproxy(args=None):  # pragma: no cover
         console_options.palette = args.palette
         console_options.palette_transparent = args.palette_transparent
         console_options.eventlog = args.eventlog
-        console_options.follow = args.follow
+        console_options.focus_follow = args.focus_follow
         console_options.intercept = args.intercept
         console_options.filter = args.filter
         console_options.no_mouse = args.no_mouse
diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py
index 750e8469..63df8307 100644
--- a/test/mitmproxy/addons/test_view.py
+++ b/test/mitmproxy/addons/test_view.py
@@ -7,6 +7,20 @@ from mitmproxy.test import taddons
 from .. import tutils
 
 
+class Options(options.Options):
+    def __init__(
+        self, *,
+        filter=None,
+        order=None,
+        order_reversed=False,
+        **kwargs
+    ):
+        self.filter = filter
+        self.order = order
+        self.order_reversed = order_reversed
+        super().__init__(**kwargs)
+
+
 def test_simple():
     v = view.View()
     f = tflow.tflow()
@@ -69,22 +83,23 @@ def test_filter():
 
 def test_order():
     v = view.View()
-    v.request(tft(method="get", start=1))
-    v.request(tft(method="put", start=2))
-    v.request(tft(method="get", start=3))
-    v.request(tft(method="put", start=4))
-    assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4]
-
-    v.set_order(view.key_request_method)
-    assert [i.request.method for i in v] == ["GET", "GET", "PUT", "PUT"]
-    v.set_reversed(True)
-    assert [i.request.method for i in v] == ["PUT", "PUT", "GET", "GET"]
+    with taddons.context(options=Options()) as tctx:
+        v.request(tft(method="get", start=1))
+        v.request(tft(method="put", start=2))
+        v.request(tft(method="get", start=3))
+        v.request(tft(method="put", start=4))
+        assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4]
 
-    v.set_order(view.key_request_start)
-    assert [i.request.timestamp_start for i in v] == [4, 3, 2, 1]
+        tctx.configure(v, order="method")
+        assert [i.request.method for i in v] == ["GET", "GET", "PUT", "PUT"]
+        v.set_reversed(True)
+        assert [i.request.method for i in v] == ["PUT", "PUT", "GET", "GET"]
+
+        tctx.configure(v, order="time")
+        assert [i.request.timestamp_start for i in v] == [4, 3, 2, 1]
 
-    v.set_reversed(False)
-    assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4]
+        v.set_reversed(False)
+        assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4]
 
 
 def test_reversed():
@@ -256,7 +271,6 @@ def test_settings():
 
     tutils.raises(KeyError, v.settings.__getitem__, f)
     v.add(f)
-    assert v.settings[f] == {}
     v.settings[f]["foo"] = "bar"
     assert v.settings[f]["foo"] == "bar"
     assert len(list(v.settings)) == 1
@@ -265,20 +279,6 @@ def test_settings():
     assert not v.settings.keys()
 
 
-class Options(options.Options):
-    def __init__(
-        self, *,
-        filter=None,
-        order=None,
-        order_reversed=False,
-        **kwargs
-    ):
-        self.filter = filter
-        self.order = order
-        self.order_reversed = order_reversed
-        super().__init__(**kwargs)
-
-
 def test_configure():
     v = view.View()
     with taddons.context(options=Options()) as tctx: