1mod base64;
8pub(crate) mod caesar;
9pub(crate) mod hex;
10pub(crate) mod inflate;
11#[cfg(feature = "decode")]
12mod javascript_static;
13mod json;
14mod limits;
15mod pipeline;
16pub(crate) mod policy;
17pub(crate) mod reverse;
18mod unicode_escape;
19mod url;
20pub(crate) mod util;
21
22pub use base64::{base64_decode, find_base64_strings, z85_decode};
23pub use base64::is_base64_candidate_byte;
29pub(crate) use base64::{
30 contains_non_padding_equals, is_standard_base64_byte, standard_base64_shape,
31};
32pub use hex::{find_hex_strings, hex_decode};
33pub(crate) use pipeline::{
34 bytecount_newlines, decoder_profile_dump, decoder_profile_reset, extract_profile_dump,
35 extract_profile_reset, splice_decoded_payload_at, with_extracted_value_spans,
36};
37pub(crate) use pipeline::{canonical_decode_order_probe_for_test, CompiledDecoderPlan};
38#[cfg(feature = "decode")]
39pub(crate) use pipeline::{decoder_admission, default_decoder_names};
40pub use pipeline::{register_decoder, try_register_decoder, DecoderRegistrationError};
41#[cfg(test)]
42pub(crate) use pipeline::{register_thread_decoder, ScopedDecoderRegistration};
43pub(crate) use util::take_hex_digits;
44
45use keyhog_core::Chunk;
46
47#[cfg(feature = "decode")]
48pub(crate) fn decode_chunk_with_policy(
49 chunk: &Chunk,
50 policy: &policy::CompiledDecodeTransformPolicy,
51 decoder_plan: &CompiledDecoderPlan,
52 max_depth: usize,
53 validate: bool,
54 deadline: Option<std::time::Instant>,
55 screen: Option<&crate::alphabet_filter::AlphabetScreen>,
56) -> Vec<Chunk> {
57 pipeline::decode_chunk_with_policy(
58 chunk,
59 policy,
60 decoder_plan,
61 max_depth,
62 validate,
63 deadline,
64 screen,
65 )
66}
67
68pub(crate) fn decode_chunk(
72 chunk: &Chunk,
73 max_depth: usize,
74 validate: bool,
75 deadline: Option<std::time::Instant>,
76 screen: Option<&crate::alphabet_filter::AlphabetScreen>,
77) -> Vec<Chunk> {
78 pipeline::decode_chunk_with_active_decoders(
79 chunk,
80 policy::bundled_compat_policy(),
81 max_depth,
82 validate,
83 deadline,
84 screen,
85 )
86}
87
88pub(crate) fn unicode_escape_decode(input: &str) -> Result<String, ()> {
89 unicode_escape::unicode_escape_decode(input)
90}
91
92#[cfg(feature = "decode")]
93pub(crate) fn quoted_printable_decode(input: &str) -> Result<String, ()> {
94 url::quoted_printable_decode(input)
95}
96
97#[cfg(feature = "decode")]
98pub(crate) fn mime_encoded_word_decode(input: &str) -> Result<String, ()> {
99 url::mime_encoded_word_decode(input)
100}
101
102#[cfg(feature = "decode")]
103pub(crate) fn octal_escape_decode(input: &str) -> Result<String, ()> {
104 url::octal_escape_decode(input)
105}
106
107pub(crate) fn extracted_value_strings_for_test(text: &str) -> Vec<String> {
108 pipeline::with_extracted_value_spans(text, |values| {
109 values.iter().map(|value| value.value.clone()).collect()
110 })
111}
112
113#[cfg(feature = "decode")]
114fn valid_html_numeric_entity_len(data: &[u8]) -> Option<usize> {
115 if !data.starts_with(b"&#") {
116 return None;
117 }
118
119 let mut index = 2usize;
120 let radix = if matches!(data.get(index), Some(b'x' | b'X')) {
121 index += 1;
122 16u32
123 } else {
124 10u32
125 };
126 let digits_start = index;
127 let mut codepoint = 0u32;
128 while index < data.len() && index - digits_start < url::MAX_NUMERIC_ENTITY_DIGITS {
129 let digit = match data[index] {
130 b'0'..=b'9' => u32::from(data[index] - b'0'),
131 b'a'..=b'f' if radix == 16 => u32::from(data[index] - b'a') + 10,
132 b'A'..=b'F' if radix == 16 => u32::from(data[index] - b'A') + 10,
133 _ => break,
134 };
135 codepoint = codepoint.checked_mul(radix)?.checked_add(digit)?;
136 index += 1;
137 }
138
139 if index == digits_start || data.get(index) != Some(&b';') {
140 return None;
141 }
142 char::from_u32(codepoint)?;
143 Some(index + 1)
144}
145
146#[cfg(feature = "decode")]
157pub(crate) fn has_decodable_payload(data: &[u8]) -> bool {
158 let mut run = 0usize;
165 let mut percent_escapes = 0usize;
166 let mut backslash_escapes = 0usize;
167 let mut html_numeric_entities = 0usize;
168 let mut has_from_char_code = false;
169 let mut has_xor_operator = false;
170 let mut i = 0usize;
171
172 while i < data.len() {
173 let b = data[i];
174
175 if b == b'^' {
176 has_xor_operator = true;
177 if has_from_char_code {
178 return true;
179 }
180 } else if b == b'f' && data[i..].starts_with(b"fromCharCode") {
181 has_from_char_code = true;
182 if has_xor_operator {
183 return true;
184 }
185 }
186
187 if b == b'%'
188 && i + 2 < data.len()
189 && data[i + 1].is_ascii_hexdigit()
190 && data[i + 2].is_ascii_hexdigit()
191 {
192 percent_escapes += 1;
193 if percent_escapes >= limits::MIN_PERCENT_ESCAPES {
194 return true;
195 }
196 run = 0;
197 i += 3;
198 continue;
199 }
200
201 if b == b'&' {
202 if let Some(entity_len) = valid_html_numeric_entity_len(&data[i..]) {
203 html_numeric_entities += 1;
204 if html_numeric_entities >= limits::MIN_HTML_NUMERIC_ENTITIES {
205 return true;
206 }
207 run = 0;
208 i += entity_len;
209 continue;
210 }
211 }
212
213 if b == b'\\' && i + 1 < data.len() {
214 match data[i + 1] {
215 b'u' if i + 5 < data.len()
216 && data[i + 2..i + 6]
217 .iter()
218 .all(|digit| digit.is_ascii_hexdigit()) =>
219 {
220 backslash_escapes += 1;
221 if backslash_escapes >= limits::MIN_BACKSLASH_ESCAPES {
222 return true;
223 }
224 run = 0;
225 i += 6;
226 continue;
227 }
228 b'x' if i + 3 < data.len()
229 && data[i + 2..i + 4]
230 .iter()
231 .all(|digit| digit.is_ascii_hexdigit()) =>
232 {
233 backslash_escapes += 1;
234 if backslash_escapes >= limits::MIN_BACKSLASH_ESCAPES {
235 return true;
236 }
237 run = 0;
238 i += 4;
239 continue;
240 }
241 b'0'..=b'7'
252 if i + 3 < data.len()
253 && (b'0'..=b'7').contains(&data[i + 2])
254 && (b'0'..=b'7').contains(&data[i + 3]) =>
255 {
256 backslash_escapes += 1;
257 if backslash_escapes >= limits::MIN_BACKSLASH_ESCAPES {
258 return true;
259 }
260 run = 0;
261 i += 4;
262 continue;
263 }
264 _ => {}
265 }
266 }
267
268 if is_base64_candidate_byte(b) {
271 run += 1;
272 if run >= limits::MIN_DECODABLE_RUN {
273 return true;
274 }
275 } else {
276 run = 0;
277 }
278 i += 1;
279 }
280 false
281}
282
283pub trait DecodeOutputSink {
289 fn push(&mut self, chunk: Chunk) -> bool;
290}
291
292impl DecodeOutputSink for Vec<Chunk> {
293 fn push(&mut self, chunk: Chunk) -> bool {
294 Vec::push(self, chunk);
295 true
296 }
297}
298
299#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
301#[error(
302 "decoder output exceeded the direct collection budget after {produced} chunks/{bytes} bytes (maximum {max_chunks} chunks/{max_bytes} bytes)"
303)]
304pub struct DecodeCollectionError {
305 pub produced: usize,
306 pub bytes: usize,
307 pub max_chunks: usize,
308 pub max_bytes: usize,
309}
310
311struct BoundedCollectSink {
312 chunks: Vec<Chunk>,
313 bytes: usize,
314 exhausted: bool,
315}
316
317impl DecodeOutputSink for BoundedCollectSink {
318 fn push(&mut self, chunk: Chunk) -> bool {
319 let Some(next_bytes) = self.bytes.checked_add(chunk.data.len()) else {
320 self.exhausted = true;
321 return false;
322 };
323 if self.chunks.len() == limits::MAX_DECODED_CHUNKS_PER_ROOT
324 || next_bytes > limits::MAX_DECODED_TOTAL_BYTES
325 {
326 self.exhausted = true;
327 return false;
328 }
329 self.bytes = next_bytes;
330 self.chunks.push(chunk);
331 true
332 }
333}
334
335pub trait Decoder: Send + Sync {
337 fn name(&self) -> &'static str;
338
339 fn version(&self) -> &'static str {
343 "1"
344 }
345
346 fn admission_sketch(&self, _chunk: &Chunk) -> DecodeAdmissionSketch {
351 DecodeAdmissionSketch::UNKNOWN
352 }
353
354 fn admission(&self, _chunk: &Chunk) -> DecodeAdmission {
361 self.admission_sketch(_chunk).admission()
362 }
363
364 fn decode_chunk_into(&self, chunk: &Chunk, sink: &mut dyn DecodeOutputSink);
370
371 fn decode_chunk(&self, chunk: &Chunk) -> Result<Vec<Chunk>, DecodeCollectionError> {
376 let mut sink = BoundedCollectSink {
377 chunks: Vec::new(),
378 bytes: 0,
379 exhausted: false,
380 };
381 self.decode_chunk_into(chunk, &mut sink);
382 if sink.exhausted {
383 return Err(DecodeCollectionError {
384 produced: sink.chunks.len(),
385 bytes: sink.bytes,
386 max_chunks: limits::MAX_DECODED_CHUNKS_PER_ROOT,
387 max_bytes: limits::MAX_DECODED_TOTAL_BYTES,
388 });
389 }
390 Ok(sink.chunks)
391 }
392}
393
394#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
400pub struct DecodeAdmissionSketch {
401 kind_mask: u32,
402 candidate_count: u16,
403 candidate_bytes: u32,
404 unknown: bool,
405}
406
407impl DecodeAdmissionSketch {
408 pub const BASE64: u32 = 1 << 0;
409 pub const HEX: u32 = 1 << 1;
410 pub const URL: u32 = 1 << 2;
411 pub const QUOTED_PRINTABLE: u32 = 1 << 3;
412 pub const HTML_NAMED_ENTITY: u32 = 1 << 4;
413 pub const HTML_NUMERIC_ENTITY: u32 = 1 << 5;
414 pub const OCTAL_ESCAPE: u32 = 1 << 6;
415 pub const MIME_ENCODED_WORD: u32 = 1 << 7;
416 pub const JSON: u32 = 1 << 8;
417 pub const UNICODE_ESCAPE: u32 = 1 << 9;
418 pub const Z85: u32 = 1 << 10;
419 pub const JAVASCRIPT_STATIC: u32 = 1 << 11;
420 pub const REVERSE: u32 = 1 << 12;
421 pub const CAESAR: u32 = 1 << 13;
422 pub const COMPRESSED_CONTAINER: u32 = 1 << 14;
424
425 pub const NONE: Self = Self {
426 kind_mask: 0,
427 candidate_count: 0,
428 candidate_bytes: 0,
429 unknown: false,
430 };
431
432 pub const UNKNOWN: Self = Self {
433 kind_mask: 0,
434 candidate_count: u16::MAX,
435 candidate_bytes: u32::MAX,
436 unknown: true,
437 };
438
439 pub const fn kind_mask(self) -> u32 {
440 self.kind_mask
441 }
442
443 pub const fn candidate_count(self) -> u16 {
444 self.candidate_count
445 }
446
447 pub const fn candidate_bytes(self) -> u32 {
448 self.candidate_bytes
449 }
450
451 pub const fn has_unknown(self) -> bool {
452 self.unknown
453 }
454
455 pub fn merge(&mut self, other: Self) {
456 self.kind_mask |= other.kind_mask;
457 self.candidate_count = self.candidate_count.saturating_add(other.candidate_count);
458 self.candidate_bytes = self.candidate_bytes.saturating_add(other.candidate_bytes);
459 self.unknown |= other.unknown;
460 if self.unknown {
461 self.candidate_count = u16::MAX;
462 self.candidate_bytes = u32::MAX;
463 }
464 }
465
466 pub(crate) fn possible(kind: u32, candidate_count: usize, candidate_bytes: usize) -> Self {
467 Self {
468 kind_mask: kind,
469 candidate_count: candidate_count.min(u16::MAX as usize) as u16,
470 candidate_bytes: candidate_bytes.min(u32::MAX as usize) as u32,
471 unknown: false,
472 }
473 }
474
475 pub(crate) const fn admission(self) -> DecodeAdmission {
476 if self.unknown {
477 DecodeAdmission::Unknown
478 } else if self.kind_mask == 0 {
479 DecodeAdmission::Impossible
480 } else {
481 DecodeAdmission::Possible
482 }
483 }
484}
485
486#[derive(Clone, Debug)]
491pub struct DecodeWorkloadPlan {
492 enabled: bool,
493 max_input_bytes: usize,
494 transforms: DecodeTransformPolicyHandle,
495 decoders: DecoderPlanHandle,
496}
497
498#[derive(Clone, Debug)]
499enum DecodeTransformPolicyHandle {
500 Bundled,
501 Compiled(std::sync::Arc<policy::CompiledDecodeTransformPolicy>),
502}
503
504#[derive(Clone, Debug)]
505enum DecoderPlanHandle {
506 Active,
507 Compiled(std::sync::Arc<CompiledDecoderPlan>),
508}
509
510impl DecodeTransformPolicyHandle {
511 fn policy(&self) -> &policy::CompiledDecodeTransformPolicy {
512 match self {
513 Self::Bundled => policy::bundled_compat_policy(),
514 Self::Compiled(policy) => policy,
515 }
516 }
517}
518
519impl PartialEq for DecodeWorkloadPlan {
520 fn eq(&self, other: &Self) -> bool {
521 self.enabled == other.enabled
522 && self.max_input_bytes == other.max_input_bytes
523 && self.transforms.policy().identity() == other.transforms.policy().identity()
524 && match (&self.decoders, &other.decoders) {
525 (DecoderPlanHandle::Active, DecoderPlanHandle::Active) => true,
526 (DecoderPlanHandle::Compiled(left), DecoderPlanHandle::Compiled(right)) => {
527 left.identity() == right.identity()
528 }
529 _ => false,
530 }
531 }
532}
533
534impl Eq for DecodeWorkloadPlan {}
535
536impl DecodeWorkloadPlan {
537 pub const fn from_limits(max_depth: usize, max_input_bytes: usize) -> Self {
543 Self {
544 enabled: cfg!(feature = "decode") && max_depth > 0,
545 max_input_bytes,
546 transforms: DecodeTransformPolicyHandle::Bundled,
547 decoders: DecoderPlanHandle::Active,
548 }
549 }
550
551 pub(crate) fn from_compiled_limits(
552 max_depth: usize,
553 max_input_bytes: usize,
554 transforms: std::sync::Arc<policy::CompiledDecodeTransformPolicy>,
555 decoders: std::sync::Arc<CompiledDecoderPlan>,
556 ) -> Self {
557 Self {
558 enabled: cfg!(feature = "decode") && max_depth > 0,
559 max_input_bytes,
560 transforms: DecodeTransformPolicyHandle::Compiled(transforms),
561 decoders: DecoderPlanHandle::Compiled(decoders),
562 }
563 }
564
565 pub const fn enabled(&self) -> bool {
566 self.enabled
567 }
568
569 pub const fn max_input_bytes(&self) -> usize {
570 self.max_input_bytes
571 }
572
573 pub fn admits(&self, chunk: &Chunk) -> bool {
574 self.enabled && chunk.data.len() <= self.max_input_bytes
575 }
576
577 pub fn sketch(&self, chunk: &Chunk) -> DecodeAdmissionSketch {
580 #[cfg(not(feature = "decode"))]
581 {
582 let _ = chunk;
584 return DecodeAdmissionSketch::NONE;
585 }
586 #[cfg(feature = "decode")]
587 if !self.admits(chunk) {
588 DecodeAdmissionSketch::NONE
589 } else {
590 match &self.decoders {
591 DecoderPlanHandle::Active => {
592 pipeline::active_decoder_admission_sketch(chunk, self.transforms.policy())
593 }
594 DecoderPlanHandle::Compiled(plan) => {
595 pipeline::decoder_admission_sketch(chunk, self.transforms.policy(), plan)
596 }
597 }
598 }
599 }
600}
601
602#[cfg(feature = "decode")]
606pub fn decode_admission_sketch(chunk: &Chunk) -> DecodeAdmissionSketch {
607 pipeline::active_decoder_admission_sketch(chunk, policy::bundled_compat_policy())
608}
609
610#[cfg(not(feature = "decode"))]
612pub fn decode_admission_sketch(_chunk: &Chunk) -> DecodeAdmissionSketch {
613 DecodeAdmissionSketch::NONE
614}
615
616#[derive(Clone, Copy, Debug, Eq, PartialEq)]
618#[non_exhaustive]
619pub enum DecodeAdmission {
620 Unknown,
622 Possible,
624 Impossible,
626}
627
628pub struct EncodedString {
630 pub value: String,
631}