From 5fc20e3e8c81d37ce551969b9b77d364c17954f0 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Tue, 15 Mar 2022 18:33:46 +0100 Subject: [PATCH] tests++ --- mitmproxy/addons/keepserving.py | 2 +- mitmproxy/addons/readfile.py | 2 +- mitmproxy/addons/script.py | 15 +- mitmproxy/command.py | 2 +- mitmproxy/contrib/tornado/__init__.py | 82 +++++ mitmproxy/contrib/tornado/asyncio.py | 282 ------------------ mitmproxy/master.py | 19 +- mitmproxy/platform/windows.py | 11 +- mitmproxy/test/taddons.py | 6 +- mitmproxy/tools/console/master.py | 6 +- mitmproxy/tools/console/statusbar.py | 3 +- mitmproxy/tools/main.py | 76 ++--- mitmproxy/tools/web/master.py | 3 +- mitmproxy/utils/asyncio_utils.py | 5 +- setup.cfg | 1 + setup.py | 2 +- test/bench/benchmark.py | 4 +- test/examples/test_examples.py | 4 +- test/mitmproxy/addons/test_script.py | 19 +- test/mitmproxy/addons/test_termlog.py | 2 + test/mitmproxy/data/addonscripts/shutdown.py | 2 +- test/mitmproxy/test_addonmanager.py | 6 +- .../tools/console/test_integration.py | 29 +- test/mitmproxy/tools/console/test_master.py | 26 -- .../mitmproxy/tools/console/test_statusbar.py | 12 +- test/mitmproxy/tools/test_cmdline.py | 5 +- test/mitmproxy/tools/test_dump.py | 6 +- test/mitmproxy/tools/web/test_app.py | 15 +- test/mitmproxy/tools/web/test_master.py | 19 -- test/mitmproxy/tservers.py | 31 -- 30 files changed, 214 insertions(+), 483 deletions(-) delete mode 100644 mitmproxy/contrib/tornado/asyncio.py delete mode 100644 test/mitmproxy/tservers.py diff --git a/mitmproxy/addons/keepserving.py b/mitmproxy/addons/keepserving.py index 161f33ff0..80cdbd925 100644 --- a/mitmproxy/addons/keepserving.py +++ b/mitmproxy/addons/keepserving.py @@ -37,4 +37,4 @@ class KeepServing: ctx.options.rfile, ] if any(opts) and not ctx.options.keepserving: - asyncio.get_event_loop().create_task(self.watch()) + asyncio.get_running_loop().create_task(self.watch()) diff --git a/mitmproxy/addons/readfile.py b/mitmproxy/addons/readfile.py index c75bc8f7d..977576e6d 100644 --- a/mitmproxy/addons/readfile.py +++ b/mitmproxy/addons/readfile.py @@ -76,7 +76,7 @@ class ReadFile: def running(self): if ctx.options.rfile: - asyncio.get_event_loop().create_task(self.doread(ctx.options.rfile)) + asyncio.get_running_loop().create_task(self.doread(ctx.options.rfile)) @command.command("readfile.reading") def reading(self) -> bool: diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py index fa2cbd34f..b09aefec9 100644 --- a/mitmproxy/addons/script.py +++ b/mitmproxy/addons/script.py @@ -14,6 +14,7 @@ from mitmproxy import command from mitmproxy import eventsequence from mitmproxy import ctx import mitmproxy.types as mtypes +from mitmproxy.utils import asyncio_utils def load_script(path: str) -> typing.Optional[types.ModuleType]: @@ -82,7 +83,10 @@ class Script: self.reloadtask = None if reload: - self.reloadtask = asyncio.ensure_future(self.watcher()) + self.reloadtask = asyncio_utils.create_task( + self.watcher(), + name=f"script watcher for {path}", + ) else: self.loadscript() @@ -107,10 +111,6 @@ class Script: ctx.master.addons.register(ns) self.ns = ns if self.ns: - # We're already running, so we have to explicitly register and - # configure the addon - if self.is_running: - ctx.master.addons.invoke_addon_sync(self.ns, hooks.RunningHook()) try: ctx.master.addons.invoke_addon_sync( self.ns, @@ -118,6 +118,9 @@ class Script: ) except exceptions.OptionsError as e: script_error_handler(self.fullpath, e, msg=str(e)) + if self.is_running: + # We're already running, so we call that on the addon now. + ctx.master.addons.invoke_addon_sync(self.ns, hooks.RunningHook()) async def watcher(self): last_mtime = 0 @@ -166,11 +169,11 @@ class ScriptLoader: mod = load_script(path) if mod: with addonmanager.safecall(): - ctx.master.addons.invoke_addon_sync(mod, hooks.RunningHook()) ctx.master.addons.invoke_addon_sync( mod, hooks.ConfigureHook(ctx.options.keys()), ) + ctx.master.addons.invoke_addon_sync(mod, hooks.RunningHook()) for f in flows: for evt in eventsequence.iterate(f): ctx.master.addons.invoke_addon_sync(mod, evt) diff --git a/mitmproxy/command.py b/mitmproxy/command.py index d67e9a424..c4389a366 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -241,7 +241,7 @@ class CommandManager: return parsed, next_params - def call(self, command_name: str, *args: typing.Sequence[typing.Any]) -> typing.Any: + def call(self, command_name: str, *args: typing.Any) -> typing.Any: """ Call a command with native arguments. May raise CommandError. """ diff --git a/mitmproxy/contrib/tornado/__init__.py b/mitmproxy/contrib/tornado/__init__.py index e69de29bb..c7000a752 100644 --- a/mitmproxy/contrib/tornado/__init__.py +++ b/mitmproxy/contrib/tornado/__init__.py @@ -0,0 +1,82 @@ +""" +SPDX-License-Identifier: Apache-2.0 + +Vendored partial copy of https://github.com/tornadoweb/tornado/blob/master/tornado/platform/asyncio.py @ e18ea03 +to fix https://github.com/tornadoweb/tornado/issues/3092. Can be removed once tornado >6.1 is out. +""" +import errno + +import select +import tornado +import tornado.platform.asyncio + + +def patch_tornado(): + if tornado.version != "6.1": + return + + def _run_select(self) -> None: + while True: + with self._select_cond: + while self._select_args is None and not self._closing_selector: + self._select_cond.wait() + if self._closing_selector: + return + assert self._select_args is not None + to_read, to_write = self._select_args + self._select_args = None + + # We use the simpler interface of the select module instead of + # the more stateful interface in the selectors module because + # this class is only intended for use on windows, where + # select.select is the only option. The selector interface + # does not have well-documented thread-safety semantics that + # we can rely on so ensuring proper synchronization would be + # tricky. + try: + # On windows, selecting on a socket for write will not + # return the socket when there is an error (but selecting + # for reads works). Also select for errors when selecting + # for writes, and merge the results. + # + # This pattern is also used in + # https://github.com/python/cpython/blob/v3.8.0/Lib/selectors.py#L312-L317 + rs, ws, xs = select.select(to_read, to_write, to_write) + ws = ws + xs + except OSError as e: + # After remove_reader or remove_writer is called, the file + # descriptor may subsequently be closed on the event loop + # thread. It's possible that this select thread hasn't + # gotten into the select system call by the time that + # happens in which case (at least on macOS), select may + # raise a "bad file descriptor" error. If we get that + # error, check and see if we're also being woken up by + # polling the waker alone. If we are, just return to the + # event loop and we'll get the updated set of file + # descriptors on the next iteration. Otherwise, raise the + # original error. + if e.errno == getattr(errno, "WSAENOTSOCK", errno.EBADF): + rs, _, _ = select.select([self._waker_r.fileno()], [], [], 0) + if rs: + ws = [] + else: + raise + else: + raise + + try: + self._real_loop.call_soon_threadsafe(self._handle_select, rs, ws) + except RuntimeError: + # "Event loop is closed". Swallow the exception for + # consistency with PollIOLoop (and logical consistency + # with the fact that we can't guarantee that an + # add_callback that completes without error will + # eventually execute). + pass + except AttributeError: + # ProactorEventLoop may raise this instead of RuntimeError + # if call_soon_threadsafe races with a call to close(). + # Swallow it too for consistency. + pass + + tornado.platform.asyncio.AddThreadSelectorEventLoop._run_select = _run_select diff --git a/mitmproxy/contrib/tornado/asyncio.py b/mitmproxy/contrib/tornado/asyncio.py deleted file mode 100644 index aaa2f2053..000000000 --- a/mitmproxy/contrib/tornado/asyncio.py +++ /dev/null @@ -1,282 +0,0 @@ -""" -SPDX-License-Identifier: Apache-2.0 - -Vendored partial copy of https://github.com/tornadoweb/tornado/blob/master/tornado/platform/asyncio.py @ e18ea03 -to fix https://github.com/tornadoweb/tornado/issues/3092. Can be removed once tornado >6.1 is out. -""" - -import asyncio -import atexit -import errno -import functools -import socket -import threading -import typing -from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union - -import select - -if typing.TYPE_CHECKING: - from typing import Set # noqa: F401 - from typing_extensions import Protocol - - - class _HasFileno(Protocol): - def fileno(self) -> int: - pass - - - _FileDescriptorLike = Union[int, _HasFileno] - -_T = TypeVar("_T") - -# Collection of selector thread event loops to shut down on exit. -_selector_loops = set() # type: Set[AddThreadSelectorEventLoop] - - -def _atexit_callback() -> None: - for loop in _selector_loops: - with loop._select_cond: - loop._closing_selector = True - loop._select_cond.notify() - try: - loop._waker_w.send(b"a") - except BlockingIOError: - pass - # If we don't join our (daemon) thread here, we may get a deadlock - # during interpreter shutdown. I don't really understand why. This - # deadlock happens every time in CI (both travis and appveyor) but - # I've never been able to reproduce locally. - loop._thread.join() - _selector_loops.clear() - - -atexit.register(_atexit_callback) - - -class AddThreadSelectorEventLoop(asyncio.AbstractEventLoop): - """Wrap an event loop to add implementations of the ``add_reader`` method family. - - Instances of this class start a second thread to run a selector. - This thread is completely hidden from the user; all callbacks are - run on the wrapped event loop's thread. - - This class is used automatically by Tornado; applications should not need - to refer to it directly. - - It is safe to wrap any event loop with this class, although it only makes sense - for event loops that do not implement the ``add_reader`` family of methods - themselves (i.e. ``WindowsProactorEventLoop``) - - Closing the ``AddThreadSelectorEventLoop`` also closes the wrapped event loop. - - """ - - # This class is a __getattribute__-based proxy. All attributes other than those - # in this set are proxied through to the underlying loop. - MY_ATTRIBUTES = { - "_consume_waker", - "_select_cond", - "_select_args", - "_closing_selector", - "_thread", - "_handle_event", - "_readers", - "_real_loop", - "_start_select", - "_run_select", - "_handle_select", - "_wake_selector", - "_waker_r", - "_waker_w", - "_writers", - "add_reader", - "add_writer", - "close", - "remove_reader", - "remove_writer", - } - - def __getattribute__(self, name: str) -> Any: - if name in AddThreadSelectorEventLoop.MY_ATTRIBUTES: - return super().__getattribute__(name) - return getattr(self._real_loop, name) - - def __init__(self, real_loop: asyncio.AbstractEventLoop) -> None: - self._real_loop = real_loop - - # Create a thread to run the select system call. We manage this thread - # manually so we can trigger a clean shutdown from an atexit hook. Note - # that due to the order of operations at shutdown, only daemon threads - # can be shut down in this way (non-daemon threads would require the - # introduction of a new hook: https://bugs.python.org/issue41962) - self._select_cond = threading.Condition() - self._select_args = ( - None - ) # type: Optional[Tuple[List[_FileDescriptorLike], List[_FileDescriptorLike]]] - self._closing_selector = False - self._thread = threading.Thread( - name="Tornado selector", - daemon=True, - target=self._run_select, - ) - self._thread.start() - # Start the select loop once the loop is started. - self._real_loop.call_soon(self._start_select) - - self._readers = {} # type: Dict[_FileDescriptorLike, Callable] - self._writers = {} # type: Dict[_FileDescriptorLike, Callable] - - # Writing to _waker_w will wake up the selector thread, which - # watches for _waker_r to be readable. - self._waker_r, self._waker_w = socket.socketpair() - self._waker_r.setblocking(False) - self._waker_w.setblocking(False) - _selector_loops.add(self) - self.add_reader(self._waker_r, self._consume_waker) - - def __del__(self) -> None: - # If the top-level application code uses asyncio interfaces to - # start and stop the event loop, no objects created in Tornado - # can get a clean shutdown notification. If we're just left to - # be GC'd, we must explicitly close our sockets to avoid - # logging warnings. - _selector_loops.discard(self) - self._waker_r.close() - self._waker_w.close() - - def close(self) -> None: - with self._select_cond: - self._closing_selector = True - self._select_cond.notify() - self._wake_selector() - self._thread.join() - _selector_loops.discard(self) - self._waker_r.close() - self._waker_w.close() - self._real_loop.close() - - def _wake_selector(self) -> None: - try: - self._waker_w.send(b"a") - except BlockingIOError: - pass - - def _consume_waker(self) -> None: - try: - self._waker_r.recv(1024) - except BlockingIOError: - pass - - def _start_select(self) -> None: - # Capture reader and writer sets here in the event loop - # thread to avoid any problems with concurrent - # modification while the select loop uses them. - with self._select_cond: - assert self._select_args is None - self._select_args = (list(self._readers.keys()), list(self._writers.keys())) - self._select_cond.notify() - - def _run_select(self) -> None: - while True: - with self._select_cond: - while self._select_args is None and not self._closing_selector: - self._select_cond.wait() - if self._closing_selector: - return - assert self._select_args is not None - to_read, to_write = self._select_args - self._select_args = None - - # We use the simpler interface of the select module instead of - # the more stateful interface in the selectors module because - # this class is only intended for use on windows, where - # select.select is the only option. The selector interface - # does not have well-documented thread-safety semantics that - # we can rely on so ensuring proper synchronization would be - # tricky. - try: - # On windows, selecting on a socket for write will not - # return the socket when there is an error (but selecting - # for reads works). Also select for errors when selecting - # for writes, and merge the results. - # - # This pattern is also used in - # https://github.com/python/cpython/blob/v3.8.0/Lib/selectors.py#L312-L317 - rs, ws, xs = select.select(to_read, to_write, to_write) - ws = ws + xs - except OSError as e: - # After remove_reader or remove_writer is called, the file - # descriptor may subsequently be closed on the event loop - # thread. It's possible that this select thread hasn't - # gotten into the select system call by the time that - # happens in which case (at least on macOS), select may - # raise a "bad file descriptor" error. If we get that - # error, check and see if we're also being woken up by - # polling the waker alone. If we are, just return to the - # event loop and we'll get the updated set of file - # descriptors on the next iteration. Otherwise, raise the - # original error. - if e.errno == getattr(errno, "WSAENOTSOCK", errno.EBADF): - rs, _, _ = select.select([self._waker_r.fileno()], [], [], 0) - if rs: - ws = [] - else: - raise - else: - raise - - try: - self._real_loop.call_soon_threadsafe(self._handle_select, rs, ws) - except RuntimeError: - # "Event loop is closed". Swallow the exception for - # consistency with PollIOLoop (and logical consistency - # with the fact that we can't guarantee that an - # add_callback that completes without error will - # eventually execute). - pass - except AttributeError: - # ProactorEventLoop may raise this instead of RuntimeError - # if call_soon_threadsafe races with a call to close(). - # Swallow it too for consistency. - pass - - def _handle_select( - self, rs: List["_FileDescriptorLike"], ws: List["_FileDescriptorLike"] - ) -> None: - for r in rs: - self._handle_event(r, self._readers) - for w in ws: - self._handle_event(w, self._writers) - self._start_select() - - def _handle_event( - self, - fd: "_FileDescriptorLike", - cb_map: Dict["_FileDescriptorLike", Callable], - ) -> None: - try: - callback = cb_map[fd] - except KeyError: - return - callback() - - def add_reader( - self, fd: "_FileDescriptorLike", callback: Callable[..., None], *args: Any - ) -> None: - self._readers[fd] = functools.partial(callback, *args) - self._wake_selector() - - def add_writer( - self, fd: "_FileDescriptorLike", callback: Callable[..., None], *args: Any - ) -> None: - self._writers[fd] = functools.partial(callback, *args) - self._wake_selector() - - def remove_reader(self, fd: "_FileDescriptorLike") -> None: - del self._readers[fd] - self._wake_selector() - - def remove_writer(self, fd: "_FileDescriptorLike") -> None: - del self._writers[fd] - self._wake_selector() diff --git a/mitmproxy/master.py b/mitmproxy/master.py index b3df5ad77..35eb3301e 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -1,5 +1,6 @@ import asyncio import traceback +from typing import Optional from mitmproxy import addonmanager, hooks from mitmproxy import command @@ -18,26 +19,32 @@ class Master: event_loop: asyncio.AbstractEventLoop - def __init__(self, opts): + def __init__(self, opts, event_loop: Optional[asyncio.AbstractEventLoop] = None): self.should_exit = asyncio.Event() self.options: options.Options = opts or options.Options() self.commands = command.CommandManager(self) self.addons = addonmanager.AddonManager(self) self.log = log.Log(self) + self.event_loop = event_loop or asyncio.get_running_loop() mitmproxy_ctx.master = self mitmproxy_ctx.log = self.log mitmproxy_ctx.options = self.options async def run(self) -> None: - self.event_loop = asyncio.get_running_loop() + old_handler = self.event_loop.get_exception_handler() self.event_loop.set_exception_handler(self._asyncio_exception_handler) - self.should_exit.clear() + try: + self.should_exit.clear() - await self.running() - await self.should_exit.wait() + # Handle scheduled tasks (configure()) first. + await asyncio.sleep(0) + await self.running() + await self.should_exit.wait() - await self.done() + await self.done() + finally: + self.event_loop.set_exception_handler(old_handler) def shutdown(self): """ diff --git a/mitmproxy/platform/windows.py b/mitmproxy/platform/windows.py index 966873e94..a9cc2c73c 100644 --- a/mitmproxy/platform/windows.py +++ b/mitmproxy/platform/windows.py @@ -17,11 +17,6 @@ import collections.abc import pydivert import pydivert.consts -if typing.TYPE_CHECKING: - class WindowsError(OSError): - @property - def winerror(self) -> int: - return 42 REDIRECT_API_HOST = "127.0.0.1" REDIRECT_API_PORT = 8085 @@ -300,7 +295,7 @@ class Redirect(threading.Thread): while True: try: packet = self.windivert.recv() - except WindowsError as e: + except OSError as e: if e.winerror == 995: return else: @@ -318,8 +313,8 @@ class Redirect(threading.Thread): """ try: return self.windivert.recv() - except WindowsError as e: - if e.winerror == 995: + except OSError as e: + if e.winerror == 995: # type: ignore return None else: raise diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index db2a4669b..09ede9489 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -21,7 +21,11 @@ class TestAddons(addonmanager.AddonManager): class RecordingMaster(mitmproxy.master.Master): def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + super().__init__(*args, **kwargs, event_loop=loop) self.addons = TestAddons(self) self.logs = [] diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py index 531917a42..f0f14065e 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -12,7 +12,7 @@ import tempfile import contextlib import threading -from mitmproxy.contrib.tornado.asyncio import AddThreadSelectorEventLoop +from tornado.platform.asyncio import AddThreadSelectorEventLoop import urwid @@ -23,6 +23,7 @@ from mitmproxy.addons import intercept from mitmproxy.addons import eventstore from mitmproxy.addons import readfile from mitmproxy.addons import view +from mitmproxy.contrib.tornado import patch_tornado from mitmproxy.tools.console import consoleaddons from mitmproxy.tools.console import defaultkeys from mitmproxy.tools.console import keymap @@ -209,8 +210,9 @@ class ConsoleMaster(master.Master): loop = asyncio.get_running_loop() if isinstance(loop, getattr(asyncio, "ProactorEventLoop", tuple())): + patch_tornado() # fix for https://bugs.python.org/issue37373 - loop = AddThreadSelectorEventLoop(loop) + loop = AddThreadSelectorEventLoop(loop) # type: ignore self.loop = urwid.MainLoop( urwid.SolidFill("x"), event_loop=urwid.AsyncioEventLoop(loop=loop), diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py index 57c5914ce..89bda0569 100644 --- a/mitmproxy/tools/console/statusbar.py +++ b/mitmproxy/tools/console/statusbar.py @@ -63,7 +63,8 @@ class ActionBar(urwid.WidgetWrap): def prep_prompt(self, p): return p.strip() + ": " - def shorten_message(self, msg, max_width): + @staticmethod + def shorten_message(msg, max_width): """ Shorten message so that it fits into a single line in the statusbar. """ diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index 0ab834e64..f9509c941 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -9,7 +9,7 @@ from mitmproxy import exceptions, master from mitmproxy import options from mitmproxy import optmanager from mitmproxy.tools import cmdline -from mitmproxy.utils import asyncio_utils, debug, arg_check +from mitmproxy.utils import debug, arg_check def assert_utf8_env(): @@ -58,48 +58,48 @@ def run( extra: Extra argument processing callable which returns a dict of options. """ - debug.register_info_dumpers() + async def main() -> master.Master: + debug.register_info_dumpers() - opts = options.Options() - master = master_cls(opts) + opts = options.Options() + master = master_cls(opts) - parser = make_parser(opts) + parser = make_parser(opts) - # To make migration from 2.x to 3.0 bearable. - if "-R" in sys.argv and sys.argv[sys.argv.index("-R") + 1].startswith("http"): - print("To use mitmproxy in reverse mode please use --mode reverse:SPEC instead") + # To make migration from 2.x to 3.0 bearable. + if "-R" in sys.argv and sys.argv[sys.argv.index("-R") + 1].startswith("http"): + print("To use mitmproxy in reverse mode please use --mode reverse:SPEC instead") - try: - args = parser.parse_args(arguments) - except SystemExit: - arg_check.check() - sys.exit(1) + try: + args = parser.parse_args(arguments) + except SystemExit: + arg_check.check() + sys.exit(1) - try: - opts.set(*args.setoptions, defer=True) - optmanager.load_paths( - opts, - os.path.join(opts.confdir, "config.yaml"), - os.path.join(opts.confdir, "config.yml"), - ) - process_options(parser, opts, args) + try: + opts.set(*args.setoptions, defer=True) + optmanager.load_paths( + opts, + os.path.join(opts.confdir, "config.yaml"), + os.path.join(opts.confdir, "config.yml"), + ) + process_options(parser, opts, args) - if args.options: - optmanager.dump_defaults(opts, sys.stdout) - sys.exit(0) - if args.commands: - master.commands.dump() - sys.exit(0) - if extra: - if args.filter_args: - master.log.info(f"Only processing flows that match \"{' & '.join(args.filter_args)}\"") - opts.update(**extra(args)) + if args.options: + optmanager.dump_defaults(opts, sys.stdout) + sys.exit(0) + if args.commands: + master.commands.dump() + sys.exit(0) + if extra: + if args.filter_args: + master.log.info(f"Only processing flows that match \"{' & '.join(args.filter_args)}\"") + opts.update(**extra(args)) - except exceptions.OptionsError as e: - print("{}: {}".format(sys.argv[0], e), file=sys.stderr) - sys.exit(1) + except exceptions.OptionsError as e: + print("{}: {}".format(sys.argv[0], e), file=sys.stderr) + sys.exit(1) - async def main(): loop = asyncio.get_running_loop() def _sigint(*_): @@ -113,10 +113,10 @@ def run( signal.signal(signal.SIGINT, _sigint) signal.signal(signal.SIGTERM, _sigterm) - return await master.run() + await master.run() + return master - asyncio.run(main()) - return master + return asyncio.run(main()) def mitmproxy(args=None) -> typing.Optional[int]: # pragma: no cover diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py index 02fe49698..201444135 100644 --- a/mitmproxy/tools/web/master.py +++ b/mitmproxy/tools/web/master.py @@ -1,6 +1,5 @@ import tornado.httpserver import tornado.ioloop -from tornado.platform.asyncio import AsyncIOMainLoop from mitmproxy import addons from mitmproxy import log @@ -11,6 +10,7 @@ from mitmproxy.addons import intercept from mitmproxy.addons import readfile from mitmproxy.addons import termlog from mitmproxy.addons import view +from mitmproxy.contrib.tornado import patch_tornado from mitmproxy.tools.web import app, webaddons, static_viewer @@ -93,6 +93,7 @@ class WebMaster(master.Master): ) async def running(self): + patch_tornado() # Register tornado with the current event loop tornado.ioloop.IOLoop.current() diff --git a/mitmproxy/utils/asyncio_utils.py b/mitmproxy/utils/asyncio_utils.py index b54d6ae8f..7c5c501ad 100644 --- a/mitmproxy/utils/asyncio_utils.py +++ b/mitmproxy/utils/asyncio_utils.py @@ -1,11 +1,8 @@ import asyncio -import concurrent.futures -import signal import sys import time -from asyncio import tasks from collections.abc import Coroutine -from typing import Awaitable, Callable, Optional, TypeVar +from typing import Optional from mitmproxy.utils import human diff --git a/setup.cfg b/setup.cfg index 22c56d67d..f3da275af 100644 --- a/setup.cfg +++ b/setup.cfg @@ -6,6 +6,7 @@ exclude = mitmproxy/contrib/*,test/mitmproxy/data/*,release/build/* addons = file,open,basestring,xrange,unicode,long,cmp [tool:pytest] +asyncio_mode = auto testpaths = test addopts = --capture=no --color=yes diff --git a/setup.py b/setup.py index 82c69fc3f..5c76b3e4e 100644 --- a/setup.py +++ b/setup.py @@ -104,7 +104,7 @@ setup( "parver>=0.1,<2.0", "pdoc>=4.0.0", "pyinstaller==4.7", - "pytest-asyncio>=0.10.0,<0.16,!=0.14", + "pytest-asyncio>=0.17.0,<0.19", "pytest-cov>=2.7.1,<3.1", "pytest-timeout>=1.3.3,<2.1", "pytest-xdist>=2.1.0,<3", diff --git a/test/bench/benchmark.py b/test/bench/benchmark.py index 96daf16cc..8b0a685e9 100644 --- a/test/bench/benchmark.py +++ b/test/bench/benchmark.py @@ -53,10 +53,10 @@ class Benchmark: def running(self): if not self.started: self.started = True - asyncio.get_event_loop().create_task(self.procs()) + asyncio.get_running_loop().create_task(self.procs()) def done(self): self.pr.dump_stats(ctx.options.benchmark_save_path + ".prof") -addons = [Benchmark()] \ No newline at end of file +addons = [Benchmark()] diff --git a/test/examples/test_examples.py b/test/examples/test_examples.py index 4a9606186..0bc7c2096 100644 --- a/test/examples/test_examples.py +++ b/test/examples/test_examples.py @@ -4,10 +4,8 @@ from mitmproxy.test import tutils from mitmproxy.test import taddons from mitmproxy.http import Headers -from ..mitmproxy import tservers - -class TestScripts(tservers.MasterTest): +class TestScripts: def test_add_header(self, tdata): with taddons.context() as tctx: a = tctx.script(tdata.path("../examples/addons/anatomy2.py")) diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index f815f1655..dc879fda1 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -194,7 +194,7 @@ class TestScriptLoader: await tctx.master.await_log("recorder response") debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ - 'recorder running', 'recorder configure', + 'recorder configure', 'recorder running', 'recorder requestheaders', 'recorder request', 'recorder responseheaders', 'recorder response' ] @@ -285,16 +285,16 @@ class TestScriptLoader: debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ 'a load', - 'a running', 'a configure', + 'a running', 'b load', - 'b running', 'b configure', + 'b running', 'c load', - 'c running', 'c configure', + 'c running', ] tctx.master.clear() @@ -330,14 +330,16 @@ class TestScriptLoader: 'b done', 'a configure', 'e load', - 'e running', 'e configure', + 'e running', ] + # stop reload tasks + tctx.configure(sc, scripts=[]) -def test_order(event_loop, tdata, capsys): + +def test_order(tdata, capsys): """Integration test: Make sure that the runtime hooks are triggered on startup in the correct order.""" - asyncio.set_event_loop(event_loop) main.mitmdump([ "-n", "-s", tdata.path("mitmproxy/data/addonscripts/recorder/recorder.py"), @@ -348,6 +350,7 @@ def test_order(event_loop, tdata, capsys): r"\('recorder', 'load', .+\n" r"\('recorder', 'configure', .+\n" r"Loading script.+shutdown.py\n" - r"\('recorder', 'running', .+\n$", + r"\('recorder', 'running', .+\n" + r"\('recorder', 'done', .+\n$", capsys.readouterr().out, ) diff --git a/test/mitmproxy/addons/test_termlog.py b/test/mitmproxy/addons/test_termlog.py index 89aea0aab..c245bd110 100644 --- a/test/mitmproxy/addons/test_termlog.py +++ b/test/mitmproxy/addons/test_termlog.py @@ -4,9 +4,11 @@ import pytest from mitmproxy.addons import termlog from mitmproxy import log from mitmproxy.test import taddons +from test.conftest import skip_windows class TestTermLog: + @skip_windows # not sure why this is suddenly necessary (03/2022) @pytest.mark.usefixtures('capfd') @pytest.mark.parametrize('outfile, expected_out, expected_err', [ (None, ['one', 'three'], ['four']), diff --git a/test/mitmproxy/data/addonscripts/shutdown.py b/test/mitmproxy/data/addonscripts/shutdown.py index f4a8f55d5..51a99b5c8 100644 --- a/test/mitmproxy/data/addonscripts/shutdown.py +++ b/test/mitmproxy/data/addonscripts/shutdown.py @@ -2,4 +2,4 @@ from mitmproxy import ctx def running(): - ctx.master.shutdown() \ No newline at end of file + ctx.master.shutdown() diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py index 8628ffb3f..0e5627df2 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -70,7 +70,7 @@ def test_command(): assert tctx.master.commands.execute("test.command") == "here" -def test_halt(): +async def test_halt(): o = options.Options() m = master.Master(o) a = addonmanager.AddonManager(m) @@ -218,7 +218,7 @@ async def test_simple(): assert ta in a -def test_load_option(): +async def test_load_option(): o = options.Options() m = master.Master(o) a = addonmanager.AddonManager(m) @@ -226,7 +226,7 @@ def test_load_option(): assert "custom_option" in m.options._options -def test_nesting(): +async def test_nesting(): o = options.Options() m = master.Master(o) a = addonmanager.AddonManager(m) diff --git a/test/mitmproxy/tools/console/test_integration.py b/test/mitmproxy/tools/console/test_integration.py index 15c156823..cdc74681c 100644 --- a/test/mitmproxy/tools/console/test_integration.py +++ b/test/mitmproxy/tools/console/test_integration.py @@ -1,3 +1,4 @@ +import asyncio import re import sys from typing import List @@ -5,7 +6,6 @@ from typing import List import pytest import mitmproxy.options -from mitmproxy import master from mitmproxy.tools.console import window from mitmproxy.tools.console.master import ConsoleMaster @@ -25,37 +25,42 @@ class ConsoleTestMaster(ConsoleMaster): for key in tokenize(input): self.window.keypress(self.ui.get_cols_rows(), key) + def screen_contents(self) -> str: + return b"\n".join(self.window.render((80, 24), True)._text_content()).decode() + @pytest.fixture -def console(monkeypatch): - monkeypatch.setattr(window.Screen, "get_cols_rows", lambda self: (120, 120)) - monkeypatch.setattr(master.Master, "run_loop", lambda *_: True) +def console(monkeypatch) -> ConsoleTestMaster: + # monkeypatch.setattr(window.Screen, "get_cols_rows", lambda self: (120, 120)) + monkeypatch.setattr(window.Screen, "start", lambda *_: True) monkeypatch.setattr(ConsoleTestMaster, "sig_call_in", lambda *_, **__: True) monkeypatch.setattr(sys.stdout, "isatty", lambda: True) - opts = mitmproxy.options.Options() - m = ConsoleTestMaster(opts) - m.run() - return m + async def make_master(): + opts = mitmproxy.options.Options() + m = ConsoleTestMaster(opts) + await m.running() + return m + return asyncio.run(make_master()) -@pytest.mark.asyncio def test_integration(tdata, console): console.type(f":view.flows.load {tdata.path('mitmproxy/data/dumpfile-7.mitm')}") console.type("") console.type("") # view second flow + assert "http://example.com/" in console.screen_contents() -@pytest.mark.asyncio def test_options_home_end(console): console.type("O") + assert "Options" in console.screen_contents() -@pytest.mark.asyncio def test_keybindings_home_end(console): console.type("K") + assert "Key Binding" in console.screen_contents() -@pytest.mark.asyncio def test_replay_count(console): console.type(":replay.server.count") + assert "Data viewer" in console.screen_contents() diff --git a/test/mitmproxy/tools/console/test_master.py b/test/mitmproxy/tools/console/test_master.py index 6df8b57f8..e69de29bb 100644 --- a/test/mitmproxy/tools/console/test_master.py +++ b/test/mitmproxy/tools/console/test_master.py @@ -1,26 +0,0 @@ -import urwid - -import pytest - -from mitmproxy import options, hooks -from mitmproxy.tools import console - -from ... import tservers - - -@pytest.mark.asyncio -class TestMaster(tservers.MasterTest): - def mkmaster(self, **opts): - o = options.Options(**opts) - m = console.master.ConsoleMaster(o) - m.addons.trigger(hooks.ConfigureHook(o.keys())) - return m - - async def test_basic(self): - m = self.mkmaster() - for i in (1, 2, 3): - try: - await self.dummy_cycle(m, 1, b"") - except urwid.ExitMainLoop: - pass - assert len(m.view) == i diff --git a/test/mitmproxy/tools/console/test_statusbar.py b/test/mitmproxy/tools/console/test_statusbar.py index 7ca78e263..2343c60a6 100644 --- a/test/mitmproxy/tools/console/test_statusbar.py +++ b/test/mitmproxy/tools/console/test_statusbar.py @@ -4,7 +4,7 @@ from mitmproxy import options from mitmproxy.tools.console import statusbar, master -def test_statusbar(monkeypatch): +async def test_statusbar(monkeypatch): o = options.Options() m = master.ConsoleMaster(o) m.options.update( @@ -48,15 +48,9 @@ def test_statusbar(monkeypatch): ("warn", "(more in eventlog)")]) ]) def test_shorten_message(message, ready_message): - o = options.Options() - m = master.ConsoleMaster(o) - ab = statusbar.ActionBar(m) - assert ab.shorten_message(message, max_width=30) == ready_message + assert statusbar.ActionBar.shorten_message(message, max_width=30) == ready_message def test_shorten_message_narrow(): - o = options.Options() - m = master.ConsoleMaster(o) - ab = statusbar.ActionBar(m) - shorten_msg = ab.shorten_message("error", max_width=4) + shorten_msg = statusbar.ActionBar.shorten_message("error", max_width=4) assert shorten_msg == [(None, "\u2026"), ("warn", "(more in eventlog)")] diff --git a/test/mitmproxy/tools/test_cmdline.py b/test/mitmproxy/tools/test_cmdline.py index d644f594d..d0b237d7e 100644 --- a/test/mitmproxy/tools/test_cmdline.py +++ b/test/mitmproxy/tools/test_cmdline.py @@ -1,7 +1,7 @@ import argparse from mitmproxy import options -from mitmproxy.tools import cmdline, web, dump, console +from mitmproxy.tools import cmdline from mitmproxy.tools import main @@ -15,20 +15,17 @@ def test_common(): def test_mitmproxy(): opts = options.Options() - console.master.ConsoleMaster(opts) ap = cmdline.mitmproxy(opts) assert ap def test_mitmdump(): opts = options.Options() - dump.DumpMaster(opts) ap = cmdline.mitmdump(opts) assert ap def test_mitmweb(): opts = options.Options() - web.master.WebMaster(opts) ap = cmdline.mitmweb(opts) assert ap diff --git a/test/mitmproxy/tools/test_dump.py b/test/mitmproxy/tools/test_dump.py index 86e76bdb1..c1e278c5c 100644 --- a/test/mitmproxy/tools/test_dump.py +++ b/test/mitmproxy/tools/test_dump.py @@ -13,21 +13,21 @@ class TestDumpMaster: m = dump.DumpMaster(o, with_termlog=False, with_dumper=False) return m - def test_has_error(self): + async def test_has_error(self): m = self.mkmaster() ent = log.LogEntry("foo", "error") m.addons.trigger(log.AddLogHook(ent)) assert m.errorcheck.has_errored @pytest.mark.parametrize("termlog", [False, True]) - def test_addons_termlog(self, termlog): + async def test_addons_termlog(self, termlog): with mock.patch('sys.stdout'): o = options.Options() m = dump.DumpMaster(o, with_termlog=termlog) assert (m.addons.get('termlog') is not None) == termlog @pytest.mark.parametrize("dumper", [False, True]) - def test_addons_dumper(self, dumper): + async def test_addons_dumper(self, dumper): with mock.patch('sys.stdout'): o = options.Options() m = dump.DumpMaster(o, with_dumper=dumper) diff --git a/test/mitmproxy/tools/web/test_app.py b/test/mitmproxy/tools/web/test_app.py index ae4ee0414..6dab4d383 100644 --- a/test/mitmproxy/tools/web/test_app.py +++ b/test/mitmproxy/tools/web/test_app.py @@ -1,4 +1,3 @@ -import asyncio import io import gzip import json @@ -76,7 +75,7 @@ def test_generate_tflow_js(tdata): ) -def test_generate_options_js(): +async def test_generate_options_js(): o = options.Options() m = webmaster.WebMaster(o) opt: optmanager._Option @@ -117,14 +116,12 @@ def test_generate_options_js(): @pytest.mark.usefixtures("no_tornado_logging", "tdata") class TestApp(tornado.testing.AsyncHTTPTestCase): - def get_new_ioloop(self): - io_loop = tornado.platform.asyncio.AsyncIOLoop() - asyncio.set_event_loop(io_loop.asyncio_loop) - return io_loop - def get_app(self): - o = options.Options(http2=False) - m = webmaster.WebMaster(o, with_termlog=False) + async def make_master(): + o = options.Options(http2=False) + return webmaster.WebMaster(o, with_termlog=False) + + m = self.io_loop.asyncio_loop.run_until_complete(make_master()) f = tflow.tflow(resp=True) f.id = "42" f.request.content = b"foo\nbar" diff --git a/test/mitmproxy/tools/web/test_master.py b/test/mitmproxy/tools/web/test_master.py index 87dc59aa9..e69de29bb 100644 --- a/test/mitmproxy/tools/web/test_master.py +++ b/test/mitmproxy/tools/web/test_master.py @@ -1,19 +0,0 @@ -import pytest - -from mitmproxy import options -from mitmproxy.tools.web import master - -from ... import tservers - - -class TestWebMaster(tservers.MasterTest): - def mkmaster(self, **opts): - o = options.Options(**opts) - return master.WebMaster(o) - - @pytest.mark.asyncio - async def test_basic(self): - m = self.mkmaster() - for i in (1, 2, 3): - await self.dummy_cycle(m, 1, b"") - assert len(m.view) == i diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py deleted file mode 100644 index bb0900198..000000000 --- a/test/mitmproxy/tservers.py +++ /dev/null @@ -1,31 +0,0 @@ -from unittest import mock - -from mitmproxy import eventsequence -from mitmproxy import io -from mitmproxy.proxy import server_hooks -from mitmproxy.test import tflow -from mitmproxy.test import tutils - - -class MasterTest: - - async def cycle(self, master, content): - f = tflow.tflow(req=tutils.treq(content=content)) - layer = mock.Mock("mitmproxy.proxy.protocol.base.Layer") - layer.client_conn = f.client_conn - await master.addons.handle_lifecycle(server_hooks.ClientConnectedHook(layer)) - for e in eventsequence.iterate(f): - await master.addons.handle_lifecycle(e) - await master.addons.handle_lifecycle(server_hooks.ClientDisconnectedHook(layer)) - return f - - async def dummy_cycle(self, master, n, content): - for i in range(n): - await self.cycle(master, content) - await master._shutdown() - - def flowfile(self, path): - with open(path, "wb") as f: - fw = io.FlowWriter(f) - t = tflow.tflow(resp=True) - fw.add(t)