From 193d76213270eb4ac34267f098757c07af4d2179 Mon Sep 17 00:00:00 2001 From: Salad Dais Date: Tue, 18 Oct 2022 22:40:15 +0000 Subject: [PATCH] Give each addon a separate addon_ctx bucket This fixes addons being able to accidentally stomp all over each others' state just because they happened to use the same name for a SessionProperty. --- addon_examples/counter.py | 2 +- addon_examples/spongecase.py | 2 +- hippolyzer/lib/proxy/addon_utils.py | 16 +++++++++++----- hippolyzer/lib/proxy/sessions.py | 6 ++++-- tests/proxy/integration/test_addons.py | 14 +++++++------- 5 files changed, 24 insertions(+), 16 deletions(-) diff --git a/addon_examples/counter.py b/addon_examples/counter.py index 8a92053..4772920 100644 --- a/addon_examples/counter.py +++ b/addon_examples/counter.py @@ -6,7 +6,7 @@ from hippolyzer.lib.proxy.sessions import Session def handle_lludp_message(session: Session, region: ProxiedRegion, message: Message): # addon_ctx will persist across addon reloads, use for storing data that # needs to survive across calls to this function - ctx = session.addon_ctx + ctx = session.addon_ctx[__name__] if message.name == "ChatFromViewer": chat = message["ChatData"]["Message"] if chat == "COUNT": diff --git a/addon_examples/spongecase.py b/addon_examples/spongecase.py index cd40c8e..892cd44 100644 --- a/addon_examples/spongecase.py +++ b/addon_examples/spongecase.py @@ -13,7 +13,7 @@ def _to_spongecase(val): def handle_lludp_message(session: Session, _region: ProxiedRegion, message: Message): - ctx = session.addon_ctx + ctx = session.addon_ctx[__name__] ctx.setdefault("spongecase", False) if message.name == "ChatFromViewer": chat = message["ChatData"]["Message"] diff --git a/hippolyzer/lib/proxy/addon_utils.py b/hippolyzer/lib/proxy/addon_utils.py index 97dc39c..f4b5a20 100644 --- a/hippolyzer/lib/proxy/addon_utils.py +++ b/hippolyzer/lib/proxy/addon_utils.py @@ -211,13 +211,17 @@ class BaseAddonProperty(abc.ABC, Generic[_T, _U]): session_manager.addon_ctx dict, without any namespacing. Can be accessed either through `AddonClass.property_name` or `addon_instance.property_name`. """ - __slots__ = ("name", "default") + __slots__ = ("name", "default", "_owner") def __init__(self, default=dataclasses.MISSING): self.default = default + self._owner = None def __set_name__(self, owner, name: str): self.name = name + # Keep track of which addon "owns" this property so that we can shove + # the data in a bucket specific to that addon name. + self._owner = owner def _make_default(self) -> _T: if self.default is not dataclasses.MISSING: @@ -235,18 +239,20 @@ class BaseAddonProperty(abc.ABC, Generic[_T, _U]): if ctx_obj is None: raise AttributeError( f"{self.__class__} {self.name} accessed outside proper context") + addon_state = ctx_obj.addon_ctx[self._owner.__name__] # Set a default if we have one, otherwise let the keyerror happen. # Maybe we should do this at addon initialization instead of on get. - if self.name not in ctx_obj.addon_ctx: + if self.name not in addon_state: default = self._make_default() if default is not dataclasses.MISSING: - ctx_obj.addon_ctx[self.name] = default + addon_state[self.name] = default else: raise AttributeError(f"{self.name} is not set") - return ctx_obj.addon_ctx[self.name] + return addon_state[self.name] def __set__(self, _obj, value: _T) -> None: - self._get_context_obj().addon_ctx[self.name] = value + addon_state = self._get_context_obj().addon_ctx[self._owner.__name__] + addon_state[self.name] = value class SessionProperty(BaseAddonProperty[_T, "Session"]): diff --git a/hippolyzer/lib/proxy/sessions.py b/hippolyzer/lib/proxy/sessions.py index 918ac1b..663a42b 100644 --- a/hippolyzer/lib/proxy/sessions.py +++ b/hippolyzer/lib/proxy/sessions.py @@ -1,5 +1,6 @@ from __future__ import annotations +import collections import dataclasses import datetime import functools @@ -43,7 +44,8 @@ class Session(BaseClientSession): self.circuit_code = circuit_code self.global_caps = {} # Bag of arbitrary data addons can use to persist data across addon reloads - self.addon_ctx = {} + # Each addon name gets its own separate dict within this dict + self.addon_ctx: Dict[str, Dict[str, Any]] = collections.defaultdict(dict) self.session_manager: SessionManager = session_manager or None self.selected: SelectionModel = SelectionModel() self.regions: List[ProxiedRegion] = [] @@ -188,7 +190,7 @@ class SessionManager: self.flow_context = HTTPFlowContext() self.asset_repo = HTTPAssetRepo() self.message_logger: Optional[BaseMessageLogger] = None - self.addon_ctx: Dict[str, Any] = {} + self.addon_ctx: Dict[str, Dict[str, Any]] = collections.defaultdict(dict) self.name_cache = ProxyNameCache() self.pending_leap_clients: List[LEAPClient] = [] diff --git a/tests/proxy/integration/test_addons.py b/tests/proxy/integration/test_addons.py index 7af1f1a..fbea5ad 100644 --- a/tests/proxy/integration/test_addons.py +++ b/tests/proxy/integration/test_addons.py @@ -88,12 +88,12 @@ class AddonIntegrationTests(BaseProxyTest): self._setup_default_circuit() self._fake_command("foobar baz") await self._wait_drained() - self.assertEqual(self.session.addon_ctx["bazquux"], "baz") + self.assertEqual(self.session.addon_ctx["MockAddon"]["bazquux"], "baz") # In session context these should be equivalent with addon_ctx.push(new_session=self.session): - self.assertEqual(self.session.addon_ctx["bazquux"], self.addon.bazquux) - self.assertEqual(self.session.addon_ctx["another"], "baz") + self.assertEqual(self.session.addon_ctx["MockAddon"]["bazquux"], self.addon.bazquux) + self.assertEqual(self.session.addon_ctx["MockAddon"]["another"], "baz") # Outside session context it should raise with self.assertRaises(AttributeError): @@ -104,7 +104,7 @@ class AddonIntegrationTests(BaseProxyTest): self.session.addon_ctx.clear() with addon_ctx.push(new_session=self.session): - # This has no default so should fail + # This has no default so it should fail with self.assertRaises(AttributeError): _something = self.addon.bazquux # This has a default @@ -144,9 +144,9 @@ class AddonIntegrationTests(BaseProxyTest): AddonManager.load_addon_from_path(str(self.parent_path), reload=True) # Wait for the init hooks to run await asyncio.sleep(0.001) - self.assertFalse("quux" in self.session_manager.addon_ctx) + self.assertFalse("quux" in self.session_manager.addon_ctx["ParentAddon"]) parent_addon_mod = AddonManager.FRESH_ADDON_MODULES['hippolyzer.user_addon_parent_addon'] self.assertEqual(0, parent_addon_mod.ParentAddon.quux) - self.assertEqual(0, self.session_manager.addon_ctx["quux"]) + self.assertEqual(0, self.session_manager.addon_ctx["ParentAddon"]["quux"]) parent_addon_mod.ParentAddon.quux = 1 - self.assertEqual(1, self.session_manager.addon_ctx["quux"]) + self.assertEqual(1, self.session_manager.addon_ctx["ParentAddon"]["quux"])