systemprompt-ai 0.43.0

//! Phrase-list heuristic safety scanner.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use async_trait::async_trait;
use systemprompt_models::wire::canonical::{CanonicalRequest, CanonicalResponse, Role};

use super::{
    Finding, PHASE_REQUEST, PHASE_REQUEST_HISTORY, PHASE_RESPONSE, SafetyScanner, Severity,
};
use crate::services::gateway::spec::HeuristicConfig;

const JAILBREAK_PHRASES: &[&str] = &[
    "ignore previous instructions",
    "ignore all previous",
    "disregard prior instructions",
    "forget your instructions",
    "act as dan",
    "developer mode enabled",
    "pretend you have no restrictions",
];

const EXCERPT_CAP: usize = 240;

#[derive(Debug, Clone)]
pub struct HeuristicScanner {
    phrases: Vec<String>,
}

impl Default for HeuristicScanner {
    fn default() -> Self {
        Self::new(&HeuristicConfig::default())
    }
}

impl HeuristicScanner {
    #[must_use]
    pub fn new(config: &HeuristicConfig) -> Self {
        Self {
            phrases: effective_phrases(config),
        }
    }
}

pub fn effective_phrases(config: &HeuristicConfig) -> Vec<String> {
    let base: Vec<String> = match (&config.phrases, config.disable_builtin) {
        (Some(list), _) => list.clone(),
        (None, true) => Vec::new(),
        (None, false) => JAILBREAK_PHRASES.iter().map(|p| (*p).to_owned()).collect(),
    };
    base.into_iter()
        .chain(config.extra_phrases.iter().cloned())
        .map(|p| p.to_ascii_lowercase())
        .filter(|p| !p.trim().is_empty())
        .collect()
}

#[async_trait]
impl SafetyScanner for HeuristicScanner {
    fn name(&self) -> &'static str {
        "heuristic"
    }

    async fn scan_request(&self, req: &CanonicalRequest) -> Vec<Finding> {
        let mut findings = Vec::new();
        if let Some(sys) = &req.system {
            scan_text(&self.phrases, PHASE_REQUEST, sys, &mut findings);
        }
        if let Some(text) = req.latest_message_text(Role::User) {
            scan_text(&self.phrases, PHASE_REQUEST, &text, &mut findings);
        }
        // Why: each leaf is its own unit — concatenating them would let two
        // unrelated strings splice into a match neither one contains.
        for leaf in req.forwarded_surface.leaves() {
            scan_text(&self.phrases, PHASE_REQUEST, &leaf.value, &mut findings);
        }
        findings
    }

    async fn scan_request_history(&self, req: &CanonicalRequest) -> Vec<Finding> {
        let mut findings = Vec::new();
        for unit in history_units(req) {
            scan_text(&self.phrases, PHASE_REQUEST_HISTORY, &unit, &mut findings);
        }
        findings
    }

    async fn scan_response_final(&self, response: &CanonicalResponse) -> Vec<Finding> {
        let mut findings = Vec::new();
        for unit in response.content_units() {
            scan_text(&self.phrases, PHASE_RESPONSE, &unit, &mut findings);
        }
        findings
    }
}

fn history_units(req: &CanonicalRequest) -> Vec<String> {
    let mut units = req.message_units();
    if req.system.is_some() && !units.is_empty() {
        units.remove(0);
    }
    if let Some(newest) = req.latest_message_text(Role::User)
        && units.last() == Some(&newest)
    {
        units.pop();
    }
    units
}

fn scan_text(phrases: &[String], phase: &'static str, text: &str, out: &mut Vec<Finding>) {
    // Why: the lowercased copy exists only for the phrase search — both
    // detectors below read `text` directly. Allocating one per leaf of a
    // multi-megabyte forwarded body was the scanner's dominant cost.
    if !phrases.is_empty() {
        let lower = text.to_ascii_lowercase();
        for phrase in phrases {
            if let Some(idx) = lower.find(phrase.as_str()) {
                let start = floor_boundary(text, idx.saturating_sub(40));
                let end = ceil_boundary(text, idx + phrase.len() + 80);
                let excerpt = text[start..end]
                    .chars()
                    .take(EXCERPT_CAP)
                    .collect::<String>();
                out.push(Finding {
                    phase,
                    severity: Severity::Medium,
                    category: "jailbreak".to_owned(),
                    excerpt: Some(excerpt),
                    scanner: "heuristic",
                });
            }
        }
    }

    if detect_email(text) {
        out.push(Finding {
            phase,
            severity: Severity::Low,
            category: "pii_email".to_owned(),
            excerpt: None,
            scanner: "heuristic",
        });
    }
    // Why: a card needs digits, and most leaves of a forwarded body — keys,
    // prose, code — have none. Cheaper than entering the scan to find out.
    if !text.bytes().any(|b| b.is_ascii_digit()) {
        return;
    }
    if detect_credit_card(text) {
        out.push(Finding {
            phase,
            severity: Severity::High,
            category: "pii_credit_card".to_owned(),
            excerpt: None,
            scanner: "heuristic",
        });
    }
}

