Merge pull request #1 from daedalus/copilot/bayes-theorem-advantage

Seed Thompson-sampling bandit with informative priors from static target profiling

Author: daedalus

README.md

@@ -20,6 +20,7 @@ Coverage-guided binary fuzzer with static target analysis, statistical novelty s
 - **TargetProfiler**: ELF static analysis at startup — extracts string constants, function boundaries, magic bytes, and input format hints
 - **Auto-populated dictionary**: interesting strings (format specifiers, error messages, keywords) and magic bytes extracted from `.rodata`
 - **Format-aware seed generation**: produces structurally meaningful initial seeds (PNG headers, text protocols, JSON, XML, HTML) based on inferred format
+- **Informative Bayesian priors**: `format_operator_priors()` seeds the Thompson-sampling bandit's Beta prior toward structure-aware operators (e.g. `png_chunk_mutate`) and dictionary operators when static analysis detects a matching format or extractable tokens, instead of always starting from the uninformative Beta(1, 1)
 - **Hot-function weighting**: seeds exercising high-branch-density functions get a proportional boost in selection
 - **Crash ETA estimation**: blends static risky density (keyword heuristic) with dynamic I(byte_position; crash) mutual information from actual executions, plus Good-Turing edge estimates and calibrated discovery rate — the MI signal strengthens as fuzzing accumulates near-miss data
 

src/fuzzer_tool/core/montecarlo.py

@@ -22,6 +22,12 @@
 
 log = logging.getLogger(__name__)
 
+# Lower clamp for Beta(alpha, beta) parameters passed to init_arm(): keeps
+# random.betavariate() numerically stable for malformed/degenerate priors
+# (e.g. a caller passing 0 or a negative value) without silently overriding
+# intentionally weak-but-valid priors above this threshold.
+MIN_BETA_PARAM = 1e-6
+
 
 class MonteCarloScheduler:
     """Thompson sampling bandit for mutation ops + CEM byte distribution.
@@ -43,6 +49,11 @@ class MonteCarloScheduler:
     """
 
     ELITE_MAX = 200
