use crate::compiler::{
build_compile_state, build_gpu_literals, build_phase2_keyword_ac, phase2_always_active_indices,
validate_compiled_pattern_detector_indices,
};
use crate::engine::{phase2_anchor, phase2_generic, scan_postprocess};
use crate::error::{Result, ScanError};
use crate::gpu_matcher_cache as gpu_cache;
use crate::scanner_config::ScannerTuningConfig;
use keyhog_core::DetectorSpec;
use std::sync::Arc;
use vyre_libs::scan::GpuLiteralSet;
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct GpuLiteralArtifact {
pub cache_key: String,
pub pattern_count: usize,
pub bytes: Vec<u8>,
pub wire_magic: [u8; 4],
pub wire_version: u32,
}
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct GpuLiteralArtifacts {
pub literal: Option<GpuLiteralArtifact>,
pub positioned_literal: Option<GpuLiteralArtifact>,
}
pub fn gpu_literal_artifact_cache_dir() -> Result<std::path::PathBuf> {
gpu_cache::gpu_matcher_cache_dir().map_err(|error| ScanError::Gpu(error.to_string()))
}
pub fn compile_gpu_literal_artifacts_default(
detectors: &[DetectorSpec],
) -> Result<GpuLiteralArtifacts> {
compile_gpu_literal_artifact_plan(detectors)
}
pub fn compile_gpu_literal_artifacts(
detectors: &[DetectorSpec],
_tuning_config: &ScannerTuningConfig,
) -> Result<GpuLiteralArtifacts> {
compile_gpu_literal_artifact_plan(detectors)
}
fn compile_gpu_literal_artifact_plan(detectors: &[DetectorSpec]) -> Result<GpuLiteralArtifacts> {
let state = build_compile_state(detectors)?;
validate_compiled_pattern_detector_indices(
&state.ac_map,
&state.phase2_patterns,
detectors.len(),
)?;
let (_, _, phase2_keywords) = build_phase2_keyword_ac(&state.phase2_patterns);
let phase2_always_active_indices = phase2_always_active_indices(&state.phase2_patterns);
let phase2_anchor_index = phase2_anchor::Phase2AnchorIndex::build(
&state.phase2_patterns,
&phase2_always_active_indices,
);
let phase2_always_anchor_literals = phase2_anchor_index
.as_ref()
.map_or(&[] as &[String], |index| index.always_anchor_literals());
let confirmed_anchor_index =
scan_postprocess::confirmed_anchor::ConfirmedAnchorIndex::build(&state.ac_map);
let confirmed_anchor_literals = confirmed_anchor_index
.as_ref()
.map_or(&[] as &[String], |index| index.anchor_literals());
let generic_keyword_literals = if detectors.iter().any(DetectorSpec::owns_entropy_policy) {
phase2_generic::keywords::GenericAssignmentKeywordPlan::compile(detectors)
.map_err(crate::error::ScanError::Config)?
.stem_literals()
.map(str::to_owned)
.collect::<Vec<_>>()
} else {
Vec::new()
};
Ok(GpuLiteralArtifacts {
literal: serialize_literal_rows(
"lit-ci",
build_gpu_literals(
&state.ac_literals,
&phase2_keywords,
phase2_always_anchor_literals,
confirmed_anchor_literals,
&generic_keyword_literals,
),
)?,
positioned_literal: None,
})
}
fn serialize_literal_rows(
cache_prefix: &'static str,
rows: Option<Arc<Vec<Vec<u8>>>>,
) -> Result<Option<GpuLiteralArtifact>> {
let Some(rows) = rows else {
return Ok(None);
};
let literal_refs: Vec<&[u8]> = rows.iter().map(Vec::as_slice).collect();
let cache_key = gpu_cache::gpu_matcher_cache_key_with_prefix(cache_prefix, &literal_refs);
let pattern_count = literal_refs.len();
let matcher = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
GpuLiteralSet::compile_case_insensitive(&literal_refs)
}))
.map_err(|panic| {
let detail = crate::error::panic_payload_detail(panic);
ScanError::Gpu(format!(
"GPU literal artifact compile panicked for cache prefix {cache_prefix} with {pattern_count} patterns: {detail}. Fix: reduce literal rows, increase VYRE's DFA budget, or shard the literal set."
))
})?;
let bytes = matcher.to_bytes().map_err(|error| {
ScanError::Gpu(format!(
"failed to serialize GPU literal artifact for cache prefix {cache_prefix} with {pattern_count} patterns: {error}. Fix: upgrade VYRE or rebuild the artifact with a compatible KeyHog binary."
))
})?;
let (wire_magic, wire_version) = literal_set_wire_header(&bytes).ok_or_else(|| {
ScanError::Gpu(format!(
"GPU literal artifact for cache prefix {cache_prefix} serialized to {} bytes, too short for VYRE's 8-byte wire envelope header. Fix: upgrade VYRE or rebuild the artifact with a compatible KeyHog binary.",
bytes.len()
))
})?;
Ok(Some(GpuLiteralArtifact {
cache_key,
pattern_count,
bytes,
wire_magic,
wire_version,
}))
}
fn literal_set_wire_header(bytes: &[u8]) -> Option<([u8; 4], u32)> {
let header = bytes.get(..8)?;
let magic = [header[0], header[1], header[2], header[3]];
let version = u32::from_le_bytes([header[4], header[5], header[6], header[7]]);
Some((magic, version))
}