// Why: phrase offsets come from an ASCII-lowercased copy, which is byte-aligned
// with `text`, but the ±40/80 excerpt padding is not — landing mid-codepoint
// would panic the scanner, and a panic here is a 500 on a customer's request.
const fn floor_boundary(text: &str, mut i: usize) -> usize {
    while i > 0 && !text.is_char_boundary(i) {
        i -= 1;
    }
    i
}

const fn ceil_boundary(text: &str, mut i: usize) -> usize {
    if i >= text.len() {
        return text.len();
    }
    while i < text.len() && !text.is_char_boundary(i) {
        i += 1;
    }
    i
}

fn detect_email(text: &str) -> bool {
    let bytes = text.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        if bytes[i] == b'@' {
            let before = bytes[..i]
                .iter()
                .rev()
                .take_while(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'+' | b'-'))
                .count();
            let after = bytes[i + 1..]
                .iter()
                .take_while(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'-'))
                .count();
            if before >= 2 && after >= 4 && bytes[i + 1..i + 1 + after].contains(&b'.') {
                return true;
            }
        }
        i += 1;
    }
    false
}

const CARD_MIN_DIGITS: usize = 13;
const CARD_MAX_DIGITS: usize = 19;

// Why: Luhn alone is a 1-in-10 coin flip, so sliding a 16-digit window over a
// long digit run all but guarantees a hit — a 40-digit hash offers 25 tries.
// Timestamps, uuids, trace ids and manifest ids were being read as cards on
// live traffic. A candidate must now be a *whole* digit run of a plausible
// card length carrying a real issuer prefix, and only then is Luhn consulted.
fn detect_credit_card(text: &str) -> bool {
    let bytes = text.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        if !bytes[i].is_ascii_digit() {
            i += 1;
            continue;
        }
        let (digits, end) = card_candidate(bytes, i);
        if is_card(&digits) {
            return true;
        }
        i = end;
    }
    false
}

// Why: a single space or hyphen *between* two digits is the grouping cards are
// written with, so it is absorbed; anything else ends the run. Without that
// stopping rule two neighbouring numbers splice into one candidate.
fn card_candidate(bytes: &[u8], start: usize) -> (Vec<u8>, usize) {
    let mut digits = Vec::new();
    let mut i = start;
    while i < bytes.len() {
        if bytes[i].is_ascii_digit() {
            digits.push(bytes[i]);
            i += 1;
            continue;
        }
        if matches!(bytes[i], b' ' | b'-') && bytes.get(i + 1).is_some_and(u8::is_ascii_digit) {
            i += 1;
            continue;
        }
        break;
    }
    (digits, i)
}

fn is_card(digits: &[u8]) -> bool {
    (CARD_MIN_DIGITS..=CARD_MAX_DIGITS).contains(&digits.len())
        && has_issuer_prefix(digits)
        && luhn(digits)
}

// Why: the assigned issuer identification ranges. Requiring one costs nothing
// against a real card and removes nine in ten of the arbitrary digit runs that
// would otherwise reach Luhn.
fn has_issuer_prefix(digits: &[u8]) -> bool {
    let Some(&second) = digits.get(1) else {
        return false;
    };
    match digits[0] {
        b'4' => true,
        b'5' => matches!(second, b'1'..=b'5'),
        b'2' => matches!(second, b'2'..=b'7'),
        b'3' => matches!(second, b'0' | b'4' | b'5' | b'6' | b'7' | b'8'),
        b'6' => digits.starts_with(b"6011") || second == b'5',
        _ => false,
    }
}

fn luhn(digits: &[u8]) -> bool {
    let mut sum = 0u32;
    for (i, b) in digits.iter().rev().enumerate() {
        let mut d = u32::from(b - b'0');
        if i % 2 == 1 {
            d *= 2;
            if d > 9 {
                d -= 9;
            }
        }
        sum += d;
    }
    sum.is_multiple_of(10)
}