tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 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::*;
ui/desktop/src/components/settings/chat/ChatSettingsSection.tsx
+7
View File
@@ -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() {
</CardContent>
</Card>
<Card className="pb-2 rounded-lg">
<CardContent className="px-2">
<SecurityToggle />
</CardContent>
</Card>
<Card className="pb-2 rounded-lg">
<CardHeader className="pb-0">
<CardTitle className="">Response Styles</CardTitle>
ui/desktop/src/components/settings/security/SecurityToggle.tsx
+93
View File
@@ -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 (
<div className="space-y-4">
<div className="flex items-center justify-between py-2 px-2 hover:bg-background-muted rounded-lg transition-all">
<div>
<h3 className="text-text-default">Enable Prompt Injection Detection</h3>
<p className="text-xs text-text-muted max-w-md mt-[2px]">
Detect and prevent potential prompt injection attacks
</p>
</div>
<div className="flex items-center">
<Switch checked={enabled} onCheckedChange={handleToggle} variant="mono" />
</div>
</div>
<div
className={`overflow-hidden transition-all duration-300 ease-in-out ${
enabled ? 'max-h-96 opacity-100' : 'max-h-0 opacity-0'
}`}
>
<div className="space-y-3 px-2 pb-2">
<div className={enabled ? '' : 'opacity-50'}>
<label
className={`text-sm font-medium ${enabled ? 'text-text-default' : 'text-text-muted'}`}
>
Detection Threshold
</label>
<p className="text-xs text-text-muted mb-2">
Higher values are more strict (0.01 = very lenient, 1.0 = maximum strict)
</p>
<input
type="number"
min={0.01}
max={1.0}
step={0.01}
value={thresholdInput}
onChange={(e) => {
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"
/>
</div>
</div>
</div>
</div>
);
};
ui/desktop/src/utils/configUtils.ts
+4
View File
@@ -14,6 +14,10 @@ export const configLabels: Record<string, string> = {
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',