+    # Declares that init_arm() accepts an informative (prior_alpha, prior_beta)
+    # override, unlike MOptScheduler/ReplicatorScheduler/EloTracker which use
+    # non-Bayesian internal representations (particle positions, population
+    # simplex, Elo ratings).
+    supports_priors = True
 
     def __init__(
         self,
@@ -79,15 +90,25 @@ def __init__(
         # Blend factor: 0.0 = pure Thompson, 1.0 = pure pairwise
         self.pairwise_blend = pairwise_blend
 
-    def init_arm(self, name: str) -> None:
-        """Register a mutation operator arm with prior (1, 1).
+    def init_arm(self, name: str, prior_alpha: float = 1.0, prior_beta: float = 1.0) -> None:
+        """Register a mutation operator arm with a Beta prior.
+
+        Defaults to the uninformative Beta(1, 1) prior. Callers with prior
+        knowledge about an operator's likely usefulness (e.g. static target
+        profiling indicating a specific file format) can pass a stronger
+        prior to bias early Thompson sampling before any evidence has been
+        observed. A no-op if the arm is already registered — the prior only
+        applies at first registration and is never overwritten by later
+        calls, matching the existing idempotent behavior of this method.
 
         Args:
             name: Name of the mutation operator.
+            prior_alpha: Prior alpha (successes + 1). Must be > 0.
+            prior_beta: Prior beta (failures + 1). Must be > 0.
         """
         if name not in self.arm_alpha:
-            self.arm_alpha[name] = 1.0
-            self.arm_beta[name] = 1.0
+            self.arm_alpha[name] = max(prior_alpha, MIN_BETA_PARAM)
+            self.arm_beta[name] = max(prior_beta, MIN_BETA_PARAM)
 
     def select_op(self, ops: list[str], prev_op: str | None = None) -> str:
         """Select mutation operator via Thompson sampling with pairwise transitions.

src/fuzzer_tool/core/target_profiler.py

@@ -612,3 +612,68 @@ def _compute_call_depths(self, profile: TargetProfile):
                 seen.add(p)
                 unique.append(p)
         profile.entry_points = unique
+
+
+# Format signature -> structure-aware mutation operators that are almost
+# certainly useful for that format (see core.mutations.FORMAT_MUTATIONS).
+_FORMAT_OPERATOR_HINTS: dict[str, tuple[str, ...]] = {
+    "png": ("png_chunk_mutate", "png_crc_fix"),
+    "jpeg": ("jpeg_chunk_mutate", "jpeg_crc_fix"),
+    "gzip": ("gzip_chunk_mutate",),
+    "bz2": ("gzip_chunk_mutate",),
+    "xz": ("gzip_chunk_mutate",),
+    "zlib": ("zlib_chunk_mutate",),
+    "riff": ("bmp_chunk_mutate",),
+}
+
+# Dictionary/token-aware mutation operators (see core.mutations.DICT_MUTATIONS).
+_DICT_OPERATORS: tuple[str, ...] = (
+    "dict_insert",
+    "dict_replace",
+    "dict_overwrite",
+    "dict_prepend",
+    "dict_append",
+    "checksum_repair",
+    "token_dup",
+)
+
+# Beta prior for operators the profile suggests are relevant: same total
+# "pseudo-observation" mass as the uninformative default (1, 1), but shifted
+# toward success so Thompson sampling favors them before real evidence
+# arrives. Weak enough that a handful of real failures will correct it.
+_BOOSTED_PRIOR: tuple[float, float] = (2.0, 1.0)
+
+
+def format_operator_priors(profile: "TargetProfile") -> dict[str, tuple[float, float]]:
+    """Derive informative Beta priors for mutation operators from a profile.
+
+    Static analysis (magic bytes, boundary markers, extracted interesting
+    strings, format signature) is prior knowledge about which
+    structure-aware mutation operators are likely to be useful *before*
+    any executions have happened. This lets the Thompson-sampling bandit
+    (:class:`fuzzer_tool.core.montecarlo.MonteCarloScheduler`) start with a
+    Beta prior biased toward those operators instead of the uninformative
+    Beta(1, 1) used for every arm by default.
+
+    Args:
+        profile: A populated :class:`TargetProfile`.
+
+    Returns:
+        Mapping of operator name to a (prior_alpha, prior_beta) override.
+        Operators not present in the mapping should use the default
+        Beta(1, 1) prior.
+    """
+    priors: dict[str, tuple[float, float]] = {}
+
+    fmt = profile.format_signature
+    if fmt:
+        for op in _FORMAT_OPERATOR_HINTS.get(fmt, ()):
+            priors[op] = _BOOSTED_PRIOR
+
+    # Boost dict-aware operators whenever static analysis extracted at least
+    # one usable token (non-empty list is truthy; empty list is falsy).
+    if profile.magic_bytes or profile.boundary_markers or profile.interesting_strings:
+        for op in _DICT_OPERATORS:
+            priors[op] = _BOOSTED_PRIOR
+
+    return priors

src/fuzzer_tool/services/fuzzer.py

@@ -544,22 +544,43 @@ def __init__(
 
             self._tracer = CrashTracer(target)
 
-        def _register_arms(scheduler):
-            """Register all mutation arms on a scheduler (mc, mopt, replicator, elo)."""
+        def _register_arms(scheduler, priors=None):
+            """Register all mutation arms on a scheduler (mc, mopt, replicator, elo).
+
+            Args:
+                scheduler: Scheduler exposing init_arm(name).
+                priors: Optional dict of operator name -> (prior_alpha,
+                    prior_beta) overrides. Only meaningful for the
+                    Beta-Bernoulli Thompson-sampling scheduler; ignored for
+                    schedulers whose init_arm() doesn't accept a prior.
+            """
+            priors = priors or {}
+
+            def _init(op):
+                prior = priors.get(op) if getattr(scheduler, "supports_priors", False) else None
+                if prior is not None and len(prior) == 2:
+                    scheduler.init_arm(op, *prior)
+                else:
+                    scheduler.init_arm(op)
+
             for op in MUTATIONS:
-                scheduler.init_arm(op)
+                _init(op)
             for op in DICT_MUTATIONS:
-                scheduler.init_arm(op)
-            scheduler.init_arm("markov_bytes")
-            scheduler.init_arm("cem_bytes")
+                _init(op)
+            _init("markov_bytes")
+            _init("cem_bytes")
             if self.grammar:
-                scheduler.init_arm("grammar_mutate")
-                scheduler.init_arm("grammar_tree_mutate")
+                _init("grammar_mutate")
+                _init("grammar_tree_mutate")
             for op in FORMAT_MUTATIONS:
-                scheduler.init_arm(op)
+                _init(op)
+
+        from fuzzer_tool.core.target_profiler import format_operator_priors
+
+        _format_priors = format_operator_priors(self._profile)
 
         if self.mc and self.mc_bandit:
-            _register_arms(self.mc)
+            _register_arms(self.mc, _format_priors)
         if self._mopt:
             _register_arms(self._mopt)
         if self._replicator:

tests/test_montecarlo.py

@@ -25,6 +25,24 @@ def test_init_arm_idempotent(self):
         mc.init_arm("bit_flip")
         assert mc.arm_alpha["bit_flip"] == 1.0
 
+    def test_init_arm_informative_prior(self):
+        mc = MonteCarloScheduler()
+        mc.init_arm("png_chunk_mutate", prior_alpha=2.0, prior_beta=1.0)
+        assert mc.arm_alpha["png_chunk_mutate"] == 2.0
+        assert mc.arm_beta["png_chunk_mutate"] == 1.0
+
+    def test_init_arm_prior_not_overwritten_on_reregister(self):
+        mc = MonteCarloScheduler()
+        mc.init_arm("bit_flip", prior_alpha=5.0, prior_beta=1.0)
+        mc.init_arm("bit_flip")  # default prior, arm already registered
+        assert mc.arm_alpha["bit_flip"] == 5.0
+
+    def test_init_arm_prior_clamped_positive(self):
+        mc = MonteCarloScheduler()
+        mc.init_arm("op", prior_alpha=0.0, prior_beta=-1.0)
+        assert mc.arm_alpha["op"] > 0
+        assert mc.arm_beta["op"] > 0
+
     def test_select_op(self):
         mc = MonteCarloScheduler()
         mc.init_arm("bit_flip")

tests/test_target_profiler.py

@@ -3,6 +3,7 @@
 import os
 import struct
 import tempfile
+import math
 
 import pytest
 
@@ -11,6 +12,8 @@
     TargetProfiler,
     TargetProfile,
     FunctionInfo,
+    _FORMAT_OPERATOR_HINTS,
+    format_operator_priors,
 )
 
 
@@ -186,3 +189,48 @@ def test_boundary_markers_are_bytes(self):
         p.boundary_markers = [b":", b"/", b"\n"]
         for bm in p.boundary_markers:
             assert isinstance(bm, bytes)
+
+
+class TestFormatOperatorPriors:
+    def test_no_hints_when_profile_empty(self):
+        p = TargetProfile()
+        assert format_operator_priors(p) == {}
+
+    def test_png_format_boosts_png_operators(self):
+        p = TargetProfile()
+        p.format_signature = "png"
+        priors = format_operator_priors(p)
+        assert priors["png_chunk_mutate"] == (2.0, 1.0)
+        assert priors["png_crc_fix"] == (2.0, 1.0)
+        assert "jpeg_chunk_mutate" not in priors
+
+    def test_unknown_format_no_format_hints(self):
+        p = TargetProfile()
+        p.format_signature = "text"
+        priors = format_operator_priors(p)
+        all_format_ops = {op for ops in _FORMAT_OPERATOR_HINTS.values() for op in ops}
+        assert not (set(priors) & all_format_ops)
+
+    def test_magic_bytes_boost_dict_operators(self):
+        p = TargetProfile()
+        p.magic_bytes = [b"\x89PNG\r\n\x1a\n"]
+        priors = format_operator_priors(p)
+        assert priors["dict_insert"] == (2.0, 1.0)
+        assert priors["checksum_repair"] == (2.0, 1.0)
+
+    def test_boundary_markers_boost_dict_operators(self):
+        p = TargetProfile()
+        p.boundary_markers = [b":"]
+        priors = format_operator_priors(p)
+        assert "dict_replace" in priors
+
+    def test_priors_are_positive_finite_floats(self):
+        p = TargetProfile()
+        p.format_signature = "gzip"
+        p.magic_bytes = [b"\x1f\x8b"]
+        priors = format_operator_priors(p)
+        for alpha, beta in priors.values():
+            assert alpha > 0
+            assert beta > 0
+            assert math.isfinite(alpha)
+            assert math.isfinite(beta)