From e6a56926568458d2bcc6ce938bf8313c5af7692c Mon Sep 17 00:00:00 2001 From: dorien-koelemeijer <62866702+dorien-koelemeijer@users.noreply.github.com> Date: Fri, 3 Oct 2025 15:01:37 +0200 Subject: [PATCH] Feat: Add prompt injection detection settings UI + update logging (#4651) --- crates/goose/src/agents/tool_execution.rs | 10 ++ crates/goose/src/security/mod.rs | 93 ++++--------------- crates/goose/src/security/scanner.rs | 9 +- .../goose/src/security/security_inspector.rs | 13 +-- crates/goose/src/tool_inspection.rs | 12 +++ .../settings/chat/ChatSettingsSection.tsx | 7 ++ .../settings/security/SecurityToggle.tsx | 93 +++++++++++++++++++ ui/desktop/src/utils/configUtils.ts | 4 + 8 files changed, 151 insertions(+), 90 deletions(-) create mode 100644 ui/desktop/src/components/settings/security/SecurityToggle.tsx diff --git a/crates/goose/src/agents/tool_execution.rs b/crates/goose/src/agents/tool_execution.rs index fa1206a2..cbcbf3e6 100644 --- a/crates/goose/src/agents/tool_execution.rs +++ b/crates/goose/src/agents/tool_execution.rs @@ -31,6 +31,7 @@ impl From>> 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; diff --git a/crates/goose/src/security/mod.rs b/crates/goose/src/security/mod.rs index 696f8db1..77d48381 100644 --- a/crates/goose/src/security/mod.rs +++ b/crates/goose/src/security/mod.rs @@ -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, - flagged_findings: Arc>>, + scanner: OnceLock, } #[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::("security") - .ok() - .and_then(|security_config| security_config.get("enabled")?.as_bool()) - .unwrap_or(false); - - tracing::debug!( - security_config = ?config.get_param::("security"), - enabled = result, - "Security configuration check completed" - ); - - result + config + .get_param::("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> { - 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, diff --git a/crates/goose/src/security/scanner.rs b/crates/goose/src/security/scanner.rs index fdaab22f..fec76b44 100644 --- a/crates/goose/src/security/scanner.rs +++ b/crates/goose/src/security/scanner.rs @@ -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::("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::("security_prompt_threshold") { + return threshold as f32; } + 0.7 // Default threshold } @@ -45,7 +43,6 @@ impl PromptInjectionScanner { ) -> Result { // 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 } diff --git a/crates/goose/src/security/security_inspector.rs b/crates/goose/src/security/security_inspector.rs index 372e6354..41a1fcad 100644 --- a/crates/goose/src/security/security_inspector.rs +++ b/crates/goose/src/security/security_inspector.rs @@ -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::("security") - .ok() - .and_then(|security_config| security_config.get("enabled")?.as_bool()) - .unwrap_or(false) + self.security_manager + .is_prompt_injection_detection_enabled() } } diff --git a/crates/goose/src/tool_inspection.rs b/crates/goose/src/tool_inspection.rs index f70157f3..56f1cdce 100644 --- a/crates/goose/src/tool_inspection.rs +++ b/crates/goose/src/tool_inspection.rs @@ -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 { + 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::*; diff --git a/ui/desktop/src/components/settings/chat/ChatSettingsSection.tsx b/ui/desktop/src/components/settings/chat/ChatSettingsSection.tsx index fff359cb..68ead2dc 100644 --- a/ui/desktop/src/components/settings/chat/ChatSettingsSection.tsx +++ b/ui/desktop/src/components/settings/chat/ChatSettingsSection.tsx @@ -2,6 +2,7 @@ import { ModeSection } from '../mode/ModeSection'; import { ToolSelectionStrategySection } from '../tool_selection_strategy/ToolSelectionStrategySection'; import SchedulerSection from '../scheduler/SchedulerSection'; import DictationSection from '../dictation/DictationSection'; +import { SecurityToggle } from '../security/SecurityToggle'; import { ResponseStylesSection } from '../response_styles/ResponseStylesSection'; import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '../../ui/card'; @@ -18,6 +19,12 @@ export default function ChatSettingsSection() { + + + + + + Response Styles diff --git a/ui/desktop/src/components/settings/security/SecurityToggle.tsx b/ui/desktop/src/components/settings/security/SecurityToggle.tsx new file mode 100644 index 00000000..e6ba583a --- /dev/null +++ b/ui/desktop/src/components/settings/security/SecurityToggle.tsx @@ -0,0 +1,93 @@ +import { useState, useEffect } from 'react'; +import { Switch } from '../../ui/switch'; +import { useConfig } from '../../ConfigContext'; + +interface SecurityConfig { + security_prompt_enabled?: boolean; + security_prompt_threshold?: number; +} + +export const SecurityToggle = () => { + const { config, upsert } = useConfig(); + + const { + security_prompt_enabled: enabled = false, + security_prompt_threshold: configThreshold = 0.7, + } = (config as SecurityConfig) ?? {}; + + const [thresholdInput, setThresholdInput] = useState(configThreshold.toString()); + + useEffect(() => { + setThresholdInput(configThreshold.toString()); + }, [configThreshold]); + + const handleToggle = async (enabled: boolean) => { + await upsert('security_prompt_enabled', enabled, false); + }; + + const handleThresholdChange = async (threshold: number) => { + const validThreshold = Math.max(0, Math.min(1, threshold)); + await upsert('security_prompt_threshold', validThreshold, false); + }; + + return ( +
+
+
+

Enable Prompt Injection Detection

+

+ Detect and prevent potential prompt injection attacks +

+
+
+ +
+
+ +
+
+
+ +

+ Higher values are more strict (0.01 = very lenient, 1.0 = maximum strict) +

+ { + setThresholdInput(e.target.value); + }} + onBlur={(e) => { + const value = parseFloat(e.target.value); + if (isNaN(value) || value < 0.01 || value > 1.0) { + // Revert to previous valid value + setThresholdInput(configThreshold.toString()); + } else { + handleThresholdChange(value); + } + }} + disabled={!enabled} + className={`w-24 px-2 py-1 text-sm border rounded ${ + enabled + ? 'border-border-default bg-background-default text-text-default' + : 'border-border-muted bg-background-muted text-text-muted cursor-not-allowed' + }`} + placeholder="0.70" + /> +
+
+
+
+ ); +}; diff --git a/ui/desktop/src/utils/configUtils.ts b/ui/desktop/src/utils/configUtils.ts index cb10a93d..5a32a56d 100644 --- a/ui/desktop/src/utils/configUtils.ts +++ b/ui/desktop/src/utils/configUtils.ts @@ -14,6 +14,10 @@ export const configLabels: Record = { GOOSE_ALLOWLIST: 'Allow List', GOOSE_RECIPE_GITHUB_REPO: 'Recipe GitHub Repo', + // security settings + security_prompt_enabled: 'Prompt Injection Detection Enabled', + security_prompt_threshold: 'Prompt Injection Detection Threshold', + // openai OPENAI_API_KEY: 'OpenAI API Key', OPENAI_HOST: 'OpenAI Host',