mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-09 02:47:45 +08:00
6511 lines
224 KiB
Rust
6511 lines
224 KiB
Rust
use std::collections::{BTreeMap, HashMap, HashSet};
|
||
use std::convert::Infallible;
|
||
use std::fmt;
|
||
use std::net::{Ipv4Addr, Ipv6Addr};
|
||
use std::sync::atomic::{AtomicU64, Ordering};
|
||
use std::sync::{Arc, LazyLock, Mutex};
|
||
use std::time::{Duration, Instant};
|
||
|
||
use aether_data_contracts::DataLayerError;
|
||
use aether_runtime_state::RuntimeState;
|
||
use chrono::NaiveDate;
|
||
use hmac::Mac;
|
||
use regex::Regex;
|
||
use serde::{Deserialize, Serialize};
|
||
use serde_json::Value;
|
||
use sha2::Digest as _;
|
||
|
||
use crate::GatewayError;
|
||
|
||
type HmacSha256 = hmac::Hmac<sha2::Sha256>;
|
||
|
||
const DEFAULT_REDACTION_TTL_SECONDS: u64 = 300;
|
||
const HMAC96_BYTES: usize = 12;
|
||
const DEFAULT_SENTINEL_NAMESPACE: &str = "AETHER";
|
||
const MAX_SENTINEL_NAMESPACE_LEN: usize = 32;
|
||
const DIRECT_RESTORE_SENTINEL_LIMIT: usize = 32;
|
||
const MAX_CACHE_SENTINEL_BYTES: usize = 128;
|
||
const MAX_CACHE_RECORD_BYTES: usize = 512;
|
||
const CHAT_PII_REDACTION_RUNTIME_CONFIG_CACHE_TTL: Duration = Duration::from_secs(5);
|
||
|
||
static EMAIL_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"(?i)[A-Z0-9._%+-]{1,64}@[A-Z0-9.-]{1,253}\.[A-Z]{2,63}")
|
||
.expect("email regex should compile")
|
||
});
|
||
static CN_MOBILE_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"(?:\+?86[- ]?)?1[3-9]\d[- ]?\d{4}[- ]?\d{4}")
|
||
.expect("cn mobile regex should compile")
|
||
});
|
||
static CN_LANDLINE_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"(?:\+?86[- ]?)?0\d{2,3}[- ]\d{7,8}(?:-\d{1,6})?")
|
||
.expect("cn landline regex should compile")
|
||
});
|
||
static E164_PHONE_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"\+[1-9]\d(?:[ -]?\d){6,13}\d").expect("e164 phone regex should compile")
|
||
});
|
||
static CN_ID_REGEX: LazyLock<Regex> =
|
||
LazyLock::new(|| Regex::new(r"(?i)\b\d{17}[\dX]\b").expect("cn id regex should compile"));
|
||
static PAYMENT_CARD_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"\b(?:\d[ -]?){12,18}\d\b").expect("payment card regex should compile")
|
||
});
|
||
static IPV4_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"\b(?:\d{1,3}\.){3}\d{1,3}\b").expect("ipv4 regex should compile")
|
||
});
|
||
static OPENAI_KEY_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"\bsk-(?:proj-)?[A-Za-z0-9_-]{20,}\b").expect("openai key regex should compile")
|
||
});
|
||
static ANTHROPIC_KEY_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"\bsk-ant-[A-Za-z0-9_-]{20,}\b").expect("anthropic key regex should compile")
|
||
});
|
||
static GITHUB_TOKEN_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"\b(?:gh[pousr]_[A-Za-z0-9_]{30,}|github_pat_[A-Za-z0-9_]{30,})\b")
|
||
.expect("github token regex should compile")
|
||
});
|
||
static SLACK_TOKEN_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"\bxox[baprs]-[A-Za-z0-9-]{20,}\b").expect("slack token regex should compile")
|
||
});
|
||
static AWS_KEY_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"\b(?:AKIA|ASIA)[0-9A-Z]{16}\b").expect("aws key regex should compile")
|
||
});
|
||
static BEARER_TOKEN_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"(?i)\bBearer\s+[A-Za-z0-9._~+/=-]{20,}")
|
||
.expect("bearer token regex should compile")
|
||
});
|
||
static JWT_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"\b[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\b")
|
||
.expect("jwt regex should compile")
|
||
});
|
||
static ACCESS_TOKEN_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r#"(?i)\baccess[_-]?token\s*[:=]\s*["']?[A-Za-z0-9._~+/=-]{20,}"#)
|
||
.expect("access token regex should compile")
|
||
});
|
||
static SECRET_KEY_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(
|
||
r#"(?i)\b(?:secret|agent[_-]?private)[_-]?key\s*[:=]\s*["']?[A-Za-z0-9._~+/=-]{20,}"#,
|
||
)
|
||
.expect("secret key regex should compile")
|
||
});
|
||
static HIGH_ENTROPY_TOKEN_REGEX: LazyLock<Regex> =
|
||
LazyLock::new(|| Regex::new(r"\b[A-Za-z0-9_-]{32,}\b").expect("api key regex should compile"));
|
||
static SENTINEL_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||
Regex::new(r"<[A-Z0-9_]+:[A-Z0-9_]+:[A-Z2-7]{20}>").expect("sentinel regex should compile")
|
||
});
|
||
|
||
#[derive(Clone, Copy, Deserialize, Eq, Hash, PartialEq, Serialize)]
|
||
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
|
||
pub(crate) enum RedactionKind {
|
||
Email,
|
||
ChinaMobile,
|
||
ChinaLandline,
|
||
Phone,
|
||
ChinaResidentId,
|
||
PaymentCard,
|
||
Ipv4,
|
||
Ipv6,
|
||
OpenAiKey,
|
||
AnthropicKey,
|
||
GitHubToken,
|
||
SlackToken,
|
||
AwsKey,
|
||
BearerToken,
|
||
Jwt,
|
||
AccessToken,
|
||
SecretKey,
|
||
ApiKey,
|
||
}
|
||
|
||
impl RedactionKind {
|
||
pub(crate) const fn label(self) -> &'static str {
|
||
match self {
|
||
Self::Email => "EMAIL",
|
||
Self::ChinaMobile | Self::ChinaLandline => "CN_PHONE",
|
||
Self::Phone => "PHONE",
|
||
Self::ChinaResidentId => "CN_ID",
|
||
Self::PaymentCard => "PAYMENT_CARD",
|
||
Self::Ipv4 => "IPV4",
|
||
Self::Ipv6 => "IPV6",
|
||
Self::OpenAiKey => "OPENAI_KEY",
|
||
Self::AnthropicKey => "ANTHROPIC_KEY",
|
||
Self::GitHubToken => "GITHUB_TOKEN",
|
||
Self::SlackToken => "SLACK_TOKEN",
|
||
Self::AwsKey => "AWS_KEY",
|
||
Self::BearerToken => "BEARER_TOKEN",
|
||
Self::Jwt => "JWT",
|
||
Self::AccessToken => "ACCESS_TOKEN",
|
||
Self::SecretKey => "SECRET_KEY",
|
||
Self::ApiKey => "API_KEY",
|
||
}
|
||
}
|
||
}
|
||
|
||
impl fmt::Debug for RedactionKind {
|
||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
formatter.write_str(self.label())
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Debug)]
|
||
struct CompiledRedactionRule {
|
||
rule_label: String,
|
||
regex: Regex,
|
||
kinds: Vec<RedactionKind>,
|
||
custom_priority: u16,
|
||
}
|
||
|
||
#[derive(Clone, Debug, Deserialize, Serialize)]
|
||
struct ChatPiiRedactionRuleConfig {
|
||
id: String,
|
||
name: String,
|
||
pattern: String,
|
||
#[serde(default)]
|
||
enabled: bool,
|
||
#[serde(default)]
|
||
features: ChatPiiRedactionRuleFeatures,
|
||
#[serde(default)]
|
||
kind: Option<String>,
|
||
#[serde(default)]
|
||
system: bool,
|
||
}
|
||
|
||
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
|
||
struct ChatPiiRedactionRuleFeatures {
|
||
#[serde(default)]
|
||
validator: Option<String>,
|
||
#[serde(flatten)]
|
||
extra: BTreeMap<String, Value>,
|
||
}
|
||
|
||
impl ChatPiiRedactionRuleConfig {
|
||
fn validator(&self) -> Option<&str> {
|
||
self.features
|
||
.validator
|
||
.as_deref()
|
||
.or(self.kind.as_deref())
|
||
.map(str::trim)
|
||
.filter(|value| !value.is_empty())
|
||
}
|
||
}
|
||
|
||
#[derive(Clone)]
|
||
pub(crate) struct RedactionSessionConfig {
|
||
hmac_key: Vec<u8>,
|
||
ttl_seconds: u64,
|
||
now_unix_secs: u64,
|
||
sentinel_namespace: String,
|
||
enabled_kinds: Option<HashSet<RedactionKind>>,
|
||
rules: Option<Arc<Vec<CompiledRedactionRule>>>,
|
||
}
|
||
|
||
impl RedactionSessionConfig {
|
||
pub(crate) fn new(hmac_key: impl Into<Vec<u8>>, ttl_seconds: u64, now_unix_secs: u64) -> Self {
|
||
Self {
|
||
hmac_key: hmac_key.into(),
|
||
ttl_seconds: ttl_seconds.max(1),
|
||
now_unix_secs,
|
||
sentinel_namespace: DEFAULT_SENTINEL_NAMESPACE.to_string(),
|
||
enabled_kinds: None,
|
||
rules: None,
|
||
}
|
||
}
|
||
|
||
pub(crate) fn default_ttl(hmac_key: impl Into<Vec<u8>>, now_unix_secs: u64) -> Self {
|
||
Self::new(hmac_key, DEFAULT_REDACTION_TTL_SECONDS, now_unix_secs)
|
||
}
|
||
|
||
fn bucket(&self) -> u64 {
|
||
self.now_unix_secs / self.ttl_seconds
|
||
}
|
||
|
||
fn expires_at_unix_secs(&self) -> u64 {
|
||
self.bucket()
|
||
.saturating_add(1)
|
||
.saturating_mul(self.ttl_seconds)
|
||
}
|
||
|
||
pub(crate) fn with_enabled_kinds(mut self, enabled_kinds: HashSet<RedactionKind>) -> Self {
|
||
self.enabled_kinds = Some(enabled_kinds);
|
||
self
|
||
}
|
||
|
||
fn with_sentinel_namespace(mut self, namespace: impl Into<String>) -> Self {
|
||
self.sentinel_namespace = normalize_sentinel_namespace_or_default(namespace.into());
|
||
self
|
||
}
|
||
|
||
fn with_rules(mut self, rules: Vec<CompiledRedactionRule>) -> Self {
|
||
self.rules = Some(Arc::new(rules));
|
||
self
|
||
}
|
||
|
||
fn kind_enabled(&self, kind: RedactionKind) -> bool {
|
||
self.enabled_kinds
|
||
.as_ref()
|
||
.is_none_or(|enabled_kinds| enabled_kinds.contains(&kind))
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||
pub(crate) enum RedactionMaskError {
|
||
/// Provider-owned reasoning text cannot be rewritten independently from
|
||
/// its encrypted continuation state. Reject the request instead of either
|
||
/// leaking detected sensitive text or corrupting the provider binding.
|
||
SensitiveOpaqueReasoningState,
|
||
}
|
||
|
||
impl From<Infallible> for RedactionMaskError {
|
||
fn from(value: Infallible) -> Self {
|
||
match value {}
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Copy, Default)]
|
||
struct RedactionScanState;
|
||
|
||
impl RedactionScanState {
|
||
fn record_scan(&mut self, _text: &str) -> Result<(), Infallible> {
|
||
Ok(())
|
||
}
|
||
|
||
fn record_detections(&mut self, _count: usize) -> Result<(), Infallible> {
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
#[derive(Default)]
|
||
struct DetectorProbe {
|
||
validator_calls: HashMap<RedactionKind, usize>,
|
||
}
|
||
|
||
impl DetectorProbe {
|
||
fn record_validator_call(&mut self, kind: RedactionKind) {
|
||
*self.validator_calls.entry(kind).or_default() += 1;
|
||
}
|
||
|
||
#[cfg(test)]
|
||
fn validator_calls(&self, kind: RedactionKind) -> usize {
|
||
self.validator_calls.get(&kind).copied().unwrap_or_default()
|
||
}
|
||
}
|
||
|
||
impl fmt::Debug for RedactionSessionConfig {
|
||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
formatter
|
||
.debug_struct("RedactionSessionConfig")
|
||
.field("hmac_key", &"<redacted>")
|
||
.field("ttl_seconds", &self.ttl_seconds)
|
||
.field("now_unix_secs", &self.now_unix_secs)
|
||
.field("sentinel_namespace", &self.sentinel_namespace)
|
||
.field(
|
||
"enabled_kind_count",
|
||
&self.enabled_kinds.as_ref().map(HashSet::len),
|
||
)
|
||
.field(
|
||
"enabled_rule_count",
|
||
&self.rules.as_ref().map(|rules| rules.len()),
|
||
)
|
||
.finish()
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Eq, Hash, PartialEq)]
|
||
struct MappingKey {
|
||
rule_label: String,
|
||
original: String,
|
||
}
|
||
|
||
#[derive(Clone)]
|
||
pub(crate) struct RedactionSession {
|
||
config: RedactionSessionConfig,
|
||
mappings: HashMap<MappingKey, RedactionMapping>,
|
||
sentinel_index: HashMap<String, MappingKey>,
|
||
collision_corpus: Vec<String>,
|
||
preserve_deepseek_opaque_reasoning_state: bool,
|
||
}
|
||
|
||
impl RedactionSession {
|
||
pub(crate) fn new(config: RedactionSessionConfig) -> Self {
|
||
Self {
|
||
config,
|
||
mappings: HashMap::new(),
|
||
sentinel_index: HashMap::new(),
|
||
collision_corpus: Vec::new(),
|
||
preserve_deepseek_opaque_reasoning_state: false,
|
||
}
|
||
}
|
||
|
||
fn apply_mask_options(&mut self, options: MaskChatRequestOptions) {
|
||
self.preserve_deepseek_opaque_reasoning_state =
|
||
options.preserve_deepseek_opaque_reasoning_state;
|
||
}
|
||
|
||
/// Updates the provider-owned reasoning replay policy after the initial
|
||
/// request has been planned and its upstream binding authenticated.
|
||
///
|
||
/// A new WebSocket chain has to redact its first event before planning, so
|
||
/// it starts with the conservative OpenAI item-id policy. Once the
|
||
/// planner has selected the provider, response restoration must use that
|
||
/// trusted binding's policy as well; otherwise a DeepSeek opaque reasoning
|
||
/// item could have a PII sentinel restored inside provider-owned state.
|
||
pub(crate) fn set_reasoning_replay_policy(
|
||
&mut self,
|
||
policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy,
|
||
) {
|
||
self.apply_mask_options(
|
||
MaskChatRequestOptions::runtime().with_reasoning_replay_policy(policy),
|
||
);
|
||
}
|
||
|
||
fn preserves_deepseek_opaque_reasoning_state(&self) -> bool {
|
||
self.preserve_deepseek_opaque_reasoning_state
|
||
}
|
||
|
||
fn set_collision_corpus(&mut self, collision_corpus: Vec<String>) {
|
||
self.collision_corpus = collision_corpus;
|
||
}
|
||
|
||
fn text_has_redaction_candidate(&self, input: &str) -> bool {
|
||
!select_non_overlapping(detect_candidates_for_session_config(
|
||
input,
|
||
&self.config,
|
||
None,
|
||
))
|
||
.is_empty()
|
||
}
|
||
|
||
pub(crate) fn redact_text(&mut self, input: &str) -> RedactedText {
|
||
self.redact_text_internal(input, None)
|
||
.expect("unlimited redaction scan should not fail")
|
||
}
|
||
|
||
fn redact_text_checked(
|
||
&mut self,
|
||
input: &str,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<RedactedText, Infallible> {
|
||
scan_state.record_scan(input)?;
|
||
self.redact_text_internal(input, Some(scan_state))
|
||
}
|
||
|
||
async fn redact_text_with_cache(
|
||
&mut self,
|
||
input: &str,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<RedactedText, RedactionMaskError> {
|
||
scan_state.record_scan(input)?;
|
||
let candidates = select_non_overlapping(detect_candidates_for_session_config(
|
||
input,
|
||
&self.config,
|
||
None,
|
||
));
|
||
scan_state.record_detections(candidates.len())?;
|
||
if candidates.is_empty() {
|
||
return Ok(RedactedText {
|
||
text: input.to_string(),
|
||
matches: Vec::new(),
|
||
});
|
||
}
|
||
|
||
let mut redacted = String::with_capacity(input.len());
|
||
let mut cursor = 0;
|
||
let mut matches = Vec::with_capacity(candidates.len());
|
||
|
||
for candidate in candidates {
|
||
redacted.push_str(&input[cursor..candidate.start]);
|
||
let sentinel = self
|
||
.sentinel_for_candidate_with_cache(input, &candidate, cache)
|
||
.await;
|
||
redacted.push_str(&sentinel);
|
||
matches.push(RedactionMatch {
|
||
rule_label: candidate.rule_label,
|
||
kind: candidate.kind,
|
||
start: candidate.start,
|
||
end: candidate.end,
|
||
original: candidate.value,
|
||
sentinel,
|
||
});
|
||
cursor = candidate.end;
|
||
}
|
||
redacted.push_str(&input[cursor..]);
|
||
|
||
Ok(RedactedText {
|
||
text: redacted,
|
||
matches,
|
||
})
|
||
}
|
||
|
||
fn redact_text_internal(
|
||
&mut self,
|
||
input: &str,
|
||
mut scan_state: Option<&mut RedactionScanState>,
|
||
) -> Result<RedactedText, Infallible> {
|
||
let candidates = select_non_overlapping(detect_candidates_for_session_config(
|
||
input,
|
||
&self.config,
|
||
None,
|
||
));
|
||
if let Some(scan_state) = scan_state.as_mut() {
|
||
scan_state.record_detections(candidates.len())?;
|
||
}
|
||
if candidates.is_empty() {
|
||
return Ok(RedactedText {
|
||
text: input.to_string(),
|
||
matches: Vec::new(),
|
||
});
|
||
}
|
||
|
||
let mut redacted = String::with_capacity(input.len());
|
||
let mut cursor = 0;
|
||
let mut matches = Vec::with_capacity(candidates.len());
|
||
|
||
for candidate in candidates {
|
||
redacted.push_str(&input[cursor..candidate.start]);
|
||
let sentinel = self.sentinel_for_candidate(input, &candidate);
|
||
redacted.push_str(&sentinel);
|
||
matches.push(RedactionMatch {
|
||
rule_label: candidate.rule_label,
|
||
kind: candidate.kind,
|
||
start: candidate.start,
|
||
end: candidate.end,
|
||
original: candidate.value,
|
||
sentinel,
|
||
});
|
||
cursor = candidate.end;
|
||
}
|
||
redacted.push_str(&input[cursor..]);
|
||
|
||
Ok(RedactedText {
|
||
text: redacted,
|
||
matches,
|
||
})
|
||
}
|
||
|
||
pub(crate) fn mapping_count(&self) -> usize {
|
||
self.mappings.len()
|
||
}
|
||
|
||
pub(crate) fn sentinel_for_original(&self, original: &str) -> Option<&str> {
|
||
self.mappings
|
||
.values()
|
||
.find(|mapping| mapping.original == original)
|
||
.map(|mapping| mapping.sentinel.as_str())
|
||
}
|
||
|
||
pub(crate) fn mappings(&self) -> impl Iterator<Item = &RedactionMapping> {
|
||
self.mappings.values()
|
||
}
|
||
|
||
fn restore_text(&self, input: &str) -> RestoredText {
|
||
if self.mapping_count() <= DIRECT_RESTORE_SENTINEL_LIMIT {
|
||
return restore_text_direct_longest_first(input, self);
|
||
}
|
||
restore_text_with_matcher(input, &SentinelMatcher::new(self))
|
||
}
|
||
|
||
fn sentinel_for_candidate(&mut self, source_text: &str, candidate: &Candidate) -> String {
|
||
let key = MappingKey {
|
||
rule_label: candidate.rule_label.clone(),
|
||
original: normalize_redaction_value(candidate.kind, &candidate.value),
|
||
};
|
||
if let Some(mapping) = self.mappings.get(&key) {
|
||
return mapping.sentinel.clone();
|
||
}
|
||
|
||
let bucket = self.config.bucket();
|
||
let mut collision_counter = 0u32;
|
||
let sentinel = loop {
|
||
let candidate_sentinel = self.build_sentinel(
|
||
&candidate.rule_label,
|
||
&key.original,
|
||
bucket,
|
||
collision_counter,
|
||
);
|
||
let collides_with_input =
|
||
self.sentinel_collides_with_original_text(source_text, &candidate_sentinel);
|
||
let collides_with_mapping = self
|
||
.sentinel_index
|
||
.get(&candidate_sentinel)
|
||
.is_some_and(|existing_key| existing_key != &key);
|
||
if !collides_with_input && !collides_with_mapping {
|
||
break candidate_sentinel;
|
||
}
|
||
collision_counter = collision_counter.saturating_add(1);
|
||
};
|
||
|
||
let mapping = RedactionMapping {
|
||
rule_label: candidate.rule_label.clone(),
|
||
kind: candidate.kind,
|
||
original: candidate.value.clone(),
|
||
normalized_value: key.original.clone(),
|
||
sentinel: sentinel.clone(),
|
||
bucket,
|
||
created_at_unix_secs: self.config.now_unix_secs,
|
||
expires_at_unix_secs: self.config.expires_at_unix_secs(),
|
||
};
|
||
self.sentinel_index.insert(sentinel.clone(), key.clone());
|
||
self.mappings.insert(key, mapping);
|
||
sentinel
|
||
}
|
||
|
||
async fn sentinel_for_candidate_with_cache(
|
||
&mut self,
|
||
source_text: &str,
|
||
candidate: &Candidate,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> String {
|
||
let key = MappingKey {
|
||
rule_label: candidate.rule_label.clone(),
|
||
original: normalize_redaction_value(candidate.kind, &candidate.value),
|
||
};
|
||
if let Some(mapping) = self.mappings.get(&key) {
|
||
return mapping.sentinel.clone();
|
||
}
|
||
|
||
let bucket = self.config.bucket();
|
||
if let Some(cache) = cache {
|
||
if let Ok(Some(sentinel)) = cache
|
||
.lookup_sentinel(
|
||
&candidate.rule_label,
|
||
&key.original,
|
||
bucket,
|
||
&self.config.sentinel_namespace,
|
||
)
|
||
.await
|
||
{
|
||
if sentinel.len() <= MAX_CACHE_SENTINEL_BYTES
|
||
&& self.cached_sentinel_usable(source_text, &key, &sentinel)
|
||
{
|
||
self.insert_mapping_for_sentinel(key, candidate, sentinel.clone(), bucket);
|
||
return sentinel;
|
||
}
|
||
}
|
||
}
|
||
|
||
let sentinel = self.sentinel_for_candidate(source_text, candidate);
|
||
if let Some(cache) = cache {
|
||
if let Some(mapping) = self.mappings.get(&key) {
|
||
let ttl_seconds = self
|
||
.config
|
||
.expires_at_unix_secs()
|
||
.saturating_sub(self.config.now_unix_secs)
|
||
.max(1);
|
||
let _ = cache
|
||
.store(
|
||
&RedactionCacheRecord::from_mapping(mapping),
|
||
&mapping.normalized_value,
|
||
ttl_seconds,
|
||
)
|
||
.await;
|
||
}
|
||
}
|
||
sentinel
|
||
}
|
||
|
||
fn cached_sentinel_usable(&self, source_text: &str, key: &MappingKey, sentinel: &str) -> bool {
|
||
if self.sentinel_collides_with_original_text(source_text, sentinel) {
|
||
return false;
|
||
}
|
||
!self
|
||
.sentinel_index
|
||
.get(sentinel)
|
||
.is_some_and(|existing_key| existing_key != key)
|
||
}
|
||
|
||
fn sentinel_collides_with_original_text(&self, source_text: &str, sentinel: &str) -> bool {
|
||
source_text.contains(sentinel)
|
||
|| self
|
||
.collision_corpus
|
||
.iter()
|
||
.any(|original_text| original_text.contains(sentinel))
|
||
}
|
||
|
||
fn insert_mapping_for_sentinel(
|
||
&mut self,
|
||
key: MappingKey,
|
||
candidate: &Candidate,
|
||
sentinel: String,
|
||
bucket: u64,
|
||
) {
|
||
let mapping = RedactionMapping {
|
||
rule_label: candidate.rule_label.clone(),
|
||
kind: candidate.kind,
|
||
original: candidate.value.clone(),
|
||
normalized_value: key.original.clone(),
|
||
sentinel: sentinel.clone(),
|
||
bucket,
|
||
created_at_unix_secs: self.config.now_unix_secs,
|
||
expires_at_unix_secs: self.config.expires_at_unix_secs(),
|
||
};
|
||
self.sentinel_index.insert(sentinel, key.clone());
|
||
self.mappings.insert(key, mapping);
|
||
}
|
||
|
||
fn build_sentinel(
|
||
&self,
|
||
rule_label: &str,
|
||
normalized_value: &str,
|
||
bucket: u64,
|
||
collision_counter: u32,
|
||
) -> String {
|
||
let mut mac = HmacSha256::new_from_slice(&self.config.hmac_key)
|
||
.expect("HMAC accepts keys of any size");
|
||
mac.update(b"aether-redaction-v1\0");
|
||
mac.update(rule_label.as_bytes());
|
||
mac.update(b"\0");
|
||
mac.update(bucket.to_string().as_bytes());
|
||
mac.update(b"\0");
|
||
mac.update(collision_counter.to_string().as_bytes());
|
||
mac.update(b"\0");
|
||
mac.update(normalized_value.as_bytes());
|
||
let digest = mac.finalize().into_bytes();
|
||
format!(
|
||
"<{}:{}:{}>",
|
||
self.config.sentinel_namespace,
|
||
rule_label,
|
||
base32_no_pad(&digest[..HMAC96_BYTES])
|
||
)
|
||
}
|
||
}
|
||
|
||
impl fmt::Debug for RedactionSession {
|
||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
let mut counts: HashMap<&str, usize> = HashMap::new();
|
||
for mapping in self.mappings.values() {
|
||
*counts.entry(mapping.rule_label.as_str()).or_default() += 1;
|
||
}
|
||
formatter
|
||
.debug_struct("RedactionSession")
|
||
.field("ttl_seconds", &self.config.ttl_seconds)
|
||
.field("bucket", &self.config.bucket())
|
||
.field("mapping_count", &self.mappings.len())
|
||
.field("type_counts", &counts)
|
||
.field(
|
||
"preserve_deepseek_opaque_reasoning_state",
|
||
&self.preserve_deepseek_opaque_reasoning_state,
|
||
)
|
||
.finish()
|
||
}
|
||
}
|
||
|
||
pub(crate) struct MaskedChatRequest {
|
||
pub(crate) body: Vec<u8>,
|
||
pub(crate) session: RedactionSession,
|
||
pub(crate) redacted: bool,
|
||
}
|
||
|
||
impl fmt::Debug for MaskedChatRequest {
|
||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
formatter
|
||
.debug_struct("MaskedChatRequest")
|
||
.field("body_len", &self.body.len())
|
||
.field("session", &self.session)
|
||
.field("redacted", &self.redacted)
|
||
.finish()
|
||
}
|
||
}
|
||
|
||
pub(crate) struct MaskedChatRequestValue {
|
||
pub(crate) body_json: Option<Value>,
|
||
pub(crate) session: RedactionSession,
|
||
pub(crate) redacted: bool,
|
||
}
|
||
|
||
impl fmt::Debug for MaskedChatRequestValue {
|
||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
formatter
|
||
.debug_struct("MaskedChatRequestValue")
|
||
.field("body_json_owned", &self.body_json.is_some())
|
||
.field("session", &self.session)
|
||
.field("redacted", &self.redacted)
|
||
.finish()
|
||
}
|
||
}
|
||
|
||
#[derive(Clone)]
|
||
pub(crate) struct CachedRequestRedaction {
|
||
pub(crate) body_json: Option<Value>,
|
||
pub(crate) session: Option<RedactionSession>,
|
||
pub(crate) redacted: bool,
|
||
}
|
||
|
||
impl CachedRequestRedaction {
|
||
pub(crate) fn unredacted() -> Self {
|
||
Self {
|
||
body_json: None,
|
||
session: None,
|
||
redacted: false,
|
||
}
|
||
}
|
||
|
||
pub(crate) fn redacted(body_json: Value, session: RedactionSession) -> Self {
|
||
Self {
|
||
body_json: Some(body_json),
|
||
session: Some(session),
|
||
redacted: true,
|
||
}
|
||
}
|
||
}
|
||
|
||
#[derive(Default, Clone)]
|
||
pub(crate) struct RedactionSessionSlot {
|
||
session: Arc<Mutex<Option<RedactionSession>>>,
|
||
sessions_by_candidate: Arc<Mutex<HashMap<String, RedactionSession>>>,
|
||
request_redactions: Arc<Mutex<HashMap<String, CachedRequestRedaction>>>,
|
||
}
|
||
|
||
impl RedactionSessionSlot {
|
||
pub(crate) fn put(&self, session: RedactionSession) {
|
||
*self
|
||
.session
|
||
.lock()
|
||
.expect("redaction session slot should lock") = Some(session);
|
||
}
|
||
|
||
pub(crate) fn put_for_candidate(&self, candidate_id: &str, session: RedactionSession) {
|
||
self.sessions_by_candidate
|
||
.lock()
|
||
.expect("redaction session candidate slot should lock")
|
||
.insert(candidate_id.to_string(), session);
|
||
}
|
||
|
||
pub(crate) fn clear(&self) {
|
||
*self
|
||
.session
|
||
.lock()
|
||
.expect("redaction session slot should lock") = None;
|
||
self.sessions_by_candidate
|
||
.lock()
|
||
.expect("redaction session candidate slot should lock")
|
||
.clear();
|
||
self.request_redactions
|
||
.lock()
|
||
.expect("redaction request cache slot should lock")
|
||
.clear();
|
||
}
|
||
|
||
pub(crate) fn take(&self) -> Option<RedactionSession> {
|
||
self.session
|
||
.lock()
|
||
.expect("redaction session slot should lock")
|
||
.take()
|
||
}
|
||
|
||
pub(crate) fn take_for_candidate(
|
||
&self,
|
||
candidate_id: Option<&str>,
|
||
) -> Option<RedactionSession> {
|
||
let Some(candidate_id) = candidate_id else {
|
||
return self.take_without_candidate_id();
|
||
};
|
||
let mut sessions_by_candidate = self
|
||
.sessions_by_candidate
|
||
.lock()
|
||
.expect("redaction session candidate slot should lock");
|
||
if let Some(session) = sessions_by_candidate.remove(candidate_id) {
|
||
return Some(session);
|
||
}
|
||
if !sessions_by_candidate.is_empty() {
|
||
return None;
|
||
}
|
||
drop(sessions_by_candidate);
|
||
self.take()
|
||
}
|
||
|
||
fn take_without_candidate_id(&self) -> Option<RedactionSession> {
|
||
let mut sessions_by_candidate = self
|
||
.sessions_by_candidate
|
||
.lock()
|
||
.expect("redaction session candidate slot should lock");
|
||
match sessions_by_candidate.len() {
|
||
0 => {
|
||
drop(sessions_by_candidate);
|
||
self.take()
|
||
}
|
||
1 => sessions_by_candidate
|
||
.drain()
|
||
.next()
|
||
.map(|(_, session)| session),
|
||
_ => None,
|
||
}
|
||
}
|
||
|
||
pub(crate) fn cached_request_redaction(&self, key: &str) -> Option<CachedRequestRedaction> {
|
||
self.request_redactions
|
||
.lock()
|
||
.expect("redaction request cache slot should lock")
|
||
.get(key)
|
||
.cloned()
|
||
}
|
||
|
||
pub(crate) fn put_cached_request_redaction(
|
||
&self,
|
||
key: impl Into<String>,
|
||
redaction: CachedRequestRedaction,
|
||
) {
|
||
self.request_redactions
|
||
.lock()
|
||
.expect("redaction request cache slot should lock")
|
||
.insert(key.into(), redaction);
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||
pub(crate) struct RedactionExecutionCandidateId(String);
|
||
|
||
impl RedactionExecutionCandidateId {
|
||
pub(crate) fn new(candidate_id: impl Into<String>) -> Self {
|
||
Self(candidate_id.into())
|
||
}
|
||
|
||
pub(crate) fn as_str(&self) -> &str {
|
||
self.0.as_str()
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Debug)]
|
||
pub(crate) struct ChatPiiRedactionRuntimeConfig {
|
||
pub(crate) enabled: bool,
|
||
rules: Vec<CompiledRedactionRule>,
|
||
pub(crate) ttl_seconds: u64,
|
||
placeholder_prefix: String,
|
||
}
|
||
|
||
impl Default for ChatPiiRedactionRuntimeConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
enabled: false,
|
||
rules: default_chat_pii_redaction_rules(),
|
||
ttl_seconds: DEFAULT_REDACTION_TTL_SECONDS,
|
||
placeholder_prefix: DEFAULT_SENTINEL_NAMESPACE.to_string(),
|
||
}
|
||
}
|
||
}
|
||
|
||
impl ChatPiiRedactionRuntimeConfig {
|
||
fn disabled() -> Self {
|
||
Self {
|
||
enabled: false,
|
||
rules: Vec::new(),
|
||
ttl_seconds: DEFAULT_REDACTION_TTL_SECONDS,
|
||
placeholder_prefix: DEFAULT_SENTINEL_NAMESPACE.to_string(),
|
||
}
|
||
}
|
||
}
|
||
|
||
#[derive(Debug, Default)]
|
||
pub(crate) struct ChatPiiRedactionRuntimeConfigCache {
|
||
value: Mutex<Option<(Instant, ChatPiiRedactionRuntimeConfig)>>,
|
||
loading_generation: Mutex<Option<u64>>,
|
||
generation: AtomicU64,
|
||
notify: tokio::sync::Notify,
|
||
}
|
||
|
||
enum ChatPiiRedactionRuntimeConfigLoadRegistration {
|
||
Leader(ChatPiiRedactionRuntimeConfigLoadGuard),
|
||
Follower,
|
||
Bypass,
|
||
}
|
||
|
||
struct ChatPiiRedactionRuntimeConfigLoadGuard {
|
||
cache: ChatPiiRedactionRuntimeConfigCacheHandle,
|
||
generation: u64,
|
||
active: bool,
|
||
}
|
||
|
||
impl ChatPiiRedactionRuntimeConfigLoadGuard {
|
||
fn generation(&self) -> u64 {
|
||
self.generation
|
||
}
|
||
}
|
||
|
||
impl Drop for ChatPiiRedactionRuntimeConfigLoadGuard {
|
||
fn drop(&mut self) {
|
||
if self.active {
|
||
self.cache.finish_load(self.generation);
|
||
}
|
||
}
|
||
}
|
||
|
||
impl ChatPiiRedactionRuntimeConfigCache {
|
||
fn get(&self) -> Option<ChatPiiRedactionRuntimeConfig> {
|
||
self.value.lock().ok().and_then(|guard| {
|
||
guard.as_ref().and_then(|(loaded_at, value)| {
|
||
(loaded_at.elapsed() <= CHAT_PII_REDACTION_RUNTIME_CONFIG_CACHE_TTL)
|
||
.then(|| value.clone())
|
||
})
|
||
})
|
||
}
|
||
|
||
fn get_stale(&self) -> Option<ChatPiiRedactionRuntimeConfig> {
|
||
self.value
|
||
.lock()
|
||
.ok()
|
||
.and_then(|guard| guard.as_ref().map(|(_, value)| value.clone()))
|
||
}
|
||
|
||
fn insert(&self, value: ChatPiiRedactionRuntimeConfig) {
|
||
if let Ok(mut guard) = self.value.lock() {
|
||
*guard = Some((Instant::now(), value));
|
||
}
|
||
}
|
||
|
||
fn insert_if_generation(&self, generation: u64, value: ChatPiiRedactionRuntimeConfig) {
|
||
if self.generation.load(Ordering::Acquire) == generation {
|
||
self.insert(value);
|
||
}
|
||
}
|
||
|
||
fn register_load(self: &Arc<Self>) -> ChatPiiRedactionRuntimeConfigLoadRegistration {
|
||
let generation = self.generation.load(Ordering::Acquire);
|
||
match self.loading_generation.lock() {
|
||
Ok(mut loading_generation) => {
|
||
if loading_generation.is_some() {
|
||
ChatPiiRedactionRuntimeConfigLoadRegistration::Follower
|
||
} else {
|
||
*loading_generation = Some(generation);
|
||
ChatPiiRedactionRuntimeConfigLoadRegistration::Leader(
|
||
ChatPiiRedactionRuntimeConfigLoadGuard {
|
||
cache: Arc::clone(self),
|
||
generation,
|
||
active: true,
|
||
},
|
||
)
|
||
}
|
||
}
|
||
Err(_) => ChatPiiRedactionRuntimeConfigLoadRegistration::Bypass,
|
||
}
|
||
}
|
||
|
||
fn notified(&self) -> tokio::sync::futures::Notified<'_> {
|
||
self.notify.notified()
|
||
}
|
||
|
||
fn finish_load(&self, generation: u64) {
|
||
let finished = self
|
||
.loading_generation
|
||
.lock()
|
||
.map(|mut loading_generation| {
|
||
if *loading_generation == Some(generation) {
|
||
*loading_generation = None;
|
||
true
|
||
} else {
|
||
false
|
||
}
|
||
})
|
||
.unwrap_or(false);
|
||
if finished {
|
||
self.notify.notify_waiters();
|
||
}
|
||
}
|
||
}
|
||
|
||
pub(crate) type ChatPiiRedactionRuntimeConfigCacheHandle = Arc<ChatPiiRedactionRuntimeConfigCache>;
|
||
|
||
pub(crate) fn new_chat_pii_redaction_runtime_config_cache(
|
||
) -> ChatPiiRedactionRuntimeConfigCacheHandle {
|
||
Arc::new(ChatPiiRedactionRuntimeConfigCache::default())
|
||
}
|
||
|
||
pub(crate) fn clear_chat_pii_redaction_runtime_config_cache(
|
||
cache: &ChatPiiRedactionRuntimeConfigCacheHandle,
|
||
) {
|
||
cache.clear();
|
||
}
|
||
|
||
impl ChatPiiRedactionRuntimeConfigCache {
|
||
fn clear(&self) {
|
||
self.generation.fetch_add(1, Ordering::AcqRel);
|
||
if let Ok(mut value) = self.value.lock() {
|
||
*value = None;
|
||
}
|
||
if let Ok(mut loading_generation) = self.loading_generation.lock() {
|
||
*loading_generation = None;
|
||
}
|
||
self.notify.notify_waiters();
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
|
||
pub(crate) struct MaskChatRequestOptions {
|
||
/// This bit is derived from the selected provider configuration, never
|
||
/// from request JSON. It permits byte-identical handling only for the
|
||
/// provider-owned, id-less continuation state used by DeepSeek's
|
||
/// Responses contract.
|
||
preserve_deepseek_opaque_reasoning_state: bool,
|
||
}
|
||
|
||
impl MaskChatRequestOptions {
|
||
pub(crate) fn runtime() -> Self {
|
||
Self::default()
|
||
}
|
||
|
||
pub(crate) fn with_reasoning_replay_policy(
|
||
mut self,
|
||
policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy,
|
||
) -> Self {
|
||
self.preserve_deepseek_opaque_reasoning_state = matches!(
|
||
policy,
|
||
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
|
||
);
|
||
self
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||
pub(crate) enum ChatPiiRedactionRequestFormat {
|
||
OpenAiChat,
|
||
OpenAiResponses,
|
||
OpenAiSearch,
|
||
ClaudeMessages,
|
||
}
|
||
|
||
impl ChatPiiRedactionRequestFormat {
|
||
pub(crate) fn from_api_format(api_format: &str) -> Option<Self> {
|
||
match api_format.trim().to_ascii_lowercase().as_str() {
|
||
"openai:chat" => Some(Self::OpenAiChat),
|
||
"openai:responses" | "openai:responses:compact" => Some(Self::OpenAiResponses),
|
||
"openai:search" => Some(Self::OpenAiSearch),
|
||
"claude:messages" => Some(Self::ClaudeMessages),
|
||
_ => None,
|
||
}
|
||
}
|
||
}
|
||
|
||
fn sanitize_redaction_rule_label(raw: &str) -> String {
|
||
let label = raw
|
||
.trim()
|
||
.chars()
|
||
.map(|ch| {
|
||
if ch.is_ascii_alphanumeric() {
|
||
ch.to_ascii_uppercase()
|
||
} else {
|
||
'_'
|
||
}
|
||
})
|
||
.collect::<String>();
|
||
let label = label.trim_matches('_').to_string();
|
||
if label.is_empty() {
|
||
"CUSTOM_RULE".to_string()
|
||
} else {
|
||
label
|
||
}
|
||
}
|
||
|
||
fn normalize_sentinel_namespace(raw: &str) -> Option<String> {
|
||
let namespace = raw.trim();
|
||
if namespace.is_empty() || namespace.len() > MAX_SENTINEL_NAMESPACE_LEN {
|
||
return None;
|
||
}
|
||
if !namespace
|
||
.chars()
|
||
.all(|ch| ch.is_ascii_alphanumeric() || ch == '_')
|
||
{
|
||
return None;
|
||
}
|
||
Some(namespace.to_ascii_uppercase())
|
||
}
|
||
|
||
fn normalize_sentinel_namespace_or_default(raw: impl AsRef<str>) -> String {
|
||
normalize_sentinel_namespace(raw.as_ref())
|
||
.unwrap_or_else(|| DEFAULT_SENTINEL_NAMESPACE.to_string())
|
||
}
|
||
|
||
fn default_chat_pii_redaction_rule_configs() -> Vec<ChatPiiRedactionRuleConfig> {
|
||
vec![
|
||
ChatPiiRedactionRuleConfig {
|
||
id: "email".to_string(),
|
||
name: "邮箱".to_string(),
|
||
pattern: r"(?i)[A-Z0-9._%+-]{1,64}@[A-Z0-9.-]{1,253}\.[A-Z]{2,63}".to_string(),
|
||
enabled: true,
|
||
features: ChatPiiRedactionRuleFeatures {
|
||
validator: Some("email".to_string()),
|
||
..Default::default()
|
||
},
|
||
kind: None,
|
||
system: true,
|
||
},
|
||
ChatPiiRedactionRuleConfig {
|
||
id: "cn_phone".to_string(),
|
||
name: "手机号".to_string(),
|
||
pattern: r"(?:\+?86[- ]?)?(?:1[3-9]\d[- ]?\d{4}[- ]?\d{4}|0\d{2,3}[- ]\d{7,8}(?:-\d{1,6})?)"
|
||
.to_string(),
|
||
enabled: true,
|
||
features: ChatPiiRedactionRuleFeatures {
|
||
validator: Some("cn_phone".to_string()),
|
||
..Default::default()
|
||
},
|
||
kind: None,
|
||
system: true,
|
||
},
|
||
ChatPiiRedactionRuleConfig {
|
||
id: "global_phone".to_string(),
|
||
name: "国际号码".to_string(),
|
||
pattern: r"\+[1-9]\d(?:[ -]?\d){6,13}\d".to_string(),
|
||
enabled: true,
|
||
features: ChatPiiRedactionRuleFeatures {
|
||
validator: Some("global_phone".to_string()),
|
||
..Default::default()
|
||
},
|
||
kind: None,
|
||
system: true,
|
||
},
|
||
ChatPiiRedactionRuleConfig {
|
||
id: "cn_id".to_string(),
|
||
name: "身份证号".to_string(),
|
||
pattern: r"(?i)\b\d{17}[\dX]\b".to_string(),
|
||
enabled: true,
|
||
features: ChatPiiRedactionRuleFeatures {
|
||
validator: Some("cn_id".to_string()),
|
||
..Default::default()
|
||
},
|
||
kind: None,
|
||
system: true,
|
||
},
|
||
ChatPiiRedactionRuleConfig {
|
||
id: "payment_card".to_string(),
|
||
name: "银行卡号".to_string(),
|
||
pattern: r"\b(?:\d[ -]?){12,18}\d\b".to_string(),
|
||
enabled: true,
|
||
features: ChatPiiRedactionRuleFeatures {
|
||
validator: Some("payment_card".to_string()),
|
||
..Default::default()
|
||
},
|
||
kind: None,
|
||
system: true,
|
||
},
|
||
ChatPiiRedactionRuleConfig {
|
||
id: "ipv4".to_string(),
|
||
name: "IPv4".to_string(),
|
||
pattern: r"\b(?:\d{1,3}\.){3}\d{1,3}\b".to_string(),
|
||
enabled: true,
|
||
features: ChatPiiRedactionRuleFeatures {
|
||
validator: Some("ipv4".to_string()),
|
||
..Default::default()
|
||
},
|
||
kind: None,
|
||
system: true,
|
||
},
|
||
ChatPiiRedactionRuleConfig {
|
||
id: "ipv6".to_string(),
|
||
name: "IPv6".to_string(),
|
||
pattern: r"\b(?:[0-9A-Fa-f]{1,4}:){2,7}[0-9A-Fa-f:.]{1,39}\b".to_string(),
|
||
enabled: true,
|
||
features: ChatPiiRedactionRuleFeatures {
|
||
validator: Some("ipv6".to_string()),
|
||
..Default::default()
|
||
},
|
||
kind: None,
|
||
system: true,
|
||
},
|
||
ChatPiiRedactionRuleConfig {
|
||
id: "api_key".to_string(),
|
||
name: "API Key".to_string(),
|
||
pattern:
|
||
r"\b(?:sk-(?:proj-)?[A-Za-z0-9_-]{20,}|sk-ant-[A-Za-z0-9_-]{20,}|(?:gh[pousr]_[A-Za-z0-9_]{30,}|github_pat_[A-Za-z0-9_]{30,})|xox[baprs]-[A-Za-z0-9-]{20,}|(?:AKIA|ASIA)[0-9A-Z]{16}|[A-Za-z0-9_-]{32,})\b"
|
||
.to_string(),
|
||
enabled: true,
|
||
features: ChatPiiRedactionRuleFeatures {
|
||
validator: Some("api_key".to_string()),
|
||
..Default::default()
|
||
},
|
||
kind: None,
|
||
system: true,
|
||
},
|
||
ChatPiiRedactionRuleConfig {
|
||
id: "access_token".to_string(),
|
||
name: "Access Token".to_string(),
|
||
pattern: r#"(?i)\baccess[_-]?token\s*[:=]\s*["']?[A-Za-z0-9._~+/=-]{20,}"#
|
||
.to_string(),
|
||
enabled: true,
|
||
features: ChatPiiRedactionRuleFeatures {
|
||
validator: Some("access_token".to_string()),
|
||
..Default::default()
|
||
},
|
||
kind: None,
|
||
system: true,
|
||
},
|
||
ChatPiiRedactionRuleConfig {
|
||
id: "secret_key".to_string(),
|
||
name: "Secret Key".to_string(),
|
||
pattern: r#"(?i)\bsecret[_-]?key\s*[:=]\s*["']?[A-Za-z0-9._~+/=-]{20,}"#
|
||
.to_string(),
|
||
enabled: true,
|
||
features: ChatPiiRedactionRuleFeatures {
|
||
validator: Some("secret_key".to_string()),
|
||
..Default::default()
|
||
},
|
||
kind: None,
|
||
system: true,
|
||
},
|
||
ChatPiiRedactionRuleConfig {
|
||
id: "bearer_token".to_string(),
|
||
name: "Bearer Token".to_string(),
|
||
pattern: r"(?i)\bBearer\s+[A-Za-z0-9._~+/=-]{20,}".to_string(),
|
||
enabled: true,
|
||
features: ChatPiiRedactionRuleFeatures {
|
||
validator: Some("bearer_token".to_string()),
|
||
..Default::default()
|
||
},
|
||
kind: None,
|
||
system: true,
|
||
},
|
||
ChatPiiRedactionRuleConfig {
|
||
id: "jwt".to_string(),
|
||
name: "JWT".to_string(),
|
||
pattern: r"\b[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\b"
|
||
.to_string(),
|
||
enabled: true,
|
||
features: ChatPiiRedactionRuleFeatures {
|
||
validator: Some("jwt".to_string()),
|
||
..Default::default()
|
||
},
|
||
kind: None,
|
||
system: true,
|
||
},
|
||
]
|
||
}
|
||
|
||
fn default_chat_pii_redaction_rules() -> Vec<CompiledRedactionRule> {
|
||
compile_chat_pii_redaction_rules(&default_chat_pii_redaction_rule_configs())
|
||
.expect("default chat pii redaction rules should compile")
|
||
}
|
||
|
||
fn compile_chat_pii_redaction_rules(
|
||
rules: &[ChatPiiRedactionRuleConfig],
|
||
) -> Result<Vec<CompiledRedactionRule>, GatewayError> {
|
||
let mut compiled = Vec::new();
|
||
for (index, rule) in rules.iter().enumerate() {
|
||
if !rule.enabled {
|
||
continue;
|
||
}
|
||
let pattern = rule.pattern.trim();
|
||
if pattern.is_empty() {
|
||
return Err(GatewayError::Internal(
|
||
"chat pii redaction rule pattern is empty".to_string(),
|
||
));
|
||
}
|
||
let regex = Regex::new(pattern).map_err(|err| {
|
||
GatewayError::Internal(format!("chat pii redaction rule regex failed: {err}"))
|
||
})?;
|
||
let kinds = match rule.validator() {
|
||
Some(validator) => {
|
||
let kinds = redaction_kinds_for_entity(validator).collect::<Vec<_>>();
|
||
if kinds.is_empty() {
|
||
return Err(GatewayError::Internal(format!(
|
||
"unsupported chat pii redaction rule validator: {validator}"
|
||
)));
|
||
}
|
||
kinds
|
||
}
|
||
None => Vec::new(),
|
||
};
|
||
compiled.push(CompiledRedactionRule {
|
||
rule_label: sanitize_redaction_rule_label(&rule.id),
|
||
regex,
|
||
kinds,
|
||
custom_priority: (1000u16).saturating_add(index as u16),
|
||
});
|
||
}
|
||
Ok(compiled)
|
||
}
|
||
|
||
fn parse_chat_pii_redaction_rules(
|
||
value: Option<&Value>,
|
||
) -> Result<Vec<CompiledRedactionRule>, GatewayError> {
|
||
let Some(value) = value else {
|
||
return Ok(default_chat_pii_redaction_rules());
|
||
};
|
||
let Some(items) = value.as_array() else {
|
||
return Err(GatewayError::Internal(
|
||
"chat pii redaction rules must be an array".to_string(),
|
||
));
|
||
};
|
||
let parsed = items
|
||
.iter()
|
||
.cloned()
|
||
.map(serde_json::from_value::<ChatPiiRedactionRuleConfig>)
|
||
.collect::<Result<Vec<_>, _>>()
|
||
.map_err(|err| {
|
||
GatewayError::Internal(format!("chat pii redaction rules invalid: {err}"))
|
||
})?;
|
||
compile_chat_pii_redaction_rules(&parsed)
|
||
}
|
||
|
||
pub(crate) async fn read_chat_pii_redaction_runtime_config(
|
||
state: &crate::AppState,
|
||
) -> Result<ChatPiiRedactionRuntimeConfig, GatewayError> {
|
||
let cache = Arc::clone(&state.chat_pii_redaction_runtime_config_cache);
|
||
if let Some(value) = cache.get() {
|
||
return Ok(value);
|
||
}
|
||
if let Some(value) = cache.get_stale() {
|
||
if let ChatPiiRedactionRuntimeConfigLoadRegistration::Leader(guard) = cache.register_load()
|
||
{
|
||
spawn_chat_pii_redaction_runtime_config_refresh(state.clone(), cache, guard);
|
||
}
|
||
return Ok(value);
|
||
}
|
||
|
||
loop {
|
||
let notified = cache.notified();
|
||
match cache.register_load() {
|
||
ChatPiiRedactionRuntimeConfigLoadRegistration::Bypass => {
|
||
let value = load_chat_pii_redaction_runtime_config(state).await?;
|
||
cache.insert(value.clone());
|
||
return Ok(value);
|
||
}
|
||
ChatPiiRedactionRuntimeConfigLoadRegistration::Follower => {
|
||
notified.await;
|
||
if let Some(value) = cache.get() {
|
||
return Ok(value);
|
||
}
|
||
}
|
||
ChatPiiRedactionRuntimeConfigLoadRegistration::Leader(_guard) => {
|
||
let generation = _guard.generation();
|
||
let value = load_chat_pii_redaction_runtime_config(state).await?;
|
||
cache.insert_if_generation(generation, value.clone());
|
||
return Ok(value);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
fn spawn_chat_pii_redaction_runtime_config_refresh(
|
||
state: crate::AppState,
|
||
cache: ChatPiiRedactionRuntimeConfigCacheHandle,
|
||
guard: ChatPiiRedactionRuntimeConfigLoadGuard,
|
||
) {
|
||
tokio::spawn(async move {
|
||
let generation = guard.generation();
|
||
match load_chat_pii_redaction_runtime_config(&state).await {
|
||
Ok(value) => cache.insert_if_generation(generation, value),
|
||
Err(err) => {
|
||
tracing::warn!(
|
||
error = ?err,
|
||
"gateway failed to refresh chat pii redaction runtime config"
|
||
);
|
||
}
|
||
}
|
||
drop(guard);
|
||
});
|
||
}
|
||
|
||
async fn load_chat_pii_redaction_runtime_config(
|
||
state: &crate::AppState,
|
||
) -> Result<ChatPiiRedactionRuntimeConfig, GatewayError> {
|
||
let enabled = state
|
||
.read_system_config_json_value("module.chat_pii_redaction.enabled")
|
||
.await?
|
||
.as_ref()
|
||
.and_then(Value::as_bool)
|
||
.unwrap_or(false);
|
||
if !enabled {
|
||
return Ok(ChatPiiRedactionRuntimeConfig::disabled());
|
||
}
|
||
let mut config = ChatPiiRedactionRuntimeConfig::default();
|
||
config.enabled = true;
|
||
config.rules = parse_chat_pii_redaction_rules(
|
||
state
|
||
.read_system_config_json_value("module.chat_pii_redaction.rules")
|
||
.await?
|
||
.as_ref(),
|
||
)?;
|
||
config.ttl_seconds = state
|
||
.read_system_config_json_value("module.chat_pii_redaction.cache_ttl_seconds")
|
||
.await?
|
||
.as_ref()
|
||
.and_then(Value::as_u64)
|
||
.unwrap_or(config.ttl_seconds)
|
||
.max(1);
|
||
config.placeholder_prefix = state
|
||
.read_system_config_json_value("module.chat_pii_redaction.placeholder_prefix")
|
||
.await?
|
||
.as_ref()
|
||
.and_then(Value::as_str)
|
||
.and_then(normalize_sentinel_namespace)
|
||
.unwrap_or(config.placeholder_prefix);
|
||
Ok(config)
|
||
}
|
||
|
||
pub(crate) fn build_redaction_session_config(
|
||
hmac_key: impl Into<Vec<u8>>,
|
||
runtime_config: &ChatPiiRedactionRuntimeConfig,
|
||
now_unix_secs: u64,
|
||
) -> RedactionSessionConfig {
|
||
RedactionSessionConfig::new(hmac_key, runtime_config.ttl_seconds, now_unix_secs)
|
||
.with_sentinel_namespace(&runtime_config.placeholder_prefix)
|
||
.with_rules(runtime_config.rules.clone())
|
||
}
|
||
|
||
pub(crate) fn mask_chat_request_json_with_options(
|
||
body: &[u8],
|
||
config: RedactionSessionConfig,
|
||
options: MaskChatRequestOptions,
|
||
) -> MaskedChatRequest {
|
||
try_mask_chat_request_json_with_options(body, config, options)
|
||
.expect("chat redaction masking should be infallible")
|
||
}
|
||
|
||
pub(crate) fn mask_chat_request_json(
|
||
body: &[u8],
|
||
config: RedactionSessionConfig,
|
||
) -> MaskedChatRequest {
|
||
mask_chat_request_json_with_options(body, config, MaskChatRequestOptions::runtime())
|
||
}
|
||
|
||
pub(crate) fn try_mask_chat_request_json_with_options(
|
||
body: &[u8],
|
||
config: RedactionSessionConfig,
|
||
options: MaskChatRequestOptions,
|
||
) -> Result<MaskedChatRequest, RedactionMaskError> {
|
||
try_mask_chat_pii_request_json_with_options(
|
||
body,
|
||
ChatPiiRedactionRequestFormat::OpenAiChat,
|
||
config,
|
||
options,
|
||
)
|
||
}
|
||
|
||
pub(crate) async fn try_mask_chat_request_json_with_cache_options(
|
||
body: &[u8],
|
||
config: RedactionSessionConfig,
|
||
options: MaskChatRequestOptions,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<MaskedChatRequest, RedactionMaskError> {
|
||
try_mask_chat_pii_request_json_with_cache_options(
|
||
body,
|
||
ChatPiiRedactionRequestFormat::OpenAiChat,
|
||
config,
|
||
options,
|
||
cache,
|
||
)
|
||
.await
|
||
}
|
||
|
||
pub(crate) fn try_mask_chat_pii_request_json_with_options(
|
||
body: &[u8],
|
||
format: ChatPiiRedactionRequestFormat,
|
||
config: RedactionSessionConfig,
|
||
options: MaskChatRequestOptions,
|
||
) -> Result<MaskedChatRequest, RedactionMaskError> {
|
||
let mut session = RedactionSession::new(config);
|
||
session.apply_mask_options(options);
|
||
let Ok(mut value) = serde_json::from_slice::<Value>(body) else {
|
||
return Ok(MaskedChatRequest {
|
||
body: body.to_vec(),
|
||
session,
|
||
redacted: false,
|
||
});
|
||
};
|
||
|
||
session.set_collision_corpus(request_collision_corpus(format, &value));
|
||
reject_sensitive_opaque_reasoning_state(format, &value, &session)?;
|
||
let mut scan_state = RedactionScanState;
|
||
let redacted = mask_request_value(format, &mut value, &mut session, &mut scan_state, options)
|
||
.unwrap_or_else(|never| match never {});
|
||
|
||
if !redacted {
|
||
return Ok(MaskedChatRequest {
|
||
body: body.to_vec(),
|
||
session,
|
||
redacted,
|
||
});
|
||
}
|
||
|
||
let body = serde_json::to_vec(&value).unwrap_or_else(|_| body.to_vec());
|
||
Ok(MaskedChatRequest {
|
||
body,
|
||
session,
|
||
redacted,
|
||
})
|
||
}
|
||
|
||
pub(crate) async fn try_mask_chat_pii_request_json_with_cache_options(
|
||
body: &[u8],
|
||
format: ChatPiiRedactionRequestFormat,
|
||
config: RedactionSessionConfig,
|
||
options: MaskChatRequestOptions,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<MaskedChatRequest, RedactionMaskError> {
|
||
let mut session = RedactionSession::new(config);
|
||
session.apply_mask_options(options);
|
||
let Ok(mut value) = serde_json::from_slice::<Value>(body) else {
|
||
return Ok(MaskedChatRequest {
|
||
body: body.to_vec(),
|
||
session,
|
||
redacted: false,
|
||
});
|
||
};
|
||
|
||
session.set_collision_corpus(request_collision_corpus(format, &value));
|
||
reject_sensitive_opaque_reasoning_state(format, &value, &session)?;
|
||
let mut scan_state = RedactionScanState;
|
||
let redacted = mask_request_value_async(
|
||
format,
|
||
&mut value,
|
||
&mut session,
|
||
&mut scan_state,
|
||
options,
|
||
cache,
|
||
)
|
||
.await?;
|
||
|
||
if !redacted {
|
||
return Ok(MaskedChatRequest {
|
||
body: body.to_vec(),
|
||
session,
|
||
redacted,
|
||
});
|
||
}
|
||
|
||
let body = serde_json::to_vec(&value).unwrap_or_else(|_| body.to_vec());
|
||
Ok(MaskedChatRequest {
|
||
body,
|
||
session,
|
||
redacted,
|
||
})
|
||
}
|
||
|
||
pub(crate) async fn try_mask_chat_pii_request_value_with_cache_options(
|
||
body_json: &Value,
|
||
format: ChatPiiRedactionRequestFormat,
|
||
config: RedactionSessionConfig,
|
||
options: MaskChatRequestOptions,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<MaskedChatRequestValue, RedactionMaskError> {
|
||
let mut session = RedactionSession::new(config);
|
||
session.apply_mask_options(options);
|
||
let mut value = body_json.clone();
|
||
|
||
session.set_collision_corpus(request_collision_corpus(format, &value));
|
||
reject_sensitive_opaque_reasoning_state(format, &value, &session)?;
|
||
let mut scan_state = RedactionScanState;
|
||
let redacted = mask_request_value_async(
|
||
format,
|
||
&mut value,
|
||
&mut session,
|
||
&mut scan_state,
|
||
options,
|
||
cache,
|
||
)
|
||
.await?;
|
||
|
||
Ok(MaskedChatRequestValue {
|
||
body_json: redacted.then_some(value),
|
||
session,
|
||
redacted,
|
||
})
|
||
}
|
||
|
||
fn request_collision_corpus(format: ChatPiiRedactionRequestFormat, value: &Value) -> Vec<String> {
|
||
match format {
|
||
ChatPiiRedactionRequestFormat::OpenAiChat => value
|
||
.get("messages")
|
||
.and_then(Value::as_array)
|
||
.map(|messages| chat_message_collision_corpus(messages))
|
||
.unwrap_or_default(),
|
||
ChatPiiRedactionRequestFormat::OpenAiResponses => openai_responses_collision_corpus(value),
|
||
ChatPiiRedactionRequestFormat::OpenAiSearch => openai_search_collision_corpus(value),
|
||
ChatPiiRedactionRequestFormat::ClaudeMessages => claude_messages_collision_corpus(value),
|
||
}
|
||
}
|
||
|
||
fn reject_sensitive_opaque_reasoning_state(
|
||
format: ChatPiiRedactionRequestFormat,
|
||
value: &Value,
|
||
session: &RedactionSession,
|
||
) -> Result<(), RedactionMaskError> {
|
||
if !session.preserves_deepseek_opaque_reasoning_state() {
|
||
return Ok(());
|
||
}
|
||
if !matches!(
|
||
format,
|
||
ChatPiiRedactionRequestFormat::OpenAiResponses
|
||
| ChatPiiRedactionRequestFormat::OpenAiSearch
|
||
) {
|
||
return Ok(());
|
||
}
|
||
let Some(items) = value.get("input").and_then(Value::as_array) else {
|
||
return Ok(());
|
||
};
|
||
if items.iter().any(|item| {
|
||
item.as_object().is_some_and(|object| {
|
||
openai_responses_object_is_opaque_reasoning_state(object)
|
||
&& opaque_reasoning_state_has_sensitive_text(item, session)
|
||
})
|
||
}) {
|
||
return Err(RedactionMaskError::SensitiveOpaqueReasoningState);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn opaque_reasoning_state_has_sensitive_text(value: &Value, session: &RedactionSession) -> bool {
|
||
let Some(object) = value.as_object() else {
|
||
return false;
|
||
};
|
||
object.iter().any(|(key, value)| {
|
||
// The provider token is ciphertext/high-entropy state. Scanning it as
|
||
// user text would reject essentially every valid continuation (the
|
||
// generic API-key detector intentionally matches high-entropy data).
|
||
// Only this exact top-level protocol field is excluded; every other
|
||
// known or future string in the opaque item is scanned fail-closed.
|
||
key != "encrypted_content" && json_value_has_redaction_candidate(value, session)
|
||
})
|
||
}
|
||
|
||
fn json_value_has_redaction_candidate(value: &Value, session: &RedactionSession) -> bool {
|
||
match value {
|
||
Value::String(text) => session.text_has_redaction_candidate(text),
|
||
Value::Array(values) => values
|
||
.iter()
|
||
.any(|value| json_value_has_redaction_candidate(value, session)),
|
||
Value::Object(values) => values
|
||
.values()
|
||
.any(|value| json_value_has_redaction_candidate(value, session)),
|
||
_ => false,
|
||
}
|
||
}
|
||
|
||
fn mask_request_value(
|
||
format: ChatPiiRedactionRequestFormat,
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
options: MaskChatRequestOptions,
|
||
) -> Result<bool, Infallible> {
|
||
match format {
|
||
ChatPiiRedactionRequestFormat::OpenAiChat => {
|
||
mask_openai_chat_request_value(value, session, scan_state, options)
|
||
}
|
||
ChatPiiRedactionRequestFormat::OpenAiResponses => {
|
||
mask_openai_responses_request_value(value, session, scan_state)
|
||
}
|
||
ChatPiiRedactionRequestFormat::OpenAiSearch => {
|
||
mask_openai_search_request_value(value, session, scan_state)
|
||
}
|
||
ChatPiiRedactionRequestFormat::ClaudeMessages => {
|
||
mask_claude_messages_request_value(value, session, scan_state)
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn mask_request_value_async(
|
||
format: ChatPiiRedactionRequestFormat,
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
options: MaskChatRequestOptions,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
match format {
|
||
ChatPiiRedactionRequestFormat::OpenAiChat => {
|
||
mask_openai_chat_request_value_async(value, session, scan_state, options, cache).await
|
||
}
|
||
ChatPiiRedactionRequestFormat::OpenAiResponses => {
|
||
mask_openai_responses_request_value_async(value, session, scan_state, cache).await
|
||
}
|
||
ChatPiiRedactionRequestFormat::OpenAiSearch => {
|
||
mask_openai_search_request_value_async(value, session, scan_state, cache).await
|
||
}
|
||
ChatPiiRedactionRequestFormat::ClaudeMessages => {
|
||
mask_claude_messages_request_value_async(value, session, scan_state, cache).await
|
||
}
|
||
}
|
||
}
|
||
|
||
fn mask_openai_chat_request_value(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
options: MaskChatRequestOptions,
|
||
) -> Result<bool, Infallible> {
|
||
let Some(messages) = value.get_mut("messages").and_then(Value::as_array_mut) else {
|
||
return Ok(false);
|
||
};
|
||
|
||
let mut redacted = false;
|
||
let mut index = 0;
|
||
while index < messages.len() {
|
||
let message_redacted = mask_chat_message_value(&mut messages[index], session, scan_state)?;
|
||
redacted |= message_redacted;
|
||
index += 1;
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
|
||
async fn mask_openai_chat_request_value_async(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
options: MaskChatRequestOptions,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
let Some(messages) = value.get_mut("messages").and_then(Value::as_array_mut) else {
|
||
return Ok(false);
|
||
};
|
||
|
||
let mut redacted = false;
|
||
let mut index = 0;
|
||
while index < messages.len() {
|
||
let message_redacted =
|
||
mask_chat_message_value_async(&mut messages[index], session, scan_state, cache).await?;
|
||
redacted |= message_redacted;
|
||
index += 1;
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
|
||
fn chat_message_collision_corpus(messages: &[Value]) -> Vec<String> {
|
||
let mut corpus = Vec::new();
|
||
for message in messages {
|
||
collect_chat_message_collision_text(message, &mut corpus);
|
||
}
|
||
corpus
|
||
}
|
||
|
||
fn collect_chat_message_collision_text(message: &Value, corpus: &mut Vec<String>) {
|
||
let Some(message) = message.as_object() else {
|
||
return;
|
||
};
|
||
if let Some(content) = message.get("content") {
|
||
collect_chat_content_collision_text(content, corpus);
|
||
}
|
||
if let Some(tool_calls) = message.get("tool_calls").and_then(Value::as_array) {
|
||
for tool_call in tool_calls {
|
||
collect_tool_call_argument_collision_text(tool_call, corpus);
|
||
}
|
||
}
|
||
}
|
||
|
||
fn collect_chat_content_collision_text(content: &Value, corpus: &mut Vec<String>) {
|
||
match content {
|
||
Value::String(text) => corpus.push(text.clone()),
|
||
Value::Array(parts) => {
|
||
for part in parts {
|
||
collect_chat_content_part_collision_text(part, corpus);
|
||
}
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
fn collect_chat_content_part_collision_text(part: &Value, corpus: &mut Vec<String>) {
|
||
let Some(part) = part.as_object() else {
|
||
return;
|
||
};
|
||
if part.get("type").and_then(Value::as_str) != Some("text") {
|
||
return;
|
||
}
|
||
if let Some(text) = part.get("text").and_then(Value::as_str) {
|
||
corpus.push(text.to_string());
|
||
}
|
||
}
|
||
|
||
fn collect_tool_call_argument_collision_text(tool_call: &Value, corpus: &mut Vec<String>) {
|
||
if let Some(arguments) = tool_call
|
||
.get("function")
|
||
.and_then(Value::as_object)
|
||
.and_then(|function| function.get("arguments"))
|
||
.and_then(Value::as_str)
|
||
{
|
||
corpus.push(arguments.to_string());
|
||
}
|
||
}
|
||
|
||
fn claude_messages_collision_corpus(value: &Value) -> Vec<String> {
|
||
let mut corpus = Vec::new();
|
||
if let Some(system) = value.get("system") {
|
||
collect_claude_text_content_collision_text(system, &mut corpus);
|
||
}
|
||
if let Some(messages) = value.get("messages").and_then(Value::as_array) {
|
||
for message in messages {
|
||
if let Some(content) = message.get("content") {
|
||
collect_claude_text_content_collision_text(content, &mut corpus);
|
||
}
|
||
}
|
||
}
|
||
corpus
|
||
}
|
||
|
||
fn collect_claude_text_content_collision_text(value: &Value, corpus: &mut Vec<String>) {
|
||
match value {
|
||
Value::String(text) => corpus.push(text.clone()),
|
||
Value::Array(parts) => {
|
||
for part in parts {
|
||
collect_claude_content_part_collision_text(part, corpus);
|
||
}
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
fn collect_claude_content_part_collision_text(part: &Value, corpus: &mut Vec<String>) {
|
||
let Some(part) = part.as_object() else {
|
||
return;
|
||
};
|
||
match part.get("type").and_then(Value::as_str) {
|
||
Some("text") => {
|
||
if let Some(text) = part.get("text").and_then(Value::as_str) {
|
||
corpus.push(text.to_string());
|
||
}
|
||
}
|
||
Some("tool_result") => {
|
||
if let Some(content) = part.get("content") {
|
||
collect_claude_text_content_collision_text(content, corpus);
|
||
}
|
||
}
|
||
Some("tool_use") => {
|
||
if let Some(input) = part.get("input").and_then(Value::as_str) {
|
||
corpus.push(input.to_string());
|
||
}
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
fn openai_responses_collision_corpus(value: &Value) -> Vec<String> {
|
||
let mut corpus = Vec::new();
|
||
if let Some(instructions) = value.get("instructions").and_then(Value::as_str) {
|
||
corpus.push(instructions.to_string());
|
||
}
|
||
if let Some(input) = value.get("input") {
|
||
collect_openai_responses_input_collision_text(input, &mut corpus);
|
||
}
|
||
corpus
|
||
}
|
||
|
||
const OPENAI_SEARCH_COMMAND_TEXT_FIELDS: [(&str, &str); 4] = [
|
||
("search_query", "q"),
|
||
("image_query", "q"),
|
||
("find", "pattern"),
|
||
("weather", "location"),
|
||
];
|
||
|
||
fn openai_search_collision_corpus(value: &Value) -> Vec<String> {
|
||
let mut corpus = openai_responses_collision_corpus(value);
|
||
let Some(commands) = value.get("commands").and_then(Value::as_object) else {
|
||
return corpus;
|
||
};
|
||
for (command, field) in OPENAI_SEARCH_COMMAND_TEXT_FIELDS {
|
||
let Some(entries) = commands.get(command).and_then(Value::as_array) else {
|
||
continue;
|
||
};
|
||
for entry in entries {
|
||
if let Some(text) = entry.get(field).and_then(Value::as_str) {
|
||
corpus.push(text.to_string());
|
||
}
|
||
}
|
||
}
|
||
corpus
|
||
}
|
||
|
||
fn collect_openai_responses_input_collision_text(value: &Value, corpus: &mut Vec<String>) {
|
||
match value {
|
||
Value::String(text) => corpus.push(text.clone()),
|
||
Value::Array(items) => {
|
||
for item in items {
|
||
collect_openai_responses_input_item_collision_text(item, corpus);
|
||
}
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
fn collect_openai_responses_input_item_collision_text(item: &Value, corpus: &mut Vec<String>) {
|
||
let Some(item) = item.as_object() else {
|
||
return;
|
||
};
|
||
if let Some(content) = item.get("content") {
|
||
collect_openai_responses_content_collision_text(content, corpus);
|
||
}
|
||
if response_textish_type(item.get("type").and_then(Value::as_str)) {
|
||
if let Some(text) = item.get("text").and_then(Value::as_str) {
|
||
corpus.push(text.to_string());
|
||
}
|
||
}
|
||
if item.get("type").and_then(Value::as_str) == Some("function_call") {
|
||
if let Some(arguments) = item.get("arguments").and_then(Value::as_str) {
|
||
corpus.push(arguments.to_string());
|
||
}
|
||
}
|
||
}
|
||
|
||
fn collect_openai_responses_content_collision_text(content: &Value, corpus: &mut Vec<String>) {
|
||
match content {
|
||
Value::String(text) => corpus.push(text.clone()),
|
||
Value::Array(parts) => {
|
||
for part in parts {
|
||
collect_openai_responses_content_part_collision_text(part, corpus);
|
||
}
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
fn collect_openai_responses_content_part_collision_text(part: &Value, corpus: &mut Vec<String>) {
|
||
let Some(part) = part.as_object() else {
|
||
return;
|
||
};
|
||
if !response_textish_type(part.get("type").and_then(Value::as_str)) {
|
||
return;
|
||
}
|
||
if let Some(text) = part.get("text").and_then(Value::as_str) {
|
||
corpus.push(text.to_string());
|
||
}
|
||
}
|
||
|
||
fn response_textish_type(raw_type: Option<&str>) -> bool {
|
||
matches!(
|
||
raw_type,
|
||
Some("text" | "input_text" | "output_text" | "summary_text" | "reasoning_text")
|
||
)
|
||
}
|
||
|
||
fn mask_chat_message_value(
|
||
message: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
let Some(message) = message.as_object_mut() else {
|
||
return Ok(false);
|
||
};
|
||
|
||
let mut redacted = false;
|
||
if let Some(content) = message.get_mut("content") {
|
||
redacted |= mask_chat_content_value(content, session, scan_state)?;
|
||
}
|
||
if let Some(tool_calls) = message.get_mut("tool_calls").and_then(Value::as_array_mut) {
|
||
for tool_call in tool_calls {
|
||
redacted |= mask_tool_call_arguments(tool_call, session, scan_state)?;
|
||
}
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
|
||
fn mask_chat_content_value(
|
||
content: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
match content {
|
||
Value::String(text) => mask_json_string(text, session, scan_state),
|
||
Value::Array(parts) => {
|
||
let mut redacted = false;
|
||
for part in parts {
|
||
redacted |= mask_chat_content_part(part, session, scan_state)?;
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
_ => Ok(false),
|
||
}
|
||
}
|
||
|
||
fn mask_chat_content_part(
|
||
part: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
let Some(part) = part.as_object_mut() else {
|
||
return Ok(false);
|
||
};
|
||
if part.get("type").and_then(Value::as_str) != Some("text") {
|
||
return Ok(false);
|
||
}
|
||
let Some(Value::String(text)) = part.get_mut("text") else {
|
||
return Ok(false);
|
||
};
|
||
mask_json_string(text, session, scan_state)
|
||
}
|
||
|
||
fn mask_tool_call_arguments(
|
||
tool_call: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
let Some(function) = tool_call.get_mut("function").and_then(Value::as_object_mut) else {
|
||
return Ok(false);
|
||
};
|
||
let Some(Value::String(arguments)) = function.get_mut("arguments") else {
|
||
return Ok(false);
|
||
};
|
||
mask_json_string(arguments, session, scan_state)
|
||
}
|
||
|
||
fn mask_claude_messages_request_value(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
let mut redacted = false;
|
||
if let Some(system) = value.get_mut("system") {
|
||
redacted |= mask_claude_text_content_value(system, session, scan_state)?;
|
||
}
|
||
if let Some(messages) = value.get_mut("messages").and_then(Value::as_array_mut) {
|
||
for message in messages {
|
||
if let Some(content) = message.get_mut("content") {
|
||
redacted |= mask_claude_text_content_value(content, session, scan_state)?;
|
||
}
|
||
}
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
|
||
fn mask_claude_text_content_value(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
match value {
|
||
Value::String(text) => mask_json_string(text, session, scan_state),
|
||
Value::Array(parts) => {
|
||
let mut redacted = false;
|
||
for part in parts {
|
||
redacted |= mask_claude_content_part(part, session, scan_state)?;
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
_ => Ok(false),
|
||
}
|
||
}
|
||
|
||
fn mask_claude_content_part(
|
||
part: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
let Some(part) = part.as_object_mut() else {
|
||
return Ok(false);
|
||
};
|
||
match part.get("type").and_then(Value::as_str) {
|
||
Some("text") => {
|
||
let Some(Value::String(text)) = part.get_mut("text") else {
|
||
return Ok(false);
|
||
};
|
||
mask_json_string(text, session, scan_state)
|
||
}
|
||
Some("tool_result") => {
|
||
let Some(content) = part.get_mut("content") else {
|
||
return Ok(false);
|
||
};
|
||
mask_claude_text_content_value(content, session, scan_state)
|
||
}
|
||
Some("tool_use") => {
|
||
let Some(Value::String(input)) = part.get_mut("input") else {
|
||
return Ok(false);
|
||
};
|
||
mask_json_string(input, session, scan_state)
|
||
}
|
||
_ => Ok(false),
|
||
}
|
||
}
|
||
|
||
fn mask_openai_responses_request_value(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
let mut redacted = false;
|
||
if let Some(Value::String(instructions)) = value.get_mut("instructions") {
|
||
redacted |= mask_json_string(instructions, session, scan_state)?;
|
||
}
|
||
if let Some(input) = value.get_mut("input") {
|
||
redacted |= mask_openai_responses_input_value(input, session, scan_state)?;
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
|
||
fn mask_openai_search_request_value(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
let mut redacted = false;
|
||
if let Some(input) = value.get_mut("input") {
|
||
redacted |= mask_openai_responses_input_value(input, session, scan_state)?;
|
||
}
|
||
let Some(commands) = value.get_mut("commands").and_then(Value::as_object_mut) else {
|
||
return Ok(redacted);
|
||
};
|
||
for (command, field) in OPENAI_SEARCH_COMMAND_TEXT_FIELDS {
|
||
let Some(entries) = commands.get_mut(command).and_then(Value::as_array_mut) else {
|
||
continue;
|
||
};
|
||
for entry in entries {
|
||
if let Some(Value::String(text)) = entry.get_mut(field) {
|
||
redacted |= mask_json_string(text, session, scan_state)?;
|
||
}
|
||
}
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
|
||
fn mask_openai_responses_input_value(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
match value {
|
||
Value::String(text) => mask_json_string(text, session, scan_state),
|
||
Value::Array(items) => {
|
||
let mut redacted = false;
|
||
for item in items {
|
||
redacted |= mask_openai_responses_input_item(item, session, scan_state)?;
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
_ => Ok(false),
|
||
}
|
||
}
|
||
|
||
fn mask_openai_responses_input_item(
|
||
item: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
let Some(item) = item.as_object_mut() else {
|
||
return Ok(false);
|
||
};
|
||
// Some Responses-compatible providers bind the clear-text reasoning
|
||
// payload to sibling opaque continuation state. The top-level validation
|
||
// rejects such an item when its text contains sensitive data; a safe item
|
||
// can then remain byte-identical for provider replay. Unbound reasoning
|
||
// text still follows the ordinary PII policy below.
|
||
if session.preserves_deepseek_opaque_reasoning_state()
|
||
&& openai_responses_object_is_opaque_reasoning_state(item)
|
||
{
|
||
return Ok(false);
|
||
}
|
||
let mut redacted = false;
|
||
if let Some(content) = item.get_mut("content") {
|
||
redacted |= mask_openai_responses_content_value(content, session, scan_state)?;
|
||
}
|
||
if response_textish_type(item.get("type").and_then(Value::as_str)) {
|
||
if let Some(Value::String(text)) = item.get_mut("text") {
|
||
redacted |= mask_json_string(text, session, scan_state)?;
|
||
}
|
||
}
|
||
if item.get("type").and_then(Value::as_str) == Some("function_call") {
|
||
if let Some(Value::String(arguments)) = item.get_mut("arguments") {
|
||
redacted |= mask_json_string(arguments, session, scan_state)?;
|
||
}
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
|
||
fn mask_openai_responses_content_value(
|
||
content: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
match content {
|
||
Value::String(text) => mask_json_string(text, session, scan_state),
|
||
Value::Array(parts) => {
|
||
let mut redacted = false;
|
||
for part in parts {
|
||
redacted |= mask_openai_responses_content_part(part, session, scan_state)?;
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
_ => Ok(false),
|
||
}
|
||
}
|
||
|
||
fn mask_openai_responses_content_part(
|
||
part: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
let Some(part) = part.as_object_mut() else {
|
||
return Ok(false);
|
||
};
|
||
if !response_textish_type(part.get("type").and_then(Value::as_str)) {
|
||
return Ok(false);
|
||
}
|
||
let Some(Value::String(text)) = part.get_mut("text") else {
|
||
return Ok(false);
|
||
};
|
||
mask_json_string(text, session, scan_state)
|
||
}
|
||
|
||
fn mask_json_string(
|
||
text: &mut String,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
) -> Result<bool, Infallible> {
|
||
let redacted = session.redact_text_checked(text, scan_state)?;
|
||
if redacted.matches.is_empty() {
|
||
return Ok(false);
|
||
}
|
||
*text = redacted.text;
|
||
Ok(true)
|
||
}
|
||
|
||
async fn mask_chat_message_value_async(
|
||
message: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
let Some(message) = message.as_object_mut() else {
|
||
return Ok(false);
|
||
};
|
||
|
||
let mut redacted = false;
|
||
if let Some(content) = message.get_mut("content") {
|
||
redacted |= mask_chat_content_value_async(content, session, scan_state, cache).await?;
|
||
}
|
||
if let Some(tool_calls) = message.get_mut("tool_calls").and_then(Value::as_array_mut) {
|
||
for tool_call in tool_calls {
|
||
redacted |=
|
||
mask_tool_call_arguments_async(tool_call, session, scan_state, cache).await?;
|
||
}
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
|
||
async fn mask_chat_content_value_async(
|
||
content: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
match content {
|
||
Value::String(text) => mask_json_string_async(text, session, scan_state, cache).await,
|
||
Value::Array(parts) => {
|
||
let mut redacted = false;
|
||
for part in parts {
|
||
redacted |= mask_chat_content_part_async(part, session, scan_state, cache).await?;
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
_ => Ok(false),
|
||
}
|
||
}
|
||
|
||
async fn mask_chat_content_part_async(
|
||
part: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
let Some(part) = part.as_object_mut() else {
|
||
return Ok(false);
|
||
};
|
||
if part.get("type").and_then(Value::as_str) != Some("text") {
|
||
return Ok(false);
|
||
}
|
||
let Some(Value::String(text)) = part.get_mut("text") else {
|
||
return Ok(false);
|
||
};
|
||
mask_json_string_async(text, session, scan_state, cache).await
|
||
}
|
||
|
||
async fn mask_tool_call_arguments_async(
|
||
tool_call: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
let Some(function) = tool_call.get_mut("function").and_then(Value::as_object_mut) else {
|
||
return Ok(false);
|
||
};
|
||
let Some(Value::String(arguments)) = function.get_mut("arguments") else {
|
||
return Ok(false);
|
||
};
|
||
mask_json_string_async(arguments, session, scan_state, cache).await
|
||
}
|
||
|
||
async fn mask_claude_messages_request_value_async(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
let mut redacted = false;
|
||
if let Some(system) = value.get_mut("system") {
|
||
redacted |=
|
||
mask_claude_text_content_value_async(system, session, scan_state, cache).await?;
|
||
}
|
||
if let Some(messages) = value.get_mut("messages").and_then(Value::as_array_mut) {
|
||
for message in messages {
|
||
if let Some(content) = message.get_mut("content") {
|
||
redacted |=
|
||
mask_claude_text_content_value_async(content, session, scan_state, cache)
|
||
.await?;
|
||
}
|
||
}
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
|
||
async fn mask_claude_text_content_value_async(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
match value {
|
||
Value::String(text) => mask_json_string_async(text, session, scan_state, cache).await,
|
||
Value::Array(parts) => {
|
||
let mut redacted = false;
|
||
for part in parts {
|
||
redacted |=
|
||
mask_claude_content_part_async(part, session, scan_state, cache).await?;
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
_ => Ok(false),
|
||
}
|
||
}
|
||
|
||
async fn mask_claude_content_part_async(
|
||
part: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
let Some(part) = part.as_object_mut() else {
|
||
return Ok(false);
|
||
};
|
||
match part.get("type").and_then(Value::as_str) {
|
||
Some("text") => {
|
||
let Some(Value::String(text)) = part.get_mut("text") else {
|
||
return Ok(false);
|
||
};
|
||
mask_json_string_async(text, session, scan_state, cache).await
|
||
}
|
||
Some("tool_result") => {
|
||
let Some(content) = part.get_mut("content") else {
|
||
return Ok(false);
|
||
};
|
||
mask_claude_tool_result_content_value_async(content, session, scan_state, cache).await
|
||
}
|
||
Some("tool_use") => {
|
||
let Some(Value::String(input)) = part.get_mut("input") else {
|
||
return Ok(false);
|
||
};
|
||
mask_json_string_async(input, session, scan_state, cache).await
|
||
}
|
||
_ => Ok(false),
|
||
}
|
||
}
|
||
|
||
async fn mask_claude_tool_result_content_value_async(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
match value {
|
||
Value::String(text) => mask_json_string_async(text, session, scan_state, cache).await,
|
||
Value::Array(parts) => {
|
||
let mut redacted = false;
|
||
for part in parts {
|
||
let Some(part) = part.as_object_mut() else {
|
||
continue;
|
||
};
|
||
if part.get("type").and_then(Value::as_str) != Some("text") {
|
||
continue;
|
||
}
|
||
if let Some(Value::String(text)) = part.get_mut("text") {
|
||
redacted |= mask_json_string_async(text, session, scan_state, cache).await?;
|
||
}
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
_ => Ok(false),
|
||
}
|
||
}
|
||
|
||
async fn mask_openai_responses_request_value_async(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
let mut redacted = false;
|
||
if let Some(Value::String(instructions)) = value.get_mut("instructions") {
|
||
redacted |= mask_json_string_async(instructions, session, scan_state, cache).await?;
|
||
}
|
||
if let Some(input) = value.get_mut("input") {
|
||
redacted |=
|
||
mask_openai_responses_input_value_async(input, session, scan_state, cache).await?;
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
|
||
async fn mask_openai_search_request_value_async(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
let mut redacted = false;
|
||
if let Some(input) = value.get_mut("input") {
|
||
redacted |=
|
||
mask_openai_responses_input_value_async(input, session, scan_state, cache).await?;
|
||
}
|
||
let Some(commands) = value.get_mut("commands").and_then(Value::as_object_mut) else {
|
||
return Ok(redacted);
|
||
};
|
||
for (command, field) in OPENAI_SEARCH_COMMAND_TEXT_FIELDS {
|
||
let Some(entries) = commands.get_mut(command).and_then(Value::as_array_mut) else {
|
||
continue;
|
||
};
|
||
for entry in entries {
|
||
if let Some(Value::String(text)) = entry.get_mut(field) {
|
||
redacted |= mask_json_string_async(text, session, scan_state, cache).await?;
|
||
}
|
||
}
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
|
||
async fn mask_openai_responses_input_value_async(
|
||
value: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
match value {
|
||
Value::String(text) => mask_json_string_async(text, session, scan_state, cache).await,
|
||
Value::Array(items) => {
|
||
let mut redacted = false;
|
||
for item in items {
|
||
redacted |=
|
||
mask_openai_responses_input_item_async(item, session, scan_state, cache)
|
||
.await?;
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
_ => Ok(false),
|
||
}
|
||
}
|
||
|
||
async fn mask_openai_responses_input_item_async(
|
||
item: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
let Some(item) = item.as_object_mut() else {
|
||
return Ok(false);
|
||
};
|
||
if session.preserves_deepseek_opaque_reasoning_state()
|
||
&& openai_responses_object_is_opaque_reasoning_state(item)
|
||
{
|
||
return Ok(false);
|
||
}
|
||
let mut redacted = false;
|
||
if let Some(content) = item.get_mut("content") {
|
||
redacted |=
|
||
mask_openai_responses_content_value_async(content, session, scan_state, cache).await?;
|
||
}
|
||
if response_textish_type(item.get("type").and_then(Value::as_str)) {
|
||
if let Some(Value::String(text)) = item.get_mut("text") {
|
||
redacted |= mask_json_string_async(text, session, scan_state, cache).await?;
|
||
}
|
||
}
|
||
if item.get("type").and_then(Value::as_str) == Some("function_call") {
|
||
if let Some(Value::String(arguments)) = item.get_mut("arguments") {
|
||
redacted |= mask_json_string_async(arguments, session, scan_state, cache).await?;
|
||
}
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
|
||
async fn mask_openai_responses_content_value_async(
|
||
content: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
match content {
|
||
Value::String(text) => mask_json_string_async(text, session, scan_state, cache).await,
|
||
Value::Array(parts) => {
|
||
let mut redacted = false;
|
||
for part in parts {
|
||
redacted |=
|
||
mask_openai_responses_content_part_async(part, session, scan_state, cache)
|
||
.await?;
|
||
}
|
||
Ok(redacted)
|
||
}
|
||
_ => Ok(false),
|
||
}
|
||
}
|
||
|
||
async fn mask_openai_responses_content_part_async(
|
||
part: &mut Value,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
let Some(part) = part.as_object_mut() else {
|
||
return Ok(false);
|
||
};
|
||
if !response_textish_type(part.get("type").and_then(Value::as_str)) {
|
||
return Ok(false);
|
||
}
|
||
let Some(Value::String(text)) = part.get_mut("text") else {
|
||
return Ok(false);
|
||
};
|
||
mask_json_string_async(text, session, scan_state, cache).await
|
||
}
|
||
|
||
async fn mask_json_string_async(
|
||
text: &mut String,
|
||
session: &mut RedactionSession,
|
||
scan_state: &mut RedactionScanState,
|
||
cache: Option<&RedisRedactionMappingCache<'_>>,
|
||
) -> Result<bool, RedactionMaskError> {
|
||
let redacted = session
|
||
.redact_text_with_cache(text, scan_state, cache)
|
||
.await?;
|
||
if redacted.matches.is_empty() {
|
||
return Ok(false);
|
||
}
|
||
*text = redacted.text;
|
||
Ok(true)
|
||
}
|
||
|
||
pub(crate) struct RestoredSyncResponseBody {
|
||
pub(crate) body: Vec<u8>,
|
||
pub(crate) restored: bool,
|
||
}
|
||
|
||
impl fmt::Debug for RestoredSyncResponseBody {
|
||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
formatter
|
||
.debug_struct("RestoredSyncResponseBody")
|
||
.field("body_len", &self.body.len())
|
||
.field("restored", &self.restored)
|
||
.finish()
|
||
}
|
||
}
|
||
|
||
struct RestoredText {
|
||
text: String,
|
||
restored: bool,
|
||
}
|
||
|
||
pub(crate) fn restore_sync_response_body(
|
||
headers: &mut BTreeMap<String, String>,
|
||
body: &[u8],
|
||
session: &RedactionSession,
|
||
) -> Result<RestoredSyncResponseBody, GatewayError> {
|
||
if session.mapping_count() == 0 {
|
||
return Ok(RestoredSyncResponseBody {
|
||
body: body.to_vec(),
|
||
restored: false,
|
||
});
|
||
}
|
||
|
||
ensure_identity_response_encoding(headers)?;
|
||
|
||
let restored = if response_body_is_json(headers, body) {
|
||
restore_json_response_body(body, session)?
|
||
} else {
|
||
restore_text_response_body(body, session)
|
||
};
|
||
if restored.restored {
|
||
set_content_length(headers, restored.body.len());
|
||
}
|
||
Ok(restored)
|
||
}
|
||
|
||
fn ensure_identity_response_encoding(
|
||
headers: &BTreeMap<String, String>,
|
||
) -> Result<(), GatewayError> {
|
||
for encoding in headers
|
||
.iter()
|
||
.filter(|(header_name, _)| header_name.eq_ignore_ascii_case("content-encoding"))
|
||
.map(|(_, value)| value)
|
||
{
|
||
if encoding.trim().is_empty() || encoding.trim().eq_ignore_ascii_case("identity") {
|
||
continue;
|
||
}
|
||
return Err(GatewayError::Internal(
|
||
"redaction response restoration does not support encoded response bodies".to_string(),
|
||
));
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn response_body_is_json(headers: &BTreeMap<String, String>, body: &[u8]) -> bool {
|
||
if header_value(headers, "content-type").is_some_and(content_type_is_json) {
|
||
return true;
|
||
}
|
||
body.iter()
|
||
.copied()
|
||
.find(|byte| !byte.is_ascii_whitespace())
|
||
.is_some_and(|byte| matches!(byte, b'{' | b'['))
|
||
}
|
||
|
||
fn content_type_is_json(content_type: &str) -> bool {
|
||
content_type
|
||
.split(';')
|
||
.next()
|
||
.map(str::trim)
|
||
.map(str::to_ascii_lowercase)
|
||
.is_some_and(|mime| mime == "application/json" || mime.ends_with("+json"))
|
||
}
|
||
|
||
fn restore_json_response_body(
|
||
body: &[u8],
|
||
session: &RedactionSession,
|
||
) -> Result<RestoredSyncResponseBody, GatewayError> {
|
||
let Ok(mut value) = serde_json::from_slice::<Value>(body) else {
|
||
return Ok(restore_text_response_body(body, session));
|
||
};
|
||
if !restore_json_strings(&mut value, session) {
|
||
return Ok(RestoredSyncResponseBody {
|
||
body: body.to_vec(),
|
||
restored: false,
|
||
});
|
||
}
|
||
let body = serde_json::to_vec(&value).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||
Ok(RestoredSyncResponseBody {
|
||
body,
|
||
restored: true,
|
||
})
|
||
}
|
||
|
||
/// 递归把 JSON 里的占位符换回真实值，只认本 `session` 记录过的映射。
|
||
///
|
||
/// 同步响应体（[`restore_sync_response_body`]）和 Responses WebSocket 的
|
||
/// provider 事件帧（`handlers::proxy::websocket::responses::redaction`）共用它，
|
||
/// 两边因此保持同一套还原语义：未映射的占位符原样保留，`type` / `model` / `id`
|
||
/// 这类协议字段虽然也被遍历，但它们不可能包含本 session 派生出的 sentinel，
|
||
/// 所以不会被改写。
|
||
///
|
||
/// Provider-owned Responses reasoning state is the exception. DeepSeek-style
|
||
/// items bind `reasoning_text` to `encrypted_content` and require both to be
|
||
/// replayed unchanged. Restoring a sentinel in only the text half would make
|
||
/// the client return a different continuation item, so those objects/events
|
||
/// stay opaque even when another response field is restored.
|
||
pub(crate) fn restore_json_strings(value: &mut Value, session: &RedactionSession) -> bool {
|
||
match value {
|
||
Value::String(text) => {
|
||
let restored = session.restore_text(text);
|
||
if !restored.restored {
|
||
return false;
|
||
}
|
||
*text = restored.text;
|
||
true
|
||
}
|
||
Value::Array(values) => {
|
||
let mut restored = false;
|
||
for value in values {
|
||
restored = restore_json_strings(value, session) || restored;
|
||
}
|
||
restored
|
||
}
|
||
Value::Object(values) => {
|
||
if session.preserves_deepseek_opaque_reasoning_state()
|
||
&& openai_responses_object_is_opaque_reasoning_state(values)
|
||
{
|
||
return false;
|
||
}
|
||
let mut restored = false;
|
||
for value in values.values_mut() {
|
||
restored = restore_json_strings(value, session) || restored;
|
||
}
|
||
restored
|
||
}
|
||
_ => false,
|
||
}
|
||
}
|
||
|
||
fn openai_responses_object_is_opaque_reasoning_state(
|
||
object: &serde_json::Map<String, Value>,
|
||
) -> bool {
|
||
let Some(item_type) = object.get("type").and_then(Value::as_str) else {
|
||
return false;
|
||
};
|
||
// Never treat a `reasoning_text` content part or delta/done event name by
|
||
// itself as proof of opaque state: ordinary OpenAI reasoning text may
|
||
// contain mask placeholders that still need client-side restoration. The
|
||
// binding evidence is the parent reasoning item carrying provider-owned
|
||
// encrypted continuation state. Returning here keeps that whole object,
|
||
// including its nested reasoning_text parts, value-for-value unchanged.
|
||
let id_is_absent_or_empty = object
|
||
.get("id")
|
||
.is_none_or(|id| id.is_null() || id.as_str().is_some_and(|value| value.trim().is_empty()));
|
||
item_type == "reasoning"
|
||
&& id_is_absent_or_empty
|
||
&& object
|
||
.get("encrypted_content")
|
||
.and_then(Value::as_str)
|
||
.is_some_and(|state| !state.trim().is_empty())
|
||
&& object
|
||
.get("content")
|
||
.and_then(Value::as_array)
|
||
.is_some_and(|content| {
|
||
content.iter().any(|part| {
|
||
part.get("type").and_then(Value::as_str) == Some("reasoning_text")
|
||
&& part.get("text").is_some_and(Value::is_string)
|
||
})
|
||
})
|
||
}
|
||
|
||
fn restore_text_response_body(body: &[u8], session: &RedactionSession) -> RestoredSyncResponseBody {
|
||
let Ok(text) = std::str::from_utf8(body) else {
|
||
return RestoredSyncResponseBody {
|
||
body: body.to_vec(),
|
||
restored: false,
|
||
};
|
||
};
|
||
let restored = session.restore_text(text);
|
||
RestoredSyncResponseBody {
|
||
body: restored.text.into_bytes(),
|
||
restored: restored.restored,
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Copy, Eq, PartialEq)]
|
||
enum StreamRestoreMode {
|
||
Sse,
|
||
Text,
|
||
}
|
||
|
||
pub(crate) struct StreamingResponseRestorer<'a> {
|
||
session: &'a RedactionSession,
|
||
matcher: SentinelMatcher<'a>,
|
||
mode: StreamRestoreMode,
|
||
max_sentinel_len: usize,
|
||
text_carry: Vec<u8>,
|
||
sse_line_buffer: Vec<u8>,
|
||
sse_json_event_lines: Vec<Vec<u8>>,
|
||
}
|
||
|
||
impl<'a> StreamingResponseRestorer<'a> {
|
||
pub(crate) fn new(
|
||
headers: &BTreeMap<String, String>,
|
||
session: &'a RedactionSession,
|
||
) -> Result<Self, GatewayError> {
|
||
if session.mapping_count() > 0 {
|
||
ensure_identity_response_encoding(headers)?;
|
||
}
|
||
let mode = if header_value(headers, "content-type").is_some_and(content_type_is_sse) {
|
||
StreamRestoreMode::Sse
|
||
} else {
|
||
StreamRestoreMode::Text
|
||
};
|
||
Ok(Self::with_mode(session, mode))
|
||
}
|
||
|
||
#[cfg(test)]
|
||
fn for_sse(session: &'a RedactionSession) -> Self {
|
||
Self::with_mode(session, StreamRestoreMode::Sse)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
fn for_text(session: &'a RedactionSession) -> Self {
|
||
Self::with_mode(session, StreamRestoreMode::Text)
|
||
}
|
||
|
||
fn with_mode(session: &'a RedactionSession, mode: StreamRestoreMode) -> Self {
|
||
let max_sentinel_len = session
|
||
.mappings()
|
||
.map(|mapping| mapping.sentinel.len())
|
||
.max()
|
||
.unwrap_or_default();
|
||
Self {
|
||
session,
|
||
matcher: SentinelMatcher::new(session),
|
||
mode,
|
||
max_sentinel_len,
|
||
text_carry: Vec::new(),
|
||
sse_line_buffer: Vec::new(),
|
||
sse_json_event_lines: Vec::new(),
|
||
}
|
||
}
|
||
|
||
pub(crate) fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, GatewayError> {
|
||
if self.max_sentinel_len == 0 {
|
||
return Ok(chunk.to_vec());
|
||
}
|
||
match self.mode {
|
||
StreamRestoreMode::Sse => self.push_sse_chunk(chunk),
|
||
StreamRestoreMode::Text => Ok(self.push_text_chunk(chunk)),
|
||
}
|
||
}
|
||
|
||
pub(crate) fn finish(&mut self) -> Result<Vec<u8>, GatewayError> {
|
||
if self.max_sentinel_len == 0 {
|
||
return Ok(Vec::new());
|
||
}
|
||
match self.mode {
|
||
StreamRestoreMode::Sse => self.finish_sse(),
|
||
StreamRestoreMode::Text => Ok(self.flush_text_carry()),
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
fn max_text_carry_len(&self) -> usize {
|
||
self.max_sentinel_len.saturating_sub(1)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
fn pending_text_carry_len(&self) -> usize {
|
||
self.text_carry.len()
|
||
}
|
||
|
||
fn push_text_chunk(&mut self, chunk: &[u8]) -> Vec<u8> {
|
||
self.text_carry.extend_from_slice(chunk);
|
||
self.restore_available_text(false)
|
||
}
|
||
|
||
fn flush_text_carry(&mut self) -> Vec<u8> {
|
||
self.restore_available_text(true)
|
||
}
|
||
|
||
fn restore_available_text(&mut self, flush: bool) -> Vec<u8> {
|
||
let scan_limit = if flush {
|
||
self.text_carry.len()
|
||
} else {
|
||
self.text_carry
|
||
.len()
|
||
.saturating_sub(self.max_sentinel_len.saturating_sub(1))
|
||
};
|
||
let mut output = Vec::with_capacity(scan_limit);
|
||
let mut index = 0;
|
||
while index < scan_limit {
|
||
if let Some(mapping) = self.matcher.matching_mapping_at(&self.text_carry, index) {
|
||
output.extend_from_slice(mapping.original.as_bytes());
|
||
index += mapping.sentinel.len();
|
||
} else {
|
||
output.push(self.text_carry[index]);
|
||
index += 1;
|
||
}
|
||
}
|
||
self.text_carry.drain(..index);
|
||
output
|
||
}
|
||
|
||
fn push_sse_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, GatewayError> {
|
||
self.sse_line_buffer.extend_from_slice(chunk);
|
||
let mut output = Vec::new();
|
||
while let Some(line_end) = self.sse_line_buffer.iter().position(|byte| *byte == b'\n') {
|
||
let line = self.sse_line_buffer.drain(..=line_end).collect::<Vec<_>>();
|
||
self.push_sse_line(line, &mut output)?;
|
||
}
|
||
Ok(output)
|
||
}
|
||
|
||
fn finish_sse(&mut self) -> Result<Vec<u8>, GatewayError> {
|
||
let mut output = Vec::new();
|
||
if !self.sse_line_buffer.is_empty() {
|
||
let line = std::mem::take(&mut self.sse_line_buffer);
|
||
self.push_sse_line(line, &mut output)?;
|
||
}
|
||
self.flush_sse_json_event(&mut output)?;
|
||
Ok(output)
|
||
}
|
||
|
||
fn push_sse_line(&mut self, line: Vec<u8>, output: &mut Vec<u8>) -> Result<(), GatewayError> {
|
||
if !self.sse_json_event_lines.is_empty() {
|
||
if sse_line_is_blank(&line) {
|
||
self.flush_sse_json_event(output)?;
|
||
output.extend_from_slice(&line);
|
||
} else {
|
||
self.sse_json_event_lines.push(line);
|
||
}
|
||
return Ok(());
|
||
}
|
||
|
||
if sse_line_is_blank(&line) {
|
||
output.extend_from_slice(&line);
|
||
return Ok(());
|
||
}
|
||
|
||
let Some(value) = sse_data_line_value(&line) else {
|
||
output.extend_from_slice(&line);
|
||
return Ok(());
|
||
};
|
||
if sse_data_value_is_done(value) {
|
||
output.extend_from_slice(&line);
|
||
return Ok(());
|
||
}
|
||
if sse_data_value_may_be_json(value) {
|
||
self.sse_json_event_lines.push(line);
|
||
return Ok(());
|
||
}
|
||
output.extend(self.restore_sse_text_data_line(&line));
|
||
Ok(())
|
||
}
|
||
|
||
fn flush_sse_json_event(&mut self, output: &mut Vec<u8>) -> Result<(), GatewayError> {
|
||
if self.sse_json_event_lines.is_empty() {
|
||
return Ok(());
|
||
}
|
||
|
||
let event_lines = std::mem::take(&mut self.sse_json_event_lines);
|
||
let Some(payload) = sse_event_data_payload(&event_lines) else {
|
||
output.extend(event_lines.into_iter().flatten());
|
||
return Ok(());
|
||
};
|
||
if payload.trim() == "[DONE]" {
|
||
output.extend(event_lines.into_iter().flatten());
|
||
return Ok(());
|
||
}
|
||
|
||
let Ok(mut value) = serde_json::from_str::<Value>(&payload) else {
|
||
for line in event_lines {
|
||
if sse_data_line_value(&line).is_some() {
|
||
output.extend(self.restore_sse_text_data_line(&line));
|
||
} else {
|
||
output.extend_from_slice(&line);
|
||
}
|
||
}
|
||
return Ok(());
|
||
};
|
||
|
||
if !restore_json_strings(&mut value, self.session) {
|
||
output.extend(event_lines.into_iter().flatten());
|
||
return Ok(());
|
||
}
|
||
|
||
let restored_payload =
|
||
serde_json::to_string(&value).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||
let mut wrote_data_line = false;
|
||
for line in event_lines {
|
||
if sse_data_line_value(&line).is_some() {
|
||
if !wrote_data_line {
|
||
output.extend(replace_sse_data_line_value(
|
||
&line,
|
||
restored_payload.as_bytes(),
|
||
));
|
||
wrote_data_line = true;
|
||
}
|
||
} else {
|
||
output.extend_from_slice(&line);
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn restore_sse_text_data_line(&self, line: &[u8]) -> Vec<u8> {
|
||
let Some(range) = sse_data_line_value_range(line) else {
|
||
return line.to_vec();
|
||
};
|
||
let (body, ending) = split_sse_line_ending(line);
|
||
let restored = restore_known_sentinels_in_bytes(&body[range.clone()], self.session);
|
||
let mut output = Vec::with_capacity(line.len());
|
||
output.extend_from_slice(&body[..range.start]);
|
||
output.extend(restored);
|
||
output.extend_from_slice(ending);
|
||
output
|
||
}
|
||
}
|
||
|
||
enum SentinelMatcher<'a> {
|
||
Direct(Vec<&'a RedactionMapping>),
|
||
Trie(SentinelTrie<'a>),
|
||
}
|
||
|
||
impl<'a> SentinelMatcher<'a> {
|
||
fn new(session: &'a RedactionSession) -> Self {
|
||
let mut mappings = session.mappings().collect::<Vec<_>>();
|
||
mappings.sort_by_key(|mapping| std::cmp::Reverse(mapping.sentinel.len()));
|
||
if mappings.len() <= DIRECT_RESTORE_SENTINEL_LIMIT {
|
||
Self::Direct(mappings)
|
||
} else {
|
||
Self::Trie(SentinelTrie::new(mappings))
|
||
}
|
||
}
|
||
|
||
fn matching_mapping_at(&self, input: &[u8], index: usize) -> Option<&'a RedactionMapping> {
|
||
match self {
|
||
Self::Direct(mappings) => mappings.iter().copied().find(|mapping| {
|
||
input
|
||
.get(index..)
|
||
.is_some_and(|remaining| remaining.starts_with(mapping.sentinel.as_bytes()))
|
||
}),
|
||
Self::Trie(trie) => trie.matching_mapping_at(input, index),
|
||
}
|
||
}
|
||
}
|
||
|
||
struct SentinelTrie<'a> {
|
||
nodes: Vec<SentinelTrieNode<'a>>,
|
||
}
|
||
|
||
struct SentinelTrieNode<'a> {
|
||
children: HashMap<u8, usize>,
|
||
mapping: Option<&'a RedactionMapping>,
|
||
}
|
||
|
||
impl<'a> SentinelTrie<'a> {
|
||
fn new(mappings: Vec<&'a RedactionMapping>) -> Self {
|
||
let mut trie = Self {
|
||
nodes: vec![SentinelTrieNode {
|
||
children: HashMap::new(),
|
||
mapping: None,
|
||
}],
|
||
};
|
||
for mapping in mappings {
|
||
trie.insert(mapping);
|
||
}
|
||
trie
|
||
}
|
||
|
||
fn insert(&mut self, mapping: &'a RedactionMapping) {
|
||
let mut node_index = 0;
|
||
for byte in mapping.sentinel.bytes() {
|
||
let next_index = if let Some(next_index) = self.nodes[node_index].children.get(&byte) {
|
||
*next_index
|
||
} else {
|
||
self.nodes.push(SentinelTrieNode {
|
||
children: HashMap::new(),
|
||
mapping: None,
|
||
});
|
||
let next_index = self.nodes.len() - 1;
|
||
self.nodes[node_index].children.insert(byte, next_index);
|
||
next_index
|
||
};
|
||
node_index = next_index;
|
||
}
|
||
self.nodes[node_index].mapping = Some(mapping);
|
||
}
|
||
|
||
fn matching_mapping_at(&self, input: &[u8], index: usize) -> Option<&'a RedactionMapping> {
|
||
let mut node_index = 0;
|
||
let mut best = None;
|
||
for byte in input.get(index..).unwrap_or_default() {
|
||
let Some(next_index) = self.nodes[node_index].children.get(byte).copied() else {
|
||
break;
|
||
};
|
||
node_index = next_index;
|
||
if let Some(mapping) = self.nodes[node_index].mapping {
|
||
best = Some(mapping);
|
||
}
|
||
}
|
||
best
|
||
}
|
||
}
|
||
|
||
fn restore_text_direct_longest_first(input: &str, session: &RedactionSession) -> RestoredText {
|
||
let mut mappings = session.mappings().collect::<Vec<_>>();
|
||
mappings.sort_by_key(|mapping| std::cmp::Reverse(mapping.sentinel.len()));
|
||
let mut text = input.to_string();
|
||
let mut restored = false;
|
||
for mapping in mappings {
|
||
if text.contains(&mapping.sentinel) {
|
||
text = text.replace(&mapping.sentinel, &mapping.original);
|
||
restored = true;
|
||
}
|
||
}
|
||
RestoredText { text, restored }
|
||
}
|
||
|
||
fn restore_text_with_matcher(input: &str, matcher: &SentinelMatcher<'_>) -> RestoredText {
|
||
let input_bytes = input.as_bytes();
|
||
let mut output = Vec::with_capacity(input_bytes.len());
|
||
let mut restored = false;
|
||
let mut index = 0;
|
||
while index < input_bytes.len() {
|
||
if let Some(mapping) = matcher.matching_mapping_at(input_bytes, index) {
|
||
output.extend_from_slice(mapping.original.as_bytes());
|
||
index += mapping.sentinel.len();
|
||
restored = true;
|
||
} else {
|
||
output.push(input_bytes[index]);
|
||
index += 1;
|
||
}
|
||
}
|
||
let text = String::from_utf8(output).unwrap_or_else(|_| input.to_string());
|
||
RestoredText { text, restored }
|
||
}
|
||
|
||
fn restore_known_sentinels_in_bytes(input: &[u8], session: &RedactionSession) -> Vec<u8> {
|
||
let matcher = SentinelMatcher::new(session);
|
||
let mut output = Vec::with_capacity(input.len());
|
||
let mut index = 0;
|
||
while index < input.len() {
|
||
if let Some(mapping) = matcher.matching_mapping_at(input, index) {
|
||
output.extend_from_slice(mapping.original.as_bytes());
|
||
index += mapping.sentinel.len();
|
||
} else {
|
||
output.push(input[index]);
|
||
index += 1;
|
||
}
|
||
}
|
||
output
|
||
}
|
||
|
||
fn content_type_is_sse(content_type: &str) -> bool {
|
||
content_type
|
||
.split(';')
|
||
.next()
|
||
.map(str::trim)
|
||
.is_some_and(|mime| mime.eq_ignore_ascii_case("text/event-stream"))
|
||
}
|
||
|
||
fn split_sse_line_ending(line: &[u8]) -> (&[u8], &[u8]) {
|
||
if line.ends_with(b"\r\n") {
|
||
(&line[..line.len() - 2], &line[line.len() - 2..])
|
||
} else if line.ends_with(b"\n") {
|
||
(&line[..line.len() - 1], &line[line.len() - 1..])
|
||
} else {
|
||
(line, &[])
|
||
}
|
||
}
|
||
|
||
fn sse_line_is_blank(line: &[u8]) -> bool {
|
||
split_sse_line_ending(line).0.is_empty()
|
||
}
|
||
|
||
fn sse_data_line_value_range(line: &[u8]) -> Option<std::ops::Range<usize>> {
|
||
let (body, _) = split_sse_line_ending(line);
|
||
if !body.starts_with(b"data:") {
|
||
return None;
|
||
}
|
||
let start = if body.get(5) == Some(&b' ') { 6 } else { 5 };
|
||
Some(start..body.len())
|
||
}
|
||
|
||
fn sse_data_line_value(line: &[u8]) -> Option<&[u8]> {
|
||
let (body, _) = split_sse_line_ending(line);
|
||
sse_data_line_value_range(line).map(|range| &body[range])
|
||
}
|
||
|
||
fn sse_data_value_is_done(value: &[u8]) -> bool {
|
||
std::str::from_utf8(value)
|
||
.map(str::trim)
|
||
.is_ok_and(|value| value == "[DONE]")
|
||
}
|
||
|
||
fn sse_data_value_may_be_json(value: &[u8]) -> bool {
|
||
value
|
||
.iter()
|
||
.copied()
|
||
.find(|byte| !byte.is_ascii_whitespace())
|
||
.is_some_and(|byte| matches!(byte, b'{' | b'['))
|
||
}
|
||
|
||
fn sse_event_data_payload(lines: &[Vec<u8>]) -> Option<String> {
|
||
let mut payload = String::new();
|
||
let mut saw_data = false;
|
||
for line in lines {
|
||
let Some(value) = sse_data_line_value(line) else {
|
||
continue;
|
||
};
|
||
if saw_data {
|
||
payload.push('\n');
|
||
}
|
||
payload.push_str(std::str::from_utf8(value).ok()?);
|
||
saw_data = true;
|
||
}
|
||
saw_data.then_some(payload)
|
||
}
|
||
|
||
fn replace_sse_data_line_value(line: &[u8], value: &[u8]) -> Vec<u8> {
|
||
let Some(range) = sse_data_line_value_range(line) else {
|
||
return line.to_vec();
|
||
};
|
||
let (body, ending) = split_sse_line_ending(line);
|
||
let mut output = Vec::with_capacity(body.len() - range.len() + value.len() + ending.len());
|
||
output.extend_from_slice(&body[..range.start]);
|
||
output.extend_from_slice(value);
|
||
output.extend_from_slice(ending);
|
||
output
|
||
}
|
||
|
||
fn header_value<'a>(headers: &'a BTreeMap<String, String>, name: &str) -> Option<&'a str> {
|
||
headers
|
||
.iter()
|
||
.find(|(header_name, _)| header_name.eq_ignore_ascii_case(name))
|
||
.map(|(_, value)| value.as_str())
|
||
}
|
||
|
||
fn set_content_length(headers: &mut BTreeMap<String, String>, body_len: usize) {
|
||
headers.retain(|name, _| !name.eq_ignore_ascii_case("content-length"));
|
||
headers.insert("content-length".to_string(), body_len.to_string());
|
||
}
|
||
|
||
pub(crate) struct RedactedText {
|
||
pub(crate) text: String,
|
||
pub(crate) matches: Vec<RedactionMatch>,
|
||
}
|
||
|
||
impl fmt::Debug for RedactedText {
|
||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
formatter
|
||
.debug_struct("RedactedText")
|
||
.field("text_len", &self.text.len())
|
||
.field("match_count", &self.matches.len())
|
||
.finish()
|
||
}
|
||
}
|
||
|
||
pub(crate) struct RedactionMatch {
|
||
pub(crate) rule_label: String,
|
||
pub(crate) kind: Option<RedactionKind>,
|
||
pub(crate) start: usize,
|
||
pub(crate) end: usize,
|
||
pub(crate) original: String,
|
||
pub(crate) sentinel: String,
|
||
}
|
||
|
||
impl fmt::Debug for RedactionMatch {
|
||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
formatter
|
||
.debug_struct("RedactionMatch")
|
||
.field("rule_label", &self.rule_label)
|
||
.field("kind", &self.kind)
|
||
.field("start", &self.start)
|
||
.field("end", &self.end)
|
||
.field("original_len", &self.original.len())
|
||
.field("sentinel", &redacted_sentinel_debug(&self.sentinel))
|
||
.finish()
|
||
}
|
||
}
|
||
|
||
#[derive(Clone)]
|
||
pub(crate) struct RedactionMapping {
|
||
pub(crate) rule_label: String,
|
||
pub(crate) kind: Option<RedactionKind>,
|
||
pub(crate) original: String,
|
||
pub(crate) normalized_value: String,
|
||
pub(crate) sentinel: String,
|
||
pub(crate) bucket: u64,
|
||
pub(crate) created_at_unix_secs: u64,
|
||
pub(crate) expires_at_unix_secs: u64,
|
||
}
|
||
|
||
impl fmt::Debug for RedactionMapping {
|
||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
formatter
|
||
.debug_struct("RedactionMapping")
|
||
.field("rule_label", &self.rule_label)
|
||
.field("kind", &self.kind)
|
||
.field("original_len", &self.original.len())
|
||
.field("normalized_value_len", &self.normalized_value.len())
|
||
.field("sentinel", &redacted_sentinel_debug(&self.sentinel))
|
||
.field("bucket", &self.bucket)
|
||
.field("created_at_unix_secs", &self.created_at_unix_secs)
|
||
.field("expires_at_unix_secs", &self.expires_at_unix_secs)
|
||
.finish()
|
||
}
|
||
}
|
||
|
||
#[derive(Deserialize, Serialize)]
|
||
pub(crate) struct RedactionCacheRecord {
|
||
pub(crate) rule_label: String,
|
||
pub(crate) kind: Option<RedactionKind>,
|
||
pub(crate) sentinel: String,
|
||
pub(crate) bucket: u64,
|
||
pub(crate) expires_at_unix_secs: u64,
|
||
}
|
||
|
||
impl RedactionCacheRecord {
|
||
pub(crate) fn from_mapping(mapping: &RedactionMapping) -> Self {
|
||
Self {
|
||
rule_label: mapping.rule_label.clone(),
|
||
kind: mapping.kind,
|
||
sentinel: mapping.sentinel.clone(),
|
||
bucket: mapping.bucket,
|
||
expires_at_unix_secs: mapping.expires_at_unix_secs,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl fmt::Debug for RedactionCacheRecord {
|
||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||
formatter
|
||
.debug_struct("RedactionCacheRecord")
|
||
.field("rule_label", &self.rule_label)
|
||
.field("kind", &self.kind)
|
||
.field("sentinel", &redacted_sentinel_debug(&self.sentinel))
|
||
.field("bucket", &self.bucket)
|
||
.field("expires_at_unix_secs", &self.expires_at_unix_secs)
|
||
.finish()
|
||
}
|
||
}
|
||
|
||
pub(crate) struct RedisRedactionMappingCache<'a> {
|
||
runtime_state: &'a RuntimeState,
|
||
key_prefix: String,
|
||
}
|
||
|
||
impl<'a> RedisRedactionMappingCache<'a> {
|
||
pub(crate) fn new(runtime_state: &'a RuntimeState) -> Self {
|
||
Self {
|
||
runtime_state,
|
||
key_prefix: "privacy:redaction:mapping".to_string(),
|
||
}
|
||
}
|
||
|
||
pub(crate) fn with_key_prefix(
|
||
runtime_state: &'a RuntimeState,
|
||
key_prefix: impl Into<String>,
|
||
) -> Self {
|
||
Self {
|
||
runtime_state,
|
||
key_prefix: key_prefix.into(),
|
||
}
|
||
}
|
||
|
||
pub(crate) async fn store(
|
||
&self,
|
||
record: &RedactionCacheRecord,
|
||
normalized_value: &str,
|
||
ttl_seconds: u64,
|
||
) -> Result<(), DataLayerError> {
|
||
if record.sentinel.len() > MAX_CACHE_SENTINEL_BYTES {
|
||
return Err(DataLayerError::InvalidInput(
|
||
"redaction cache sentinel is too large".to_string(),
|
||
));
|
||
}
|
||
let value = serde_json::to_string(record)
|
||
.map_err(|err| DataLayerError::InvalidInput(err.to_string()))?;
|
||
if value.len() > MAX_CACHE_RECORD_BYTES {
|
||
return Err(DataLayerError::InvalidInput(
|
||
"redaction cache record is too large".to_string(),
|
||
));
|
||
}
|
||
let ttl_seconds = ttl_seconds.max(1);
|
||
self.runtime_state
|
||
.kv_set(
|
||
&self.forward_cache_key(&record.rule_label, normalized_value, record.bucket),
|
||
record.sentinel.clone(),
|
||
Some(Duration::from_secs(ttl_seconds)),
|
||
)
|
||
.await?;
|
||
self.runtime_state
|
||
.kv_set(
|
||
&self.reverse_cache_key(&record.sentinel),
|
||
value,
|
||
Some(Duration::from_secs(ttl_seconds)),
|
||
)
|
||
.await?;
|
||
Ok(())
|
||
}
|
||
|
||
pub(crate) async fn lookup_sentinel(
|
||
&self,
|
||
rule_label: &str,
|
||
normalized_value: &str,
|
||
bucket: u64,
|
||
namespace: &str,
|
||
) -> Result<Option<String>, DataLayerError> {
|
||
let sentinel = self
|
||
.get_string(&self.forward_cache_key(rule_label, normalized_value, bucket))
|
||
.await?;
|
||
Ok(sentinel.filter(|value| {
|
||
value.len() <= MAX_CACHE_SENTINEL_BYTES
|
||
&& sentinel_matches_rule_label(namespace, rule_label, value)
|
||
}))
|
||
}
|
||
|
||
pub(crate) async fn load(
|
||
&self,
|
||
sentinel: &str,
|
||
) -> Result<Option<RedactionCacheRecord>, DataLayerError> {
|
||
let raw = self.get_string(&self.reverse_cache_key(sentinel)).await?;
|
||
raw.map(|value| {
|
||
serde_json::from_str::<RedactionCacheRecord>(&value)
|
||
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))
|
||
})
|
||
.transpose()
|
||
}
|
||
|
||
async fn get_string(&self, key: &str) -> Result<Option<String>, DataLayerError> {
|
||
self.runtime_state.kv_get(key).await
|
||
}
|
||
|
||
fn forward_cache_key(&self, rule_label: &str, normalized_value: &str, bucket: u64) -> String {
|
||
format!(
|
||
"{}:forward:{}:{}:{}",
|
||
self.key_prefix,
|
||
rule_label,
|
||
bucket,
|
||
normalized_value_cache_digest(normalized_value)
|
||
)
|
||
}
|
||
|
||
fn reverse_cache_key(&self, sentinel: &str) -> String {
|
||
format!("{}:reverse:{}", self.key_prefix, sentinel)
|
||
}
|
||
}
|
||
|
||
#[derive(Clone)]
|
||
struct Candidate {
|
||
rule_label: String,
|
||
kind: Option<RedactionKind>,
|
||
start: usize,
|
||
end: usize,
|
||
value: String,
|
||
priority: u16,
|
||
}
|
||
|
||
impl Candidate {
|
||
fn new(
|
||
rule_label: impl Into<String>,
|
||
kind: Option<RedactionKind>,
|
||
start: usize,
|
||
end: usize,
|
||
value: &str,
|
||
priority: u16,
|
||
) -> Self {
|
||
Self {
|
||
rule_label: rule_label.into(),
|
||
kind,
|
||
start,
|
||
end,
|
||
value: value.to_string(),
|
||
priority,
|
||
}
|
||
}
|
||
|
||
fn len(&self) -> usize {
|
||
self.end.saturating_sub(self.start)
|
||
}
|
||
}
|
||
|
||
fn detect_candidates(input: &str) -> Vec<Candidate> {
|
||
detect_candidates_with_probe(input, None)
|
||
}
|
||
|
||
fn detect_candidates_with_probe(
|
||
input: &str,
|
||
mut probe: Option<&mut DetectorProbe>,
|
||
) -> Vec<Candidate> {
|
||
// Segment-level scan-bypass cache is intentionally deferred until profiling shows regex scanning is the bottleneck.
|
||
let mut candidates = Vec::new();
|
||
|
||
if input.contains('@') {
|
||
push_regex_candidates(
|
||
input,
|
||
&EMAIL_REGEX,
|
||
RedactionKind::Email,
|
||
10,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
|value| value.contains('@'),
|
||
);
|
||
}
|
||
|
||
let digit_count = input.chars().filter(|ch| ch.is_ascii_digit()).count();
|
||
if digit_count >= 8 {
|
||
push_regex_candidates(
|
||
input,
|
||
&CN_MOBILE_REGEX,
|
||
RedactionKind::ChinaMobile,
|
||
20,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
is_valid_cn_mobile,
|
||
);
|
||
push_regex_candidates(
|
||
input,
|
||
&CN_LANDLINE_REGEX,
|
||
RedactionKind::ChinaLandline,
|
||
21,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
is_valid_cn_landline,
|
||
);
|
||
push_regex_candidates(
|
||
input,
|
||
&E164_PHONE_REGEX,
|
||
RedactionKind::Phone,
|
||
30,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
is_valid_e164_phone,
|
||
);
|
||
push_regex_candidates(
|
||
input,
|
||
&CN_ID_REGEX,
|
||
RedactionKind::ChinaResidentId,
|
||
40,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
is_valid_cn_resident_id,
|
||
);
|
||
push_regex_candidates(
|
||
input,
|
||
&PAYMENT_CARD_REGEX,
|
||
RedactionKind::PaymentCard,
|
||
50,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
is_valid_payment_card,
|
||
);
|
||
}
|
||
|
||
if input.contains('.') {
|
||
push_regex_candidates(
|
||
input,
|
||
&IPV4_REGEX,
|
||
RedactionKind::Ipv4,
|
||
60,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
|value| value.parse::<Ipv4Addr>().is_ok(),
|
||
);
|
||
push_regex_candidates(
|
||
input,
|
||
&JWT_REGEX,
|
||
RedactionKind::Jwt,
|
||
120,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
is_valid_jwt_like,
|
||
);
|
||
}
|
||
|
||
if input.contains(':') {
|
||
candidates.extend(detect_ipv6_candidates(input));
|
||
}
|
||
|
||
if input.contains("sk-") {
|
||
push_regex_candidates(
|
||
input,
|
||
&OPENAI_KEY_REGEX,
|
||
RedactionKind::OpenAiKey,
|
||
80,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
|_| true,
|
||
);
|
||
}
|
||
if input.contains("sk-ant-") {
|
||
push_regex_candidates(
|
||
input,
|
||
&ANTHROPIC_KEY_REGEX,
|
||
RedactionKind::AnthropicKey,
|
||
70,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
|_| true,
|
||
);
|
||
}
|
||
if input.contains("ghp_")
|
||
|| input.contains("gho_")
|
||
|| input.contains("ghu_")
|
||
|| input.contains("ghs_")
|
||
|| input.contains("ghr_")
|
||
|| input.contains("github_pat_")
|
||
{
|
||
push_regex_candidates(
|
||
input,
|
||
&GITHUB_TOKEN_REGEX,
|
||
RedactionKind::GitHubToken,
|
||
90,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
|_| true,
|
||
);
|
||
}
|
||
if input.contains("xox") {
|
||
push_regex_candidates(
|
||
input,
|
||
&SLACK_TOKEN_REGEX,
|
||
RedactionKind::SlackToken,
|
||
100,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
|_| true,
|
||
);
|
||
}
|
||
if input.contains("AKIA") || input.contains("ASIA") {
|
||
push_regex_candidates(
|
||
input,
|
||
&AWS_KEY_REGEX,
|
||
RedactionKind::AwsKey,
|
||
110,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
|_| true,
|
||
);
|
||
}
|
||
if input.contains("Bearer ") || input.contains("bearer ") {
|
||
push_regex_candidates(
|
||
input,
|
||
&BEARER_TOKEN_REGEX,
|
||
RedactionKind::BearerToken,
|
||
115,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
is_valid_bearer_token,
|
||
);
|
||
}
|
||
if input.contains("access_token")
|
||
|| input.contains("access-token")
|
||
|| input.contains("AccessToken")
|
||
{
|
||
push_regex_candidates(
|
||
input,
|
||
&ACCESS_TOKEN_REGEX,
|
||
RedactionKind::AccessToken,
|
||
116,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
is_valid_named_token,
|
||
);
|
||
}
|
||
if input.contains("secret_key")
|
||
|| input.contains("secret-key")
|
||
|| input.contains("SecretKey")
|
||
|| input.contains("agent_private_key")
|
||
|| input.contains("agent-private-key")
|
||
|| input.contains("agentPrivateKey")
|
||
{
|
||
push_regex_candidates(
|
||
input,
|
||
&SECRET_KEY_REGEX,
|
||
RedactionKind::SecretKey,
|
||
117,
|
||
&mut candidates,
|
||
probe.as_deref_mut(),
|
||
is_valid_named_token,
|
||
);
|
||
}
|
||
|
||
if input_may_contain_high_entropy_token(input) {
|
||
push_regex_candidates(
|
||
input,
|
||
&HIGH_ENTROPY_TOKEN_REGEX,
|
||
RedactionKind::ApiKey,
|
||
200,
|
||
&mut candidates,
|
||
probe,
|
||
is_strict_high_entropy_token,
|
||
);
|
||
}
|
||
|
||
candidates.retain(|candidate| !looks_like_existing_sentinel(&candidate.value));
|
||
candidates
|
||
}
|
||
|
||
fn detect_candidates_for_session_config(
|
||
input: &str,
|
||
config: &RedactionSessionConfig,
|
||
probe: Option<&mut DetectorProbe>,
|
||
) -> Vec<Candidate> {
|
||
if let Some(rules) = config.rules.as_ref() {
|
||
return detect_candidates_from_compiled_rules(input, rules, probe);
|
||
}
|
||
detect_candidates(input)
|
||
.into_iter()
|
||
.filter(|candidate| candidate.kind.is_some_and(|kind| config.kind_enabled(kind)))
|
||
.collect()
|
||
}
|
||
|
||
fn detect_candidates_from_compiled_rules(
|
||
input: &str,
|
||
rules: &[CompiledRedactionRule],
|
||
mut probe: Option<&mut DetectorProbe>,
|
||
) -> Vec<Candidate> {
|
||
let mut candidates = Vec::new();
|
||
for rule in rules {
|
||
if rule.kinds.is_empty() {
|
||
push_compiled_rule_candidates(input, rule, &mut candidates, probe.as_deref_mut());
|
||
continue;
|
||
}
|
||
for kind in &rule.kinds {
|
||
push_compiled_rule_candidate(input, rule, *kind, &mut candidates, probe.as_deref_mut());
|
||
}
|
||
}
|
||
candidates.retain(|candidate| !looks_like_existing_sentinel(&candidate.value));
|
||
candidates
|
||
}
|
||
|
||
fn push_compiled_rule_candidates(
|
||
input: &str,
|
||
rule: &CompiledRedactionRule,
|
||
candidates: &mut Vec<Candidate>,
|
||
mut probe: Option<&mut DetectorProbe>,
|
||
) {
|
||
for matched in rule.regex.find_iter(input) {
|
||
let value = matched.as_str();
|
||
if !has_token_boundary(input, matched.start(), matched.end()) {
|
||
continue;
|
||
}
|
||
if probe.is_some() {
|
||
// Custom rules have no validator call accounting.
|
||
candidates.push(Candidate::new(
|
||
rule.rule_label.clone(),
|
||
None,
|
||
matched.start(),
|
||
matched.end(),
|
||
value,
|
||
rule.custom_priority,
|
||
));
|
||
continue;
|
||
}
|
||
candidates.push(Candidate::new(
|
||
rule.rule_label.clone(),
|
||
None,
|
||
matched.start(),
|
||
matched.end(),
|
||
value,
|
||
rule.custom_priority,
|
||
));
|
||
}
|
||
}
|
||
|
||
fn push_compiled_rule_candidate(
|
||
input: &str,
|
||
rule: &CompiledRedactionRule,
|
||
kind: RedactionKind,
|
||
candidates: &mut Vec<Candidate>,
|
||
mut probe: Option<&mut DetectorProbe>,
|
||
) {
|
||
for matched in rule.regex.find_iter(input) {
|
||
let value = matched.as_str();
|
||
if !has_token_boundary(input, matched.start(), matched.end()) {
|
||
continue;
|
||
}
|
||
if let Some(probe) = probe.as_deref_mut() {
|
||
probe.record_validator_call(kind);
|
||
}
|
||
if redaction_candidate_is_valid(kind, value) {
|
||
candidates.push(Candidate::new(
|
||
rule.rule_label.clone(),
|
||
Some(kind),
|
||
matched.start(),
|
||
matched.end(),
|
||
value,
|
||
candidate_priority(Some(kind), rule.custom_priority),
|
||
));
|
||
}
|
||
}
|
||
}
|
||
|
||
fn push_regex_candidates(
|
||
input: &str,
|
||
regex: &Regex,
|
||
kind: RedactionKind,
|
||
priority: u16,
|
||
candidates: &mut Vec<Candidate>,
|
||
mut probe: Option<&mut DetectorProbe>,
|
||
validator: impl Fn(&str) -> bool,
|
||
) {
|
||
for matched in regex.find_iter(input) {
|
||
let value = matched.as_str();
|
||
if !has_token_boundary(input, matched.start(), matched.end()) {
|
||
continue;
|
||
}
|
||
if let Some(probe) = probe.as_deref_mut() {
|
||
probe.record_validator_call(kind);
|
||
}
|
||
if validator(value) {
|
||
candidates.push(Candidate::new(
|
||
kind.label(),
|
||
Some(kind),
|
||
matched.start(),
|
||
matched.end(),
|
||
value,
|
||
priority,
|
||
));
|
||
}
|
||
}
|
||
}
|
||
|
||
fn detect_ipv6_candidates(input: &str) -> Vec<Candidate> {
|
||
let mut candidates = Vec::new();
|
||
let mut start = None;
|
||
for (index, ch) in input.char_indices() {
|
||
if ch.is_ascii_hexdigit() || ch == ':' || ch == '.' {
|
||
start.get_or_insert(index);
|
||
continue;
|
||
}
|
||
if let Some(token_start) = start.take() {
|
||
push_ipv6_candidate(input, token_start, index, &mut candidates);
|
||
}
|
||
}
|
||
if let Some(token_start) = start {
|
||
push_ipv6_candidate(input, token_start, input.len(), &mut candidates);
|
||
}
|
||
candidates
|
||
}
|
||
|
||
fn push_ipv6_candidate(input: &str, start: usize, end: usize, candidates: &mut Vec<Candidate>) {
|
||
let value = &input[start..end];
|
||
if value.contains(':') && value.parse::<Ipv6Addr>().is_ok() {
|
||
candidates.push(Candidate::new(
|
||
RedactionKind::Ipv6.label(),
|
||
Some(RedactionKind::Ipv6),
|
||
start,
|
||
end,
|
||
value,
|
||
65,
|
||
));
|
||
}
|
||
}
|
||
|
||
fn select_non_overlapping(mut candidates: Vec<Candidate>) -> Vec<Candidate> {
|
||
candidates.sort_by(|left, right| {
|
||
right
|
||
.len()
|
||
.cmp(&left.len())
|
||
.then_with(|| left.priority.cmp(&right.priority))
|
||
.then_with(|| left.start.cmp(&right.start))
|
||
});
|
||
let mut occupied = HashSet::new();
|
||
let mut selected = Vec::new();
|
||
for candidate in candidates {
|
||
if (candidate.start..candidate.end).any(|index| occupied.contains(&index)) {
|
||
continue;
|
||
}
|
||
occupied.extend(candidate.start..candidate.end);
|
||
selected.push(candidate);
|
||
}
|
||
selected.sort_by_key(|candidate| candidate.start);
|
||
selected
|
||
}
|
||
|
||
fn has_token_boundary(input: &str, start: usize, end: usize) -> bool {
|
||
!has_sensitive_token_before(input, start) && !has_sensitive_token_after(input, end)
|
||
}
|
||
|
||
fn is_sensitive_token_char(ch: char) -> bool {
|
||
ch.is_ascii_alphanumeric() || matches!(ch, '_' | '-' | '.' | '@')
|
||
}
|
||
|
||
fn has_sensitive_token_before(input: &str, start: usize) -> bool {
|
||
let Some(before) = input[..start].chars().next_back() else {
|
||
return false;
|
||
};
|
||
if before != '.' {
|
||
return is_sensitive_token_char(before);
|
||
}
|
||
input[..start - before.len_utf8()]
|
||
.chars()
|
||
.next_back()
|
||
.is_some_and(is_sensitive_token_char)
|
||
}
|
||
|
||
fn has_sensitive_token_after(input: &str, end: usize) -> bool {
|
||
let Some(after) = input[end..].chars().next() else {
|
||
return false;
|
||
};
|
||
if after != '.' {
|
||
return is_sensitive_token_char(after);
|
||
}
|
||
input[end + after.len_utf8()..]
|
||
.chars()
|
||
.next()
|
||
.is_some_and(is_sensitive_token_char)
|
||
}
|
||
|
||
fn input_may_contain_high_entropy_token(input: &str) -> bool {
|
||
let mut token_len = 0usize;
|
||
let mut class_mask = 0u8;
|
||
for ch in input.chars() {
|
||
let class = high_entropy_token_char_class(ch);
|
||
if class == 0 {
|
||
if token_len >= 32 && class_mask.count_ones() >= 3 {
|
||
return true;
|
||
}
|
||
token_len = 0;
|
||
class_mask = 0;
|
||
continue;
|
||
}
|
||
token_len += 1;
|
||
class_mask |= class;
|
||
}
|
||
token_len >= 32 && class_mask.count_ones() >= 3
|
||
}
|
||
|
||
fn high_entropy_token_char_class(ch: char) -> u8 {
|
||
match ch {
|
||
'a'..='z' => 0b0001,
|
||
'A'..='Z' => 0b0010,
|
||
'0'..='9' => 0b0100,
|
||
'_' | '-' => 0b1000,
|
||
_ => 0,
|
||
}
|
||
}
|
||
|
||
fn digits_only(value: &str) -> String {
|
||
value.chars().filter(|ch| ch.is_ascii_digit()).collect()
|
||
}
|
||
|
||
fn is_valid_cn_mobile(value: &str) -> bool {
|
||
let mut digits = digits_only(value);
|
||
if digits.len() == 13 && digits.starts_with("86") {
|
||
digits = digits[2..].to_string();
|
||
}
|
||
digits.len() == 11
|
||
&& digits.starts_with('1')
|
||
&& digits
|
||
.as_bytes()
|
||
.get(1)
|
||
.is_some_and(|digit| matches!(digit, b'3'..=b'9'))
|
||
}
|
||
|
||
fn is_valid_cn_landline(value: &str) -> bool {
|
||
if !value.contains('-') {
|
||
return false;
|
||
}
|
||
let mut digits = digits_only(value);
|
||
if digits.len() >= 12 && digits.starts_with("86") {
|
||
digits = digits[2..].to_string();
|
||
}
|
||
digits.starts_with('0') && (10..=17).contains(&digits.len())
|
||
}
|
||
|
||
fn is_valid_e164_phone(value: &str) -> bool {
|
||
if !value.starts_with('+') {
|
||
return false;
|
||
}
|
||
let digits = digits_only(value);
|
||
(8..=15).contains(&digits.len()) && !digits.starts_with('0')
|
||
}
|
||
|
||
fn is_valid_cn_resident_id(value: &str) -> bool {
|
||
if value.len() != 18 || !value[..17].chars().all(|ch| ch.is_ascii_digit()) {
|
||
return false;
|
||
}
|
||
let birth = &value[6..14];
|
||
let year = birth[0..4].parse::<i32>().ok();
|
||
let month = birth[4..6].parse::<u32>().ok();
|
||
let day = birth[6..8].parse::<u32>().ok();
|
||
if year
|
||
.zip(month)
|
||
.zip(day)
|
||
.and_then(|((year, month), day)| NaiveDate::from_ymd_opt(year, month, day))
|
||
.is_none()
|
||
{
|
||
return false;
|
||
}
|
||
|
||
const WEIGHTS: [u32; 17] = [7, 9, 10, 5, 8, 4, 2, 1, 6, 3, 7, 9, 10, 5, 8, 4, 2];
|
||
const CHECK_CODES: [char; 11] = ['1', '0', 'X', '9', '8', '7', '6', '5', '4', '3', '2'];
|
||
let checksum = value[..17]
|
||
.chars()
|
||
.zip(WEIGHTS)
|
||
.map(|(digit, weight)| digit.to_digit(10).unwrap_or_default() * weight)
|
||
.sum::<u32>();
|
||
let expected = CHECK_CODES[(checksum % 11) as usize];
|
||
value
|
||
.chars()
|
||
.nth(17)
|
||
.is_some_and(|actual| actual.to_ascii_uppercase() == expected)
|
||
}
|
||
|
||
fn is_valid_payment_card(value: &str) -> bool {
|
||
let digits = digits_only(value);
|
||
if !(13..=19).contains(&digits.len())
|
||
|| digits
|
||
.chars()
|
||
.all(|ch| ch == digits.chars().next().unwrap_or_default())
|
||
{
|
||
return false;
|
||
}
|
||
if !matches!(digits.as_bytes().first(), Some(b'3' | b'4' | b'5' | b'6')) {
|
||
return false;
|
||
}
|
||
luhn_valid(&digits)
|
||
}
|
||
|
||
fn luhn_valid(digits: &str) -> bool {
|
||
let mut sum = 0u32;
|
||
let mut double = false;
|
||
for digit in digits.chars().rev().filter_map(|ch| ch.to_digit(10)) {
|
||
let mut value = digit;
|
||
if double {
|
||
value *= 2;
|
||
if value > 9 {
|
||
value -= 9;
|
||
}
|
||
}
|
||
sum += value;
|
||
double = !double;
|
||
}
|
||
sum.is_multiple_of(10)
|
||
}
|
||
|
||
fn is_valid_bearer_token(value: &str) -> bool {
|
||
value
|
||
.split_once(char::is_whitespace)
|
||
.map(|(_, token)| token.trim().len() >= 20)
|
||
.unwrap_or(false)
|
||
}
|
||
|
||
fn is_valid_named_token(value: &str) -> bool {
|
||
value
|
||
.split_once([':', '='])
|
||
.map(|(_, token)| {
|
||
let token = token.trim().trim_matches(['\'', '"']);
|
||
token.len() >= 20 && token.chars().any(|ch| ch.is_ascii_digit())
|
||
})
|
||
.unwrap_or(false)
|
||
}
|
||
|
||
fn is_valid_jwt_like(value: &str) -> bool {
|
||
let parts = value.split('.').collect::<Vec<_>>();
|
||
parts.len() == 3
|
||
&& parts.iter().all(|part| {
|
||
part.len() >= 10
|
||
&& part
|
||
.chars()
|
||
.all(|ch| ch.is_ascii_alphanumeric() || ch == '_' || ch == '-')
|
||
})
|
||
}
|
||
|
||
fn is_strict_high_entropy_token(value: &str) -> bool {
|
||
if value.len() < 32 || looks_like_existing_sentinel(value) {
|
||
return false;
|
||
}
|
||
let mut classes = 0;
|
||
classes += value.chars().any(|ch| ch.is_ascii_lowercase()) as u8;
|
||
classes += value.chars().any(|ch| ch.is_ascii_uppercase()) as u8;
|
||
classes += value.chars().any(|ch| ch.is_ascii_digit()) as u8;
|
||
classes += value.chars().any(|ch| matches!(ch, '_' | '-')) as u8;
|
||
classes >= 3 && shannon_entropy(value) >= 4.0
|
||
}
|
||
|
||
fn shannon_entropy(value: &str) -> f64 {
|
||
let mut counts = HashMap::<u8, usize>::new();
|
||
for byte in value.bytes() {
|
||
*counts.entry(byte).or_default() += 1;
|
||
}
|
||
let len = value.len() as f64;
|
||
counts
|
||
.values()
|
||
.map(|count| {
|
||
let probability = *count as f64 / len;
|
||
-probability * probability.log2()
|
||
})
|
||
.sum()
|
||
}
|
||
|
||
fn looks_like_existing_sentinel(value: &str) -> bool {
|
||
SENTINEL_REGEX.is_match(value)
|
||
}
|
||
|
||
fn sentinel_matches_rule_label(namespace: &str, rule_label: &str, sentinel: &str) -> bool {
|
||
if !full_sentinel_match(sentinel) {
|
||
return false;
|
||
}
|
||
sentinel
|
||
.strip_prefix('<')
|
||
.and_then(|value| value.strip_suffix('>'))
|
||
.and_then(|value| {
|
||
let mut parts = value.split(':');
|
||
let actual_namespace = parts.next()?;
|
||
let actual_rule = parts.next()?;
|
||
let digest = parts.next()?;
|
||
(parts.next().is_none()).then_some((actual_namespace, actual_rule, digest))
|
||
})
|
||
.is_some_and(|(actual_namespace, actual_rule, _)| {
|
||
actual_namespace == namespace && actual_rule == rule_label
|
||
})
|
||
}
|
||
|
||
fn full_sentinel_match(value: &str) -> bool {
|
||
SENTINEL_REGEX
|
||
.find(value)
|
||
.is_some_and(|matched| matched.start() == 0 && matched.end() == value.len())
|
||
}
|
||
|
||
fn normalize_redaction_value(kind: Option<RedactionKind>, value: &str) -> String {
|
||
match kind {
|
||
Some(RedactionKind::Email) => value.trim().to_ascii_lowercase(),
|
||
Some(RedactionKind::ChinaMobile)
|
||
| Some(RedactionKind::ChinaLandline)
|
||
| Some(RedactionKind::Phone)
|
||
| Some(RedactionKind::PaymentCard) => digits_only(value),
|
||
Some(RedactionKind::AccessToken) | Some(RedactionKind::SecretKey) => value
|
||
.split_once([':', '='])
|
||
.map(|(_, token)| token.trim().trim_matches(['\'', '"']).to_string())
|
||
.unwrap_or_else(|| value.trim().to_string()),
|
||
Some(RedactionKind::BearerToken) => value
|
||
.split_once(char::is_whitespace)
|
||
.map(|(_, token)| token.trim().to_string())
|
||
.unwrap_or_else(|| value.trim().to_string()),
|
||
_ => value.trim().to_string(),
|
||
}
|
||
}
|
||
|
||
fn candidate_priority(kind: Option<RedactionKind>, fallback_priority: u16) -> u16 {
|
||
match kind {
|
||
Some(RedactionKind::Email) => 10,
|
||
Some(RedactionKind::ChinaMobile) => 20,
|
||
Some(RedactionKind::ChinaLandline) => 21,
|
||
Some(RedactionKind::Phone) => 30,
|
||
Some(RedactionKind::ChinaResidentId) => 40,
|
||
Some(RedactionKind::PaymentCard) => 50,
|
||
Some(RedactionKind::Ipv4) => 60,
|
||
Some(RedactionKind::Ipv6) => 65,
|
||
Some(RedactionKind::AnthropicKey) => 70,
|
||
Some(RedactionKind::OpenAiKey) => 80,
|
||
Some(RedactionKind::GitHubToken) => 90,
|
||
Some(RedactionKind::SlackToken) => 100,
|
||
Some(RedactionKind::AwsKey) => 110,
|
||
Some(RedactionKind::BearerToken) => 115,
|
||
Some(RedactionKind::AccessToken) => 116,
|
||
Some(RedactionKind::SecretKey) => 117,
|
||
Some(RedactionKind::Jwt) => 120,
|
||
Some(RedactionKind::ApiKey) => 200,
|
||
None => fallback_priority,
|
||
}
|
||
}
|
||
|
||
fn redaction_candidate_is_valid(kind: RedactionKind, value: &str) -> bool {
|
||
match kind {
|
||
RedactionKind::Email => value.contains('@'),
|
||
RedactionKind::ChinaMobile => is_valid_cn_mobile(value),
|
||
RedactionKind::ChinaLandline => is_valid_cn_landline(value),
|
||
RedactionKind::Phone => is_valid_e164_phone(value),
|
||
RedactionKind::ChinaResidentId => is_valid_cn_resident_id(value),
|
||
RedactionKind::PaymentCard => is_valid_payment_card(value),
|
||
RedactionKind::Ipv4 => value.parse::<Ipv4Addr>().is_ok(),
|
||
RedactionKind::Ipv6 => value.parse::<Ipv6Addr>().is_ok(),
|
||
RedactionKind::OpenAiKey
|
||
| RedactionKind::AnthropicKey
|
||
| RedactionKind::GitHubToken
|
||
| RedactionKind::SlackToken
|
||
| RedactionKind::AwsKey => true,
|
||
RedactionKind::BearerToken => is_valid_bearer_token(value),
|
||
RedactionKind::Jwt => is_valid_jwt_like(value),
|
||
RedactionKind::AccessToken => is_valid_named_token(value),
|
||
RedactionKind::SecretKey => is_valid_named_token(value),
|
||
RedactionKind::ApiKey => is_strict_high_entropy_token(value),
|
||
}
|
||
}
|
||
|
||
fn normalized_value_cache_digest(normalized_value: &str) -> String {
|
||
let digest = sha2::Sha256::digest(normalized_value.as_bytes());
|
||
base32_no_pad(&digest[..HMAC96_BYTES])
|
||
}
|
||
|
||
fn parse_enabled_redaction_kinds(value: Option<&Value>) -> HashSet<RedactionKind> {
|
||
let Some(values) = value.and_then(Value::as_array) else {
|
||
return default_enabled_redaction_kinds();
|
||
};
|
||
values
|
||
.iter()
|
||
.filter_map(Value::as_str)
|
||
.flat_map(redaction_kinds_for_entity)
|
||
.collect::<HashSet<_>>()
|
||
}
|
||
|
||
fn default_enabled_redaction_kinds() -> HashSet<RedactionKind> {
|
||
[
|
||
RedactionKind::Email,
|
||
RedactionKind::ChinaMobile,
|
||
RedactionKind::ChinaLandline,
|
||
RedactionKind::Phone,
|
||
RedactionKind::ChinaResidentId,
|
||
RedactionKind::PaymentCard,
|
||
RedactionKind::Ipv4,
|
||
RedactionKind::Ipv6,
|
||
RedactionKind::OpenAiKey,
|
||
RedactionKind::AnthropicKey,
|
||
RedactionKind::GitHubToken,
|
||
RedactionKind::SlackToken,
|
||
RedactionKind::AwsKey,
|
||
RedactionKind::ApiKey,
|
||
RedactionKind::AccessToken,
|
||
RedactionKind::SecretKey,
|
||
RedactionKind::BearerToken,
|
||
RedactionKind::Jwt,
|
||
]
|
||
.into_iter()
|
||
.collect()
|
||
}
|
||
|
||
fn redaction_kinds_for_entity(entity: &str) -> impl Iterator<Item = RedactionKind> {
|
||
match entity.trim().to_ascii_lowercase().as_str() {
|
||
"email" => vec![RedactionKind::Email],
|
||
"cn_phone" => vec![RedactionKind::ChinaMobile, RedactionKind::ChinaLandline],
|
||
"global_phone" => vec![RedactionKind::Phone],
|
||
"cn_id" => vec![RedactionKind::ChinaResidentId],
|
||
"payment_card" => vec![RedactionKind::PaymentCard],
|
||
"ipv4" => vec![RedactionKind::Ipv4],
|
||
"ipv6" => vec![RedactionKind::Ipv6],
|
||
"api_key" => vec![
|
||
RedactionKind::OpenAiKey,
|
||
RedactionKind::AnthropicKey,
|
||
RedactionKind::GitHubToken,
|
||
RedactionKind::SlackToken,
|
||
RedactionKind::AwsKey,
|
||
RedactionKind::ApiKey,
|
||
],
|
||
"access_token" => vec![RedactionKind::AccessToken],
|
||
"secret_key" => vec![RedactionKind::SecretKey],
|
||
"bearer_token" => vec![RedactionKind::BearerToken],
|
||
"jwt" => vec![RedactionKind::Jwt],
|
||
_ => Vec::new(),
|
||
}
|
||
.into_iter()
|
||
}
|
||
|
||
fn base32_no_pad(bytes: &[u8]) -> String {
|
||
const ALPHABET: &[u8; 32] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
|
||
let mut output = String::with_capacity((bytes.len() * 8).div_ceil(5));
|
||
let mut bit_buffer = 0u32;
|
||
let mut bit_count = 0u8;
|
||
|
||
for byte in bytes {
|
||
bit_buffer = (bit_buffer << 8) | u32::from(*byte);
|
||
bit_count += 8;
|
||
while bit_count >= 5 {
|
||
bit_count -= 5;
|
||
let index = ((bit_buffer >> bit_count) & 0x1f) as usize;
|
||
output.push(ALPHABET[index] as char);
|
||
}
|
||
}
|
||
|
||
if bit_count > 0 {
|
||
let index = ((bit_buffer << (5 - bit_count)) & 0x1f) as usize;
|
||
output.push(ALPHABET[index] as char);
|
||
}
|
||
|
||
output
|
||
}
|
||
|
||
fn redacted_sentinel_debug(sentinel: &str) -> String {
|
||
let (namespace, kind) = sentinel
|
||
.strip_prefix('<')
|
||
.and_then(|rest| rest.split_once(':'))
|
||
.map(|(namespace, rest)| {
|
||
let kind = rest
|
||
.split_once(':')
|
||
.map(|(kind, _)| kind)
|
||
.unwrap_or("UNKNOWN");
|
||
(namespace, kind)
|
||
})
|
||
.unwrap_or((DEFAULT_SENTINEL_NAMESPACE, "UNKNOWN"));
|
||
format!("<{namespace}:{kind}:redacted>")
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::{
|
||
build_redaction_session_config, detect_candidates_with_probe, mask_chat_request_json,
|
||
mask_chat_request_json_with_options, parse_chat_pii_redaction_rules, restore_json_strings,
|
||
restore_sync_response_body, try_mask_chat_pii_request_json_with_options,
|
||
try_mask_chat_pii_request_value_with_cache_options,
|
||
try_mask_chat_request_json_with_cache_options, try_mask_chat_request_json_with_options,
|
||
ChatPiiRedactionRequestFormat, ChatPiiRedactionRuntimeConfig, DetectorProbe, MappingKey,
|
||
MaskChatRequestOptions, RedactionKind, RedactionMapping, RedactionMaskError,
|
||
RedactionSession, RedactionSessionConfig, RedactionSessionSlot, RedisRedactionMappingCache,
|
||
SentinelMatcher, StreamingResponseRestorer,
|
||
};
|
||
use std::collections::BTreeMap;
|
||
use std::time::Duration;
|
||
|
||
use aether_runtime_state::{RedisClientConfig, RuntimeState};
|
||
use aether_test_support::ManagedRedisServer;
|
||
use serde_json::{json, Value};
|
||
|
||
fn assert_debug_surface_hides_values(debug: &str, originals: &[&str], sentinels: &[String]) {
|
||
for original in originals {
|
||
assert!(
|
||
!debug.contains(original),
|
||
"debug leaked original {original}"
|
||
);
|
||
}
|
||
for sentinel in sentinels {
|
||
assert!(
|
||
!debug.contains(sentinel),
|
||
"debug leaked sentinel {sentinel}"
|
||
);
|
||
}
|
||
}
|
||
|
||
fn source_chunk_before_struct<'a>(source: &'a str, marker: &str) -> &'a str {
|
||
let marker_start = source.find(marker).expect("struct marker should exist");
|
||
let chunk_start = marker_start.saturating_sub(220);
|
||
&source[chunk_start..marker_start]
|
||
}
|
||
|
||
fn assert_not_serde_derived(source: &str, marker: &str) {
|
||
let chunk = source_chunk_before_struct(source, marker);
|
||
assert!(
|
||
!chunk.contains("Serialize"),
|
||
"{marker} must not derive Serialize"
|
||
);
|
||
assert!(
|
||
!chunk.contains("Deserialize"),
|
||
"{marker} must not derive Deserialize"
|
||
);
|
||
}
|
||
|
||
fn session_at(now_unix_secs: u64) -> RedactionSession {
|
||
RedactionSession::new(RedactionSessionConfig::new(
|
||
b"redaction-test-key".to_vec(),
|
||
300,
|
||
now_unix_secs,
|
||
))
|
||
}
|
||
|
||
fn assert_base32_sentinel(sentinel: &str) {
|
||
let body = sentinel
|
||
.strip_prefix("<AETHER:")
|
||
.and_then(|rest| rest.strip_suffix('>'))
|
||
.expect("sentinel should use Aether wrapper");
|
||
let (_, payload) = body
|
||
.split_once(':')
|
||
.expect("sentinel should include type and payload");
|
||
assert_eq!(payload.len(), 20, "HMAC96 base32 payload is 20 chars");
|
||
assert!(payload
|
||
.chars()
|
||
.all(|ch| matches!(ch, 'A'..='Z' | '2'..='7')));
|
||
assert_ne!(payload.len(), 24, "old hex HMAC96 payload must not be used");
|
||
}
|
||
|
||
fn assert_namespaced_base32_sentinel(sentinel: &str, namespace: &str, rule_label: &str) {
|
||
let body = sentinel
|
||
.strip_prefix('<')
|
||
.and_then(|rest| rest.strip_suffix('>'))
|
||
.expect("sentinel should use wrapper");
|
||
let parts = body.split(':').collect::<Vec<_>>();
|
||
assert_eq!(parts.len(), 3);
|
||
assert_eq!(parts[0], namespace);
|
||
assert_eq!(parts[1], rule_label);
|
||
assert_eq!(parts[2].len(), 20, "HMAC96 base32 payload is 20 chars");
|
||
assert!(parts[2]
|
||
.chars()
|
||
.all(|ch| matches!(ch, 'A'..='Z' | '2'..='7')));
|
||
}
|
||
|
||
fn session_with_response_mapping(original: &str, sentinel: &str) -> RedactionSession {
|
||
let mut session = session_at(600);
|
||
let key = MappingKey {
|
||
rule_label: "EMAIL".to_string(),
|
||
original: original.to_string(),
|
||
};
|
||
session
|
||
.sentinel_index
|
||
.insert(sentinel.to_string(), key.clone());
|
||
session.mappings.insert(
|
||
key,
|
||
RedactionMapping {
|
||
rule_label: "EMAIL".to_string(),
|
||
kind: Some(RedactionKind::Email),
|
||
original: original.to_string(),
|
||
normalized_value: original.to_string(),
|
||
sentinel: sentinel.to_string(),
|
||
bucket: 2,
|
||
created_at_unix_secs: 600,
|
||
expires_at_unix_secs: 900,
|
||
},
|
||
);
|
||
session
|
||
}
|
||
|
||
async fn start_managed_redis_or_skip() -> Option<ManagedRedisServer> {
|
||
match ManagedRedisServer::start().await {
|
||
Ok(server) => Some(server),
|
||
Err(err) if err.to_string().contains("No such file or directory") => None,
|
||
Err(err) => panic!("redis server should start: {err}"),
|
||
}
|
||
}
|
||
|
||
async fn redis_cache_runtime_state(redis_url: &str, key_prefix: &str) -> RuntimeState {
|
||
RuntimeState::redis(
|
||
RedisClientConfig {
|
||
url: redis_url.to_string(),
|
||
key_prefix: Some(key_prefix.to_string()),
|
||
},
|
||
Some(1_000),
|
||
)
|
||
.await
|
||
.expect("redis runtime state should build")
|
||
}
|
||
|
||
fn email_only_runtime_config() -> ChatPiiRedactionRuntimeConfig {
|
||
ChatPiiRedactionRuntimeConfig {
|
||
enabled: true,
|
||
rules: parse_chat_pii_redaction_rules(Some(&json!([{
|
||
"id": "email",
|
||
"name": "邮箱",
|
||
"pattern": "(?i)[A-Z0-9._%+-]{1,64}@[A-Z0-9.-]{1,253}\\.[A-Z]{2,63}",
|
||
"enabled": true,
|
||
"kind": "email",
|
||
"system": true
|
||
}])))
|
||
.expect("email rule should compile"),
|
||
..ChatPiiRedactionRuntimeConfig::default()
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_session_detects_reuses_rolls_over_and_avoids_collisions() {
|
||
let input = concat!(
|
||
"email [email protected] again [email protected] ",
|
||
"mobile 13800138000 landline 010-12345678 phone +14155552671 ",
|
||
"id 11010519491231002X card 4111 1111 1111 1111 ",
|
||
"ips 192.168.0.1 2001:db8::1 ",
|
||
"keys sk-proj-AbCdEfGhIjKlMnOpQrStUvWxYz123456 sk-ant-api03-AbCdEfGhIjKlMnOpQrStUvWxYz1234567890 ",
|
||
"ghp_abcdefghijklmnopqrstuvwxyzABCDEFGHIJ ",
|
||
"xo", "xb-123456789012-123456789012-AbCdEfGhIjKlMnOpQrStUvWx ",
|
||
"AKIAIOSFODNN7EXAMPLE Bearer abcdefghijklmnopqrstuvwxyzABCDEF1234567890 ",
|
||
"access_token=accessValueABCDEF1234567890abcdef secret_key=secretValueABCDEF1234567890abcdef ",
|
||
"eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJpc3MiOiJ0ZXN0IiwiZXhwIjoyMDAwMDAwMDAwfQ.signature123456789 ",
|
||
"q8HjK2LmN9PqR4StU7VwX0YzA3BcD6Ef"
|
||
);
|
||
let originals = [
|
||
"[email protected]",
|
||
"13800138000",
|
||
"010-12345678",
|
||
"+14155552671",
|
||
"11010519491231002X",
|
||
"4111 1111 1111 1111",
|
||
"192.168.0.1",
|
||
"2001:db8::1",
|
||
"sk-proj-AbCdEfGhIjKlMnOpQrStUvWxYz123456",
|
||
"sk-ant-api03-AbCdEfGhIjKlMnOpQrStUvWxYz1234567890",
|
||
"ghp_abcdefghijklmnopqrstuvwxyzABCDEFGHIJ",
|
||
concat!("xo", "xb-123456789012-123456789012-AbCdEfGhIjKlMnOpQrStUvWx"),
|
||
"AKIAIOSFODNN7EXAMPLE",
|
||
"Bearer abcdefghijklmnopqrstuvwxyzABCDEF1234567890",
|
||
"access_token=accessValueABCDEF1234567890abcdef",
|
||
"secret_key=secretValueABCDEF1234567890abcdef",
|
||
"eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJpc3MiOiJ0ZXN0IiwiZXhwIjoyMDAwMDAwMDAwfQ.signature123456789",
|
||
"q8HjK2LmN9PqR4StU7VwX0YzA3BcD6Ef",
|
||
];
|
||
|
||
let mut session = session_at(600);
|
||
let redacted = session.redact_text(input);
|
||
|
||
for original in originals {
|
||
assert!(!redacted.text.contains(original), "leaked {original}");
|
||
}
|
||
assert!(redacted.text.contains("<AETHER:EMAIL:"));
|
||
assert!(redacted.text.contains("<AETHER:CN_PHONE:"));
|
||
assert!(redacted.text.contains("<AETHER:ACCESS_TOKEN:"));
|
||
assert!(redacted.text.contains("<AETHER:SECRET_KEY:"));
|
||
assert!(redacted
|
||
.matches
|
||
.iter()
|
||
.all(|matched| matched.sentinel.starts_with("<AETHER:")
|
||
&& matched.sentinel.ends_with('>')));
|
||
for matched in &redacted.matches {
|
||
assert_base32_sentinel(&matched.sentinel);
|
||
}
|
||
assert!(redacted.matches.iter().any(|matched| {
|
||
matched.kind == Some(RedactionKind::ChinaMobile)
|
||
&& matched.sentinel.contains(":CN_PHONE:")
|
||
}));
|
||
assert!(redacted.matches.iter().any(|matched| {
|
||
matched.kind == Some(RedactionKind::ChinaLandline)
|
||
&& matched.sentinel.contains(":CN_PHONE:")
|
||
}));
|
||
assert!(redacted
|
||
.matches
|
||
.iter()
|
||
.any(|matched| matched.kind == Some(RedactionKind::AccessToken)));
|
||
assert!(redacted
|
||
.matches
|
||
.iter()
|
||
.any(|matched| matched.kind == Some(RedactionKind::SecretKey)));
|
||
assert_eq!(session.mapping_count(), originals.len());
|
||
|
||
let first_email = session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("email sentinel should exist")
|
||
.to_string();
|
||
assert_base32_sentinel(&first_email);
|
||
assert_eq!(redacted.text.matches(&first_email).count(), 2);
|
||
|
||
let debug = format!("{session:?}");
|
||
assert!(debug.contains("mapping_count"));
|
||
assert!(!debug.contains("[email protected]"));
|
||
assert!(!debug.contains(&first_email));
|
||
assert!(!debug.contains("redaction-test-key"));
|
||
|
||
let mut same_bucket = session_at(899);
|
||
same_bucket.redact_text("[email protected]");
|
||
assert_eq!(
|
||
same_bucket.sentinel_for_original("[email protected]"),
|
||
Some(first_email.as_str())
|
||
);
|
||
|
||
let mut next_bucket = session_at(900);
|
||
next_bucket.redact_text("[email protected]");
|
||
assert_ne!(
|
||
next_bucket.sentinel_for_original("[email protected]"),
|
||
Some(first_email.as_str())
|
||
);
|
||
|
||
let colliding_literal = first_email;
|
||
let mut collision_session = session_at(600);
|
||
collision_session.redact_text(&format!(
|
||
"literal {colliding_literal} value [email protected]"
|
||
));
|
||
assert_ne!(
|
||
collision_session.sentinel_for_original("[email protected]"),
|
||
Some(colliding_literal.as_str())
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_hides_agent_identity_private_keys() {
|
||
let private_key =
|
||
"agent_private_key=MC4CAQAwBQYDK2VwBCIEIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
|
||
let mut session = session_at(601);
|
||
|
||
let redacted = session.redact_text(private_key);
|
||
|
||
assert!(!redacted.text.contains(private_key));
|
||
assert!(redacted.text.contains("<AETHER:SECRET_KEY:"));
|
||
assert!(redacted
|
||
.matches
|
||
.iter()
|
||
.any(|matched| matched.kind == Some(RedactionKind::SecretKey)));
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_session_uses_configured_sentinel_namespace() {
|
||
let mut session = RedactionSession::new(
|
||
RedactionSessionConfig::new(b"redaction-test-key".to_vec(), 300, 600)
|
||
.with_sentinel_namespace("vendor_safe"),
|
||
);
|
||
|
||
let redacted = session.redact_text("contact [email protected]");
|
||
let sentinel = session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("email sentinel should exist");
|
||
|
||
assert_namespaced_base32_sentinel(sentinel, "VENDOR_SAFE", "EMAIL");
|
||
assert!(redacted.text.contains(sentinel));
|
||
assert!(!redacted.text.contains("<AETHER:"));
|
||
assert_eq!(
|
||
session.restore_text(&redacted.text).text,
|
||
"contact [email protected]"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_no_mapping_leak() {
|
||
let originals = [
|
||
"[email protected]",
|
||
"sk-proj-AbCdEfGhIjKlMnOpQrStUvWxYz123456",
|
||
"access_token=accessValueABCDEF1234567890abcdef",
|
||
];
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [{
|
||
"role": "user",
|
||
"content": format!(
|
||
"Email {} key {} token {}",
|
||
originals[0], originals[1], originals[2]
|
||
),
|
||
}]
|
||
});
|
||
let masked = mask_chat_request_json(
|
||
&serde_json::to_vec(&request).expect("request should serialize"),
|
||
test_config(),
|
||
);
|
||
assert!(masked.redacted);
|
||
let sentinels = originals
|
||
.iter()
|
||
.map(|original| {
|
||
masked
|
||
.session
|
||
.sentinel_for_original(original)
|
||
.expect("sentinel should exist")
|
||
.to_string()
|
||
})
|
||
.collect::<Vec<_>>();
|
||
|
||
for debug in [
|
||
format!("{:?}", masked.session),
|
||
format!("{:?}", masked),
|
||
format!(
|
||
"{:?}",
|
||
RedactionSessionConfig::new(originals[1].as_bytes().to_vec(), 300, 600)
|
||
),
|
||
] {
|
||
assert_debug_surface_hides_values(&debug, &originals, &sentinels);
|
||
}
|
||
|
||
let mut session = session_at(600);
|
||
let redacted = session.redact_text(&format!(
|
||
"Email {} key {} token {}",
|
||
originals[0], originals[1], originals[2]
|
||
));
|
||
let redacted_debug = format!("{redacted:?}");
|
||
assert_debug_surface_hides_values(&redacted_debug, &originals, &sentinels);
|
||
for matched in &redacted.matches {
|
||
let matched_debug = format!("{matched:?}");
|
||
assert_debug_surface_hides_values(&matched_debug, &originals, &sentinels);
|
||
assert!(matched_debug.contains("<AETHER:"));
|
||
assert!(matched_debug.contains(":redacted>"));
|
||
}
|
||
for mapping in session.mappings() {
|
||
let mapping_debug = format!("{mapping:?}");
|
||
assert_debug_surface_hides_values(&mapping_debug, &originals, &sentinels);
|
||
assert!(mapping_debug.contains("original_len"));
|
||
assert!(mapping_debug.contains("normalized_value_len"));
|
||
assert!(mapping_debug.contains(":redacted>"));
|
||
}
|
||
|
||
let privacy_source = include_str!("mod.rs");
|
||
for marker in [
|
||
"pub(crate) struct RedactionSession {",
|
||
"pub(crate) struct RedactionSessionSlot {",
|
||
"pub(crate) struct MaskedChatRequest {",
|
||
"pub(crate) struct RedactionMatch {",
|
||
"pub(crate) struct RedactionMapping {",
|
||
] {
|
||
assert_not_serde_derived(privacy_source, marker);
|
||
}
|
||
let session_serialize_impl = format!("impl {} for {}", "Serialize", "RedactionSession");
|
||
let mapping_serialize_impl = format!("impl {} for {}", "Serialize", "RedactionMapping");
|
||
assert!(!privacy_source.contains(&session_serialize_impl));
|
||
assert!(!privacy_source.contains(&mapping_serialize_impl));
|
||
|
||
let report_context_source = include_str!("../ai_serving/planner/report_context.rs");
|
||
for forbidden in [
|
||
"RedactionSession",
|
||
"RedactionMapping",
|
||
"RedactionSessionSlot",
|
||
"redaction_session",
|
||
"redaction_mapping",
|
||
"sentinel_index",
|
||
] {
|
||
assert!(
|
||
!report_context_source.contains(forbidden),
|
||
"report context source must not expose {forbidden}"
|
||
);
|
||
}
|
||
assert!(report_context_source.contains("original_request_body_json"));
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_false_positives_are_not_detected() {
|
||
let input = concat!(
|
||
"Order AETHER-20240502-123456 date 2026-05-02 decimal 12345.67890 ",
|
||
"short ids 123456 9876543210 invalid card 4111111111111112 ",
|
||
"short sk word sk-test low entropy aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa ",
|
||
"name Alice Zhang address 北京市朝阳区建国路88号"
|
||
);
|
||
let mut session = session_at(600);
|
||
let redacted = session.redact_text(input);
|
||
|
||
assert_eq!(redacted.text, input);
|
||
assert!(redacted.matches.is_empty());
|
||
assert_eq!(session.mapping_count(), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_request_masks_chat_message_text_and_tool_arguments() {
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"metadata": {
|
||
"owner": "[email protected]"
|
||
},
|
||
"temperature": 0.2,
|
||
"tools": [{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "lookup_customer",
|
||
"description": "schema owner [email protected] stays visible",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"email": {"type": "string", "description": "example [email protected]"}
|
||
}
|
||
}
|
||
}
|
||
}],
|
||
"messages": [
|
||
{
|
||
"role": "system",
|
||
"content": "Contact [email protected] with quote \"ok\", path C:\\tmp, and newline\nend."
|
||
},
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": "Phone +14155552671 should be masked."},
|
||
{"type": "image_url", "text": "image text [email protected] must stay."},
|
||
{"type": "input_text", "text": "input text [email protected] must stay."},
|
||
{"type": "text", "text": 42}
|
||
]
|
||
},
|
||
{
|
||
"role": "assistant",
|
||
"content": null,
|
||
"tool_calls": [{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "lookup_customer",
|
||
"arguments": "{\"phone\":\"13800138000\",\"secret\":\"secret_key=secretValueABCDEF1234567890abcdef\"}"
|
||
}
|
||
}]
|
||
},
|
||
{
|
||
"role": "tool",
|
||
"tool_call_id": "call_1",
|
||
"content": "Tool result access_token=accessValueABCDEF1234567890abcdef"
|
||
}
|
||
]
|
||
});
|
||
let raw = serde_json::to_vec(&request).expect("request should serialize");
|
||
|
||
let masked = mask_chat_request_json(&raw, test_config());
|
||
|
||
assert!(masked.redacted);
|
||
assert_eq!(masked.session.mapping_count(), 5);
|
||
let masked_json: serde_json::Value =
|
||
serde_json::from_slice(&masked.body).expect("masked request should stay valid JSON");
|
||
|
||
assert_eq!(masked_json["model"], "gpt-5");
|
||
assert_eq!(masked_json["metadata"]["owner"], "[email protected]");
|
||
assert_eq!(masked_json["temperature"], 0.2);
|
||
assert_eq!(
|
||
masked_json["tools"][0]["function"]["description"],
|
||
"schema owner [email protected] stays visible"
|
||
);
|
||
assert_eq!(
|
||
masked_json["tools"][0]["function"]["parameters"]["properties"]["email"]["description"],
|
||
"example [email protected]"
|
||
);
|
||
|
||
let system_content = masked_json["messages"][0]["content"]
|
||
.as_str()
|
||
.expect("system content should remain a string");
|
||
assert!(!system_content.contains("[email protected]"));
|
||
assert!(system_content.contains("\"ok\""));
|
||
assert!(system_content.contains("C:\\tmp"));
|
||
assert!(system_content.contains("newline\nend"));
|
||
assert!(system_content.contains(
|
||
masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("email sentinel should exist")
|
||
));
|
||
|
||
let text_part = masked_json["messages"][1]["content"][0]["text"]
|
||
.as_str()
|
||
.expect("text part should remain a string");
|
||
assert!(!text_part.contains("+14155552671"));
|
||
assert!(text_part.contains(
|
||
masked
|
||
.session
|
||
.sentinel_for_original("+14155552671")
|
||
.expect("phone sentinel should exist")
|
||
));
|
||
assert_eq!(
|
||
masked_json["messages"][1]["content"][1]["text"],
|
||
"image text [email protected] must stay."
|
||
);
|
||
assert_eq!(
|
||
masked_json["messages"][1]["content"][2]["text"],
|
||
"input text [email protected] must stay."
|
||
);
|
||
|
||
let arguments = masked_json["messages"][2]["tool_calls"][0]["function"]["arguments"]
|
||
.as_str()
|
||
.expect("tool call arguments should remain a string");
|
||
assert!(!arguments.contains("13800138000"));
|
||
assert!(!arguments.contains("secretValueABCDEF1234567890abcdef"));
|
||
assert!(arguments.contains(
|
||
masked
|
||
.session
|
||
.sentinel_for_original("13800138000")
|
||
.expect("mobile sentinel should exist")
|
||
));
|
||
assert!(arguments.contains(
|
||
masked
|
||
.session
|
||
.sentinel_for_original("secret_key=secretValueABCDEF1234567890abcdef")
|
||
.expect("secret sentinel should exist")
|
||
));
|
||
|
||
let tool_content = masked_json["messages"][3]["content"]
|
||
.as_str()
|
||
.expect("tool content should remain a string");
|
||
assert!(!tool_content.contains("accessValueABCDEF1234567890abcdef"));
|
||
assert!(tool_content.contains(
|
||
masked
|
||
.session
|
||
.sentinel_for_original("access_token=accessValueABCDEF1234567890abcdef")
|
||
.expect("access token sentinel should exist")
|
||
));
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_request_masks_claude_messages_system_and_content() {
|
||
let request = json!({
|
||
"model": "claude-sonnet-4",
|
||
"system": [
|
||
{"type": "text", "text": "System owner [email protected]"},
|
||
{"type": "image", "source": {"type": "base64", "data": "ignored"}}
|
||
],
|
||
"metadata": {"owner": "[email protected]"},
|
||
"messages": [{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": "Call 13800138000 before noon."},
|
||
{"type": "tool_result", "content": "token access_token=accessValueABCDEF1234567890abcdef"},
|
||
{"type": "image", "text": "image [email protected] must stay"}
|
||
]
|
||
}]
|
||
});
|
||
let raw = serde_json::to_vec(&request).expect("request should serialize");
|
||
|
||
let masked = try_mask_chat_pii_request_json_with_options(
|
||
&raw,
|
||
ChatPiiRedactionRequestFormat::ClaudeMessages,
|
||
test_config(),
|
||
MaskChatRequestOptions::runtime(),
|
||
)
|
||
.expect("claude messages request should mask");
|
||
|
||
assert!(masked.redacted);
|
||
assert_eq!(masked.session.mapping_count(), 3);
|
||
let masked_json: serde_json::Value =
|
||
serde_json::from_slice(&masked.body).expect("masked request should stay valid JSON");
|
||
assert_eq!(masked_json["metadata"]["owner"], "[email protected]");
|
||
assert!(!masked_json["system"][0]["text"]
|
||
.as_str()
|
||
.expect("system text should remain a string")
|
||
.contains("[email protected]"));
|
||
assert!(!masked_json["messages"][0]["content"][0]["text"]
|
||
.as_str()
|
||
.expect("message text should remain a string")
|
||
.contains("13800138000"));
|
||
assert!(!masked_json["messages"][0]["content"][1]["content"]
|
||
.as_str()
|
||
.expect("tool result should remain a string")
|
||
.contains("accessValueABCDEF1234567890abcdef"));
|
||
assert_eq!(
|
||
masked_json["messages"][0]["content"][2]["text"],
|
||
"image [email protected] must stay"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_request_masks_email_before_sentence_period() {
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [{
|
||
"role": "user",
|
||
"content": "Contact [email protected]."
|
||
}]
|
||
});
|
||
let raw = serde_json::to_vec(&request).expect("request should serialize");
|
||
|
||
let masked = try_mask_chat_pii_request_json_with_options(
|
||
&raw,
|
||
ChatPiiRedactionRequestFormat::OpenAiChat,
|
||
test_config(),
|
||
MaskChatRequestOptions::runtime(),
|
||
)
|
||
.expect("chat request should mask");
|
||
|
||
assert!(masked.redacted);
|
||
let masked_json: serde_json::Value =
|
||
serde_json::from_slice(&masked.body).expect("masked request should stay valid JSON");
|
||
let content = masked_json["messages"][0]["content"]
|
||
.as_str()
|
||
.expect("content should stay text");
|
||
assert!(content.contains("<AETHER:EMAIL:"));
|
||
assert!(content.ends_with('.'));
|
||
assert!(!content.contains("[email protected]"));
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_request_masks_openai_responses_input_and_function_arguments() {
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"instructions": "Route replies for [email protected] contact.",
|
||
"metadata": {"owner": "[email protected]"},
|
||
"input": [
|
||
{
|
||
"type": "message",
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "input_text", "text": "Phone +14155552671"},
|
||
{"type": "input_image", "text": "image [email protected] must stay"}
|
||
]
|
||
},
|
||
{
|
||
"type": "function_call",
|
||
"name": "lookup_customer",
|
||
"arguments": "{\"secret\":\"secret_key=secretValueABCDEF1234567890abcdef\"}"
|
||
}
|
||
]
|
||
});
|
||
let raw = serde_json::to_vec(&request).expect("request should serialize");
|
||
|
||
let masked = try_mask_chat_pii_request_json_with_options(
|
||
&raw,
|
||
ChatPiiRedactionRequestFormat::OpenAiResponses,
|
||
test_config(),
|
||
MaskChatRequestOptions::runtime(),
|
||
)
|
||
.expect("responses request should mask");
|
||
|
||
assert!(masked.redacted);
|
||
assert!(masked.session.mapping_count() >= 2);
|
||
let masked_json: serde_json::Value =
|
||
serde_json::from_slice(&masked.body).expect("masked request should stay valid JSON");
|
||
assert_eq!(masked_json["metadata"]["owner"], "[email protected]");
|
||
let instructions = masked_json["instructions"]
|
||
.as_str()
|
||
.expect("instructions should remain a string");
|
||
assert!(!instructions.contains("[email protected]"));
|
||
assert!(!masked_json["input"][0]["content"][0]["text"]
|
||
.as_str()
|
||
.expect("input text should remain a string")
|
||
.contains("+14155552671"));
|
||
assert_eq!(
|
||
masked_json["input"][0]["content"][1]["text"],
|
||
"image [email protected] must stay"
|
||
);
|
||
assert!(!masked_json["input"][1]["arguments"]
|
||
.as_str()
|
||
.expect("arguments should remain a string")
|
||
.contains("secretValueABCDEF1234567890abcdef"));
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_preserves_opaque_reasoning_state_while_masking_user_input() {
|
||
let request = json!({
|
||
"model": "deepseek-reasoner",
|
||
"input": [
|
||
{
|
||
"type": "reasoning",
|
||
"encrypted_content": "opaque-state-must-remain-byte-identical",
|
||
"future_state": {"version": 2},
|
||
"content": [
|
||
{
|
||
"type": "reasoning_text",
|
||
"text": "Provider reasoning state must remain byte-identical"
|
||
},
|
||
{
|
||
"type": "summary_text",
|
||
"text": "Bound future text must also remain byte-identical"
|
||
}
|
||
]
|
||
},
|
||
{
|
||
"type": "message",
|
||
"role": "user",
|
||
"content": [{
|
||
"type": "input_text",
|
||
"text": "Please contact [email protected]"
|
||
}]
|
||
}
|
||
]
|
||
});
|
||
let raw = serde_json::to_vec(&request).expect("request should serialize");
|
||
|
||
let masked = try_mask_chat_pii_request_json_with_options(
|
||
&raw,
|
||
ChatPiiRedactionRequestFormat::OpenAiResponses,
|
||
test_config(),
|
||
deepseek_opaque_replay_options(),
|
||
)
|
||
.expect("reasoning request should mask");
|
||
|
||
assert!(masked.redacted);
|
||
let masked_json: serde_json::Value =
|
||
serde_json::from_slice(&masked.body).expect("masked request should stay valid JSON");
|
||
assert_eq!(
|
||
masked_json["input"][0]["encrypted_content"],
|
||
"opaque-state-must-remain-byte-identical"
|
||
);
|
||
assert_eq!(
|
||
masked_json["input"][0]["future_state"],
|
||
json!({"version": 2})
|
||
);
|
||
assert_eq!(
|
||
masked_json["input"][0]["content"][0]["text"],
|
||
"Provider reasoning state must remain byte-identical"
|
||
);
|
||
assert_eq!(
|
||
masked_json["input"][0]["content"][1]["text"],
|
||
"Bound future text must also remain byte-identical",
|
||
"the complete provider-bound item must remain unchanged"
|
||
);
|
||
let user_text = masked_json["input"][1]["content"][0]["text"]
|
||
.as_str()
|
||
.expect("user text should remain a string");
|
||
assert!(!user_text.contains("[email protected]"));
|
||
assert!(user_text.contains("<AETHER:EMAIL:"));
|
||
|
||
let sentinel = masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("masked user email should have a sentinel");
|
||
let provider_response = serde_json::to_vec(&json!({
|
||
"output": [
|
||
{
|
||
"type": "reasoning",
|
||
"encrypted_content": "provider-bound-state",
|
||
"future_state": {"version": 3},
|
||
"content": [{
|
||
"type": "reasoning_text",
|
||
"text": format!("opaque {sentinel}")
|
||
}]
|
||
},
|
||
{
|
||
"type": "message",
|
||
"content": [{
|
||
"type": "output_text",
|
||
"text": format!("visible {sentinel}")
|
||
}]
|
||
}
|
||
]
|
||
}))
|
||
.expect("provider response should serialize");
|
||
let mut headers =
|
||
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||
let restored =
|
||
restore_sync_response_body(&mut headers, provider_response.as_slice(), &masked.session)
|
||
.expect("provider response should restore");
|
||
let restored_json: serde_json::Value =
|
||
serde_json::from_slice(&restored.body).expect("restored response should remain JSON");
|
||
let nested_reasoning_text = restored_json["output"][0]["content"][0].clone();
|
||
assert_eq!(
|
||
restored_json["output"][0]["content"][0]["text"],
|
||
format!("opaque {sentinel}"),
|
||
"provider-bound reasoning text must remain byte-identical"
|
||
);
|
||
assert_eq!(
|
||
restored_json["output"][0]["encrypted_content"],
|
||
"provider-bound-state"
|
||
);
|
||
assert_eq!(
|
||
restored_json["output"][0]["future_state"],
|
||
json!({"version": 3})
|
||
);
|
||
assert_eq!(
|
||
restored_json["output"][1]["content"][0]["text"], "visible [email protected]",
|
||
"ordinary assistant output should still be restored"
|
||
);
|
||
assert_eq!(
|
||
nested_reasoning_text,
|
||
json!({
|
||
"type": "reasoning_text",
|
||
"text": format!("opaque {sentinel}")
|
||
}),
|
||
"opaque reasoning parts must retain all original fields and values"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_rejects_sensitive_provider_bound_reasoning_state() {
|
||
let request = json!({
|
||
"model": "deepseek-reasoner",
|
||
"input": [{
|
||
"type": "reasoning",
|
||
"encrypted_content": "opaque-state-must-remain-byte-identical",
|
||
"content": [{
|
||
"type": "reasoning_text",
|
||
"text": "Provider state contains [email protected]"
|
||
}]
|
||
}]
|
||
});
|
||
let raw = serde_json::to_vec(&request).expect("request should serialize");
|
||
|
||
let error = try_mask_chat_pii_request_json_with_options(
|
||
&raw,
|
||
ChatPiiRedactionRequestFormat::OpenAiResponses,
|
||
test_config(),
|
||
deepseek_opaque_replay_options(),
|
||
)
|
||
.expect_err("sensitive opaque state must fail closed");
|
||
|
||
assert_eq!(error, RedactionMaskError::SensitiveOpaqueReasoningState);
|
||
}
|
||
|
||
#[test]
|
||
fn forged_opaque_reasoning_shape_cannot_opt_out_of_pii_masking() {
|
||
let request = json!({
|
||
"model": "client-selected-model",
|
||
"input": [{
|
||
"type": "reasoning",
|
||
"encrypted_content": "client-forged-opaque-state",
|
||
"content": [{
|
||
"type": "reasoning_text",
|
||
"text": "Send this to [email protected]"
|
||
}]
|
||
}]
|
||
});
|
||
let raw = serde_json::to_vec(&request).expect("request should serialize");
|
||
|
||
let masked = try_mask_chat_pii_request_json_with_options(
|
||
&raw,
|
||
ChatPiiRedactionRequestFormat::OpenAiResponses,
|
||
test_config(),
|
||
MaskChatRequestOptions::runtime(),
|
||
)
|
||
.expect("an untrusted reasoning shape should use ordinary masking");
|
||
|
||
assert!(masked.redacted);
|
||
let masked_json: Value =
|
||
serde_json::from_slice(&masked.body).expect("masked request should remain JSON");
|
||
let text = masked_json["input"][0]["content"][0]["text"]
|
||
.as_str()
|
||
.expect("reasoning text should remain a string");
|
||
assert!(!text.contains("[email protected]"));
|
||
assert!(text.contains("<AETHER:EMAIL:"));
|
||
assert_eq!(
|
||
masked_json["input"][0]["encrypted_content"],
|
||
"client-forged-opaque-state"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn cached_pii_redaction_rejects_sensitive_provider_bound_reasoning_state() {
|
||
let request = json!({
|
||
"model": "deepseek-reasoner",
|
||
"input": [{
|
||
"type": "reasoning",
|
||
"encrypted_content": "opaque-state-must-remain-byte-identical",
|
||
"future_state": {
|
||
"nested": {"future_owner": "[email protected]"}
|
||
},
|
||
"content": [{
|
||
"type": "reasoning_text",
|
||
"text": "Provider state remains byte-identical"
|
||
}]
|
||
}]
|
||
});
|
||
|
||
let error = try_mask_chat_pii_request_value_with_cache_options(
|
||
&request,
|
||
ChatPiiRedactionRequestFormat::OpenAiResponses,
|
||
test_config(),
|
||
deepseek_opaque_replay_options(),
|
||
None,
|
||
)
|
||
.await
|
||
.expect_err("sensitive future text in opaque state must fail closed");
|
||
|
||
assert_eq!(error, RedactionMaskError::SensitiveOpaqueReasoningState);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_restore_still_restores_unbound_reasoning_text_events() {
|
||
let masked = try_mask_chat_pii_request_json_with_options(
|
||
br#"{"model":"gpt","input":"[email protected]"}"#,
|
||
ChatPiiRedactionRequestFormat::OpenAiResponses,
|
||
test_config(),
|
||
MaskChatRequestOptions::runtime(),
|
||
)
|
||
.expect("input should mask");
|
||
let sentinel = masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("masked email should have a sentinel");
|
||
let mut event = json!({
|
||
"type": "response.reasoning_text.delta",
|
||
"delta": format!("thinking about {sentinel}")
|
||
});
|
||
|
||
assert!(restore_json_strings(&mut event, &masked.session));
|
||
assert_eq!(
|
||
event["delta"], "thinking about [email protected]",
|
||
"an event name alone must not be treated as opaque bound state"
|
||
);
|
||
|
||
let mut unbound_part = json!({
|
||
"type": "reasoning_text",
|
||
"text": format!("thinking about {sentinel}")
|
||
});
|
||
assert!(restore_json_strings(&mut unbound_part, &masked.session));
|
||
assert_eq!(
|
||
unbound_part["text"], "thinking about [email protected]",
|
||
"a reasoning_text part without encrypted parent state must still restore"
|
||
);
|
||
|
||
let unbound_request = json!({
|
||
"model": "gpt",
|
||
"input": [{
|
||
"type": "reasoning",
|
||
"content": [{
|
||
"type": "reasoning_text",
|
||
"text": "unbound [email protected]"
|
||
}]
|
||
}]
|
||
});
|
||
let unbound_raw =
|
||
serde_json::to_vec(&unbound_request).expect("unbound request should serialize");
|
||
let unbound_masked = try_mask_chat_pii_request_json_with_options(
|
||
&unbound_raw,
|
||
ChatPiiRedactionRequestFormat::OpenAiResponses,
|
||
test_config(),
|
||
MaskChatRequestOptions::runtime(),
|
||
)
|
||
.expect("unbound reasoning request should mask");
|
||
let unbound_masked_json: Value = serde_json::from_slice(&unbound_masked.body)
|
||
.expect("masked unbound request should remain JSON");
|
||
let masked_reasoning_text = unbound_masked_json["input"][0]["content"][0]["text"]
|
||
.as_str()
|
||
.expect("unbound reasoning text should remain a string");
|
||
assert!(!masked_reasoning_text.contains("[email protected]"));
|
||
assert!(masked_reasoning_text.contains("<AETHER:EMAIL:"));
|
||
|
||
let mut invalid_opaque_item = json!({
|
||
"type": "reasoning",
|
||
"encrypted_content": "non-empty-but-unbound",
|
||
"content": [{
|
||
"type": "summary_text",
|
||
"text": format!("thinking about {sentinel}")
|
||
}]
|
||
});
|
||
assert!(restore_json_strings(
|
||
&mut invalid_opaque_item,
|
||
&masked.session
|
||
));
|
||
assert_eq!(
|
||
invalid_opaque_item["content"][0]["text"], "thinking about [email protected]",
|
||
"encrypted state without a reasoning_text part must not suppress restoration"
|
||
);
|
||
|
||
let mut identified_reasoning_item = json!({
|
||
"type": "reasoning",
|
||
"id": "rs_provider_123",
|
||
"encrypted_content": "provider-state",
|
||
"content": [{
|
||
"type": "reasoning_text",
|
||
"text": format!("thinking about {sentinel}")
|
||
}]
|
||
});
|
||
assert!(restore_json_strings(
|
||
&mut identified_reasoning_item,
|
||
&masked.session
|
||
));
|
||
assert_eq!(
|
||
identified_reasoning_item["content"][0]["text"], "thinking about [email protected]",
|
||
"only id-less DeepSeek-style state may bypass ordinary restoration"
|
||
);
|
||
|
||
let identified_request = json!({
|
||
"model": "gpt",
|
||
"input": [{
|
||
"type": "reasoning",
|
||
"id": "rs_provider_123",
|
||
"encrypted_content": "provider-state",
|
||
"content": [{
|
||
"type": "reasoning_text",
|
||
"text": "identified [email protected]"
|
||
}]
|
||
}]
|
||
});
|
||
let identified_masked = try_mask_chat_pii_request_json_with_options(
|
||
&serde_json::to_vec(&identified_request).expect("request should serialize"),
|
||
ChatPiiRedactionRequestFormat::OpenAiResponses,
|
||
test_config(),
|
||
MaskChatRequestOptions::runtime(),
|
||
)
|
||
.expect("identified reasoning request should mask");
|
||
let identified_masked_json: Value = serde_json::from_slice(&identified_masked.body)
|
||
.expect("identified reasoning request should remain JSON");
|
||
let identified_text = identified_masked_json["input"][0]["content"][0]["text"]
|
||
.as_str()
|
||
.expect("identified reasoning text should remain a string");
|
||
assert!(!identified_text.contains("[email protected]"));
|
||
assert!(identified_text.contains("<AETHER:EMAIL:"));
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_request_masks_openai_search_text_fields() {
|
||
let request = json!({
|
||
"id": "session-1",
|
||
"model": "gpt-5.6",
|
||
"input": "Find [email protected]",
|
||
"commands": {
|
||
"search_query": [{"q": "Phone +14155552671"}],
|
||
"image_query": [{"q": "Image for [email protected]"}],
|
||
"find": [{"ref_id": "https://example.com/[email protected]", "pattern": "secret_key=secretValueABCDEF1234567890abcdef"}],
|
||
"weather": [{"location": "Contact [email protected]"}],
|
||
"open": [{"ref_id": "https://example.com/[email protected]"}]
|
||
}
|
||
});
|
||
let raw = serde_json::to_vec(&request).expect("request should serialize");
|
||
|
||
let masked = try_mask_chat_pii_request_json_with_options(
|
||
&raw,
|
||
ChatPiiRedactionRequestFormat::OpenAiSearch,
|
||
test_config(),
|
||
MaskChatRequestOptions::runtime(),
|
||
)
|
||
.expect("search request should mask");
|
||
let masked_json: Value =
|
||
serde_json::from_slice(&masked.body).expect("masked request should parse");
|
||
|
||
assert!(masked.redacted);
|
||
assert!(!masked_json["input"]
|
||
.as_str()
|
||
.unwrap()
|
||
.contains("[email protected]"));
|
||
assert!(!masked_json["commands"]["search_query"][0]["q"]
|
||
.as_str()
|
||
.unwrap()
|
||
.contains("+14155552671"));
|
||
assert!(!masked_json["commands"]["image_query"][0]["q"]
|
||
.as_str()
|
||
.unwrap()
|
||
.contains("[email protected]"));
|
||
assert!(!masked_json["commands"]["find"][0]["pattern"]
|
||
.as_str()
|
||
.unwrap()
|
||
.contains("secretValueABCDEF1234567890abcdef"));
|
||
assert!(!masked_json["commands"]["weather"][0]["location"]
|
||
.as_str()
|
||
.unwrap()
|
||
.contains("[email protected]"));
|
||
assert_eq!(
|
||
masked_json["commands"]["open"][0]["ref_id"],
|
||
"https://example.com/[email protected]"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn pii_redaction_async_request_masks_openai_search_text_fields() {
|
||
let request = json!({
|
||
"id": "session-1",
|
||
"model": "gpt-5.6",
|
||
"input": "Find [email protected]",
|
||
"commands": {
|
||
"search_query": [{"q": "Phone +14155552671"}],
|
||
"find": [{"ref_id": "turn0search0", "pattern": "[email protected]"}]
|
||
}
|
||
});
|
||
|
||
let masked = try_mask_chat_pii_request_value_with_cache_options(
|
||
&request,
|
||
ChatPiiRedactionRequestFormat::OpenAiSearch,
|
||
test_config(),
|
||
MaskChatRequestOptions::runtime(),
|
||
None,
|
||
)
|
||
.await
|
||
.expect("search request should mask");
|
||
let masked_json = masked.body_json.expect("masked body should be present");
|
||
|
||
assert!(masked.redacted);
|
||
assert!(!masked_json["input"]
|
||
.as_str()
|
||
.unwrap()
|
||
.contains("[email protected]"));
|
||
assert!(!masked_json["commands"]["search_query"][0]["q"]
|
||
.as_str()
|
||
.unwrap()
|
||
.contains("+14155552671"));
|
||
assert!(!masked_json["commands"]["find"][0]["pattern"]
|
||
.as_str()
|
||
.unwrap()
|
||
.contains("[email protected]"));
|
||
assert_eq!(masked_json["commands"]["find"][0]["ref_id"], "turn0search0");
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_request_avoids_cross_message_and_tool_argument_sentinel_collisions() {
|
||
let mut probe = session_at(600);
|
||
probe.redact_text("Contact [email protected]");
|
||
let colliding_literal = probe
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("probe sentinel should exist")
|
||
.to_string();
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [
|
||
{"role": "user", "content": format!("literal {colliding_literal} stays unknown")},
|
||
{
|
||
"role": "assistant",
|
||
"content": null,
|
||
"tool_calls": [{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "lookup_contact",
|
||
"arguments": "{\"email\":\"[email protected]\"}"
|
||
}
|
||
}]
|
||
}
|
||
]
|
||
});
|
||
|
||
let masked = mask_chat_request_json(
|
||
&serde_json::to_vec(&request).expect("request should serialize"),
|
||
test_config(),
|
||
);
|
||
|
||
let current_sentinel = masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("current request sentinel should exist")
|
||
.to_string();
|
||
assert_ne!(current_sentinel, colliding_literal);
|
||
let masked_json: serde_json::Value =
|
||
serde_json::from_slice(&masked.body).expect("masked request should parse");
|
||
assert_eq!(
|
||
masked_json["messages"][0]["content"],
|
||
format!("literal {colliding_literal} stays unknown")
|
||
);
|
||
let arguments = masked_json["messages"][1]["tool_calls"][0]["function"]["arguments"]
|
||
.as_str()
|
||
.expect("arguments should be text");
|
||
assert!(!arguments.contains("[email protected]"));
|
||
assert!(arguments.contains(¤t_sentinel));
|
||
|
||
let mut headers =
|
||
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||
let restored = restore_sync_response_body(
|
||
&mut headers,
|
||
serde_json::to_vec(&json!({
|
||
"message": format!("literal {colliding_literal} current {current_sentinel}")
|
||
}))
|
||
.expect("response should serialize")
|
||
.as_slice(),
|
||
&masked.session,
|
||
)
|
||
.expect("restore should succeed");
|
||
let restored_json: serde_json::Value =
|
||
serde_json::from_slice(&restored.body).expect("restored body should parse");
|
||
assert_eq!(
|
||
restored_json["message"],
|
||
format!("literal {colliding_literal} current [email protected]")
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_request_noops_for_invalid_absent_or_malformed_messages() {
|
||
for raw in [
|
||
b"{not json".to_vec(),
|
||
br#"{"model":"gpt-5","metadata":{"owner":"[email protected]"}}"#.to_vec(),
|
||
br#"{"model":"gpt-5","messages":{"role":"user","content":"[email protected]"}}"#
|
||
.to_vec(),
|
||
] {
|
||
let masked = mask_chat_request_json(&raw, test_config());
|
||
|
||
assert!(!masked.redacted);
|
||
assert_eq!(masked.body, raw);
|
||
assert_eq!(masked.session.mapping_count(), 0);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_request_noops_when_chat_messages_are_clean() {
|
||
let raw = br#"{"model":"gpt-5","messages":[{"role":"user","content":"hello"}],"metadata":{"trace":"keep"}}"#;
|
||
|
||
let masked = mask_chat_request_json(raw, test_config());
|
||
|
||
assert!(!masked.redacted);
|
||
assert_eq!(masked.body, raw);
|
||
assert_eq!(masked.session.mapping_count(), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_sync_restore_json_strings_reserialize_safely_and_updates_content_length() {
|
||
let sentinel = "<AETHER:EMAIL:ABCDEFGHIJKLMNOPQRST>";
|
||
let original = "[email protected] said \"ok\" from C:\\tmp\\a\nnext";
|
||
let session = session_with_response_mapping(original, sentinel);
|
||
let body = serde_json::to_vec(&json!({
|
||
"choices": [{
|
||
"message": {
|
||
"role": "assistant",
|
||
"content": format!("Restored: {sentinel}"),
|
||
}
|
||
}]
|
||
}))
|
||
.expect("response should serialize");
|
||
let mut headers = BTreeMap::from([
|
||
("content-type".to_string(), "application/json".to_string()),
|
||
("content-length".to_string(), body.len().to_string()),
|
||
]);
|
||
|
||
let restored = restore_sync_response_body(&mut headers, &body, &session)
|
||
.expect("restore should succeed");
|
||
|
||
assert!(restored.restored);
|
||
let restored_len = restored.body.len().to_string();
|
||
assert_eq!(
|
||
headers.get("content-length").map(String::as_str),
|
||
Some(restored_len.as_str())
|
||
);
|
||
let restored_json: serde_json::Value =
|
||
serde_json::from_slice(&restored.body).expect("restored body should stay valid JSON");
|
||
assert_eq!(
|
||
restored_json["choices"][0]["message"]["content"],
|
||
format!("Restored: {original}")
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_sync_restore_assistant_content_and_tool_call_arguments() {
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [
|
||
{"role": "user", "content": "Email [email protected]"},
|
||
{
|
||
"role": "assistant",
|
||
"content": null,
|
||
"tool_calls": [{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "lookup_customer",
|
||
"arguments": "{\"phone\":\"13800138000\"}"
|
||
}
|
||
}]
|
||
}
|
||
]
|
||
});
|
||
let masked = mask_chat_request_json(
|
||
&serde_json::to_vec(&request).expect("request should serialize"),
|
||
test_config(),
|
||
);
|
||
let email_sentinel = masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("email sentinel should exist");
|
||
let phone_sentinel = masked
|
||
.session
|
||
.sentinel_for_original("13800138000")
|
||
.expect("phone sentinel should exist");
|
||
let body = serde_json::to_vec(&json!({
|
||
"choices": [{
|
||
"message": {
|
||
"role": "assistant",
|
||
"content": format!("Customer {email_sentinel}"),
|
||
"tool_calls": [{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "lookup_customer",
|
||
"arguments": format!("{{\"phone\":\"{phone_sentinel}\"}}")
|
||
}
|
||
}]
|
||
}
|
||
}]
|
||
}))
|
||
.expect("response should serialize");
|
||
let mut headers = BTreeMap::from([(
|
||
"content-type".to_string(),
|
||
"application/json; charset=utf-8".to_string(),
|
||
)]);
|
||
|
||
let restored = restore_sync_response_body(&mut headers, &body, &masked.session)
|
||
.expect("restore should succeed");
|
||
|
||
let restored_json: serde_json::Value =
|
||
serde_json::from_slice(&restored.body).expect("restored body should stay valid JSON");
|
||
assert_eq!(
|
||
restored_json["choices"][0]["message"]["content"],
|
||
"Customer [email protected]"
|
||
);
|
||
assert_eq!(
|
||
restored_json["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"],
|
||
"{\"phone\":\"13800138000\"}"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_sync_restore_unknown_sentinel_looking_strings_remain_unchanged() {
|
||
let known_sentinel = "<AETHER:EMAIL:ABCDEFGHIJKLMNOPQRST>";
|
||
let unknown_sentinel = "<AETHER:EMAIL:TSRQPONMLKJIHGFEDCBA>";
|
||
let session = session_with_response_mapping("[email protected]", known_sentinel);
|
||
let body = serde_json::to_vec(&json!({
|
||
"message": format!("known {known_sentinel} unknown {unknown_sentinel}"),
|
||
}))
|
||
.expect("response should serialize");
|
||
let mut headers =
|
||
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||
|
||
let restored = restore_sync_response_body(&mut headers, &body, &session)
|
||
.expect("restore should succeed");
|
||
|
||
let restored_json: serde_json::Value =
|
||
serde_json::from_slice(&restored.body).expect("restored body should stay valid JSON");
|
||
assert_eq!(
|
||
restored_json["message"],
|
||
format!("known [email protected] unknown {unknown_sentinel}")
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_sync_restore_text_body_replaces_current_request_sentinels() {
|
||
let known_sentinel = "<AETHER:EMAIL:ABCDEFGHIJKLMNOPQRST>";
|
||
let unknown_sentinel = "<AETHER:EMAIL:TSRQPONMLKJIHGFEDCBA>";
|
||
let session = session_with_response_mapping("[email protected]", known_sentinel);
|
||
let body = format!("plain {known_sentinel} and {unknown_sentinel}").into_bytes();
|
||
let mut headers = BTreeMap::from([
|
||
("content-type".to_string(), "text/plain".to_string()),
|
||
("Content-Length".to_string(), body.len().to_string()),
|
||
]);
|
||
|
||
let restored = restore_sync_response_body(&mut headers, &body, &session)
|
||
.expect("restore should succeed");
|
||
|
||
assert_eq!(
|
||
std::str::from_utf8(&restored.body).expect("restored text should be utf8"),
|
||
format!("plain [email protected] and {unknown_sentinel}")
|
||
);
|
||
let restored_len = restored.body.len().to_string();
|
||
assert_eq!(headers.get("Content-Length"), None);
|
||
assert_eq!(
|
||
headers.get("content-length").map(String::as_str),
|
||
Some(restored_len.as_str())
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_compressed_response_safe_error() {
|
||
let sentinel = "<AETHER:EMAIL:ABCDEFGHIJKLMNOPQRST>";
|
||
let session = session_with_response_mapping("[email protected]", sentinel);
|
||
let body = format!("compressed bytes with {sentinel}").into_bytes();
|
||
let mut headers = BTreeMap::from([
|
||
("Content-Encoding".to_string(), "identity".to_string()),
|
||
("content-encoding".to_string(), "gzip".to_string()),
|
||
("content-length".to_string(), body.len().to_string()),
|
||
]);
|
||
|
||
let err = restore_sync_response_body(&mut headers, &body, &session)
|
||
.expect_err("compressed response should fail safely");
|
||
|
||
let message = format!("{err:?}");
|
||
assert!(message.contains("encoded response bodies"));
|
||
assert!(!message.contains("[email protected]"));
|
||
assert!(!message.contains(sentinel));
|
||
assert_eq!(
|
||
headers.get("Content-Encoding").map(String::as_str),
|
||
Some("identity")
|
||
);
|
||
assert_eq!(
|
||
headers.get("content-encoding").map(String::as_str),
|
||
Some("gzip")
|
||
);
|
||
let original_len = body.len().to_string();
|
||
assert_eq!(
|
||
headers.get("content-length").map(String::as_str),
|
||
Some(original_len.as_str())
|
||
);
|
||
|
||
let stream_err = match StreamingResponseRestorer::new(&headers, &session) {
|
||
Ok(_) => panic!("compressed stream response should fail safely"),
|
||
Err(err) => err,
|
||
};
|
||
let stream_message = format!("{stream_err:?}");
|
||
assert!(stream_message.contains("encoded response bodies"));
|
||
assert!(!stream_message.contains("[email protected]"));
|
||
assert!(!stream_message.contains(sentinel));
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_stream_restore_raw_text_split_sentinel_across_three_chunks() {
|
||
let known_sentinel = "<AETHER:EMAIL:ABCDEFGHIJKLMNOPQRST>";
|
||
let unknown_sentinel = "<AETHER:EMAIL:TSRQPONMLKJIHGFEDCBA>";
|
||
let session = session_with_response_mapping("[email protected]", known_sentinel);
|
||
let mut restorer = StreamingResponseRestorer::for_text(&session);
|
||
let chunks = [
|
||
b"prefix <AET".as_slice(),
|
||
b"HER:EMAIL:ABC".as_slice(),
|
||
b"DEFGHIJKLMNOPQRST> middle ".as_slice(),
|
||
unknown_sentinel.as_bytes(),
|
||
];
|
||
|
||
let mut output = Vec::new();
|
||
for chunk in chunks {
|
||
output.extend(restorer.push_chunk(chunk).expect("restore should succeed"));
|
||
assert!(restorer.pending_text_carry_len() <= restorer.max_text_carry_len());
|
||
}
|
||
output.extend(restorer.finish().expect("finish should succeed"));
|
||
|
||
assert_eq!(
|
||
String::from_utf8(output).expect("restored stream should be utf8"),
|
||
format!("prefix [email protected] middle {unknown_sentinel}")
|
||
);
|
||
assert_eq!(restorer.pending_text_carry_len(), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_stream_restore_bounded_carry_and_final_flush() {
|
||
let sentinel = "<AETHER:EMAIL:ABCDEFGHIJKLMNOPQRST>";
|
||
let session = session_with_response_mapping("[email protected]", sentinel);
|
||
let mut restorer = StreamingResponseRestorer::for_text(&session);
|
||
let chunk = vec![b'x'; sentinel.len() * 3];
|
||
|
||
let first = restorer.push_chunk(&chunk).expect("restore should succeed");
|
||
|
||
assert!(!first.is_empty());
|
||
assert!(restorer.pending_text_carry_len() <= restorer.max_text_carry_len());
|
||
let flushed = restorer.finish().expect("finish should succeed");
|
||
assert_eq!(first.len() + flushed.len(), chunk.len());
|
||
assert_eq!(restorer.pending_text_carry_len(), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_stream_restore_sse_json_split_sentinel_across_three_chunks() {
|
||
let sentinel = "<AETHER:EMAIL:ABCDEFGHIJKLMNOPQRST>";
|
||
let session = session_with_response_mapping("[email protected]", sentinel);
|
||
let mut restorer = StreamingResponseRestorer::for_sse(&session);
|
||
let chunks = [
|
||
b"data: {\"choices\":[{\"delta\":{\"content\":\"hello <AET".as_slice(),
|
||
b"HER:EMAIL:ABC".as_slice(),
|
||
b"DEFGHIJKLMNOPQRST>\"}}]}\n\n".as_slice(),
|
||
];
|
||
|
||
let mut output = Vec::new();
|
||
for chunk in chunks {
|
||
output.extend(restorer.push_chunk(chunk).expect("restore should succeed"));
|
||
}
|
||
output.extend(restorer.finish().expect("finish should succeed"));
|
||
let output_text = String::from_utf8(output).expect("restored stream should be utf8");
|
||
|
||
assert!(output_text.contains("hello [email protected]"));
|
||
assert!(!output_text.contains(sentinel));
|
||
assert!(output_text.ends_with("\n\n"));
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_stream_restore_sse_tool_call_argument_delta_json_strings() {
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [
|
||
{"role": "user", "content": "Email [email protected]"},
|
||
{
|
||
"role": "assistant",
|
||
"content": null,
|
||
"tool_calls": [{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "lookup_customer",
|
||
"arguments": "{\"phone\":\"13800138000\"}"
|
||
}
|
||
}]
|
||
}
|
||
]
|
||
});
|
||
let masked = mask_chat_request_json(
|
||
&serde_json::to_vec(&request).expect("request should serialize"),
|
||
test_config(),
|
||
);
|
||
let email_sentinel = masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("email sentinel should exist");
|
||
let phone_sentinel = masked
|
||
.session
|
||
.sentinel_for_original("13800138000")
|
||
.expect("phone sentinel should exist");
|
||
let event = format!(
|
||
"data: {}\n\n",
|
||
serde_json::to_string(&json!({
|
||
"choices": [{
|
||
"delta": {
|
||
"content": format!("Customer {email_sentinel}"),
|
||
"tool_calls": [{
|
||
"index": 0,
|
||
"function": {
|
||
"arguments": format!("{{\"phone\":\"{phone_sentinel}\"}}")
|
||
}
|
||
}]
|
||
}
|
||
}]
|
||
}))
|
||
.expect("event should serialize")
|
||
);
|
||
let mut restorer = StreamingResponseRestorer::for_sse(&masked.session);
|
||
|
||
let output = restorer
|
||
.push_chunk(event.as_bytes())
|
||
.expect("restore should succeed");
|
||
|
||
let output_text = String::from_utf8(output).expect("restored stream should be utf8");
|
||
assert!(output_text.contains("Customer [email protected]"));
|
||
assert!(output_text.contains("\\\"phone\\\":\\\"13800138000\\\""));
|
||
assert!(!output_text.contains(email_sentinel));
|
||
assert!(!output_text.contains(phone_sentinel));
|
||
assert!(restorer.finish().expect("finish should succeed").is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_stream_restore_sse_preserves_comments_blank_crlf_multiple_data_and_done() {
|
||
let sentinel = "<AETHER:EMAIL:ABCDEFGHIJKLMNOPQRST>";
|
||
let session = session_with_response_mapping("[email protected]", sentinel);
|
||
let mut restorer = StreamingResponseRestorer::for_sse(&session);
|
||
let input = format!(
|
||
": keep-alive\r\n\r\nevent: message\r\ndata: plain {sentinel}\r\ndata: second line\r\n\r\ndata: [DONE]\r\n\r\n"
|
||
);
|
||
|
||
let mut output = restorer
|
||
.push_chunk(input.as_bytes())
|
||
.expect("restore should succeed");
|
||
output.extend(restorer.finish().expect("finish should succeed"));
|
||
|
||
assert_eq!(
|
||
String::from_utf8(output).expect("restored stream should be utf8"),
|
||
concat!(
|
||
": keep-alive\r\n",
|
||
"\r\n",
|
||
"event: message\r\n",
|
||
"data: plain [email protected]\r\n",
|
||
"data: second line\r\n",
|
||
"\r\n",
|
||
"data: [DONE]\r\n",
|
||
"\r\n"
|
||
)
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn proxy_pii_redaction_rule_selection_controls_email_detector() {
|
||
let config = ChatPiiRedactionRuntimeConfig {
|
||
enabled: true,
|
||
rules: parse_chat_pii_redaction_rules(Some(&json!([{
|
||
"id": "email",
|
||
"name": "邮箱",
|
||
"pattern": "(?i)[A-Z0-9._%+-]{1,64}@[A-Z0-9.-]{1,253}\\.[A-Z]{2,63}",
|
||
"enabled": true,
|
||
"kind": "email",
|
||
"system": true
|
||
}])))
|
||
.expect("email rule should compile"),
|
||
..ChatPiiRedactionRuntimeConfig::default()
|
||
};
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [{"role": "user", "content": "Email [email protected] phone +14155552671"}]
|
||
});
|
||
|
||
let masked = mask_chat_request_json_with_options(
|
||
&serde_json::to_vec(&request).expect("request should serialize"),
|
||
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
|
||
MaskChatRequestOptions::runtime(),
|
||
);
|
||
|
||
let masked_json: serde_json::Value =
|
||
serde_json::from_slice(&masked.body).expect("masked request should parse");
|
||
let content = masked_json["messages"][0]["content"]
|
||
.as_str()
|
||
.expect("content should stay string");
|
||
assert!(!content.contains("[email protected]"));
|
||
assert!(content.contains("+14155552671"));
|
||
assert!(content.contains("<AETHER:EMAIL:"));
|
||
assert!(!content.contains("<AETHER:PHONE:"));
|
||
}
|
||
|
||
#[test]
|
||
fn proxy_pii_redaction_disabled_rules_disable_all_detectors() {
|
||
let config = ChatPiiRedactionRuntimeConfig {
|
||
enabled: true,
|
||
rules: parse_chat_pii_redaction_rules(Some(&json!([{
|
||
"id": "email",
|
||
"name": "邮箱",
|
||
"pattern": "(?i)[A-Z0-9._%+-]{1,64}@[A-Z0-9.-]{1,253}\\.[A-Z]{2,63}",
|
||
"enabled": false,
|
||
"kind": "email",
|
||
"system": true
|
||
}])))
|
||
.expect("disabled email rule should compile"),
|
||
..ChatPiiRedactionRuntimeConfig::default()
|
||
};
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [{"role": "user", "content": "Email [email protected] phone +14155552671"}]
|
||
});
|
||
|
||
let masked = mask_chat_request_json_with_options(
|
||
&serde_json::to_vec(&request).expect("request should serialize"),
|
||
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
|
||
MaskChatRequestOptions::runtime(),
|
||
);
|
||
|
||
assert!(!masked.redacted);
|
||
let masked_json: serde_json::Value =
|
||
serde_json::from_slice(&masked.body).expect("masked request should parse");
|
||
assert_eq!(masked_json, request);
|
||
assert!(masked_json.to_string().contains("[email protected]"));
|
||
assert!(!masked_json.to_string().contains("<AETHER:"));
|
||
}
|
||
|
||
#[test]
|
||
fn redaction_session_slot_uses_single_candidate_when_response_header_missing() {
|
||
let slot = RedactionSessionSlot::default();
|
||
let sentinel = "<AETHER:EMAIL:SINGLECANDIDATE001>";
|
||
slot.put_for_candidate(
|
||
"candidate-1",
|
||
session_with_response_mapping("[email protected]", sentinel),
|
||
);
|
||
|
||
let session = slot
|
||
.take_for_candidate(None)
|
||
.expect("single candidate session should be used without header");
|
||
|
||
assert_eq!(
|
||
session.sentinel_for_original("[email protected]"),
|
||
Some(sentinel)
|
||
);
|
||
assert!(slot.take_for_candidate(Some("candidate-1")).is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn redaction_session_slot_does_not_guess_when_response_header_missing_for_multiple_candidates()
|
||
{
|
||
let slot = RedactionSessionSlot::default();
|
||
let alice_sentinel = "<AETHER:EMAIL:MULTICANDIDATE001>";
|
||
let bob_sentinel = "<AETHER:EMAIL:MULTICANDIDATE002>";
|
||
slot.put_for_candidate(
|
||
"candidate-1",
|
||
session_with_response_mapping("[email protected]", alice_sentinel),
|
||
);
|
||
slot.put_for_candidate(
|
||
"candidate-2",
|
||
session_with_response_mapping("[email protected]", bob_sentinel),
|
||
);
|
||
|
||
assert!(slot.take_for_candidate(None).is_none());
|
||
|
||
let session = slot
|
||
.take_for_candidate(Some("candidate-2"))
|
||
.expect("specific candidate session should remain available");
|
||
assert_eq!(
|
||
session.sentinel_for_original("[email protected]"),
|
||
Some(bob_sentinel)
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn redaction_session_slot_does_not_fallback_to_other_candidate_on_header_miss() {
|
||
let slot = RedactionSessionSlot::default();
|
||
slot.put_for_candidate(
|
||
"candidate-1",
|
||
session_with_response_mapping("[email protected]", "<AETHER:EMAIL:HEADERMISS001>"),
|
||
);
|
||
|
||
assert!(slot.take_for_candidate(Some("candidate-2")).is_none());
|
||
assert!(slot.take_for_candidate(Some("candidate-1")).is_some());
|
||
}
|
||
|
||
#[test]
|
||
fn proxy_pii_redaction_provider_bound_request_uses_sentinels_without_prompt_notice() {
|
||
let config = ChatPiiRedactionRuntimeConfig::default();
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [
|
||
{"role": "system", "content": "be brief"},
|
||
{"role": "user", "content": "Contact [email protected]"},
|
||
{"role": "assistant", "content": "ok"}
|
||
]
|
||
});
|
||
|
||
let masked = mask_chat_request_json_with_options(
|
||
&serde_json::to_vec(&request).expect("request should serialize"),
|
||
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
|
||
MaskChatRequestOptions::runtime(),
|
||
);
|
||
|
||
assert!(masked.redacted);
|
||
let masked_json: serde_json::Value =
|
||
serde_json::from_slice(&masked.body).expect("masked request should parse");
|
||
let messages = masked_json["messages"]
|
||
.as_array()
|
||
.expect("messages should be an array");
|
||
assert_eq!(messages.len(), 3);
|
||
assert_eq!(messages[0]["role"], "system");
|
||
assert!(messages[1..]
|
||
.iter()
|
||
.all(|message| message["role"].as_str() != Some("system")));
|
||
assert_eq!(messages[1]["role"], "user");
|
||
let content = messages[1]["content"]
|
||
.as_str()
|
||
.expect("user content should be text");
|
||
assert!(!content.contains("[email protected]"));
|
||
assert!(content.contains("<AETHER:EMAIL:"));
|
||
assert_eq!(messages[2]["role"], "assistant");
|
||
}
|
||
|
||
#[test]
|
||
fn proxy_pii_redaction_sync_and_stream_client_responses_restore_current_sentinels() {
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [{"role": "user", "content": "Email [email protected]"}]
|
||
});
|
||
let masked = mask_chat_request_json(
|
||
&serde_json::to_vec(&request).expect("request should serialize"),
|
||
test_config(),
|
||
);
|
||
let sentinel = masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("email sentinel should exist");
|
||
let sync_body = serde_json::to_vec(&json!({
|
||
"choices": [{"message": {"role": "assistant", "content": format!("hello {sentinel}")}}]
|
||
}))
|
||
.expect("response should serialize");
|
||
let mut headers =
|
||
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||
|
||
let restored = restore_sync_response_body(&mut headers, &sync_body, &masked.session)
|
||
.expect("sync restore should succeed");
|
||
let restored_json: serde_json::Value =
|
||
serde_json::from_slice(&restored.body).expect("restored response should parse");
|
||
assert_eq!(
|
||
restored_json["choices"][0]["message"]["content"],
|
||
"hello [email protected]"
|
||
);
|
||
|
||
let sse =
|
||
format!("data: {{\"choices\":[{{\"delta\":{{\"content\":\"{sentinel}\"}}}}]}}\n\n");
|
||
let mut restorer = StreamingResponseRestorer::for_sse(&masked.session);
|
||
let mut output = restorer
|
||
.push_chunk(sse.as_bytes())
|
||
.expect("stream restore should succeed");
|
||
output.extend(restorer.finish().expect("stream finish should succeed"));
|
||
let output = String::from_utf8(output).expect("stream should be utf8");
|
||
assert!(output.contains("[email protected]"));
|
||
assert!(!output.contains(sentinel));
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_performance_prefilters_skip_impossible_expensive_validators() {
|
||
let mut clean_probe = DetectorProbe::default();
|
||
let clean_candidates = detect_candidates_with_probe(
|
||
"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz",
|
||
Some(&mut clean_probe),
|
||
);
|
||
|
||
assert!(clean_candidates.is_empty());
|
||
for kind in [
|
||
RedactionKind::ChinaResidentId,
|
||
RedactionKind::PaymentCard,
|
||
RedactionKind::Jwt,
|
||
RedactionKind::ApiKey,
|
||
] {
|
||
assert_eq!(clean_probe.validator_calls(kind), 0);
|
||
}
|
||
|
||
let mut token_probe = DetectorProbe::default();
|
||
let token_candidates = detect_candidates_with_probe(
|
||
"token q8HjK2LmN9PqR4StU7VwX0YzA3BcD6Ef",
|
||
Some(&mut token_probe),
|
||
);
|
||
|
||
assert_eq!(token_probe.validator_calls(RedactionKind::ApiKey), 1);
|
||
assert!(token_candidates
|
||
.iter()
|
||
.any(|candidate| candidate.kind == Some(RedactionKind::ApiKey)));
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_accepts_text_and_detection_counts_above_previous_caps() {
|
||
let large_request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [{
|
||
"role": "user",
|
||
"content": format!("[email protected] {}", "x".repeat(2 * 1024 * 1024 + 1))
|
||
}]
|
||
});
|
||
let large_masked = try_mask_chat_request_json_with_options(
|
||
&serde_json::to_vec(&large_request).expect("request should serialize"),
|
||
test_config(),
|
||
MaskChatRequestOptions::runtime(),
|
||
)
|
||
.expect("text above the previous scan cap should be redacted");
|
||
assert!(large_masked.redacted);
|
||
assert_eq!(large_masked.session.mapping_count(), 1);
|
||
|
||
let dense_content = (0..1025)
|
||
.map(|index| format!("user{index}@example.com"))
|
||
.collect::<Vec<_>>()
|
||
.join(" ");
|
||
let dense_request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [{"role": "user", "content": dense_content}]
|
||
});
|
||
let dense_masked = try_mask_chat_request_json_with_options(
|
||
&serde_json::to_vec(&dense_request).expect("request should serialize"),
|
||
test_config(),
|
||
MaskChatRequestOptions::runtime(),
|
||
)
|
||
.expect("detections above the previous cap should be redacted");
|
||
assert!(dense_masked.redacted);
|
||
assert_eq!(dense_masked.session.mapping_count(), 1025);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn pii_redaction_performance_multiturn_cache_reuses_and_rolls_over_sentinels() {
|
||
let Some(redis) = start_managed_redis_or_skip().await else {
|
||
return;
|
||
};
|
||
let runtime_state =
|
||
redis_cache_runtime_state(redis.redis_url(), "pii_redaction_performance_cache").await;
|
||
let cache = RedisRedactionMappingCache::new(&runtime_state);
|
||
let config = email_only_runtime_config();
|
||
|
||
let first_request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [{"role": "user", "content": "Contact [email protected]"}],
|
||
"provider_context": "provider-a"
|
||
});
|
||
let first_masked = try_mask_chat_request_json_with_cache_options(
|
||
&serde_json::to_vec(&first_request).expect("request should serialize"),
|
||
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
|
||
MaskChatRequestOptions::runtime(),
|
||
Some(&cache),
|
||
)
|
||
.await
|
||
.expect("first masking should succeed");
|
||
let first_sentinel = first_masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("first sentinel should exist")
|
||
.to_string();
|
||
|
||
let second_request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [{"role": "user", "content": "Contact [email protected] and [email protected]"}],
|
||
"provider_context": "provider-b"
|
||
});
|
||
let second_masked = try_mask_chat_request_json_with_cache_options(
|
||
&serde_json::to_vec(&second_request).expect("request should serialize"),
|
||
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 899),
|
||
MaskChatRequestOptions::runtime(),
|
||
Some(&cache),
|
||
)
|
||
.await
|
||
.expect("second masking should succeed");
|
||
assert_eq!(
|
||
second_masked
|
||
.session
|
||
.sentinel_for_original("[email protected]"),
|
||
Some(first_sentinel.as_str())
|
||
);
|
||
assert!(second_masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.is_some());
|
||
let second_masked_json: serde_json::Value = serde_json::from_slice(&second_masked.body)
|
||
.expect("second masked request should parse");
|
||
assert_eq!(second_masked_json["provider_context"], "provider-b");
|
||
let second_content = second_masked_json["messages"][0]["content"]
|
||
.as_str()
|
||
.expect("second content should be text");
|
||
assert!(second_content.contains(&first_sentinel));
|
||
assert!(!second_content.contains("[email protected]"));
|
||
assert!(!second_content.contains("[email protected]"));
|
||
let mut headers =
|
||
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||
let restored = restore_sync_response_body(
|
||
&mut headers,
|
||
serde_json::to_vec(&json!({"message": format!("cached {first_sentinel}")}))
|
||
.expect("response should serialize")
|
||
.as_slice(),
|
||
&second_masked.session,
|
||
)
|
||
.expect("restore should succeed");
|
||
let restored_json: serde_json::Value =
|
||
serde_json::from_slice(&restored.body).expect("restored body should parse");
|
||
assert_eq!(restored_json["message"], "cached [email protected]");
|
||
|
||
let rolled_masked = try_mask_chat_request_json_with_cache_options(
|
||
&serde_json::to_vec(&second_request).expect("request should serialize"),
|
||
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 900),
|
||
MaskChatRequestOptions::runtime(),
|
||
Some(&cache),
|
||
)
|
||
.await
|
||
.expect("rolled masking should succeed");
|
||
let rolled_sentinel = rolled_masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("rolled sentinel should exist")
|
||
.to_string();
|
||
assert_ne!(rolled_sentinel, first_sentinel);
|
||
let mut rolled_headers =
|
||
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||
let rolled_restored = restore_sync_response_body(
|
||
&mut rolled_headers,
|
||
serde_json::to_vec(&json!({"message": format!("rolled {rolled_sentinel}")}))
|
||
.expect("response should serialize")
|
||
.as_slice(),
|
||
&rolled_masked.session,
|
||
)
|
||
.expect("rolled restore should succeed");
|
||
let rolled_restored_json: serde_json::Value = serde_json::from_slice(&rolled_restored.body)
|
||
.expect("rolled restored body should parse");
|
||
assert_eq!(rolled_restored_json["message"], "rolled [email protected]");
|
||
|
||
let keys = runtime_state
|
||
.scan_keys("privacy:redaction:mapping:*", 100)
|
||
.await
|
||
.expect("keys should read");
|
||
assert!(!keys.is_empty());
|
||
for key in keys {
|
||
assert!(!key.contains("[email protected]"));
|
||
let raw_key = runtime_state.strip_namespace(&key).to_string();
|
||
let ttl = runtime_state
|
||
.kv_ttl_seconds(&raw_key)
|
||
.await
|
||
.expect("ttl should read");
|
||
assert!(
|
||
ttl.unwrap_or_default() > 0,
|
||
"redaction cache key {key} should have ttl"
|
||
);
|
||
}
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn pii_redaction_performance_cache_hit_avoids_cached_sentinel_literal_collision() {
|
||
let Some(redis) = start_managed_redis_or_skip().await else {
|
||
return;
|
||
};
|
||
let runtime_state = redis_cache_runtime_state(
|
||
redis.redis_url(),
|
||
"pii_redaction_performance_cache_collision",
|
||
)
|
||
.await;
|
||
let cache = RedisRedactionMappingCache::new(&runtime_state);
|
||
let config = email_only_runtime_config();
|
||
|
||
let first_request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [{"role": "user", "content": "Contact [email protected]"}]
|
||
});
|
||
let first_masked = try_mask_chat_request_json_with_cache_options(
|
||
&serde_json::to_vec(&first_request).expect("request should serialize"),
|
||
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
|
||
MaskChatRequestOptions::runtime(),
|
||
Some(&cache),
|
||
)
|
||
.await
|
||
.expect("first masking should succeed");
|
||
let cached_sentinel = first_masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("cached sentinel should exist")
|
||
.to_string();
|
||
|
||
let colliding_request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": format!("literal {cached_sentinel} stays unknown")
|
||
},
|
||
{
|
||
"role": "assistant",
|
||
"content": null,
|
||
"tool_calls": [{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "lookup_contact",
|
||
"arguments": "{\"email\":\"[email protected]\"}"
|
||
}
|
||
}]
|
||
}
|
||
]
|
||
});
|
||
let second_masked = try_mask_chat_request_json_with_cache_options(
|
||
&serde_json::to_vec(&colliding_request).expect("request should serialize"),
|
||
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 899),
|
||
MaskChatRequestOptions::runtime(),
|
||
Some(&cache),
|
||
)
|
||
.await
|
||
.expect("collision masking should succeed");
|
||
let current_sentinel = second_masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("current request sentinel should exist")
|
||
.to_string();
|
||
assert_ne!(current_sentinel, cached_sentinel);
|
||
|
||
let masked_json: serde_json::Value =
|
||
serde_json::from_slice(&second_masked.body).expect("masked body should parse");
|
||
let content = masked_json["messages"][0]["content"]
|
||
.as_str()
|
||
.expect("content should be text");
|
||
assert_eq!(content.matches(&cached_sentinel).count(), 1);
|
||
assert!(!content.contains("[email protected]"));
|
||
let arguments = masked_json["messages"][1]["tool_calls"][0]["function"]["arguments"]
|
||
.as_str()
|
||
.expect("arguments should be text");
|
||
assert_eq!(arguments.matches(¤t_sentinel).count(), 1);
|
||
assert!(!arguments.contains("[email protected]"));
|
||
|
||
let mut headers =
|
||
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||
let restored = restore_sync_response_body(
|
||
&mut headers,
|
||
serde_json::to_vec(&json!({
|
||
"message": format!("literal {cached_sentinel} and {current_sentinel}")
|
||
}))
|
||
.expect("response should serialize")
|
||
.as_slice(),
|
||
&second_masked.session,
|
||
)
|
||
.expect("restore should succeed");
|
||
let restored_json: serde_json::Value =
|
||
serde_json::from_slice(&restored.body).expect("restored body should parse");
|
||
assert_eq!(
|
||
restored_json["message"],
|
||
format!("literal {cached_sentinel} and [email protected]")
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn pii_redaction_performance_ignores_malformed_cached_sentinel() {
|
||
let Some(redis) = start_managed_redis_or_skip().await else {
|
||
return;
|
||
};
|
||
let runtime_state =
|
||
redis_cache_runtime_state(redis.redis_url(), "pii_redaction_performance_bad_cache")
|
||
.await;
|
||
let cache = RedisRedactionMappingCache::new(&runtime_state);
|
||
let config = email_only_runtime_config();
|
||
runtime_state
|
||
.kv_set(
|
||
&cache.forward_cache_key("EMAIL", "[email protected]", 2),
|
||
"<AETHER:PHONE:ABCDEFGHIJKLMNOPQRST>",
|
||
Some(Duration::from_secs(300)),
|
||
)
|
||
.await
|
||
.expect("malformed cache value should write");
|
||
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [{"role": "user", "content": "Contact [email protected]"}]
|
||
});
|
||
let masked = try_mask_chat_request_json_with_cache_options(
|
||
&serde_json::to_vec(&request).expect("request should serialize"),
|
||
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
|
||
MaskChatRequestOptions::runtime(),
|
||
Some(&cache),
|
||
)
|
||
.await
|
||
.expect("masking should ignore malformed cache value");
|
||
|
||
let sentinel = masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("email sentinel should exist");
|
||
assert!(sentinel.starts_with("<AETHER:EMAIL:"));
|
||
assert_ne!(sentinel, "<AETHER:PHONE:ABCDEFGHIJKLMNOPQRST>");
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn pii_redaction_performance_ignores_cached_sentinel_with_wrong_namespace() {
|
||
let Some(redis) = start_managed_redis_or_skip().await else {
|
||
return;
|
||
};
|
||
let runtime_state = redis_cache_runtime_state(
|
||
redis.redis_url(),
|
||
"pii_redaction_performance_wrong_namespace",
|
||
)
|
||
.await;
|
||
let cache = RedisRedactionMappingCache::new(&runtime_state);
|
||
let config = ChatPiiRedactionRuntimeConfig {
|
||
placeholder_prefix: "SAFE".to_string(),
|
||
..email_only_runtime_config()
|
||
};
|
||
runtime_state
|
||
.kv_set(
|
||
&cache.forward_cache_key("EMAIL", "[email protected]", 2),
|
||
"<AETHER:EMAIL:ABCDEFGHIJKLMNOPQRST>",
|
||
Some(Duration::from_secs(300)),
|
||
)
|
||
.await
|
||
.expect("wrong namespace cache value should write");
|
||
|
||
let request = json!({
|
||
"model": "gpt-5",
|
||
"messages": [{"role": "user", "content": "Contact [email protected]"}]
|
||
});
|
||
let masked = try_mask_chat_request_json_with_cache_options(
|
||
&serde_json::to_vec(&request).expect("request should serialize"),
|
||
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
|
||
MaskChatRequestOptions::runtime(),
|
||
Some(&cache),
|
||
)
|
||
.await
|
||
.expect("masking should ignore wrong namespace cache value");
|
||
|
||
let sentinel = masked
|
||
.session
|
||
.sentinel_for_original("[email protected]")
|
||
.expect("email sentinel should exist");
|
||
assert!(sentinel.starts_with("<SAFE:EMAIL:"));
|
||
assert_ne!(sentinel, "<AETHER:EMAIL:ABCDEFGHIJKLMNOPQRST>");
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_performance_restore_uses_large_mapping_matcher() {
|
||
let mut session = session_at(600);
|
||
for index in 0..32 {
|
||
let original = format!("user{index}@example.com");
|
||
let sentinel = format!("<AETHER:EMAIL:{index:0>20}>");
|
||
let key = MappingKey {
|
||
rule_label: "EMAIL".to_string(),
|
||
original: original.clone(),
|
||
};
|
||
session.sentinel_index.insert(sentinel.clone(), key.clone());
|
||
session.mappings.insert(
|
||
key,
|
||
RedactionMapping {
|
||
rule_label: "EMAIL".to_string(),
|
||
kind: Some(RedactionKind::Email),
|
||
original,
|
||
normalized_value: format!("user{index}@example.com"),
|
||
sentinel,
|
||
bucket: 2,
|
||
created_at_unix_secs: 600,
|
||
expires_at_unix_secs: 900,
|
||
},
|
||
);
|
||
}
|
||
assert!(matches!(
|
||
SentinelMatcher::new(&session),
|
||
SentinelMatcher::Direct(_)
|
||
));
|
||
|
||
for index in 32..40 {
|
||
let original = format!("user{index}@example.com");
|
||
let sentinel = format!("<AETHER:EMAIL:{index:0>20}>");
|
||
let key = MappingKey {
|
||
rule_label: "EMAIL".to_string(),
|
||
original: original.clone(),
|
||
};
|
||
session.sentinel_index.insert(sentinel.clone(), key.clone());
|
||
session.mappings.insert(
|
||
key,
|
||
RedactionMapping {
|
||
rule_label: "EMAIL".to_string(),
|
||
kind: Some(RedactionKind::Email),
|
||
original,
|
||
normalized_value: format!("user{index}@example.com"),
|
||
sentinel,
|
||
bucket: 2,
|
||
created_at_unix_secs: 600,
|
||
expires_at_unix_secs: 900,
|
||
},
|
||
);
|
||
}
|
||
assert!(matches!(
|
||
SentinelMatcher::new(&session),
|
||
SentinelMatcher::Trie(_)
|
||
));
|
||
|
||
let text = "before <AETHER:EMAIL:00000000000000000039> middle <AETHER:EMAIL:00000000000000000001> after";
|
||
let restored = session.restore_text(text);
|
||
assert!(restored.restored);
|
||
assert_eq!(
|
||
restored.text,
|
||
"before [email protected] middle [email protected] after"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn pii_redaction_performance_documents_segment_scan_bypass_deferral() {
|
||
let source = include_str!("mod.rs");
|
||
assert!(source.contains("Segment-level scan-bypass cache is intentionally deferred"));
|
||
}
|
||
|
||
fn test_config() -> RedactionSessionConfig {
|
||
RedactionSessionConfig::new(b"redaction-test-key".to_vec(), 300, 600)
|
||
}
|
||
|
||
fn deepseek_opaque_replay_options() -> MaskChatRequestOptions {
|
||
MaskChatRequestOptions::runtime().with_reasoning_replay_policy(
|
||
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque,
|
||
)
|
||
}
|
||
}
|