diff --git a/secator/runners/_base.py b/secator/runners/_base.py index 38454a0df..f5e99dff7 100644 --- a/secator/runners/_base.py +++ b/secator/runners/_base.py @@ -178,6 +178,7 @@ def __init__(self, config, inputs=[], results=[], run_opts={}, hooks={}, validat self.resolved_hooks = {name: [] for name in HOOKS + getattr(self, 'hooks', [])} self.debug('registering hooks', obj=list(self.resolved_hooks.keys()), sub='init') self.register_hooks(hooks) + self._apply_context_drivers() # Validators self.resolved_validators = {name: [] for name in VALIDATORS + getattr(self, 'validators', [])} @@ -464,46 +465,50 @@ def run(self): """ return list(self.__iter__()) - def __getstate__(self): - """Custom pickle: strip hook functions so dynamically-loaded modules - (e.g. secator.hooks.cockpit) don't cause ModuleNotFoundError on workers. - Driver names in context['drivers'] are used to re-load hooks in __setstate__. + def _apply_context_drivers(self): + """Register hooks for drivers named in ``context['drivers']``. + + ``context['drivers']`` is the cross-process source of truth for driver + hooks. Loading them at construction (rather than on unpickle) means every + runner gets its driver hooks exactly once, wherever it is built: the + initial dispatch, a chunk task rebuilt by ``run_command`` via + ``task_cls(targets, **opts)``, and a chord callback runner. Because + ``register_hooks()`` is idempotent, drivers also supplied explicitly via + ``self._hooks`` (CLI-resolved hooks or library callers passing + ``hooks=HOOKS``) are not registered twice. + + This replaces the former ``__getstate__``/``__setstate__`` pair: with the + driver modules discovered at worker startup (``celery.py`` ``IN_WORKER``), + hook functions now pickle/unpickle natively by qualified name, so runners + no longer need to strip hooks on pickle and rebuild them on every unpickle + (which, under ``replace``/chord synchronization, re-registered hooks + O(chunks) times and flooded ``SECATOR_DEBUG=runner`` logs). """ - state = self.__dict__.copy() - state['_hooks'] = {} - state['resolved_hooks'] = {name: [] for name in state['resolved_hooks']} - return state - - def __setstate__(self, state): - """Custom unpickle: restore runner state then re-register hooks.""" - self.__dict__.update(state) drivers = self.context.get('drivers', []) - if drivers: - from secator.loader import discover_external_drivers, order_drivers - - discover_external_drivers() - # Order by canonical priority so authoritative backends (e.g. mongodb) - # register their hooks before relay drivers (e.g. api). Hook lists are - # concatenated in driver order, so this decides hook execution order. - drivers = order_drivers(drivers) + if not drivers: + return + from secator.loader import discover_external_drivers, order_drivers + + discover_external_drivers() + # Order by canonical priority so authoritative backends (e.g. mongodb) + # register their hooks before relay drivers (e.g. api). Hook lists are + # concatenated in driver order, so this decides hook execution order. + drivers = order_drivers(drivers) hooks_list = [] for driver in drivers: driver_hooks = import_dynamic(f'secator.hooks.{driver}', 'HOOKS') if driver_hooks: hooks_list.append(driver_hooks) - merged_hooks = {} - if hooks_list: - from secator.utils import deep_merge_dicts + if not hooks_list: + return + from secator.utils import deep_merge_dicts - merged_hooks = deep_merge_dicts(*hooks_list) + merged_hooks = deep_merge_dicts(*hooks_list) # Driver HOOKS dicts are keyed by base runner class (Scan/Workflow/Task). A task # runner's class is its command subclass (e.g. ``whois``), never the base ``Task``, # so register_hooks()' exact ``hooks.get(self.__class__)`` lookup would miss the - # ``Task`` entry and the task's on_end hook would never re-register on unpickle. - # That left chunk-parent tasks — the only tasks pickled into a chord callback — - # stuck in RUNNING because mark_runner_completed() ran zero on_end hooks. Flatten - # to the base runner type's hooks first (same convention as Workflow handing - # ``self._hooks.get(Task)`` to its task signatures) so they restore in the worker. + # ``Task`` entry. Flatten to the base runner type's hooks first (same convention as + # Workflow handing ``self._hooks.get(Task)`` to its task signatures). from secator.runners import Scan, Task, Workflow base_cls = {'scan': Scan, 'workflow': Workflow, 'task': Task}.get(self.config.type) @@ -1046,24 +1051,33 @@ def run_validators(self, validator_type, *args, error=True, sub='validators'): def register_hooks(self, hooks): """Register hooks. + Idempotent: a hook already present in ``resolved_hooks[key]`` is skipped + (and not re-logged). This lets the same driver be supplied via both + ``self._hooks`` (e.g. library callers passing ``hooks=HOOKS``) and + ``context['drivers']`` without registering — or logging — it twice. + Args: hooks (dict[str, List[Callable]]): List of hooks to register. """ for key in self.resolved_hooks: + registered = self.resolved_hooks[key] + # Register class + derived class hooks class_hook = getattr(self, key, None) - if class_hook: + if class_hook and class_hook not in registered: fun = self.get_func_path(class_hook) self.debug('hook registered', obj={'name': key, 'fun': fun}, sub='init') - self.resolved_hooks[key].append(class_hook) + registered.append(class_hook) - # Register user hooks - user_hooks = hooks.get(self.__class__, {}).get(key, []) + # Register user hooks (copy so we never mutate a caller's/shared list) + user_hooks = list(hooks.get(self.__class__, {}).get(key, [])) user_hooks.extend(hooks.get(key, [])) for hook in user_hooks: + if hook in registered: + continue fun = self.get_func_path(hook) self.debug('hook registered', obj={'name': key, 'fun': fun}, sub='init') - self.resolved_hooks[key].extend(user_hooks) + registered.append(hook) def register_validators(self, validators): """Register validators. diff --git a/tests/unit/test_celery.py b/tests/unit/test_celery.py index 2382a0278..1051d43c3 100644 --- a/tests/unit/test_celery.py +++ b/tests/unit/test_celery.py @@ -344,14 +344,18 @@ def test_run_scan_celery_task(self): class TestRunnerPickle(unittest.TestCase): """Test that Runner objects with dynamic driver hooks can be pickled/unpickled.""" - def test_runner_pickle_survives_missing_module_on_worker(self): - """Reproduces the original bug: unpickling a runner whose hooks reference a - dynamically-loaded module that does not exist on the worker side. - Without the __getstate__/__setstate__ fix this raises ModuleNotFoundError.""" + def test_runner_pickle_survives_with_discovered_module(self): + """Runners pickle/unpickle natively (no __getstate__/__setstate__). + + Driver hook functions reference dynamically-loaded modules. The worker + discovers external drivers at startup (celery.py IN_WORKER), so those + modules are in sys.modules and hook functions resolve on unpickle. The + hook survives the round-trip and unpickling does NOT re-register hooks + (which, under replace()/chord synchronization, previously flooded logs).""" import pickle import types import sys - from secator.runners import Workflow + from secator.runners import Runner, Workflow from secator.loader import get_configs_by_type workflows = get_configs_by_type('workflow') @@ -360,7 +364,7 @@ def test_runner_pickle_survives_missing_module_on_worker(self): config = workflows[0] - # 1) CLI side: load the driver into sys.modules (as the loader does) + # Driver module present in sys.modules, as on a worker post-discovery fake_module = types.ModuleType('secator.hooks.testdriver') def on_start(runner, *args): @@ -371,23 +375,31 @@ def on_start(runner, *args): fake_module.on_start = on_start sys.modules['secator.hooks.testdriver'] = fake_module - hooks = {Workflow: {'on_start': [on_start]}} - runner = Workflow(config, inputs=['example.com'], run_opts={'dry_run': True}, hooks=hooks, context={}) - - # Pickle while the module is available (simulates the CLI/sender side) - pickled = pickle.dumps(runner) - - # 2) Worker side: remove the module to simulate it not being installed - del sys.modules['secator.hooks.testdriver'] - - # Without the fix, unpickling would raise: - # ModuleNotFoundError: No module named 'secator.hooks.testdriver' - # With the fix, __getstate__ strips hooks so the bytes contain no reference - # to the dynamic module and unpickling succeeds. - restored = pickle.loads(pickled) - self.assertEqual(restored.name, runner.name) - # Hooks were stripped; dynamic hook not re-registered (driver not in context['drivers']) - self.assertNotIn(on_start, restored.resolved_hooks.get('on_start', [])) + try: + hooks = {Workflow: {'on_start': [on_start]}} + runner = Workflow(config, inputs=['example.com'], run_opts={'dry_run': True}, hooks=hooks, context={}) + self.assertIn(on_start, runner.resolved_hooks.get('on_start', [])) + + # Unpickling must NOT call register_hooks (native pickling restores state) + calls = {'n': 0} + orig = Runner.register_hooks + + def counting(self, h): + calls['n'] += 1 + return orig(self, h) + + Runner.register_hooks = counting + try: + restored = pickle.loads(pickle.dumps(runner)) + finally: + Runner.register_hooks = orig + + self.assertEqual(calls['n'], 0, 'unpickle must not re-register hooks') + self.assertEqual(restored.name, runner.name) + # The dynamically-referenced hook survives natively + self.assertIn(on_start, restored.resolved_hooks.get('on_start', [])) + finally: + del sys.modules['secator.hooks.testdriver'] def test_runner_pickle_restores_hooks_from_context_drivers(self): """Unpickling a Runner re-registers hooks from context['drivers'].""" @@ -426,7 +438,7 @@ def on_end(runner, *args): pickled = pickle.dumps(runner) restored = pickle.loads(pickled) self.assertEqual(restored.name, runner.name) - # __setstate__ re-registers on_end from fakedriver via context['drivers'] + # on_end is loaded at init from context['drivers'] and survives pickling natively self.assertIn(on_end, restored.resolved_hooks.get('on_end', [])) finally: diff --git a/tests/unit/test_runner_hooks.py b/tests/unit/test_runner_hooks.py new file mode 100644 index 000000000..33a07ccf3 --- /dev/null +++ b/tests/unit/test_runner_hooks.py @@ -0,0 +1,117 @@ +import pickle +import sys +import tempfile +import unittest +from pathlib import Path + +from secator.config import CONFIG +from secator.loader import discover_external_drivers +from secator.runners._base import Runner + +# A minimal external driver, as a user would drop into ~/.secator/templates/. +CUSTOM_DRIVER = ''' +from secator.runners import Task + + +def cd_on_item(self, item): + return item + + +def cd_on_end(self, *args, **kwargs): + return None + + +HOOKS = {Task: {'on_item': [cd_on_item], 'on_end': [cd_on_end]}} +''' + + +def lib_on_item(self, item): + """Module-level hook (picklable by qualified name), for the library-mode test.""" + return item + + +class TestRunnerHooks(unittest.TestCase): + """Driver hooks are registered once at construction and survive pickling + natively (no __getstate__/__setstate__), so unpickling never re-registers. + """ + + @classmethod + def setUpClass(cls): + cls._tmpdir = tempfile.TemporaryDirectory() + tmp = Path(cls._tmpdir.name) + (tmp / 'mydriver.py').write_text(CUSTOM_DRIVER) + cls._orig_templates = CONFIG.dirs.templates + CONFIG.dirs.templates = tmp + # discover_external_drivers is @cache'd; clear it so it re-scans our tmp dir + # (an earlier test/import may have populated the cache with the real dir) + discover_external_drivers.cache_clear() + discover_external_drivers() + + @classmethod + def tearDownClass(cls): + CONFIG.dirs.templates = cls._orig_templates + cls._tmpdir.cleanup() + sys.modules.pop('secator.hooks.mydriver', None) + # reset the cache so later tests re-discover against the restored dir + discover_external_drivers.cache_clear() + + def _build_task(self, **kwargs): + # Bind the class straight from sys.modules so its identity matches what + # pickle resolves (other suite tests may reload the task module, which + # would otherwise make pickle reject the "not the same object" class). + import importlib + mod = importlib.import_module('secator.tasks.httpx') + return mod.httpx(['http://localhost'], enable_hooks=False, dry_run=True, **kwargs) + + @staticmethod + def _hook_names(task, event): + return [h.__name__ for h in task.resolved_hooks[event]] + + def _count_register_hooks(self, fn): + calls = {'n': 0} + orig = Runner.register_hooks + + def wrapped(self, hooks): + calls['n'] += 1 + return orig(self, hooks) + + Runner.register_hooks = wrapped + try: + fn() + finally: + Runner.register_hooks = orig + return calls['n'] + + def test_driver_hooks_loaded_at_init(self): + # context['drivers'] hooks must be present right after construction (no pickle) + task = self._build_task(context={'drivers': ['mydriver']}) + self.assertIn('cd_on_item', self._hook_names(task, 'on_item')) + self.assertIn('cd_on_end', self._hook_names(task, 'on_end')) + + def test_custom_driver_hook_survives_pickle_without_reregistration(self): + task = self._build_task(context={'drivers': ['mydriver']}) + # unpickling must NOT call register_hooks (this is what flooded the logs) + reg = self._count_register_hooks(lambda: pickle.loads(pickle.dumps(task))) + self.assertEqual(reg, 0, 'unpickle must not re-register hooks') + # and the hook must still be there + callable + back = pickle.loads(pickle.dumps(task)) + self.assertIn('cd_on_item', self._hook_names(back, 'on_item')) + self.assertTrue(callable(back.resolved_hooks['on_item'][0])) + + def test_registration_is_idempotent(self): + task = self._build_task(context={'drivers': ['mydriver']}) + before = len(task.resolved_hooks['on_item']) + # re-apply the same context drivers -> no duplicate registration + task._apply_context_drivers() + self.assertEqual(len(task.resolved_hooks['on_item']), before) + + def test_library_mode_explicit_hooks_survive_pickle(self): + # Library callers pass hooks explicitly and may set NO context['drivers']. + task = self._build_task(hooks={'on_item': [lib_on_item]}) + self.assertIn('lib_on_item', self._hook_names(task, 'on_item')) + back = pickle.loads(pickle.dumps(task)) + self.assertIn('lib_on_item', self._hook_names(back, 'on_item')) + + +if __name__ == '__main__': + unittest.main()