keyhog_scanner/decode/pipeline/
registry.rs

1use crate::decode::base64::{Base64Decoder, Z85Decoder};
2use crate::decode::caesar::CaesarDecoder;
3use crate::decode::hex::HexDecoder;
4#[cfg(feature = "decode")]
5use crate::decode::javascript_static::JavaScriptStaticDecoder;
6use crate::decode::json::JsonDecoder;
7use crate::decode::reverse::ReverseDecoder;
8use crate::decode::url::{
9    HtmlNamedEntityDecoder, HtmlNumericEntityDecoder, MimeEncodedWordDecoder, OctalEscapeDecoder,
10    QuotedPrintableDecoder, UnicodeEscapeDecoder, UrlDecoder,
11};
12#[cfg(feature = "decode")]
13use crate::decode::DecodeAdmission;
14#[cfg(any(feature = "decode", test))]
15use crate::decode::DecodeAdmissionSketch;
16use crate::decode::Decoder;
17use parking_lot::RwLock;
18#[cfg(test)]
19use std::cell::RefCell;
20use std::sync::Arc;
21
22// The active decoder set is stored behind one shared `Arc<Vec<..>>`. Scanner
23// construction captures this Arc in its immutable execution plan. Standalone
24// compatibility helpers clone one Arc per call. Registration is copy-on-write,
25// so existing scanners and in-flight compatibility calls keep their snapshot.
26static DECODERS: std::sync::OnceLock<RwLock<DecoderRegistryState>> = std::sync::OnceLock::new();
27
28struct DecoderRegistryState {
29    decoders: Arc<Vec<RegisteredDecoder>>,
30    compatibility_failure: Option<DecoderRegistrationError>,
31}
32
33#[derive(Clone)]
34pub(crate) enum RegisteredDecoder {
35    Shared(Arc<dyn Decoder>),
36    Reverse,
37    Caesar,
38}
39
40impl RegisteredDecoder {
41    pub(crate) fn name(&self) -> &'static str {
42        match self {
43            Self::Shared(decoder) => decoder.name(),
44            Self::Reverse => "reverse",
45            Self::Caesar => "caesar",
46        }
47    }
48
49    fn version(&self) -> &'static str {
50        match self {
51            Self::Shared(decoder) => decoder.version(),
52            Self::Reverse | Self::Caesar => "detector-policy-1",
53        }
54    }
55
56    #[cfg(feature = "decode")]
57    fn admission(
58        &self,
59        chunk: &keyhog_core::Chunk,
60        policy: &super::super::policy::CompiledDecodeTransformPolicy,
61    ) -> DecodeAdmission {
62        match self {
63            Self::Shared(decoder) => decoder.admission(chunk),
64            Self::Reverse => ReverseDecoder
65                .admission_sketch_with_policy(chunk, policy)
66                .admission(),
67            Self::Caesar => CaesarDecoder
68                .admission_sketch_with_policy(chunk, policy)
69                .admission(),
70        }
71    }
72
73    #[cfg(any(feature = "decode", test))]
74    fn admission_sketch(
75        &self,
76        chunk: &keyhog_core::Chunk,
77        policy: &super::super::policy::CompiledDecodeTransformPolicy,
78    ) -> DecodeAdmissionSketch {
79        match self {
80            Self::Shared(decoder) => decoder.admission_sketch(chunk),
81            Self::Reverse => ReverseDecoder.admission_sketch_with_policy(chunk, policy),
82            Self::Caesar => CaesarDecoder.admission_sketch_with_policy(chunk, policy),
83        }
84    }
85
86    pub(super) fn decode_chunk(
87        &self,
88        chunk: &keyhog_core::Chunk,
89        policy: &super::super::policy::CompiledDecodeTransformPolicy,
90    ) -> Vec<keyhog_core::Chunk> {
91        match self {
92            Self::Shared(decoder) => decoder.decode_chunk(chunk),
93            Self::Reverse => ReverseDecoder.decode_chunk_with_policy(chunk, policy),
94            Self::Caesar => CaesarDecoder.decode_chunk_with_policy(chunk, policy),
95        }
96    }
97}
98
99#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
100pub enum DecoderRegistrationError {
101    #[error("decoder name must be non-empty ASCII without whitespace")]
102    InvalidName,
103    #[error("decoder {name:?} version must be non-empty ASCII without whitespace")]
104    InvalidVersion { name: &'static str },
105    #[error("decoder name {0:?} is already registered")]
106    DuplicateName(&'static str),
107}
108
109#[derive(Clone)]
110pub(crate) struct CompiledDecoderPlan {
111    decoders: Arc<Vec<RegisteredDecoder>>,
112    identity: u64,
113}
114
115impl std::fmt::Debug for CompiledDecoderPlan {
116    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
117        formatter
118            .debug_struct("CompiledDecoderPlan")
119            .field("decoder_count", &self.decoders.len())
120            .field("identity", &self.identity)
121            .finish()
122    }
123}
124
125impl CompiledDecoderPlan {
126    pub(crate) fn snapshot() -> Result<Self, DecoderRegistrationError> {
127        let decoders = snapshot_decoders()?;
128        let mut names = std::collections::HashSet::with_capacity(decoders.len());
129        let mut hasher = blake3::Hasher::new();
130        hasher.update(b"keyhog-compiled-decoder-plan-v1\0");
131        for decoder in decoders.iter() {
132            let name = decoder.name();
133            let version = decoder.version();
134            validate_descriptor(name, version)?;
135            if !names.insert(name) {
136                return Err(DecoderRegistrationError::DuplicateName(name));
137            }
138            hash_descriptor(&mut hasher, name, version);
139        }
140        let mut bytes = [0u8; 8];
141        bytes.copy_from_slice(&hasher.finalize().as_bytes()[..8]);
142        Ok(Self {
143            decoders,
144            identity: u64::from_le_bytes(bytes),
145        })
146    }
147
148    pub(crate) fn identity(&self) -> u64 {
149        self.identity
150    }
151
152    #[cfg(feature = "decode")]
153    pub(crate) fn decoders(&self) -> &[RegisteredDecoder] {
154        &self.decoders
155    }
156}
157
158fn validate_descriptor(
159    name: &'static str,
160    version: &'static str,
161) -> Result<(), DecoderRegistrationError> {
162    if name.is_empty() || !name.is_ascii() || name.bytes().any(|byte| byte.is_ascii_whitespace()) {
163        return Err(DecoderRegistrationError::InvalidName);
164    }
165    if version.is_empty()
166        || !version.is_ascii()
167        || version.bytes().any(|byte| byte.is_ascii_whitespace())
168    {
169        return Err(DecoderRegistrationError::InvalidVersion { name });
170    }
171    Ok(())
172}
173
174fn hash_descriptor(hasher: &mut blake3::Hasher, name: &str, version: &str) {
175    for value in [name.as_bytes(), version.as_bytes()] {
176        hasher.update(&(value.len() as u64).to_le_bytes());
177        hasher.update(value);
178    }
179}
180
181#[cfg(test)]
182thread_local! {
183    static THREAD_DECODERS: RefCell<Vec<Arc<dyn Decoder>>> = RefCell::new(Vec::new());
184}
185
186/// Per-decoder wall-time profiler (measurement only). Enabled by the single
187/// scanner profiler switch (`keyhog scan --profile`) so profiling has one runtime
188/// owner instead of one env knob per pass. Records which decoder dominates
189/// decode generation. Zero-cost unset.
190pub(super) fn profile_enabled() -> bool {
191    crate::scan_profile::enabled()
192}
193
194/// Fixed number of per-decoder profiler slots. The `DECODER_NS` / `DECODER_PRODUCED`
195/// accumulators and every index/clamp into them share this one capacity. There
196/// are 14 default decoders today, so the cap carries headroom; a decoder past
197/// slot `MAX_PROFILED_DECODERS` is simply not profiled (`record_decoder_run`
198/// drops it), the `decoder_registry_within_profiler_capacity` gap test guards
199/// the default set against silently outgrowing this.
200const MAX_PROFILED_DECODERS: usize = 16;
201
202static DECODER_NS: [std::sync::atomic::AtomicU64; MAX_PROFILED_DECODERS] = {
203    const Z: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
204    [Z; MAX_PROFILED_DECODERS]
205};
206
207/// Sub-chunks EMITTED per decoder (pre-dedup/screen). The sub-chunk COUNT - not
208/// gen time - is what drives the dominant decode-rescan + per-sub-chunk fixed
209/// phase-1 cost, so this isolates which decoders to gate with a sound prune.
210static DECODER_PRODUCED: [std::sync::atomic::AtomicU64; MAX_PROFILED_DECODERS] = {
211    const Z: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
212    [Z; MAX_PROFILED_DECODERS]
213};
214
215pub(super) fn record_decoder_run(
216    decoder_index: usize,
217    elapsed: std::time::Duration,
218    produced: usize,
219) {
220    if decoder_index >= MAX_PROFILED_DECODERS {
221        return;
222    }
223    use std::sync::atomic::Ordering::Relaxed;
224    DECODER_NS[decoder_index].fetch_add(elapsed.as_nanos() as u64, Relaxed);
225    DECODER_PRODUCED[decoder_index].fetch_add(produced as u64, Relaxed);
226}
227
228/// Print and reset the accumulated per-decoder times (paired with registry
229/// names). Folded into the unified scanner profile dump.
230pub(crate) fn decoder_profile_dump() {
231    use std::sync::atomic::Ordering::Relaxed;
232    let decoders = active_decoders();
233    let names: Vec<&str> = decoders.iter().map(|d| d.name()).collect();
234    let mut rows: Vec<(String, f64)> = (0..names.len().min(MAX_PROFILED_DECODERS))
235        .map(|i| {
236            (
237                names[i].to_string(),
238                DECODER_NS[i].swap(0, Relaxed) as f64 / 1e6,
239            )
240        })
241        .collect();
242    rows.sort_by(|a, b| b.1.total_cmp(&a.1));
243    let total: f64 = rows.iter().map(|r| r.1).sum();
244    let mut prod: Vec<(String, u64)> = (0..names.len().min(MAX_PROFILED_DECODERS))
245        .map(|i| (names[i].to_string(), DECODER_PRODUCED[i].swap(0, Relaxed)))
246        .collect();
247    prod.sort_by(|a, b| b.1.cmp(&a.1));
248    let prod_total: u64 = prod.iter().map(|r| r.1).sum();
249    if total == 0.0 && prod_total == 0 {
250        return;
251    }
252    eprintln!("=== per-decoder decode_chunk time ===");
253    for (name, ms) in &rows {
254        let pct = if total > 0.0 { 100.0 * ms / total } else { 0.0 };
255        eprintln!("  {name:<18}: {ms:>8.1} ms ({pct:>5.1}%)");
256    }
257    eprintln!("  {:<18}: {total:>8.1} ms", "TOTAL");
258    eprintln!("=== per-decoder sub-chunks EMITTED (pre-dedup/screen) ===");
259    for (name, n) in &prod {
260        let pct = if prod_total > 0 {
261            100.0 * *n as f64 / prod_total as f64
262        } else {
263            0.0
264        };
265        eprintln!("  {name:<18}: {n:>8} ({pct:>5.1}%)");
266    }
267    eprintln!("  {:<18}: {prod_total:>8}", "TOTAL");
268}
269
270pub(crate) fn decoder_profile_reset() {
271    use std::sync::atomic::Ordering::Relaxed;
272    for slot in &DECODER_NS {
273        slot.store(0, Relaxed);
274    }
275    for slot in &DECODER_PRODUCED {
276        slot.store(0, Relaxed);
277    }
278}
279
280fn default_decoders() -> Vec<RegisteredDecoder> {
281    vec![
282        RegisteredDecoder::Shared(Arc::new(Base64Decoder)),
283        RegisteredDecoder::Shared(Arc::new(HexDecoder)),
284        RegisteredDecoder::Shared(Arc::new(UrlDecoder)),
285        RegisteredDecoder::Shared(Arc::new(QuotedPrintableDecoder)),
286        RegisteredDecoder::Shared(Arc::new(HtmlNamedEntityDecoder)),
287        RegisteredDecoder::Shared(Arc::new(HtmlNumericEntityDecoder)),
288        RegisteredDecoder::Shared(Arc::new(OctalEscapeDecoder)),
289        RegisteredDecoder::Shared(Arc::new(MimeEncodedWordDecoder)),
290        // JSON unescape - strips `\"` / `\\` / `\n` style escapes inside JSON
291        // string values so credentials stored as JSON-encoded fields survive
292        // into the scanner.
293        RegisteredDecoder::Shared(Arc::new(JsonDecoder)),
294        RegisteredDecoder::Shared(Arc::new(UnicodeEscapeDecoder)),
295        RegisteredDecoder::Shared(Arc::new(Z85Decoder)),
296        // Bounded, side-effect-free JavaScript constant recovery. Keep it after
297        // representation decoders and before the asymmetric evasion decoders.
298        #[cfg(feature = "decode")]
299        RegisteredDecoder::Shared(Arc::new(JavaScriptStaticDecoder)),
300        RegisteredDecoder::Reverse,
301        RegisteredDecoder::Caesar,
302    ]
303}
304
305/// The `name()` of each default decoder, in registration order. This is the
306/// canonical decode-pipeline composition, the order is load-bearing (the
307/// `reverse` and `caesar` decoders deliberately run last, after the structural
308/// decoders), and is pinned by `decoder_registry_default_order` so a reorder
309/// or addition can't silently shift the pipeline.
310#[cfg(feature = "decode")]
311pub(crate) fn default_decoder_names() -> Vec<&'static str> {
312    default_decoders().iter().map(|d| d.name()).collect()
313}
314
315/// Aggregate decoder-owned admission proofs for one root chunk.
316///
317/// Candidate extraction is primed once so built-in predicates that use the
318/// shared extractor do not each allocate and rescan independently. Any custom
319/// decoder that keeps the trait default returns `Unknown`, which is preserved
320/// unless another decoder already proves the chunk is `Possible`.
321#[cfg(feature = "decode")]
322pub(crate) fn decoder_admission(
323    chunk: &keyhog_core::Chunk,
324    policy: &super::super::policy::CompiledDecodeTransformPolicy,
325    plan: &CompiledDecoderPlan,
326) -> DecodeAdmission {
327    super::extractor::clear_shared_candidates();
328    super::extractor::prime_shared_candidates(&chunk.data);
329
330    let mut aggregate = DecodeAdmission::Impossible;
331    for decoder in plan.decoders() {
332        match decoder.admission(chunk, policy) {
333            DecodeAdmission::Possible => {
334                aggregate = DecodeAdmission::Possible;
335                break;
336            }
337            DecodeAdmission::Unknown => aggregate = DecodeAdmission::Unknown,
338            DecodeAdmission::Impossible => {}
339        }
340    }
341
342    super::extractor::clear_shared_candidates();
343    aggregate
344}
345
346#[cfg(any(feature = "decode", test))]
347pub(crate) fn decoder_admission_sketch(
348    chunk: &keyhog_core::Chunk,
349    policy: &super::super::policy::CompiledDecodeTransformPolicy,
350    plan: &CompiledDecoderPlan,
351) -> DecodeAdmissionSketch {
352    decoder_admission_sketch_with_decoders(chunk, policy, plan.decoders())
353}
354
355#[cfg(any(feature = "decode", test))]
356pub(crate) fn active_decoder_admission_sketch(
357    chunk: &keyhog_core::Chunk,
358    policy: &super::super::policy::CompiledDecodeTransformPolicy,
359) -> DecodeAdmissionSketch {
360    let decoders = active_decoders();
361    decoder_admission_sketch_with_decoders(chunk, policy, &decoders)
362}
363
364#[cfg(any(feature = "decode", test))]
365fn decoder_admission_sketch_with_decoders(
366    chunk: &keyhog_core::Chunk,
367    policy: &super::super::policy::CompiledDecodeTransformPolicy,
368    decoders: &[RegisteredDecoder],
369) -> DecodeAdmissionSketch {
370    super::extractor::clear_shared_candidates();
371    super::extractor::prime_shared_candidates(&chunk.data);
372
373    let mut aggregate = DecodeAdmissionSketch::NONE;
374    for decoder in decoders {
375        aggregate.merge(decoder.admission_sketch(chunk, policy));
376    }
377
378    super::extractor::clear_shared_candidates();
379    aggregate
380}
381
382fn decoder_registry() -> &'static RwLock<DecoderRegistryState> {
383    DECODERS.get_or_init(|| {
384        RwLock::new(DecoderRegistryState {
385            decoders: Arc::new(default_decoders()),
386            compatibility_failure: None,
387        })
388    })
389}
390
391#[cfg(not(test))]
392pub(super) fn active_decoders() -> Arc<Vec<RegisteredDecoder>> {
393    // One `Arc` clone (a single atomic increment) instead of deep-cloning the
394    // decoder Vec on every `decode_chunk`. Callers only iterate, so the shared
395    // snapshot suffices.
396    Arc::clone(&decoder_registry().read().decoders)
397}
398
399#[cfg(test)]
400pub(super) fn active_decoders() -> Arc<Vec<RegisteredDecoder>> {
401    let base = Arc::clone(&decoder_registry().read().decoders);
402    THREAD_DECODERS.with(|thread_decoders| {
403        let thread = thread_decoders.borrow();
404        if thread.is_empty() {
405            // Common case: no per-test decoder registered, hand back the shared
406            // snapshot with no allocation, matching the non-test fast path.
407            base
408        } else {
409            let mut combined = (*base).clone();
410            combined.extend(thread.iter().cloned().map(RegisteredDecoder::Shared));
411            Arc::new(combined)
412        }
413    })
414}
415
416#[cfg(not(test))]
417fn snapshot_decoders() -> Result<Arc<Vec<RegisteredDecoder>>, DecoderRegistrationError> {
418    let registry = decoder_registry().read();
419    if let Some(error) = registry.compatibility_failure.clone() {
420        Err(error)
421    } else {
422        Ok(Arc::clone(&registry.decoders))
423    }
424}
425
426#[cfg(test)]
427fn snapshot_decoders() -> Result<Arc<Vec<RegisteredDecoder>>, DecoderRegistrationError> {
428    let base = {
429        let registry = decoder_registry().read();
430        if let Some(error) = registry.compatibility_failure.clone() {
431            return Err(error);
432        }
433        Arc::clone(&registry.decoders)
434    };
435    THREAD_DECODERS.with(|thread_decoders| {
436        let thread = thread_decoders.borrow();
437        if thread.is_empty() {
438            Ok(base)
439        } else {
440            let mut combined = (*base).clone();
441            combined.extend(thread.iter().cloned().map(RegisteredDecoder::Shared));
442            Ok(Arc::new(combined))
443        }
444    })
445}
446
447/// Register a custom decoder for scanners compiled afterward.
448///
449/// Use [`try_register_decoder`] when the caller can handle a registration
450/// error. This compatibility entry point records any error, and later scanner
451/// compilation returns it. Existing compiled scanners retain their immutable
452/// decoder plan.
453pub fn register_decoder(decoder: Box<dyn Decoder>) {
454    if let Err(error) = register_decoder_inner(decoder, true) {
455        tracing::error!(%error, "decoder registration failed; later scanner compilation will fail");
456    }
457}
458
459/// Register a custom decoder and return descriptor or collision errors.
460///
461/// The name and version must be non-empty ASCII without whitespace. A name can
462/// be registered only once. Existing compiled scanners retain their immutable
463/// decoder plan.
464pub fn try_register_decoder(decoder: Box<dyn Decoder>) -> Result<(), DecoderRegistrationError> {
465    register_decoder_inner(decoder, false)
466}
467
468fn register_decoder_inner(
469    decoder: Box<dyn Decoder>,
470    record_failure: bool,
471) -> Result<(), DecoderRegistrationError> {
472    let decoder_name = decoder.name();
473    let mut guard = decoder_registry().write();
474    if let Some(error) = guard.compatibility_failure.clone() {
475        return Err(error);
476    }
477    let result = validate_descriptor(decoder_name, decoder.version()).and_then(|()| {
478        if guard
479            .decoders
480            .iter()
481            .any(|existing| existing.name() == decoder_name)
482        {
483            Err(DecoderRegistrationError::DuplicateName(decoder_name))
484        } else {
485            Ok(())
486        }
487    });
488    if let Err(error) = result {
489        if record_failure && guard.compatibility_failure.is_none() {
490            guard.compatibility_failure = Some(error.clone());
491        }
492        return Err(error);
493    }
494    // Copy-on-write: publish a fresh snapshot so any `active_decoders()` Arc
495    // already handed out keeps its consistent view. Registration happens at
496    // startup / test setup, never on the decode hot path, so this one-time Vec
497    // clone is not a concern.
498    let mut next = (*guard.decoders).clone();
499    next.push(RegisteredDecoder::Shared(Arc::from(decoder)));
500    guard.decoders = Arc::new(next);
501    Ok(())
502}
503
504#[cfg(test)]
505pub(crate) struct ScopedDecoderRegistration {
506    name: &'static str,
507    active: bool,
508}
509
510#[cfg(test)]
511impl Drop for ScopedDecoderRegistration {
512    fn drop(&mut self) {
513        if !self.active {
514            return;
515        }
516        THREAD_DECODERS.with(|thread_decoders| {
517            let mut decoders = thread_decoders.borrow_mut();
518            if let Some(index) = decoders
519                .iter()
520                .rposition(|decoder| decoder.name() == self.name)
521            {
522                decoders.remove(index);
523            }
524        });
525    }
526}
527
528#[cfg(test)]
529pub(crate) fn register_thread_decoder(decoder: Box<dyn Decoder>) -> ScopedDecoderRegistration {
530    let decoder_name = decoder.name();
531    let has_global_duplicate = decoder_registry()
532        .read()
533        .decoders
534        .iter()
535        .any(|existing| existing.name() == decoder_name);
536    if has_global_duplicate {
537        tracing::warn!(
538            decoder = decoder_name,
539            "register_thread_decoder called with a duplicate global decoder name; decoder ignored"
540        );
541        return ScopedDecoderRegistration {
542            name: decoder_name,
543            active: false,
544        };
545    }
546
547    let mut inserted = false;
548    let decoder = Arc::from(decoder);
549    THREAD_DECODERS.with(|thread_decoders| {
550        let mut decoders = thread_decoders.borrow_mut();
551        if decoders.iter().any(|existing| existing.name() == decoder_name) {
552            tracing::warn!(
553                decoder = decoder_name,
554                "register_thread_decoder called with a duplicate thread decoder name; decoder ignored"
555            );
556            return;
557        }
558        decoders.push(decoder);
559        inserted = true;
560    });
561    ScopedDecoderRegistration {
562        name: decoder_name,
563        active: inserted,
564    }
565}
566
567#[cfg(test)]
568#[path = "../../../tests/unit/decode_admission_soundness.rs"]
569mod admission_soundness_tests;
keyhog_scanner/decode/pipeline/registry.rs

keyhog_scanner/decode/pipeline/
registry.rs
