use super::withhold_rules::names::{
self, AUDIT_FIELDS, Action, Class, Leaf, PAIR_VALUE, Pair, SCHEMA_KEYWORDS,
SCHEMA_VALUE_KEYWORDS, SUBSCHEMA_KEYWORDS, action, classify, header_action, header_class,
header_item_action, is_header_map, member_class,
};
use super::withhold_rules::values::{
has_private_key, looks_like_json_text, model_output_withheld, shows_credential, text_withheld,
unprinted_credentials, unreadable_json_strings,
};
use super::withhold_rules::view::{self, Json, Kind};
use crate::error::{SecretSet, WITHHELD};
use crate::json::OValue;
use scm::constants::{AI_GATEWAY_SECRET_FIELDS, AIGatewaySecretOperation as Op};
impl Json for OValue {
fn kind(&self) -> Kind {
match self {
Self::Null => Kind::Null,
Self::Bool(_) => Kind::Bool,
Self::Number(_) => Kind::Number,
Self::String(_) => Kind::String,
Self::Array(_) => Kind::Array,
Self::Object(_) => Kind::Object,
}
}
fn text(&self) -> Option<&str> {
self.as_str()
}
fn items(&self) -> Vec<&Self> {
self.as_array()
.map(|items| items.iter().collect())
.unwrap_or_default()
}
fn members(&self) -> Vec<(&str, &Self)> {
self.as_object()
.map(|entries| {
entries
.iter()
.map(|(name, member)| (name.as_str(), member))
.collect()
})
.unwrap_or_default()
}
}
#[cfg(test)]
impl super::withhold_rules::check::Document for OValue {
fn document(text: &str) -> Option<Self> {
json_document(text)
}
}
type Path = &'static [&'static str];
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Policy {
pub operations: &'static [Op],
pub paths: &'static [Path],
}
impl Policy {
fn rule_paths(&self) -> impl Iterator<Item = Path> + '_ {
sdk_paths(self.operations).chain(self.paths.iter().copied())
}
}
fn sdk_paths(operations: &[Op]) -> impl Iterator<Item = Path> + '_ {
AI_GATEWAY_SECRET_FIELDS
.iter()
.filter(|(operation, _)| operations.contains(operation))
.flat_map(|(_, rules)| rules.request.iter().chain(rules.response))
.map(|rule| rule.path)
}
pub const PLAIN: Policy = Policy {
operations: &[],
paths: &[],
};
pub const FILTER_BOUNDARIES: Policy = Policy {
operations: &[Op::TelemetryFilterBoundaries],
paths: &[],
};
pub const API_KEYS: Policy = Policy {
operations: &[
Op::ApiKeysCreateService,
Op::ApiKeysCreateUser,
Op::ApiKeysRotateService,
Op::ApiKeysRotateUser,
],
paths: &[],
};
pub const PROVIDERS: Policy = Policy {
operations: &[Op::ProvidersGet],
paths: &[],
};
pub const INTEGRATIONS: Policy = Policy {
operations: &[Op::IntegrationsCreate, Op::IntegrationsUpdate],
paths: &[],
};
pub const MCP: Policy = Policy {
operations: &[Op::McpIntegrationsCreate, Op::McpIntegrationsUpdate],
paths: &[],
};
pub const AUTH_SETTINGS: Policy = Policy {
operations: &[
Op::OrganisationsGetAuthSettings,
Op::OrganisationsUpdateAuthSettings,
],
paths: &[&["data", "auth_settings", "client_secret"]],
};
pub const DEPLOYMENTS: Policy = Policy {
operations: &[Op::DeploymentsCreate, Op::DeploymentsUpdate],
paths: &[],
};
pub const PLUGINS: Policy = Policy {
operations: &[Op::PluginsCreate],
paths: &[],
};
pub const ORGANISATION_INFO: Policy = Policy {
operations: &[Op::OrganisationsGetInfo],
paths: &[],
};
const POLICIES: [&Policy; 10] = [
&PLAIN,
&FILTER_BOUNDARIES,
&API_KEYS,
&PROVIDERS,
&INTEGRATIONS,
&MCP,
&AUTH_SETTINGS,
&DEPLOYMENTS,
&PLUGINS,
&ORGANISATION_INFO,
];
#[cfg(test)]
const AUDIT_ONLY: [Op; 5] = [
Op::SecretReferencesList,
Op::SecretReferencesCreate,
Op::SecretReferencesGet,
Op::SecretReferencesUpdate,
Op::SecretReferencesDelete,
];
const NOT_AUDITED: [Op; 1] = [Op::TelemetryFilterBoundaries];
fn audit_paths() -> impl Iterator<Item = Path> {
AI_GATEWAY_SECRET_FIELDS
.iter()
.filter(|(operation, _)| !NOT_AUDITED.contains(operation))
.flat_map(|(_, rules)| rules.request.iter().chain(rules.response))
.map(|rule| rule.path)
.chain(
POLICIES
.iter()
.flat_map(|policy| policy.paths.iter().copied()),
)
}
pub fn apply(doc: &mut OValue, policy: &Policy, secrets: &SecretSet) {
let mut withheld = Vec::new();
let paths: Vec<Path> = policy.rule_paths().collect();
if !paths.is_empty() {
let mut apply_paths = |target: &mut OValue| {
for path in &paths {
at_path(target, path, Paths::Resource, &mut withheld);
}
};
match &mut *doc {
OValue::Array(rows) => rows.iter_mut().for_each(&mut apply_paths),
other => apply_paths(other),
}
}
names_and_urls(doc, &mut withheld);
register(withheld, secrets);
crate::secrets::withhold(doc, secrets);
}
pub fn apply_model_output(doc: &mut OValue, secrets: &SecretSet) {
value_rules(doc);
crate::secrets::withhold(doc, secrets);
}
fn value_rules(value: &mut OValue) {
match value {
OValue::Object(entries) => entries
.iter_mut()
.for_each(|(_, member)| value_rules(member)),
OValue::Array(items) => items.iter_mut().for_each(value_rules),
OValue::String(text) => {
if let Some(shown) = model_output_withheld(text) {
*text = shown;
}
}
_ => {}
}
}
fn in_place(text: &mut String, withheld: &mut Vec<String>) {
if let Some((shown, parts)) = text_withheld(text) {
withheld.extend(parts);
*text = shown;
}
}
pub fn apply_audit(rows: &mut OValue, reveal: bool, secrets: &SecretSet) {
let mut withheld = Vec::new();
if let OValue::Array(records) = rows {
for record in records {
let OValue::Object(entries) = record else {
continue;
};
for (name, member) in entries.iter_mut() {
if AUDIT_FIELDS.contains(&name.as_str()) {
let headers = name == "request_headers";
audit_field(member, headers, reveal, &mut withheld);
}
}
}
}
register(withheld, secrets);
apply(rows, &PLAIN, secrets);
}
fn audit_field(member: &mut OValue, headers: bool, reveal: bool, withheld: &mut Vec<String>) {
let document = match &*member {
OValue::String(text) => json_document(text),
OValue::Object(_) | OValue::Array(_) => Some(member.clone()),
_ => None,
};
if let Some(mut document) = document {
audit_document(&mut document, headers, withheld);
if reveal {
if !matches!(member, OValue::String(_)) {
*member = document;
return;
}
if let Ok(text) = crate::json::to_compact(&document) {
*member = OValue::String(text);
return;
}
}
}
match &*member {
OValue::Null => return,
OValue::String(text) if text == WITHHELD => return,
OValue::String(text) => {
if json_document(text).is_none() {
withheld.extend(unprinted_credentials(text));
}
}
OValue::Number(number) => {
withheld.extend(names::registered_number(Leaf::Inner, &number.to_string()));
}
_ => {}
}
*member = OValue::String(WITHHELD.to_owned());
}
type Walk = fn(&mut OValue, &mut Vec<String>);
fn audit_document(doc: &mut OValue, headers: bool, withheld: &mut Vec<String>) {
if headers {
header_map(doc, audit_walk, withheld);
return;
}
for path in audit_paths() {
let mut at = |target: &mut OValue| match target {
OValue::Array(rows) => rows
.iter_mut()
.for_each(|row| at_path(row, path, Paths::Audit, withheld)),
other => at_path(other, path, Paths::Audit, withheld),
};
at(doc);
if let OValue::Object(entries) = doc
&& let Some((_, data)) = entries.iter_mut().find(|(name, _)| name == "data")
{
at(data);
}
}
audit_walk(doc, withheld);
}
fn audit_walk(value: &mut OValue, withheld: &mut Vec<String>) {
match value {
OValue::Object(entries) => {
let pair = pair_of(entries);
let located = located(entries);
for (name, member) in entries.iter_mut() {
match &pair {
Some(pair) if pair.names_member(name, member.as_str()) => {
audit_walk(member, withheld);
}
Some(pair) if name == PAIR_VALUE => take(
header_action(pair.class(), view::shape(member)),
member,
audit_walk,
withheld,
),
_ => withhold_named(name, located, member, audit_walk, withheld),
}
}
}
OValue::Array(items) => items.iter_mut().for_each(|item| audit_walk(item, withheld)),
OValue::String(_) => withhold_text(
value,
|doc, withheld| audit_document(doc, false, withheld),
withheld,
),
_ => {}
}
}
fn take(action: Action, member: &mut OValue, walk: Walk, withheld: &mut Vec<String>) {
match action {
Action::Read => walk(member, withheld),
Action::Headers => header_map(member, walk, withheld),
Action::Keep => {}
Action::Withhold(leaf) => withhold_value(leaf, member, withheld),
}
}
fn withhold_named(
name: &str,
located: bool,
member: &mut OValue,
walk: Walk,
withheld: &mut Vec<String>,
) {
let action = action(member_class(name, located), view::shape(member));
take(action, member, walk, withheld);
}
fn pair_of(entries: &[(String, OValue)]) -> Option<Pair> {
Pair::of(
entries
.iter()
.map(|(name, member)| (name.as_str(), member.as_str())),
)
}
fn located(entries: &[(String, OValue)]) -> bool {
names::is_secret_location(entries.iter().map(|(name, _)| name.as_str()))
}
fn withhold_text(value: &mut OValue, inner: Walk, withheld: &mut Vec<String>) {
if within_json_text(value, inner, withheld) {
return;
}
let OValue::String(text) = value else {
return;
};
if has_private_key(text) {
register_leaf(Leaf::Credential, text, withheld);
blank(value);
} else {
in_place(text, withheld);
}
}
fn within_json_text(
value: &mut OValue,
inner: impl FnOnce(&mut OValue, &mut Vec<String>),
withheld: &mut Vec<String>,
) -> bool {
let Some(text) = value.as_str() else {
return false;
};
let Some(mut parsed) = json_document(text) else {
if !looks_like_json_text(text) {
return false;
}
withheld.extend(unreadable_json_strings(text));
blank(value);
return true;
};
let received = parsed.clone();
inner(&mut parsed, withheld);
if parsed != received {
match crate::json::to_compact(&parsed) {
Ok(compact) => *value = OValue::String(compact),
Err(_) => blank(value),
}
}
true
}
fn header_map(value: &mut OValue, walk: Walk, withheld: &mut Vec<String>) {
match value {
OValue::Object(entries) => {
let pair = pair_of(entries);
for (name, member) in entries.iter_mut() {
let class = match &pair {
Some(pair) if pair.names_member(name, member.as_str()) => {
walk(member, withheld);
continue;
}
Some(pair) if name == PAIR_VALUE => pair.class(),
_ => header_class(name),
};
take(
header_action(class, view::shape(member)),
member,
walk,
withheld,
);
}
}
OValue::Array(items) => {
for item in items {
if matches!(item, OValue::Object(_)) {
header_map(item, walk, withheld);
} else {
let action = header_item_action(item.as_str(), view::shape(item));
take(action, item, walk, withheld);
}
}
}
OValue::String(_) => {
let inner =
|doc: &mut OValue, withheld: &mut Vec<String>| header_map(doc, walk, withheld);
if !within_json_text(value, inner, withheld) {
withhold_value(Leaf::Generic, value, withheld);
}
}
_ => {}
}
}
fn header_leaf(class: Class) -> Leaf {
match class {
Class::Count | Class::Credential | Class::Member => Leaf::Credential,
_ => Leaf::Generic,
}
}
fn register_leaf(leaf: Leaf, text: &str, withheld: &mut Vec<String>) {
withheld.extend(names::registered(leaf, text));
}
fn withhold_value(leaf: Leaf, value: &mut OValue, withheld: &mut Vec<String>) {
match value {
OValue::Null | OValue::Bool(_) => return,
OValue::String(text) if text == WITHHELD => return,
OValue::String(text) => {
json_text_credentials(text, withheld);
register_leaf(leaf, text, withheld);
}
OValue::Number(number) => {
withheld.extend(names::registered_number(leaf, &number.to_string()));
}
OValue::Object(_) | OValue::Array(_) => credentials_inside(value, withheld),
}
*value = OValue::String(WITHHELD.to_owned());
}
fn blank(value: &mut OValue) {
if matches!(
value,
OValue::String(_) | OValue::Object(_) | OValue::Array(_)
) {
*value = OValue::String(WITHHELD.to_owned());
}
}
fn credentials_inside(doc: &mut OValue, withheld: &mut Vec<String>) {
names_and_urls(doc, withheld);
inner_leaves(doc, false, withheld);
}
fn inner_leaves(value: &OValue, describing: bool, withheld: &mut Vec<String>) {
match value {
OValue::String(text) if !text.contains(WITHHELD) => {
let leaf = if describing {
Leaf::Describing
} else {
Leaf::Inner
};
register_leaf(leaf, text, withheld);
}
OValue::Array(items) => items
.iter()
.for_each(|item| inner_leaves(item, describing, withheld)),
OValue::Object(entries) => {
let schema = view::is_schema(value);
for (name, member) in entries {
let describing = describing || (schema && SCHEMA_KEYWORDS.contains(&name.as_str()));
inner_leaves(member, describing, withheld);
}
}
_ => {}
}
}
fn withhold_described_credentials(value: &mut OValue) {
match value {
OValue::String(text) if shows_credential(text) => blank(value),
OValue::Array(items) => items.iter_mut().for_each(withhold_described_credentials),
OValue::Object(entries) => entries
.iter_mut()
.for_each(|(_, member)| withhold_described_credentials(member)),
_ => {}
}
}
fn json_document(text: &str) -> Option<OValue> {
if !text.trim_start().starts_with(['{', '[']) {
return None;
}
crate::json::parse(text)
.ok()
.filter(|parsed| matches!(parsed, OValue::Object(_) | OValue::Array(_)))
}
fn json_text_credentials(text: &str, withheld: &mut Vec<String>) {
if let Some(mut parsed) = json_document(text) {
credentials_inside(&mut parsed, withheld);
} else if looks_like_json_text(text) {
withheld.extend(unreadable_json_strings(text));
}
}
fn register(withheld: Vec<String>, secrets: &SecretSet) {
for value in withheld.iter().filter(|value| names::registrable(value)) {
secrets.add(value);
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Paths {
Resource,
Audit,
}
fn at_path(value: &mut OValue, path: &[&str], paths: Paths, withheld: &mut Vec<String>) {
let Some((first, rest)) = path.split_first() else {
return;
};
match value {
OValue::Object(entries) => {
for (name, member) in entries.iter_mut() {
if *first != "*" && name != first {
continue;
}
if rest.is_empty() {
withhold_value(path_leaf(name, *first == "*", paths), member, withheld);
} else if rest == ["*"] && is_header_map(name) {
header_values(member, withheld);
} else {
at_path(member, rest, paths, withheld);
}
}
}
OValue::Array(items) if *first == "*" => {
for item in items {
if rest.is_empty() {
withhold_value(Leaf::Path, item, withheld);
} else {
at_path(item, rest, paths, withheld);
}
}
}
_ => {}
}
}
fn path_leaf(name: &str, wildcard: bool, paths: Paths) -> Leaf {
match classify(name) {
Class::Count | Class::Credential | Class::Member => Leaf::Path,
Class::Generic if paths == Paths::Resource => Leaf::Path,
Class::Generic => Leaf::Generic,
Class::Plain | Class::Fingerprint | Class::HeaderMap if wildcard => Leaf::Inner,
Class::Plain | Class::Fingerprint | Class::HeaderMap => Leaf::Path,
}
}
fn header_values(map: &mut OValue, withheld: &mut Vec<String>) {
match map {
OValue::Object(entries) => {
for (name, member) in entries.iter_mut() {
withhold_value(header_leaf(header_class(name)), member, withheld);
}
}
OValue::Array(items) => {
for item in items {
let pair = match item {
OValue::Object(entries) => pair_of(entries),
_ => None,
};
let Some(pair) = pair else {
withhold_value(header_leaf(header_class("")), item, withheld);
continue;
};
if let OValue::Object(entries) = item {
for (name, member) in entries.iter_mut() {
if name == PAIR_VALUE {
withhold_value(header_leaf(pair.class()), member, withheld);
}
}
}
blank(item);
}
}
_ => {}
}
}
fn names_and_urls(value: &mut OValue, withheld: &mut Vec<String>) {
match value {
OValue::Object(entries) => {
let pair = pair_of(entries);
let located = located(entries);
let owned_modes: Vec<String> = names::mode_values(
entries
.iter()
.map(|(name, member)| (name.as_str(), member.as_str())),
)
.into_iter()
.map(str::to_owned)
.collect();
let modes: Vec<&str> = owned_modes.iter().map(String::as_str).collect();
for (name, member) in entries.iter_mut() {
let mode = member
.as_str()
.filter(|text| names::repeats_a_mode(text, &modes))
.map(str::to_owned);
let before = withheld.len();
match &pair {
Some(pair) if pair.names_member(name, member.as_str()) => {
names_and_urls(member, withheld);
}
Some(pair) if name == PAIR_VALUE => take(
action(pair.class_outside_headers(), view::shape(member)),
member,
names_and_urls,
withheld,
),
_ if name == "properties" && matches!(member, OValue::Object(_)) => {
property_map(member, withheld);
}
_ => withhold_named(name, located, member, names_and_urls, withheld),
}
if let Some(mode) = mode {
let added = withheld.split_off(before);
withheld.extend(added.into_iter().filter(|found| *found != mode));
}
}
}
OValue::Array(items) => items
.iter_mut()
.for_each(|item| names_and_urls(item, withheld)),
OValue::String(_) => withhold_text(value, names_and_urls, withheld),
_ => {}
}
}
fn property_map(map: &mut OValue, withheld: &mut Vec<String>) {
let OValue::Object(entries) = map else {
return;
};
for (name, member) in entries.iter_mut() {
if view::is_schema(member) {
schema_property(name, member, withheld);
} else {
withhold_named(name, false, member, names_and_urls, withheld);
}
}
}
fn schema_property(name: &str, schema: &mut OValue, withheld: &mut Vec<String>) {
let OValue::Object(entries) = schema else {
return;
};
let credential = classify(name) > Class::Plain;
for (keyword, member) in entries.iter_mut() {
let keyword = keyword.as_str();
if keyword == "properties" && matches!(member, OValue::Object(_)) {
property_map(member, withheld);
} else if SUBSCHEMA_KEYWORDS.contains(&keyword) {
match member {
OValue::Object(_) => schema_property(name, member, withheld),
OValue::Array(items) => {
for item in items {
if matches!(item, OValue::Object(_)) {
schema_property(name, item, withheld);
} else {
instance_member(name, keyword, item, withheld);
}
}
}
other => instance_member(name, keyword, other, withheld),
}
} else if SCHEMA_KEYWORDS.contains(&keyword) {
names_and_urls(member, withheld);
if credential {
withhold_described_credentials(member);
}
} else if SCHEMA_VALUE_KEYWORDS.contains(&keyword) && view::is_empty(&*member) {
} else {
instance_member(name, keyword, member, withheld);
}
}
}
fn instance_member(property: &str, own: &str, member: &mut OValue, withheld: &mut Vec<String>) {
let name = if classify(own) >= classify(property) {
own
} else {
property
};
withhold_named(name, false, member, names_and_urls, withheld);
}
#[cfg(test)]
mod tests;