Improve command handling testing

hippolyzer/lib/proxy/addons.py

@@ -73,9 +73,11 @@ class AddonManager:
     _SUBPROCESS: bool = False
     _REPL_TASK: Optional[asyncio.Task] = None
     _HOT_RELOADING_STACK: Set[str] = set()
+    _SWALLOW_ADDON_EXCEPTIONS: bool = True
 
     @classmethod
-    def init(cls, addon_script_paths, session_manager, addon_objects=None, subprocess=False):
+    def init(cls, addon_script_paths, session_manager, addon_objects=None, subprocess=False,
+             swallow_addon_exceptions=True):
         addon_objects = addon_objects or []
         cls.BASE_ADDON_SPECS.clear()
         cls.FRESH_ADDON_MODULES.clear()
@@ -85,6 +87,7 @@ class AddonManager:
         # Are we a child process from `multiprocessing`?
         cls._SUBPROCESS = subprocess
         cls._HOT_RELOADING_STACK.clear()
+        cls._SWALLOW_ADDON_EXCEPTIONS = swallow_addon_exceptions
         for path in addon_script_paths:
             # defer reloading until we've collected them all
             cls.load_addon_from_path(path, reload=False)
@@ -369,6 +372,8 @@ class AddonManager:
             return hook_func(*args, **kwargs)
         except:
             logging.exception("Exploded in %r's %s hook" % (addon, hook_name))
+            if not cls._SWALLOW_ADDON_EXCEPTIONS:
+                raise
 
     @classmethod
     def _show_message(cls, text: str, session: Optional[Session] = None):
@@ -377,7 +382,8 @@ class AddonManager:
         try:
             show_message(text, session=session)
         except:
-            pass
+            if not cls._SWALLOW_ADDON_EXCEPTIONS:
+                raise
 
     @classmethod
     def _show_error(cls, text: str, session: Optional[Session] = None):
@@ -389,13 +395,15 @@ class AddonManager:
         cls._reload_addons()
         if message.name == "ChatFromViewer" and "ChatData" in message:
             if message["ChatData"]["Channel"] == cls.COMMAND_CHANNEL:
+                region.circuit.drop_message(message)
                 with addon_ctx.push(session, region):
                     try:
                         cls._handle_command(session, region, message["ChatData"]["Message"])
                     except Exception as e:
                         LOG.exception(f'Failed while handling command {message["ChatData"]["Message"]}')
                         cls._show_message(str(e), session)
-                    region.circuit.drop_message(message)
+                        if not cls._SWALLOW_ADDON_EXCEPTIONS:
+                            raise
                     return True
         if message.name == "ChatFromSimulator" and "ChatData" in message:
             chat: str = message["ChatData"]["Message"]
@@ -416,6 +424,8 @@ class AddonManager:
                         return True
                 except:
                     LOG.exception(f"Failed while handling command {chat!r}")
+                    if not cls._SWALLOW_ADDON_EXCEPTIONS:
+                        raise
 
         with addon_ctx.push(session, region):
             return cls._call_all_addon_hooks("handle_lludp_message", session, region, message)

hippolyzer/lib/proxy/tests/integration/__init__.py

@@ -51,6 +51,7 @@ class BaseIntegrationTest(unittest.IsolatedAsyncioTestCase):
         self.session_manager.claim_session(self.session.id)
         self.session.open_circuit(self.client_addr, self.region_addr,
                                   self.protocol.transport)
+        self.session.main_region = self.session.regions[-1]
 
     def _msg_to_datagram(self, msg: ProxiedMessage, src, dst, direction, socks_header=True):
         serialized = self.serializer.serialize(msg)

hippolyzer/lib/proxy/tests/integration/test_addons.py

@@ -2,7 +2,12 @@ from __future__ import annotations
 
 from hippolyzer.lib.base.message.message import Block
 from hippolyzer.lib.proxy import addon_ctx
-from hippolyzer.lib.proxy.addon_utils import BaseAddon, SessionProperty
+from hippolyzer.lib.proxy.addon_utils import (
+    BaseAddon,
+    SessionProperty,
+    send_chat,
+    show_message,
+)
 from hippolyzer.lib.proxy.addons import AddonManager
 from hippolyzer.lib.proxy.commands import handle_command
 from hippolyzer.lib.proxy.message import ProxiedMessage
@@ -20,13 +25,18 @@ class MockAddon(BaseAddon):
     async def foobar(self, _session: Session, _region: ProxiedRegion, bar: str):
         self.bazquux = bar
         self.another = bar
+        send_chat(bar)
+        show_message(bar)
 
 
 class AddonIntegrationTests(BaseIntegrationTest):
     def setUp(self) -> None:
         super().setUp()
         self.addon = MockAddon()
-        AddonManager.init([], self.session_manager, [self.addon])
+        AddonManager.init([], self.session_manager, [self.addon], swallow_addon_exceptions=False)
+
+    def tearDown(self) -> None:
+        AddonManager.shutdown()
 
     def _fake_command(self, command: str) -> None:
         msg = ProxiedMessage(
@@ -63,3 +73,6 @@ class AddonIntegrationTests(BaseIntegrationTest):
                 _something = self.addon.bazquux
             # This has a default
             self.assertEqual(self.addon.another, "default")
+        # Should have sent out the two injected packets for inbound and outbound chat
+        # But not the original chatfromviewer from our command.
+        self.assertEqual(len(self.transport.packets), 2)