use std::borrow::Cow;
use std::sync::LazyLock;
use std::sync::atomic::{AtomicBool, Ordering};
pub const PREFIX: &str = "DRE_SECRET_";
pub const MASK: &str = "*****";
pub const MASKED_SOURCE_ERROR: &str =
"source error details omitted because the request contains a masked secret";
const MIN_LEN: usize = 4;
static ENABLED: AtomicBool = AtomicBool::new(true);
fn eligible(value: &str) -> bool {
value.chars().count() >= MIN_LEN
}
static SECRETS: LazyLock<Vec<String>> = LazyLock::new(|| {
let mut v: Vec<String> = std::env::vars()
.filter(|(k, v)| k.starts_with(PREFIX) && eligible(v))
.map(|(_, v)| v)
.collect();
v.sort_by_key(|s| std::cmp::Reverse(s.len()));
v.dedup();
v
});
pub fn set_enabled(on: bool) {
ENABLED.store(on, Ordering::Relaxed);
}
pub(crate) fn enabled() -> bool {
ENABLED.load(Ordering::Relaxed)
}
fn mask_with<'a>(s: &'a str, secrets: &[String]) -> Cow<'a, str> {
if !secrets.iter().any(|x| s.contains(x.as_str())) {
return Cow::Borrowed(s);
}
let mut out = s.to_string();
for x in secrets {
out = out.replace(x.as_str(), MASK);
}
Cow::Owned(out)
}
pub fn mask(s: &str) -> Cow<'_, str> {
if !ENABLED.load(Ordering::Relaxed) {
return Cow::Borrowed(s);
}
mask_with(s, &SECRETS)
}
pub fn contains_secret(s: &str) -> bool {
ENABLED.load(Ordering::Relaxed) && SECRETS.iter().any(|secret| s.contains(secret))
}
fn mask_source_error_with<'a>(error: &'a str, sensitive: bool, secrets: &[String]) -> Cow<'a, str> {
if sensitive {
Cow::Borrowed(MASKED_SOURCE_ERROR)
} else {
mask_with(error, secrets)
}
}
pub fn mask_source_error(error: &str, sensitive: bool) -> Cow<'_, str> {
if !ENABLED.load(Ordering::Relaxed) {
return Cow::Borrowed(error);
}
mask_source_error_with(error, sensitive, &SECRETS)
}
fn unique_masked_names_with(names: &[String], secrets: &[String]) -> Vec<String> {
use std::collections::BTreeSet;
let masked: Vec<String> = names
.iter()
.map(|name| mask_with(name, secrets).into_owned())
.collect();
if names == masked {
return names.to_vec();
}
let mut used = BTreeSet::new();
let mut output = vec![String::new(); names.len()];
let mut pending = Vec::new();
for (index, (original, protected)) in names.iter().zip(&masked).enumerate() {
if original == protected {
used.insert(original.clone());
output[index].clone_from(original);
} else {
pending.push((protected.clone(), original.clone(), index));
}
}
pending.sort();
for (base, _, index) in pending {
let mut candidate = base.clone();
let mut suffix = 2;
while !used.insert(candidate.clone()) {
candidate = format!("{base}#{suffix}");
suffix += 1;
}
output[index] = candidate;
}
output
}
fn redact_json_with(v: &mut serde_json::Value, secrets: &[String]) {
use serde_json::Value;
match v {
Value::String(s) => {
if let Cow::Owned(m) = mask_with(s, secrets) {
*s = m;
}
}
Value::Array(a) => a.iter_mut().for_each(|v| redact_json_with(v, secrets)),
Value::Object(o) => {
let entries = std::mem::take(o);
let keys: Vec<String> = entries.keys().cloned().collect();
let protected_keys = unique_masked_names_with(&keys, secrets);
for ((_, mut val), protected_key) in entries.into_iter().zip(protected_keys) {
redact_json_with(&mut val, secrets);
o.insert(protected_key, val);
}
}
_ => {}
}
}
pub fn redact_json(v: &mut serde_json::Value) {
if ENABLED.load(Ordering::Relaxed) {
redact_json_with(v, &SECRETS);
}
}
fn redacted<T: serde::Serialize>(v: &T) -> serde_json::Result<serde_json::Value> {
let mut v = serde_json::to_value(v)?;
redact_json(&mut v);
Ok(v)
}
pub fn to_json_pretty<T: serde::Serialize>(v: &T) -> serde_json::Result<String> {
serde_json::to_string_pretty(&redacted(v)?)
}
pub fn to_json_line<T: serde::Serialize>(v: &T) -> serde_json::Result<String> {
serde_json::to_string(&redacted(v)?)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn secrets(values: &[&str]) -> Vec<String> {
let mut values: Vec<String> = values.iter().map(|v| (*v).to_string()).collect();
values.sort_by_key(|s| std::cmp::Reverse(s.len()));
values
}
#[test]
fn redacts_nested_values_longest_first() {
let secrets = secrets(&["nested-secret", "nested-secret-and-more"]);
let mut value = json!({"items": ["nested-secret-and-more", {"value": "x nested-secret y"}]});
redact_json_with(&mut value, &secrets);
assert_eq!(value, json!({"items": ["*****", {"value": "x ***** y"}]}));
}
#[test]
fn preserves_entries_when_masked_keys_collide() {
let secrets = secrets(&["alpha-secret-key", "bravo-secret-key"]);
let mut value = json!({
"alpha-secret-key": 1,
"bravo-secret-key": 2,
"*****": 3,
"*****#2": 4,
});
redact_json_with(&mut value, &secrets);
assert_eq!(
value,
json!({
"*****": 3,
"*****#2": 4,
"*****#3": 1,
"*****#4": 2,
})
);
}
#[test]
fn masking_one_name_does_not_rename_duplicate_literal_names() {
let secrets = secrets(&["secret-column"]);
let names = ["a", "a", "secret-column"].map(str::to_string);
assert_eq!(unique_masked_names_with(&names, &secrets), ["a", "a", "*****"]);
}
#[test]
fn suppresses_transformed_source_errors_only_for_secret_requests() {
let secrets = secrets(&["long-secret-value-that-a-source-truncated"]);
let error = "LINE 1: select * from long-secret-value-that-a-source...";
assert_eq!(mask_source_error_with(error, true, &secrets), MASKED_SOURCE_ERROR);
assert_eq!(
mask_source_error_with("ordinary syntax error", false, &secrets),
"ordinary syntax error"
);
}
#[test]
fn minimum_length_counts_unicode_characters() {
assert!(!eligible("ééé"));
assert!(eligible("éééé"));
}
}