vyre 0.4.0

GPU compute intermediate representation with a standard operation library
Documentation
//! Catalog entry for `detect_jwt`.

use crate::ops::security_detection::detector_support::{spans, ByteSpan, DetectionError};

/// Embedded operation spec formerly stored in metadata/spec.toml.
pub const SPEC_TOML: &str = r#"schema_version = 1
id = "security_detection.detect_jwt"
archetype = "match-bytes-pattern"
display_name = "Detect JWT"
summary = "Returns offset-length spans for three-part base64url JWT-like tokens."
category = "C"

[intrinsic]
wgsl = "security_detection_detect_jwt"

[signature]
inputs = ["Bytes", "Bytes"]
output = "Bytes"

laws = []
equivalence_classes = ["three_segments", "short_segment", "bad_alphabet", "t47_cap"]
workgroup_size = [64, 1, 1]
tags = ["security-detection", "jwt", "secret-scan", "t47"]
fixtures_dir = "fixtures/"
"#;

/// Embedded reference vectors formerly stored in fixtures/reference-vectors.toml.
pub const REFERENCE_VECTORS_TOML: &str = r#"[[case]]
name = "positive_jwt"
input = "bearer eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0In0.c2lnbmF0dXJl"
expected_spans = [{ offset = 7, len = 60 }]

[[case]]
name = "negative_two_parts"
input = "bearer eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0In0"
expected_spans = []
"#;

/// WGSL lowering source for this detector.
pub mod lowering {
    /// Return the detector-specific WGSL source.
    #[must_use]
    pub const fn source() -> &'static str {
        r#"struct Params {
    input_len: u32,
    max_spans: u32,
    _pad0: u32,
    _pad1: u32,
}

struct SpanOutput {
    count: atomic<u32>,
    data: array<u32>,
}

@group(0) @binding(0) var<storage, read> input: array<u32>;
@group(0) @binding(1) var<storage, read_write> output: SpanOutput;
@group(0) @binding(2) var<uniform> params: Params;

fn is_alpha(byte: u32) -> bool {
    return (byte >= 65u && byte <= 90u) || (byte >= 97u && byte <= 122u);
}

fn is_digit(byte: u32) -> bool {
    return byte >= 48u && byte <= 57u;
}

fn is_boundary(byte: u32) -> bool {
    return !(is_alpha(byte) || is_digit(byte) || byte == 95u || byte == 45u);
}

fn is_base64url(byte: u32) -> bool {
    return is_alpha(byte) || is_digit(byte) || byte == 45u || byte == 95u;
}

fn parse_segment(start: u32) -> u32 {
    var pos = start;
    loop {
        if (pos >= params.input_len || !is_base64url(input[pos])) {
            break;
        }
        pos = pos + 1u;
    }
    if (pos - start < 8u) {
        return 0xffffffffu;
    }
    return pos;
}

fn emit_span(offset: u32, len: u32) {
    let slot = atomicAdd(&output.count, 1u);
    if (slot < params.max_spans) {
        output.data[slot * 2u] = offset;
        output.data[slot * 2u + 1u] = len;
    }
}

@compute @workgroup_size(64)
fn security_detection_detect_jwt(@builtin(global_invocation_id) gid: vec3<u32>) {
    let start = gid.x;
    if (start >= params.input_len) {
        return;
    }
    let before_ok = start == 0u || is_boundary(input[start - 1u]);
    if (!before_ok) {
        return;
    }
    var pos = parse_segment(start);
    if (pos == 0xffffffffu || pos >= params.input_len || input[pos] != 46u) {
        return;
    }
    pos = parse_segment(pos + 1u);
    if (pos == 0xffffffffu || pos >= params.input_len || input[pos] != 46u) {
        return;
    }
    pos = parse_segment(pos + 1u);
    if (pos == 0xffffffffu) {
        return;
    }
    let after_ok = pos >= params.input_len || is_boundary(input[pos]);
    if (after_ok) {
        emit_span(start, pos - start);
    }
}
"#
    }
}

/// Return JWT-like base64url.base64url.base64url spans.
///
/// # Errors
///
/// Returns `Fix: ...` when input exceeds 64 MiB.
pub fn detect_jwt(input: &[u8]) -> Result<Vec<ByteSpan>, DetectionError> {
    spans::jwt_spans(input)
}

/// Compatibility surface for the previous generated implementation module.
pub mod implementation {
    pub use super::detect_jwt;
    /// Compatibility module for callers that used the generated kernel path.
    pub mod kernel {
        pub use super::super::detect_jwt;
    }

    /// Compatibility module for callers that used the generated lowering path.
    pub mod lowering {
        /// Compatibility module for callers that used `implementation::lowering::wgsl`.
        pub mod wgsl {
            pub use super::super::super::lowering::source;
        }
    }
}