From 0f7f0d05340318d3ca4bf2ccddce30eede29adf7 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Thu, 7 Mar 2024 21:41:26 +0100
Subject: [PATCH] fix display of error messages on early shutdown (#6719)

fix #6707
fix #6716
---
 .gitignore                               |  1 +
 CHANGELOG.md                             |  2 ++
 mitmproxy/addons/errorcheck.py           | 11 +++++++++--
 mitmproxy/master.py                      | 20 +++++++++++++++-----
 test/mitmproxy/addons/test_errorcheck.py |  5 +++--
 5 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/.gitignore b/.gitignore
index 2a6f5eb37..a5dfafcc2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,6 +2,7 @@
 MANIFEST
 **/tmp
 /venv*
+/.venv*
 *.py[cdo]
 *.swp
 *.swo
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5a9cf993b..6634b3d7e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,8 @@
 
 ## Unreleased: mitmproxy next
 
+* Fix a bug where errors during startup would not be displayed when running mitmproxy.
+  ([#6719](https://github.com/mitmproxy/mitmproxy/pull/6719), @mhils)
 * Use newer cryptography APIs to avoid CryptographyDeprecationWarnings.
   This bumps the minimum required version to cryptography 42.0.
   ([#6718](https://github.com/mitmproxy/mitmproxy/pull/6718), @mhils)
diff --git a/mitmproxy/addons/errorcheck.py b/mitmproxy/addons/errorcheck.py
index 9b6eff66a..c46f92035 100644
--- a/mitmproxy/addons/errorcheck.py
+++ b/mitmproxy/addons/errorcheck.py
@@ -3,6 +3,8 @@ import logging
 import sys
 
 from mitmproxy import log
+from mitmproxy.contrib import click as miniclick
+from mitmproxy.utils import vt_codes
 
 
 class ErrorCheck:
@@ -29,8 +31,13 @@ class ErrorCheck:
         if self.logger.has_errored:
             plural = "s" if len(self.logger.has_errored) > 1 else ""
             if self.repeat_errors_on_stderr:
-                msg = "\n".join(self.logger.format(r) for r in self.logger.has_errored)
-                print(f"Error{plural} logged during startup:\n{msg}", file=sys.stderr)
+                message = f"Error{plural} logged during startup:"
+                if vt_codes.ensure_supported(sys.stderr):  # pragma: no cover
+                    message = miniclick.style(message, fg="red")
+                details = "\n".join(
+                    self.logger.format(r) for r in self.logger.has_errored
+                )
+                print(f"{message}\n{details}", file=sys.stderr)
             else:
                 print(
                     f"Error{plural} logged during startup, exiting...", file=sys.stderr
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 69740280c..7ed6fde5c 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -58,6 +58,7 @@ class Master:
         ):
             self.should_exit.clear()
 
+            # Can we exit before even bringing up servers?
             if ec := self.addons.get("errorcheck"):
                 await ec.shutdown_if_errored()
             if ps := self.addons.get("proxyserver"):
@@ -69,14 +70,23 @@ class Master:
                     ],
                     return_when=asyncio.FIRST_COMPLETED,
                 )
-            await self.running()
-            if ec := self.addons.get("errorcheck"):
-                await ec.shutdown_if_errored()
-                ec.finish()
+                if self.should_exit.is_set():
+                    return
+                # Did bringing up servers fail?
+                if ec := self.addons.get("errorcheck"):
+                    await ec.shutdown_if_errored()
+
             try:
+                await self.running()
+                # Any errors in the final part of startup?
+                if ec := self.addons.get("errorcheck"):
+                    await ec.shutdown_if_errored()
+                    ec.finish()
+
                 await self.should_exit.wait()
             finally:
-                # .wait might be cancelled (e.g. by sys.exit)
+                # if running() was called, we also always want to call done().
+                # .wait might be cancelled (e.g. by sys.exit), so  this needs to be in a finally block.
                 await self.done()
 
     def shutdown(self):
diff --git a/test/mitmproxy/addons/test_errorcheck.py b/test/mitmproxy/addons/test_errorcheck.py
index a9ab75ad5..bc71f6aaf 100644
--- a/test/mitmproxy/addons/test_errorcheck.py
+++ b/test/mitmproxy/addons/test_errorcheck.py
@@ -6,10 +6,11 @@ from mitmproxy.addons.errorcheck import ErrorCheck
 from mitmproxy.tools import main
 
 
-def test_errorcheck(tdata, capsys):
+@pytest.mark.parametrize("run_main", [main.mitmdump, main.mitmproxy])
+def test_errorcheck(tdata, capsys, run_main):
     """Integration test: Make sure that we catch errors on startup an exit."""
     with pytest.raises(SystemExit):
-        main.mitmdump(
+        run_main(
             [
                 "-n",
                 "-s",