fix: log until the very end (mitmproxy#6541)

CHANGELOG.md

@@ -11,6 +11,8 @@
   ([#6513](https://github.com/mitmproxy/mitmproxy/pull/6513), @dsphper)
 * DNS resolution is now exempted from `--ignore-hosts` in WireGuard Mode.
   ([#6513](https://github.com/mitmproxy/mitmproxy/pull/6513), @dsphper)
+* Fix a bug where logging was stopped prematurely during shutdown.
+  ([#6541](https://github.com/mitmproxy/mitmproxy/pull/6541), @mhils)
 * For plaintext traffic, `--ignore-hosts` now also takes HTTP/1 host headers into account.
   ([#6513](https://github.com/mitmproxy/mitmproxy/pull/6513), @dsphper)
 * Fix empty cookie attributes being set to `Key=` instead of `Key`

mitmproxy/addons/termlog.py

@@ -27,16 +27,10 @@ def configure(self, updated):
         if "termlog_verbosity" in updated:
             self.logger.setLevel(ctx.options.termlog_verbosity.upper())
 
-    def done(self):
-        t = self._teardown()
-        try:
-            # try to delay teardown a bit.
-            self._teardown_task = asyncio.create_task(t)
-        except RuntimeError:
-            # no event loop, we're in a test.
-            asyncio.run(t)
-
-    async def _teardown(self):
+    def uninstall(self) -> None:
+        # uninstall the log dumper.
+        # This happens at the very very end after done() is completed,
+        # because we don't want to uninstall while other addons are still logging.
         self.logger.uninstall()
 
 

mitmproxy/ctx.py

@@ -1,7 +1,12 @@
-import mitmproxy.log
-import mitmproxy.master
-import mitmproxy.options
+from __future__ import annotations
 
-log: "mitmproxy.log.Log"
-master: "mitmproxy.master.Master"
-options: "mitmproxy.options.Options"
+import typing
+
+if typing.TYPE_CHECKING:
+    import mitmproxy.log
+    import mitmproxy.master
+    import mitmproxy.options
+
+log: mitmproxy.log.Log
+master: mitmproxy.master.Master
+options: mitmproxy.options.Options

mitmproxy/log.py

@@ -2,14 +2,17 @@
 
 import logging
 import os
+import typing
 import warnings
 from dataclasses import dataclass
 
 from mitmproxy import hooks
-from mitmproxy import master
 from mitmproxy.contrib import click as miniclick
 from mitmproxy.utils import human
 
+if typing.TYPE_CHECKING:
+    from mitmproxy import master
+
 ALERT = logging.INFO + 1
 """
 The ALERT logging level has the same urgency as info, but

mitmproxy/master.py

@@ -2,6 +2,7 @@
 import logging
 
 from . import ctx as mitmproxy_ctx
+from .addons import termlog
 from .proxy.mode_specs import ReverseMode
 from mitmproxy import addonmanager
 from mitmproxy import command
@@ -20,16 +21,22 @@ class Master:
     """
 
     event_loop: asyncio.AbstractEventLoop
+    _termlog_addon: termlog.TermLog | None = None
 
     def __init__(
         self,
         opts: options.Options,
         event_loop: asyncio.AbstractEventLoop | None = None,
+        with_termlog: bool = False,
     ):
         self.options: options.Options = opts or options.Options()
         self.commands = command.CommandManager(self)
         self.addons = addonmanager.AddonManager(self)
 
+        if with_termlog:
+            self._termlog_addon = termlog.TermLog()
+            self.addons.add(self._termlog_addon)
+
         self.log = log.Log(self)  # deprecated, do not use.
         self._legacy_log_events = log.LegacyLogEvents(self)
         self._legacy_log_events.install()
@@ -85,6 +92,8 @@ async def running(self) -> None:
     async def done(self) -> None:
         await self.addons.trigger_event(hooks.DoneHook())
         self._legacy_log_events.uninstall()
+        if self._termlog_addon is not None:
+            self._termlog_addon.uninstall()
 
     def _asyncio_exception_handler(self, loop, context) -> None:
         try:

mitmproxy/proxy/mode_servers.py

@@ -28,6 +28,7 @@
 from typing import ClassVar
 from typing import Generic
 from typing import get_args
+from typing import TYPE_CHECKING
 from typing import TypeVar
 
 import mitmproxy_rs
@@ -36,7 +37,6 @@
 from mitmproxy import flow
 from mitmproxy import platform
 from mitmproxy.connection import Address
-from mitmproxy.master import Master
 from mitmproxy.net import local_ip
 from mitmproxy.net import udp
 from mitmproxy.proxy import commands
@@ -52,6 +52,9 @@
 else:
     from typing import Self
 
+if TYPE_CHECKING:
+    from mitmproxy.master import Master
+
 logger = logging.getLogger(__name__)
 
 

mitmproxy/tools/dump.py

@@ -5,7 +5,6 @@
 from mitmproxy.addons import errorcheck
 from mitmproxy.addons import keepserving
 from mitmproxy.addons import readfile
-from mitmproxy.addons import termlog
 
 
 class DumpMaster(master.Master):
@@ -16,9 +15,7 @@ def __init__(
         with_termlog=True,
         with_dumper=True,
     ) -> None:
-        super().__init__(options, event_loop=loop)
-        if with_termlog:
-            self.addons.add(termlog.TermLog())
+        super().__init__(options, event_loop=loop, with_termlog=with_termlog)
         self.addons.add(*addons.default_addons())
         if with_dumper:
             self.addons.add(dumper.Dumper())

mitmproxy/tools/web/master.py

@@ -14,7 +14,6 @@
 from mitmproxy.addons import eventstore
 from mitmproxy.addons import intercept
 from mitmproxy.addons import readfile
-from mitmproxy.addons import termlog
 from mitmproxy.addons import view
 from mitmproxy.addons.proxyserver import Proxyserver
 from mitmproxy.tools.web import app
@@ -26,7 +25,7 @@
 
 class WebMaster(master.Master):
     def __init__(self, opts: options.Options, with_termlog: bool = True):
-        super().__init__(opts)
+        super().__init__(opts, with_termlog=with_termlog)
         self.view = view.View()
         self.view.sig_view_add.connect(self._sig_view_add)
         self.view.sig_view_remove.connect(self._sig_view_remove)
@@ -39,8 +38,6 @@ def __init__(self, opts: options.Options, with_termlog: bool = True):
 
         self.options.changed.connect(self._sig_options_update)
 
-        if with_termlog:
-            self.addons.add(termlog.TermLog())
         self.addons.add(*addons.default_addons())
         self.addons.add(
             webaddons.WebAddon(),

test/mitmproxy/addons/test_termlog.py

@@ -1,4 +1,3 @@
-import asyncio
 import builtins
 import io
 import logging
@@ -16,14 +15,6 @@ def ensure_cleanup():
     assert not any(isinstance(x, termlog.TermLogHandler) for x in logging.root.handlers)
 
 
-async def test_delayed_teardown():
-    t = termlog.TermLog()
-    t.done()
-    assert t.logger in logging.root.handlers
-    await asyncio.sleep(0)
-    assert t.logger not in logging.root.handlers
-
-
 def test_output(capsys):
     logging.getLogger().setLevel(logging.DEBUG)
     t = termlog.TermLog()
@@ -39,7 +30,7 @@ def test_output(capsys):
     assert "two" not in out
     assert "three" in out
     assert "four" in out
-    t.done()
+    t.uninstall()
 
 
 async def test_styling(monkeypatch) -> None:
@@ -52,7 +43,7 @@ async def test_styling(monkeypatch) -> None:
         logging.warning("hello")
 
     assert "\x1b[33mhello\x1b[0m" in f.getvalue()
-    t.done()
+    t.uninstall()
 
 
 async def test_cannot_print(monkeypatch) -> None:
@@ -69,4 +60,4 @@ def _raise(*args, **kwargs):
 
         assert exc_info.value.args[0] == 1
 
-    t.done()
+    t.uninstall()