Files
galaxy/app/src/ai/blocklist/block/secret_redaction.rs
T

709 lines
31 KiB
Rust

use std::collections::HashMap;
use itertools::Itertools;
use similar::DiffableStr;
use warpui::elements::{MouseStateHandle, PartialClickableElement, SecretRange};
use warpui::platform::Cursor;
use crate::ai::agent::{AIAgentOutput, AIAgentTextSection, AgentOutputText};
use crate::terminal::model::secrets::{SecretLevel, REGEX_LEVEL_METADATA, SECRETS_REGEX};
use super::{AIBlockAction, TextLocation};
pub const SECRET_REDACTION_REPLACEMENT_CHARACTER: &str = "*";
/// Returns the ranges of detected secrets in the given text.
pub(crate) fn find_secrets_in_text(text: &str) -> Vec<SecretRange> {
find_secrets_in_text_with_levels(text)
.into_iter()
.map(|(range, _level)| range)
.collect()
}
/// Returns the ranges of detected secrets in the given text along with their SecretLevel.
pub(crate) fn find_secrets_in_text_with_levels(text: &str) -> Vec<(SecretRange, SecretLevel)> {
// Combine all regex patterns into a single regex pattern with non-capturing groups, for efficiency.
// Note that we purposely use regex::Regex instead of RegexDFAs since we are working a Text (containing
// a normal String) rather than the Grid (where text is in Cells with 1 character each).
let regex = SECRETS_REGEX.read();
let metadata = REGEX_LEVEL_METADATA.read();
let mut secret_ranges = vec![];
let mut byte_to_char_index = vec![0; text.len() + 1]; // Map byte index to char index
// Track the current character index while iterating through the string.
let mut char_index = 0;
for (byte_index, _) in text.char_indices() {
byte_to_char_index[byte_index] = char_index;
char_index += 1;
}
byte_to_char_index[text.len()] = char_index; // Map the last byte to the last character index
// Iterate over the text once, finding all matches against secret regex. Map the byte ranges
// to character ranges and store them.
for mat in regex.find_iter(text) {
let start_byte = mat.start();
let end_byte = mat.end();
let start_char = byte_to_char_index[start_byte];
let end_char = byte_to_char_index[end_byte];
// Determine which pattern matched by getting the pattern ID and map via counts
let pattern_id = mat.pattern().as_usize();
let total_patterns = metadata.enterprise_count + metadata.user_count;
if pattern_id >= total_patterns {
log::error!("Secret level not found for pattern ID {pattern_id}");
continue;
}
let secret_level = if pattern_id < metadata.enterprise_count {
SecretLevel::Enterprise
} else {
SecretLevel::User
};
secret_ranges.push((
SecretRange {
char_range: start_char..end_char,
byte_range: start_byte..end_byte,
},
secret_level,
));
}
// Merge overlapping ranges, preserving the highest priority SecretLevel
merge_sorted_ranges_with_levels(secret_ranges)
}
/// Merges overlapping ranges while preserving the highest priority SecretLevel
fn merge_sorted_ranges_with_levels(
ranges: Vec<(SecretRange, SecretLevel)>,
) -> Vec<(SecretRange, SecretLevel)> {
if ranges.is_empty() {
return ranges;
}
let mut merged_ranges = vec![];
let mut current_range = ranges[0].0.clone();
let mut current_level = ranges[0].1;
for (range, level) in ranges.into_iter().skip(1) {
// We can merge based on character ranges since non-overlapping character ranges result in non-overlapping byte ranges.
if range.char_range.start <= current_range.char_range.end {
// Extend the current range to include the overlapping range.
current_range.extend_range_end(&range);
// Keep the highest priority level
if level.priority() > current_level.priority() {
current_level = level;
}
} else {
// No overlap, push the current range and move to the next.
merged_ranges.push((current_range, current_level));
current_range = range;
current_level = level;
}
}
// Add the last range.
merged_ranges.push((current_range, current_level));
merged_ranges
}
#[derive(Debug, Eq, PartialEq)]
pub struct SecretLocation {
pub secret_range: SecretRange,
pub location: TextLocation,
}
#[derive(Clone, Debug)]
pub struct Secret {
pub secret: String,
pub is_obfuscated: bool,
pub mouse_state: MouseStateHandle,
pub secret_level: SecretLevel,
}
#[derive(Default, Debug)]
pub struct DetectedSecretsInTextLocation {
pub detected_secrets: HashMap<SecretRange, Secret>,
}
#[derive(Default, Debug)]
pub struct SecretRedactionState {
/// Last byte index of text we've scanned for secret redaction (avoid re-scanning in the context
/// of streaming output text). Note we want to redact secrets WHILE streaming to avoid any full secrets
/// ever being shown on the screen! This applies to the last output step we've currently received.
last_scanned_secret_redaction_byte_index: usize,
/// Buffer to hold the last word that's been scanned, since we may need to combine it with the next
/// tokens we're receiving (since a secret could be split apart across streaming chunks). Words are defined
/// to be separated by whitespace.
last_word_to_rescan_for_redaction: String,
/// Buffer to hold the current line. This is important, as markdown parsing could potentially mutate the entire line.
current_line_for_redaction: String,
/// Keeps track of the last step index we've run secret redaction scanning on, while streaming output.
last_text_section_index_scanned_for_redaction: usize,
last_line_index_scanned_for_redaction: usize,
detected_secrets: HashMap<TextLocation, DetectedSecretsInTextLocation>,
currently_hovered_secret_location: Option<SecretLocation>,
// This is separate from currently_hovered_secret_location because after clicking
// on a secret to open the tooltip, this secret should remain highlighted and the tooltip in place
// even if we hover over other secrets.
secret_location_open_tooltip: Option<SecretLocation>,
}
impl SecretRedactionState {
pub fn open_tooltip_location(&self) -> Option<&SecretLocation> {
self.secret_location_open_tooltip.as_ref()
}
pub fn hovered_location(&self) -> Option<&SecretLocation> {
self.currently_hovered_secret_location.as_ref()
}
pub fn has_open_tooltip(&self, location: &TextLocation, range: &SecretRange) -> bool {
self.open_tooltip_location()
.is_some_and(|tooltip_location| {
tooltip_location.location == *location && tooltip_location.secret_range == *range
})
}
pub fn is_hovered(&self, location: &TextLocation, range: &SecretRange) -> bool {
self.hovered_location().is_some_and(|tooltip_location| {
tooltip_location.location == *location && tooltip_location.secret_range == *range
})
}
pub fn reset(&mut self) {
self.last_text_section_index_scanned_for_redaction = 0;
self.last_scanned_secret_redaction_byte_index = 0;
self.last_word_to_rescan_for_redaction = Default::default();
}
/// Clears secret redaction state for `user_query` locations.
///
/// A bit of an edge case, but this is required when the user accepts a 'suggest new conversation' action
/// for an existing query. This query becomes the first query in a new conversation, and we prefix '/agent'
/// to all initial user queries, so detected secret ranges (if any) become stale.
pub fn clear_user_query_locations(&mut self) {
self.detected_secrets
.retain(|location, _| !matches!(location, TextLocation::Query { .. }));
if self
.currently_hovered_secret_location
.as_ref()
.is_some_and(|location| matches!(location.location, TextLocation::Query { .. }))
{
self.currently_hovered_secret_location = None;
}
if self
.secret_location_open_tooltip
.as_ref()
.is_some_and(|location| matches!(location.location, TextLocation::Query { .. }))
{
self.secret_location_open_tooltip = None;
}
}
pub fn show_secret_tooltip(
&mut self,
location: &TextLocation,
secret_range: &SecretRange,
) -> Option<&mut Secret> {
self.secret_location_open_tooltip = Some(SecretLocation {
secret_range: secret_range.clone(),
location: *location,
});
self.get_secret_mut(location, secret_range)
}
pub fn dismiss_tooltip(&mut self) {
self.secret_location_open_tooltip = None;
}
pub fn set_obfuscated(
&mut self,
location: &TextLocation,
secret_range: &SecretRange,
is_obfuscated: bool,
) {
if let Some(hoverable_secret_mut) = self.get_secret_mut(location, secret_range) {
hoverable_secret_mut.is_obfuscated = is_obfuscated;
}
}
pub fn set_hover_state_for_secret(
&mut self,
location: &TextLocation,
secret_range: &SecretRange,
is_hovering: bool,
) {
if is_hovering {
self.currently_hovered_secret_location = Some(SecretLocation {
secret_range: secret_range.clone(),
location: *location,
});
} else if self.currently_hovered_secret_location.as_ref().is_some_and(
|currently_hovered_secret| {
currently_hovered_secret.secret_range == *secret_range
&& currently_hovered_secret.location == *location
},
) {
self.currently_hovered_secret_location = None;
}
}
pub fn secrets_for_location(
&self,
location: &TextLocation,
) -> Option<&DetectedSecretsInTextLocation> {
self.detected_secrets.get(location)
}
fn get_secret_mut(
&mut self,
location: &TextLocation,
secret_range: &SecretRange,
) -> Option<&mut Secret> {
self.detected_secrets
.get_mut(location)
.and_then(|detected_location| detected_location.detected_secrets.get_mut(secret_range))
}
pub fn run_redaction_for_location(
&mut self,
text: &str,
location: TextLocation,
should_obfuscate: bool,
) {
// Detect secrets in user's query.
let secret_ranges_with_levels = find_secrets_in_text_with_levels(text);
for (secret_range, secret_level) in secret_ranges_with_levels {
if let Some(secret_text) =
text.get(secret_range.byte_range.start..secret_range.byte_range.end)
{
self.detected_secrets
.entry(location)
.or_default()
.detected_secrets
.insert(
secret_range,
Secret {
secret: secret_text.to_string(),
is_obfuscated: should_obfuscate,
mouse_state: Default::default(),
secret_level,
},
);
}
}
}
pub fn run_incremental_redaction_on_partial_output(
&mut self,
output: &AIAgentOutput,
should_obfuscate: bool,
) {
// Steps are sequentially streamed, hence we always check the last step.
// Important: all the *lines* in the context are markdown lines, which could be rendered as multiple lines on screen.
if let Some((section_index, text_section)) = output
.all_text()
.flat_map(|text| text.sections.iter())
.enumerate()
.last()
{
let line_index = match text_section {
AIAgentTextSection::PlainText {
text:
AgentOutputText {
formatted_lines: Some(text),
..
},
} => text
.lines()
.iter()
.enumerate()
.next_back()
.map_or(0, |(index, _)| index),
_ => 0,
};
let mut start_of_last_word_byte_index;
if section_index == self.last_text_section_index_scanned_for_redaction
&& line_index == self.last_line_index_scanned_for_redaction
{
// The boundary between what we're done scanning and what we still need to scan.
start_of_last_word_byte_index = self.last_scanned_secret_redaction_byte_index
- self.last_word_to_rescan_for_redaction.len();
self.detected_secrets
// We remove all secrets that are beyond the cutoff boundary where we start rescanning,
// specifically to avoid having duplicate secrets from the last word buffer.
.retain(|location, secrets| {
if let TextLocation::Output {
section_index: current_section_index,
line_index: current_line_index,
} = location
{
// Only clear secrets from the last step we scanned.
if *current_section_index == section_index
&& *current_line_index == line_index
{
secrets.detected_secrets.retain(|secret_range, _| {
secret_range.byte_range.start < start_of_last_word_byte_index
&& secret_range.byte_range.end
<= start_of_last_word_byte_index
});
}
!secrets.detected_secrets.is_empty()
} else {
true
}
});
} else {
// Addition needs to happen before subtraction to prevent an intermediate usize value
// from causing a numeric overflow and panicking, as Rust arithmetic operators are left-associative.
// TODO: Investigate why, sometimes, the section index is smaller than the last section index scanned.
let num_sections_to_scan = (section_index + 1)
.saturating_sub(self.last_text_section_index_scanned_for_redaction);
for (rerun_section_index, rerun_section) in output
.all_text()
.flat_map(|text| text.sections.iter())
.enumerate()
.skip(self.last_text_section_index_scanned_for_redaction)
.take(num_sections_to_scan)
{
if rerun_section_index == self.last_text_section_index_scanned_for_redaction {
// First step: rerun secret detection on the step which the secret detection last got ran.
if let AIAgentTextSection::PlainText { text } = rerun_section {
let end_line_index = if self
.last_text_section_index_scanned_for_redaction
== section_index
{
// Only step index changed, we're still in the same step. Run secret detection up to the current line.
line_index
} else {
// We're in a new step - run secret detection up to and include the last line.
text.formatted_lines
.as_ref()
.map(|lines| lines.lines().len())
.unwrap_or(0)
};
// Starting on the last detected line (rerun, in case markdown parsing altered the text)
// TODO: Optimization: only run secret detection on the newly part of the line if the parser didn't alter it
for rerun_line_index in
self.last_line_index_scanned_for_redaction..end_line_index
{
self.rerun_secret_detection_for_output_line(
rerun_section_index,
rerun_line_index,
text,
)
}
} else {
self.detected_secrets.remove(&TextLocation::Output {
section_index: rerun_section_index,
line_index: 0,
});
}
} else if rerun_section_index == section_index {
// Last step: rerun secret detection up until the current line.
// This part is already handled in the previous condition if we're in the same step.
if let AIAgentTextSection::PlainText { text } = rerun_section {
for rerun_line_index in 0..line_index {
self.rerun_secret_detection_for_output_line(
rerun_section_index,
rerun_line_index,
text,
)
}
}
// No need to clear secrets in an else case - these lines has not been scanned for secrets yet.
} else {
// Middle sections: run secret detection on all lines.
if let AIAgentTextSection::PlainText { text } = rerun_section {
for rerun_line_index in 0..text
.formatted_lines
.as_ref()
.map(|lines| lines.lines().len())
.unwrap_or(0)
{
self.rerun_secret_detection_for_output_line(
rerun_section_index,
rerun_line_index,
text,
)
}
}
// No need to clear secrets in an else case - these lines has not been scanned for secrets yet.
}
}
self.last_text_section_index_scanned_for_redaction = section_index;
self.last_line_index_scanned_for_redaction = line_index;
self.last_scanned_secret_redaction_byte_index = 0;
self.last_word_to_rescan_for_redaction.clear();
self.current_line_for_redaction.clear();
start_of_last_word_byte_index = 0;
}
if let AIAgentTextSection::PlainText { text } = &text_section {
let text = match &text.formatted_lines {
Some(text) => text.lines().iter().last().map(|line| line.raw_text()),
_ => None,
};
if let Some(text) = text {
// Trim the trailing newline as the parser automatically
// adds a newline to the end of the text.
let text = if text.ends_with_newline() {
&text[..text.len() - 1]
} else {
text
};
if text.len() >= self.current_line_for_redaction.len()
&& text.starts_with(&self.current_line_for_redaction)
{
// If the current line is a prefix of the new text, we can just append the new text.
self.current_line_for_redaction
.push_str(&text[self.current_line_for_redaction.len()..]);
} else {
// If the current line is not a prefix of the new text, we need to clear the current line and start over.
self.current_line_for_redaction.clear();
self.current_line_for_redaction.push_str(text);
self.last_scanned_secret_redaction_byte_index = 0;
self.last_word_to_rescan_for_redaction.clear();
start_of_last_word_byte_index = 0;
self.detected_secrets.retain(|location, _| {
if let TextLocation::Output {
section_index: cur_step_index,
line_index: cur_line_index,
} = location
{
// Clear all secrets of the current line in the current step.
!(*cur_step_index == section_index && *cur_line_index == line_index)
} else {
true
}
});
}
// Combine the last word in last word buffer with the new text to handle
// the case where a secret is split across streaming chunks.
let combined_text = format!(
"{}{}",
&self.last_word_to_rescan_for_redaction,
&text[self.last_scanned_secret_redaction_byte_index..]
);
let secret_ranges_with_levels =
find_secrets_in_text_with_levels(&combined_text);
for (secret_range, secret_level) in secret_ranges_with_levels {
// Adjust the ranges to map correctly within the new text.
let adjusted_byte_start =
start_of_last_word_byte_index + secret_range.byte_range.start;
let adjusted_byte_end =
start_of_last_word_byte_index + secret_range.byte_range.end;
let adjusted_char_start = text[..adjusted_byte_start].chars().count();
let adjusted_char_end = text[..adjusted_byte_end].chars().count();
let adjusted_secret_range = SecretRange {
char_range: adjusted_char_start..adjusted_char_end,
byte_range: adjusted_byte_start..adjusted_byte_end,
};
if let Some(secret_text) = text.get(adjusted_byte_start..adjusted_byte_end)
{
self.detected_secrets
.entry(TextLocation::Output {
section_index,
line_index,
})
.or_default()
.detected_secrets
.insert(
adjusted_secret_range,
Secret {
secret: secret_text.to_string(),
is_obfuscated: should_obfuscate,
mouse_state: Default::default(),
secret_level,
},
);
}
}
// Update the last scanned position to the end of the current text.
self.last_scanned_secret_redaction_byte_index = text.len();
// Extract and store the last word (whitespace-separated) in the last word buffer.
if let Some(last_space_byte_index) =
combined_text.rfind(|c: char| c.is_whitespace())
{
let slice_with_space = &combined_text[last_space_byte_index..];
let space_offset = slice_with_space
.chars()
.next()
.expect("The whitespace character should be present")
.len_utf8();
self.last_word_to_rescan_for_redaction =
combined_text[last_space_byte_index + space_offset..].to_string();
} else {
// If no whitespace is found, store the entire combined text as the prefix.
self.last_word_to_rescan_for_redaction = combined_text.clone();
}
}
}
}
}
pub fn run_redaction_on_complete_output(&mut self, output: &AIAgentOutput) {
// Delete all output secrets as we'll be rescanning the entire output.
self.detected_secrets
.retain(|location, _| !matches!(location, TextLocation::Output { .. }));
for (section_index, section) in output
.all_text()
.flat_map(|text| text.sections.iter())
.enumerate()
{
if let AIAgentTextSection::PlainText { text } = section {
let texts = match &text.formatted_lines {
Some(text) => text.lines().iter().map(|line| line.raw_text()).collect(),
_ => vec![text.text()],
};
for (line_index, text) in texts.iter().enumerate() {
let secret_ranges_with_levels = find_secrets_in_text_with_levels(text);
for (secret_range, secret_level) in secret_ranges_with_levels {
if let Some(secret_text) =
text.get(secret_range.byte_range.start..secret_range.byte_range.end)
{
self.detected_secrets
.entry(TextLocation::Output {
section_index,
line_index,
})
.or_default()
.detected_secrets
.insert(
secret_range,
Secret {
secret: secret_text.to_string(),
is_obfuscated: true,
mouse_state: Default::default(),
secret_level,
},
);
}
}
}
}
}
}
fn rerun_secret_detection_for_output_line(
&mut self,
section_index: usize,
line_index: usize,
text: &AgentOutputText,
) {
let text_line = text
.formatted_lines
.as_ref()
.and_then(|lines| lines.lines().get(line_index).map(|line| line.raw_text()));
let entry_location = TextLocation::Output {
section_index,
line_index,
};
self.detected_secrets.remove(&entry_location);
if let Some(line_text) = text_line {
let secret_ranges_with_levels = find_secrets_in_text_with_levels(line_text);
for (secret_range, secret_level) in secret_ranges_with_levels {
// No adjustment is needed - we're redoing this entire line
if let Some(secret_text) = line_text.get(secret_range.byte_range.clone()) {
self.detected_secrets
.entry(entry_location)
.or_default()
.detected_secrets
.insert(
secret_range,
Secret {
secret: secret_text.to_string(),
is_obfuscated: true,
mouse_state: Default::default(),
secret_level,
},
);
}
}
}
}
}
pub(crate) fn redact_secrets_in_element<T: PartialClickableElement>(
mut element: T,
detected_secrets: &DetectedSecretsInTextLocation,
location: TextLocation,
should_hide: bool,
) -> T {
// Collect the secrets into a Vec, so we can reverse sort them by starting byte position.
let secrets: std::iter::Rev<std::vec::IntoIter<(SecretRange, Secret)>> = detected_secrets
.detected_secrets
.iter()
.map(|(range, hoverable)| (range.clone(), hoverable.clone()))
.sorted_by_key(|(detected_secret_range, _)| detected_secret_range.byte_range.start)
.rev();
// Process the secrets in reverse order to avoid issues where we replace multibyte characters with single-width asterisks, changing
// indices of subsequent secrets.
for (detected_secret_range, hoverable_secret) in secrets {
let detected_secret_range_click_clone = detected_secret_range.clone();
element = element.with_clickable_char_range(
detected_secret_range_click_clone.char_range.clone(),
move |_modifiers, ctx, _app| {
ctx.dispatch_typed_action(AIBlockAction::OpenSecretTooltip {
secret_range: detected_secret_range_click_clone.clone(),
location,
});
},
);
if hoverable_secret.is_obfuscated && should_hide {
element.replace_text_range(
detected_secret_range.clone(),
SECRET_REDACTION_REPLACEMENT_CHARACTER
.repeat(
// Use character range length here! Even wide characters should be replaced by a single *, in rich content
// block secret redaction e.g. 码1234 -> *1234 not **1234.
detected_secret_range.char_range.end
- detected_secret_range.char_range.start,
)
.into(),
);
}
let detected_secret_range_hover_clone = detected_secret_range.clone();
element = element.with_hoverable_char_range(
detected_secret_range_hover_clone.char_range.clone(),
hoverable_secret.mouse_state.clone(),
Some(Cursor::PointingHand),
move |is_hovering, ctx, _app| {
ctx.dispatch_typed_action(AIBlockAction::ChangedHoverOnSecret {
secret_range: detected_secret_range_hover_clone.clone(),
location,
is_hovering,
})
},
);
}
element
}
#[cfg(test)]
#[path = "secret_redaction_test.rs"]
mod test;