From bda9c4ea3acab80293796a479a0d59766fa3b3e8 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Sun, 10 Dec 2023 02:53:33 +0700 Subject: [PATCH] fix: log until the very end (#6541) --- CHANGELOG.md | 2 ++ mitmproxy/addons/termlog.py | 14 ++++---------- mitmproxy/ctx.py | 17 +++++++++++------ mitmproxy/log.py | 5 ++++- mitmproxy/master.py | 9 +++++++++ mitmproxy/proxy/mode_servers.py | 5 ++++- mitmproxy/tools/dump.py | 5 +---- mitmproxy/tools/web/master.py | 5 +---- test/mitmproxy/addons/test_termlog.py | 15 +++------------ 9 files changed, 39 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 14f237559f..7db955d1f3 100644 --- a/CHANGELOG.md +++ b/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` diff --git a/mitmproxy/addons/termlog.py b/mitmproxy/addons/termlog.py index f6d0038643..0bf1b63848 100644 --- a/mitmproxy/addons/termlog.py +++ b/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() diff --git a/mitmproxy/ctx.py b/mitmproxy/ctx.py index 2ce9c7c232..8380a100be 100644 --- a/mitmproxy/ctx.py +++ b/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 diff --git a/mitmproxy/log.py b/mitmproxy/log.py index 221b12a932..5eaa324ad4 100644 --- a/mitmproxy/log.py +++ b/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 diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 58da37bd17..aa92df9749 100644 --- a/mitmproxy/master.py +++ b/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: diff --git a/mitmproxy/proxy/mode_servers.py b/mitmproxy/proxy/mode_servers.py index 6e57a5f77c..5ada94ed11 100644 --- a/mitmproxy/proxy/mode_servers.py +++ b/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__) diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py index 1f6d62f4db..2f815b46be 100644 --- a/mitmproxy/tools/dump.py +++ b/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()) diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py index 520cf4d4ea..e80ca45688 100644 --- a/mitmproxy/tools/web/master.py +++ b/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(), diff --git a/test/mitmproxy/addons/test_termlog.py b/test/mitmproxy/addons/test_termlog.py index 0ebf3601a0..5ee5517a64 100644 --- a/test/mitmproxy/addons/test_termlog.py +++ b/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()