tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

Feat: Add prompt injection detection settings UI + update logging (#4651)

Browse Source
This commit is contained in:
dorien-koelemeijer
2025-10-03 15:01:37 +02:00
committed by GitHub
parent c96476fc24
commit e6a5692656
8 changed files with 151 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/tool_execution.rs
+10
View File
@@ -31,6 +31,7 @@ impl From<ToolResult<Vec<Content>>> for ToolCallResult {
use super::agent::{tool_stream, ToolStream};
use crate::agents::Agent;
use crate::conversation::message::{Message, ToolRequest};
use crate::tool_inspection::get_security_finding_id_from_results;
pub const DECLINED_RESPONSE: &str = "The user has declined to run this tool. \
DO NOT attempt to call this tool again. \
@@ -79,6 +80,15 @@ impl Agent {
let mut rx = self.confirmation_rx.lock().await;
while let Some((req_id, confirmation)) = rx.recv().await {
if req_id == request.id {
// Log user decision if this was a security alert
if let Some(finding_id) = get_security_finding_id_from_results(&request.id, inspection_results) {
tracing::info!(
"🔒 User security decision: {:?} for finding ID: {}",
confirmation.permission,
finding_id
);
}
if confirmation.permission == Permission::AllowOnce || confirmation.permission == Permission::AlwaysAllow {
// Clone tool_call to avoid moving it
let (req_id, tool_result) = self.dispatch_tool_call(tool_call.clone(), request.id.clone(), cancellation_token.clone(), &None).await;
crates/goose/src/security/mod.rs
+20 -73
View File
@@ -6,15 +6,11 @@ use crate::conversation::message::{Message, ToolRequest};
use crate::permission::permission_judge::PermissionCheckResult;
use anyhow::Result;
use scanner::PromptInjectionScanner;
use std::collections::{hash_map::DefaultHasher, HashSet};
use std::hash::{Hash, Hasher};
use std::sync::{Arc, Mutex};
use std::sync::OnceLock;
use uuid::Uuid;
/// Simple security manager for the POC
/// Focuses on tool call analysis with conversation context
pub struct SecurityManager {
scanner: Option<PromptInjectionScanner>,
flagged_findings: Arc<Mutex<HashSet<String>>>,
scanner: OnceLock<PromptInjectionScanner>,
}
#[derive(Debug, Clone)]
@@ -29,43 +25,19 @@ pub struct SecurityResult {
impl SecurityManager {
pub fn new() -> Self {
// Initialize scanner based on config
let should_enable = Self::should_enable_security();
let scanner = if should_enable {
tracing::info!("Security scanner initialized and enabled");
Some(PromptInjectionScanner::new())
} else {
tracing::debug!("Security scanning disabled via configuration");
None
};
Self {
scanner,
flagged_findings: Arc::new(Mutex::new(HashSet::new())),
scanner: OnceLock::new(),
}
}
/// Check if security should be enabled based on config
fn should_enable_security() -> bool {
// Check config file for security settings
/// Check if prompt injection security is enabled
pub fn is_prompt_injection_detection_enabled(&self) -> bool {
use crate::config::Config;
let config = Config::global();
// Try to get security.enabled from config
let result = config
.get_param::<serde_json::Value>("security")
.ok()
.and_then(|security_config| security_config.get("enabled")?.as_bool())
.unwrap_or(false);
tracing::debug!(
security_config = ?config.get_param::<serde_json::Value>("security"),
enabled = result,
"Security configuration check completed"
);
result
config
.get_param::<bool>("security_prompt_enabled")
.unwrap_or(false)
}
/// New method for tool inspection framework - works directly with tool requests
@@ -74,11 +46,15 @@ impl SecurityManager {
tool_requests: &[ToolRequest],
messages: &[Message],
) -> Result<Vec<SecurityResult>> {
let Some(scanner) = &self.scanner else {
// Security disabled, return empty results
if !self.is_prompt_injection_detection_enabled() {
tracing::debug!("🔓 Security scanning disabled - returning empty results");
return Ok(vec![]);
};
}
let scanner = self.scanner.get_or_init(|| {
tracing::info!("Security scanner initialized and enabled");
PromptInjectionScanner::new()
});
let mut results = Vec::new();
@@ -88,8 +64,7 @@ impl SecurityManager {
messages.len()
);
// Only analyze CURRENT tool requests, not historical ones from conversation
// This prevents re-flagging the same malicious content from previous messages
// Analyze each tool request
for (i, tool_request) in tool_requests.iter().enumerate() {
if let Ok(tool_call) = &tool_request.tool_call {
tracing::info!(
@@ -100,44 +75,16 @@ impl SecurityManager {
"🔍 Starting security analysis for current tool call"
);
// Analyze only the current tool call content, not the entire conversation history
// This prevents re-analyzing and re-flagging historical malicious content
let analysis_result = scanner
.analyze_tool_call_with_context(tool_call, &[]) // Pass empty messages to avoid historical analysis
.analyze_tool_call_with_context(tool_call, messages)
.await?;
// Get threshold from config - only flag things above threshold
let config_threshold = scanner.get_threshold_from_config();
if analysis_result.is_malicious && analysis_result.confidence > config_threshold {
// Generate a unique finding ID based on normalized tool call content
// This ensures the same malicious content always gets the same finding ID
// regardless of JSON formatting or tool request ID variations
let normalized_content = format!(
"{}:{}",
tool_call.name,
serde_json::to_string(&tool_call.arguments).unwrap_or_default()
);
let mut hasher = DefaultHasher::new();
normalized_content.hash(&mut hasher);
let content_hash = hasher.finish();
let finding_id = format!("SEC-{:016x}", content_hash);
// Check if we've already flagged this exact finding before
let mut flagged_set = self.flagged_findings.lock().unwrap();
if flagged_set.contains(&finding_id) {
tracing::debug!(
tool_name = %tool_call.name,
tool_request_id = %tool_request.id,
finding_id = %finding_id,
"🔄 Skipping already flagged security finding - preventing re-flagging"
);
continue;
}
// Mark this finding as flagged
flagged_set.insert(finding_id.clone());
drop(flagged_set); // Release the lock
// Generate a globally unique finding ID for each security finding
let finding_id = format!("SEC-{}", Uuid::new_v4().simple());
tracing::warn!(
tool_name = %tool_call.name,
crates/goose/src/security/scanner.rs
+3 -6
View File
@@ -27,12 +27,10 @@ impl PromptInjectionScanner {
use crate::config::Config;
let config = Config::global();
// Get security config and extract threshold
if let Ok(security_value) = config.get_param::<serde_json::Value>("security") {
if let Some(threshold) = security_value.get("threshold").and_then(|t| t.as_f64()) {
return threshold as f32;
}
if let Ok(threshold) = config.get_param::<f64>("security_prompt_threshold") {
return threshold as f32;
}
0.7 // Default threshold
}
@@ -45,7 +43,6 @@ impl PromptInjectionScanner {
) -> Result<ScanResult> {
// For Phase 1, focus on tool call content analysis
// Phase 2 will add conversation context analysis
let tool_content = self.extract_tool_content(tool_call);
self.scan_for_dangerous_patterns(&tool_content).await
}
crates/goose/src/security/security_inspector.rs
+2 -11
View File
@@ -75,8 +75,6 @@ impl ToolInspector for SecurityInspector {
let inspection_results = security_results
.into_iter()
.map(|security_result| {
// Extract the tool request ID from the security result's context
// The SecurityManager should provide this information
let tool_request_id = security_result.tool_request_id.clone();
self.convert_security_result(&security_result, tool_request_id)
})
@@ -86,15 +84,8 @@ impl ToolInspector for SecurityInspector {
}
fn is_enabled(&self) -> bool {
// Check if security is enabled in config
use crate::config::Config;
let config = Config::global();
config
.get_param::<serde_json::Value>("security")
.ok()
.and_then(|security_config| security_config.get("enabled")?.as_bool())
.unwrap_or(false)
self.security_manager
.is_prompt_injection_detection_enabled()
}
}
crates/goose/src/tool_inspection.rs
+12
View File
@@ -267,6 +267,18 @@ pub fn apply_inspection_results_to_permissions(
permission_result
}
pub fn get_security_finding_id_from_results(
tool_request_id: &str,
inspection_results: &[InspectionResult],
) -> Option<String> {
inspection_results
.iter()
.find(|result| {
result.tool_request_id == tool_request_id && result.inspector_name == "security"
})
.and_then(|result| result.finding_id.clone())
}
#[cfg(test)]
mod tests {
use super::*;