From 9e645d1fdbd07fb6735b39379011a66ef6050906 Mon Sep 17 00:00:00 2001 From: hart_s3 Date: Tue, 14 Jul 2026 14:37:11 +0200 Subject: [PATCH] Add prompt injection detection framework and UI alert dialog --- .../Agents/AgentRetrievalContextValidation.cs | 16 +- .../Assistants/AssistantBase.razor.cs | 19 ++ .../DocumentAnalysisAssistant.razor.cs | 2 +- .../Assistants/I18N/allTexts.lua | 51 +++ .../AssistantPromptOptimizer.razor.cs | 5 + .../SlideBuilder/SlideAssistant.razor.cs | 7 +- app/MindWork AI Studio/Chat/ContentText.cs | 8 +- .../Components/ReadFileContent.razor.cs | 6 + .../Components/ReadWebContent.razor.cs | 9 + .../Dialogs/DocumentCheckDialog.razor.cs | 6 + .../Dialogs/PromptInjectionAlertDialog.razor | 50 +++ .../PromptInjectionAlertDialog.razor.cs | 17 + .../Dialogs/Settings/SettingsDialogChat.razor | 4 + .../Layout/MainLayout.razor.cs | 33 +- .../Plugins/configuration/plugin.lua | 8 + app/MindWork AI Studio/Program.cs | 5 +- .../Settings/DataModel/DataChat.cs | 10 + .../Tools/AIJobs/AIJobService.cs | 7 + app/MindWork AI Studio/Tools/Event.cs | 5 + .../Tools/PluginSystem/PluginConfiguration.cs | 2 + .../PluginSystem/PluginFactory.Loading.cs | 6 + .../Tools/RAG/IRetrievalContextExtensions.cs | 78 +++-- .../Security/PromptInjectionAlertMessage.cs | 3 + .../PromptInjectionBlockedException.cs | 6 + .../Tools/Security/PromptInjectionFinding.cs | 3 + .../Security/PromptInjectionGuardService.cs | 44 +++ .../Tools/Security/PromptInjectionPatterns.cs | 141 +++++++++ .../Security/PromptInjectionScanResult.cs | 12 + .../Tools/Security/PromptInjectionScanner.cs | 293 ++++++++++++++++++ .../Tools/Security/PromptInjectionSource.cs | 12 + app/MindWork AI Studio/Tools/UserFile.cs | 4 +- .../wwwroot/changelog/v26.7.3.md | 1 + 32 files changed, 836 insertions(+), 37 deletions(-) create mode 100644 app/MindWork AI Studio/Dialogs/PromptInjectionAlertDialog.razor create mode 100644 app/MindWork AI Studio/Dialogs/PromptInjectionAlertDialog.razor.cs create mode 100644 app/MindWork AI Studio/Tools/Security/PromptInjectionAlertMessage.cs create mode 100644 app/MindWork AI Studio/Tools/Security/PromptInjectionBlockedException.cs create mode 100644 app/MindWork AI Studio/Tools/Security/PromptInjectionFinding.cs create mode 100644 app/MindWork AI Studio/Tools/Security/PromptInjectionGuardService.cs create mode 100644 app/MindWork AI Studio/Tools/Security/PromptInjectionPatterns.cs create mode 100644 app/MindWork AI Studio/Tools/Security/PromptInjectionScanResult.cs create mode 100644 app/MindWork AI Studio/Tools/Security/PromptInjectionScanner.cs create mode 100644 app/MindWork AI Studio/Tools/Security/PromptInjectionSource.cs diff --git a/app/MindWork AI Studio/Agents/AgentRetrievalContextValidation.cs b/app/MindWork AI Studio/Agents/AgentRetrievalContextValidation.cs index ee2437d9..ee409fa7 100644 --- a/app/MindWork AI Studio/Agents/AgentRetrievalContextValidation.cs +++ b/app/MindWork AI Studio/Agents/AgentRetrievalContextValidation.cs @@ -4,6 +4,7 @@ using AIStudio.Chat; using AIStudio.Provider; using AIStudio.Settings; using AIStudio.Tools.RAG; +using AIStudio.Tools.Security; using AIStudio.Tools.Services; namespace AIStudio.Agents; @@ -237,7 +238,20 @@ public sealed class AgentRetrievalContextValidation (ILogger(); - var markdownRetrievalContext = await retrievalContext.AsMarkdown(token: token); + string markdownRetrievalContext; + try + { + markdownRetrievalContext = await retrievalContext.AsMarkdown(token: token); + } + catch (PromptInjectionBlockedException exception) + { + logger.LogWarning( + exception, + "Blocked retrieval context '{DataSourceName}' at '{Path}' before validation.", + retrievalContext.DataSourceName, + retrievalContext.Path); + return new(false, "The retrieval context was blocked due to suspected prompt injection.", 1.0f, retrievalContext); + } additionalData.Add("retrievalContext", markdownRetrievalContext); // diff --git a/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs b/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs index 577ce61f..7411e1a9 100644 --- a/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs +++ b/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs @@ -4,6 +4,7 @@ using AIStudio.Settings; using AIStudio.Dialogs.Settings; using AIStudio.Tools.AIJobs; using AIStudio.Tools.AssistantSessions; +using AIStudio.Tools.Security; using AIStudio.Tools.Services; using Microsoft.AspNetCore.Components; @@ -263,6 +264,12 @@ public abstract partial class AssistantBase : AssistantLowerBase wher this.Logger.LogError(e, "The provider request failed for assistant '{AssistantTitle}'. Status={StatusCode}, Reason='{ReasonPhrase}', Body='{ResponseBody}'", this.Title, e.StatusCode, e.ReasonPhrase, e.ResponseBody); await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.CloudOff, e.UserMessage)); } + catch (PromptInjectionBlockedException e) + { + sessionStatus = AssistantSessionStatus.FAILED; + errorMessage = e.Message; + this.Logger.LogWarning(e, "Blocked prompt-injection content during assistant session '{AssistantTitle}'.", this.Title); + } catch (Exception e) { sessionStatus = AssistantSessionStatus.FAILED; @@ -451,6 +458,18 @@ public abstract partial class AssistantBase : AssistantLowerBase wher return string.Empty; } + catch (PromptInjectionBlockedException e) + { + this.Logger.LogWarning(e, "Blocked prompt-injection content before sending assistant request for '{AssistantTitle}'.", this.Title); + + if (this.ResultingContentBlock is not null && string.IsNullOrWhiteSpace(aiText.Text)) + { + this.ChatThread?.Blocks.Remove(this.ResultingContentBlock); + this.ResultingContentBlock = null; + } + + return string.Empty; + } finally { this.IsProcessing = this.assistantSessionId is not null && (this.AssistantSessionService.TryGetSnapshot(this.assistantSessionKey)?.IsActive ?? false); diff --git a/app/MindWork AI Studio/Assistants/DocumentAnalysis/DocumentAnalysisAssistant.razor.cs b/app/MindWork AI Studio/Assistants/DocumentAnalysis/DocumentAnalysisAssistant.razor.cs index 436c5c4d..500081b0 100644 --- a/app/MindWork AI Studio/Assistants/DocumentAnalysis/DocumentAnalysisAssistant.razor.cs +++ b/app/MindWork AI Studio/Assistants/DocumentAnalysis/DocumentAnalysisAssistant.razor.cs @@ -705,7 +705,7 @@ public partial class DocumentAnalysisAssistant : AssistantBaseCore { + [Inject] + private IDialogService DialogService { get; init; } = null!; + protected override Tools.Components Component => Tools.Components.SLIDE_BUILDER_ASSISTANT; protected override string Title => T("Slide Planner Assistant"); @@ -382,7 +387,7 @@ public partial class SlideAssistant : AssistantBaseCore(); + var safeFileContent = await guardService.EnsureSafeForLlmAsync( + await Program.RUST_SERVICE.ReadArbitraryFileData(document.FilePath, int.MaxValue), + PromptInjectionSource.ChatAttachment(document.FilePath)); sb.AppendLine(); sb.AppendLine("---------------------------------------"); sb.AppendLine($"File path: {document.FilePath}"); sb.AppendLine("File content:"); sb.AppendLine("````"); - sb.AppendLine(await Program.RUST_SERVICE.ReadArbitraryFileData(document.FilePath, int.MaxValue)); + sb.AppendLine(safeFileContent); sb.AppendLine("````"); } diff --git a/app/MindWork AI Studio/Components/ReadFileContent.razor.cs b/app/MindWork AI Studio/Components/ReadFileContent.razor.cs index c301a541..f48b3df4 100644 --- a/app/MindWork AI Studio/Components/ReadFileContent.razor.cs +++ b/app/MindWork AI Studio/Components/ReadFileContent.razor.cs @@ -1,4 +1,5 @@ using AIStudio.Tools.Rust; +using AIStudio.Tools.Security; using AIStudio.Tools.Services; using AIStudio.Tools.Validation; @@ -192,6 +193,11 @@ public partial class ReadFileContent : MSGComponentBase this.Logger.LogInformation("Successfully loaded file content: {FilePath}", filePath); return true; } + catch (PromptInjectionBlockedException) + { + this.Logger.LogWarning("Blocked suspected prompt injection while loading file content: {FilePath}", filePath); + return false; + } catch (Exception ex) { this.Logger.LogError(ex, "Failed to load file content: {FilePath}", filePath); diff --git a/app/MindWork AI Studio/Components/ReadWebContent.razor.cs b/app/MindWork AI Studio/Components/ReadWebContent.razor.cs index 53a5e616..3dade2d0 100644 --- a/app/MindWork AI Studio/Components/ReadWebContent.razor.cs +++ b/app/MindWork AI Studio/Components/ReadWebContent.razor.cs @@ -1,5 +1,6 @@ using AIStudio.Agents; using AIStudio.Chat; +using AIStudio.Tools.Security; using Microsoft.AspNetCore.Components; @@ -13,6 +14,9 @@ public partial class ReadWebContent : MSGComponentBase [Inject] private AgentTextContentCleaner AgentTextContentCleaner { get; init; } = null!; + [Inject] + private PromptInjectionGuardService PromptInjectionGuardService { get; init; } = null!; + [Parameter] public string Content { get; set; } = string.Empty; @@ -87,6 +91,7 @@ public partial class ReadWebContent : MSGComponentBase this.processStep = this.process[ReadWebContentSteps.PARSING]; this.StateHasChanged(); markdown = this.HTMLParser.ParseToMarkdown(html); + markdown = await this.PromptInjectionGuardService.EnsureSafeForLlmAsync(markdown, PromptInjectionSource.WebContent(this.providedURL)); if (this.PreselectContentCleanerAgent && this.providerSettings != AIStudio.Settings.Provider.NONE) { @@ -120,6 +125,10 @@ public partial class ReadWebContent : MSGComponentBase this.StateHasChanged(); } } + catch (PromptInjectionBlockedException) + { + markdown = string.Empty; + } catch { if (this.AgentIsRunning) diff --git a/app/MindWork AI Studio/Dialogs/DocumentCheckDialog.razor.cs b/app/MindWork AI Studio/Dialogs/DocumentCheckDialog.razor.cs index 4bf306f1..aaa88b73 100644 --- a/app/MindWork AI Studio/Dialogs/DocumentCheckDialog.razor.cs +++ b/app/MindWork AI Studio/Dialogs/DocumentCheckDialog.razor.cs @@ -1,5 +1,6 @@ using AIStudio.Chat; using AIStudio.Components; +using AIStudio.Tools.Security; using AIStudio.Tools.Services; using Microsoft.AspNetCore.Components; @@ -42,6 +43,11 @@ public partial class DocumentCheckDialog : MSGComponentBase this.FileContent = fileContent; } } + catch (PromptInjectionBlockedException exception) + { + this.Logger.LogWarning(exception, "Blocked suspected prompt injection while previewing '{FilePath}'", this.Document?.FilePath); + this.FileContent = string.Empty; + } catch (Exception ex) { this.Logger.LogError(ex, "Failed to load file content from '{FilePath}'", this.Document); diff --git a/app/MindWork AI Studio/Dialogs/PromptInjectionAlertDialog.razor b/app/MindWork AI Studio/Dialogs/PromptInjectionAlertDialog.razor new file mode 100644 index 00000000..e7dbe490 --- /dev/null +++ b/app/MindWork AI Studio/Dialogs/PromptInjectionAlertDialog.razor @@ -0,0 +1,50 @@ +@using AIStudio.Tools.Security +@inherits MSGComponentBase + + + + + + @T("Prompt Injection Detected") + + + + @if (this.Result is not null) + { + + @T("AI Studio blocked this content before it reached a model or agent.") + + + + @T("Source kind"): @this.Result.Source.Kind + + + @T("Source"): @this.Result.Source.Label + + + + @T("Detected signals") + + + @foreach (var finding in this.Result.Findings) + { + + @($"{finding.Category} / {finding.DetectionStage}: {finding.Snippet}") + + } + + + + @T("More information"): + + @PromptInjectionGuardService.WIKI_URL + + + } + + + + @T("Close") + + + \ No newline at end of file diff --git a/app/MindWork AI Studio/Dialogs/PromptInjectionAlertDialog.razor.cs b/app/MindWork AI Studio/Dialogs/PromptInjectionAlertDialog.razor.cs new file mode 100644 index 00000000..0ddd0c1b --- /dev/null +++ b/app/MindWork AI Studio/Dialogs/PromptInjectionAlertDialog.razor.cs @@ -0,0 +1,17 @@ +using AIStudio.Components; +using AIStudio.Tools.Security; + +using Microsoft.AspNetCore.Components; + +namespace AIStudio.Dialogs; + +public partial class PromptInjectionAlertDialog : MSGComponentBase +{ + [CascadingParameter] + private IMudDialogInstance MudDialog { get; set; } = null!; + + [Parameter] + public PromptInjectionScanResult Result { get; set; } = null!; + + private void Close() => this.MudDialog.Close(); +} \ No newline at end of file diff --git a/app/MindWork AI Studio/Dialogs/Settings/SettingsDialogChat.razor b/app/MindWork AI Studio/Dialogs/Settings/SettingsDialogChat.razor index f80fa857..a8d35a8d 100644 --- a/app/MindWork AI Studio/Dialogs/Settings/SettingsDialogChat.razor +++ b/app/MindWork AI Studio/Dialogs/Settings/SettingsDialogChat.razor @@ -14,6 +14,10 @@ + + + + diff --git a/app/MindWork AI Studio/Layout/MainLayout.razor.cs b/app/MindWork AI Studio/Layout/MainLayout.razor.cs index b7f9aae1..439ca0db 100644 --- a/app/MindWork AI Studio/Layout/MainLayout.razor.cs +++ b/app/MindWork AI Studio/Layout/MainLayout.razor.cs @@ -4,6 +4,7 @@ using AIStudio.Settings.DataModel; using AIStudio.Tools.AIJobs; using AIStudio.Tools.AssistantSessions; using AIStudio.Tools.PluginSystem; +using AIStudio.Tools.Security; using AIStudio.Tools.Rust; using AIStudio.Tools.Services; @@ -64,6 +65,7 @@ public partial class MainLayout : LayoutComponentBase, IMessageBusReceiver, ILan private bool startupCompleted; private bool settingsWriteProtectionWarningShown; private readonly SemaphoreSlim mandatoryInfoDialogSemaphore = new(1, 1); + private readonly SemaphoreSlim promptInjectionDialogSemaphore = new(1, 1); private IReadOnlyCollection navItems = []; @@ -104,7 +106,7 @@ public partial class MainLayout : LayoutComponentBase, IMessageBusReceiver, ILan this.MessageBus.ApplyFilters(this, [], [ Event.UPDATE_AVAILABLE, Event.CONFIGURATION_CHANGED, Event.COLOR_THEME_CHANGED, Event.SHOW_ERROR, - Event.SHOW_WARNING, Event.SHOW_SUCCESS, Event.STARTUP_PLUGIN_SYSTEM, Event.PLUGINS_RELOADED, + Event.SHOW_WARNING, Event.SHOW_SUCCESS, Event.SHOW_PROMPT_INJECTION_ALERT, Event.STARTUP_PLUGIN_SYSTEM, Event.PLUGINS_RELOADED, Event.INSTALL_UPDATE, Event.STARTUP_COMPLETED, Event.AI_JOB_CHANGED, Event.AI_JOB_FINISHED, Event.CHAT_GENERATION_CHANGED, Event.ASSISTANT_SESSION_CHANGED, Event.ASSISTANT_SESSION_FINISHED, ]); @@ -244,6 +246,12 @@ public partial class MainLayout : LayoutComponentBase, IMessageBusReceiver, ILan break; + case Event.SHOW_PROMPT_INJECTION_ALERT: + if (data is PromptInjectionAlertMessage promptInjectionAlert) + await this.ShowPromptInjectionAlertAsync(promptInjectionAlert); + + break; + case Event.SHOW_ERROR: if (data is DataErrorMessage error) error.Show(this.Snackbar); @@ -332,6 +340,29 @@ public partial class MainLayout : LayoutComponentBase, IMessageBusReceiver, ILan }); } + private async Task ShowPromptInjectionAlertAsync(PromptInjectionAlertMessage alert) + { + await this.promptInjectionDialogSemaphore.WaitAsync(); + try + { + var dialogParameters = new DialogParameters + { + { x => x.Result, alert.Result }, + }; + + var dialogReference = await this.DialogService.ShowAsync( + T("Prompt Injection Detected"), + dialogParameters, + DialogOptions.BLOCKING_FULLSCREEN); + + await dialogReference.Result; + } + finally + { + this.promptInjectionDialogSemaphore.Release(); + } + } + public Task ProcessMessageWithResult(ComponentBase? sendingComponent, Event triggeredEvent, TPayload? data) { return Task.FromResult(default); diff --git a/app/MindWork AI Studio/Plugins/configuration/plugin.lua b/app/MindWork AI Studio/Plugins/configuration/plugin.lua index 23a128d0..7b367245 100644 --- a/app/MindWork AI Studio/Plugins/configuration/plugin.lua +++ b/app/MindWork AI Studio/Plugins/configuration/plugin.lua @@ -255,6 +255,14 @@ CONFIG["SETTINGS"] = {} -- This must be enabled for the chat-specific provider, profile, and chat template to take effect. -- CONFIG["SETTINGS"]["DataChat.PreselectOptions"] = true -- +-- Configure prompt-injection protection for external content such as webpages, +-- retrieved context, and file attachments before it is sent to an LLM. +-- CONFIG["SETTINGS"]["DataChat.EnablePromptInjectionProtection"] = true +-- +-- Configure whether AI Studio shows a blocking explanation dialog when +-- suspicious content is detected and blocked. +-- CONFIG["SETTINGS"]["DataChat.ShowPromptInjectionAlert"] = true +-- -- Configure the preselected provider for chats. -- It must be one of the provider IDs defined in CONFIG["LLM_PROVIDERS"]. -- CONFIG["SETTINGS"]["DataChat.PreselectedProvider"] = "00000000-0000-0000-0000-000000000000" diff --git a/app/MindWork AI Studio/Program.cs b/app/MindWork AI Studio/Program.cs index b37b729b..0d7263db 100644 --- a/app/MindWork AI Studio/Program.cs +++ b/app/MindWork AI Studio/Program.cs @@ -6,6 +6,7 @@ using AIStudio.Tools.AIJobs; using AIStudio.Tools.AssistantSessions; using AIStudio.Tools.PluginSystem; using AIStudio.Tools.PluginSystem.Assistants; +using AIStudio.Tools.Security; using AIStudio.Tools.Services; using Microsoft.AspNetCore.Server.Kestrel.Core; @@ -129,6 +130,8 @@ internal sealed class Program builder.Services.AddSingleton(rust); builder.Services.AddMudMarkdownClipboardService(); builder.Services.AddSingleton(); + builder.Services.AddSingleton(); + builder.Services.AddSingleton(); builder.Services.AddSingleton(); builder.Services.AddSingleton(); builder.Services.AddSingleton(); @@ -251,4 +254,4 @@ internal sealed class Program PluginFactory.Dispose(); programLogger.LogInformation("The AI Studio server was stopped."); } -} +} \ No newline at end of file diff --git a/app/MindWork AI Studio/Settings/DataModel/DataChat.cs b/app/MindWork AI Studio/Settings/DataModel/DataChat.cs index 67b3b313..9556e988 100644 --- a/app/MindWork AI Studio/Settings/DataModel/DataChat.cs +++ b/app/MindWork AI Studio/Settings/DataModel/DataChat.cs @@ -93,6 +93,16 @@ public sealed class DataChat(Expression>? configSelection = } } + /// + /// Whether prompt-injection protection is enabled for external and attached content. + /// + public bool EnablePromptInjectionProtection { get; set; } = ManagedConfiguration.Register(configSelection, n => n.EnablePromptInjectionProtection, true); + + /// + /// Whether an alert dialog should be shown when prompt-injection content is blocked. + /// + public bool ShowPromptInjectionAlert { get; set; } = ManagedConfiguration.Register(configSelection, n => n.ShowPromptInjectionAlert, true); + /// /// Should we show the latest message after loading? When false, we show the first (aka oldest) message. /// diff --git a/app/MindWork AI Studio/Tools/AIJobs/AIJobService.cs b/app/MindWork AI Studio/Tools/AIJobs/AIJobService.cs index 7619b6f7..af6b45a0 100644 --- a/app/MindWork AI Studio/Tools/AIJobs/AIJobService.cs +++ b/app/MindWork AI Studio/Tools/AIJobs/AIJobService.cs @@ -5,6 +5,7 @@ using AIStudio.Provider; using AIStudio.Settings; using AIStudio.Tools.PluginSystem; using AIStudio.Tools.RAG.RAGProcesses; +using AIStudio.Tools.Security; namespace AIStudio.Tools.AIJobs; @@ -266,6 +267,12 @@ public sealed class AIJobService( await this.CompleteChatGenerationAsync(state, AIJobStatus.FAILED, e.UserMessage); await MessageBus.INSTANCE.SendError(new(Icons.Material.Filled.CloudOff, e.UserMessage)); } + catch (PromptInjectionBlockedException e) + { + logger.LogWarning(e, "Blocked prompt-injection content during chat generation job '{JobId}'.", state.Snapshot.JobId); + RemoveEmptyAIResponse(state); + await this.CompleteChatGenerationAsync(state, AIJobStatus.FAILED, e.Message); + } catch (Exception e) { logger.LogError(e, "The chat generation job '{JobId}' failed.", state.Snapshot.JobId); diff --git a/app/MindWork AI Studio/Tools/Event.cs b/app/MindWork AI Studio/Tools/Event.cs index dbc737e4..fe99097d 100644 --- a/app/MindWork AI Studio/Tools/Event.cs +++ b/app/MindWork AI Studio/Tools/Event.cs @@ -73,6 +73,11 @@ public enum Event /// SHOW_SUCCESS, + /// + /// Requests display of a prompt-injection alert dialog. + /// + SHOW_PROMPT_INJECTION_ALERT, + /// /// Carries an event received from the Tauri runtime. /// diff --git a/app/MindWork AI Studio/Tools/PluginSystem/PluginConfiguration.cs b/app/MindWork AI Studio/Tools/PluginSystem/PluginConfiguration.cs index 1574f8e2..5309b01d 100644 --- a/app/MindWork AI Studio/Tools/PluginSystem/PluginConfiguration.cs +++ b/app/MindWork AI Studio/Tools/PluginSystem/PluginConfiguration.cs @@ -269,6 +269,8 @@ public sealed class PluginConfiguration(bool isInternal, LuaState state, PluginT ManagedConfiguration.TryProcessConfiguration(x => x.Chat, x => x.PreselectedDataSourcesAutomaticValidation, this.Id, settingsTable, dryRun); ManagedConfiguration.TryProcessConfiguration(x => x.Chat, x => x.PreselectedDataSourceIds, this.Id, settingsTable, dryRun); ManagedConfiguration.TryProcessConfiguration(x => x.Chat, x => x.SendToChatDataSourceBehavior, this.Id, settingsTable, dryRun); + ManagedConfiguration.TryProcessConfiguration(x => x.Chat, x => x.EnablePromptInjectionProtection, this.Id, settingsTable, dryRun); + ManagedConfiguration.TryProcessConfiguration(x => x.Chat, x => x.ShowPromptInjectionAlert, this.Id, settingsTable, dryRun); // Config: transcription provider? ManagedConfiguration.TryProcessConfiguration(x => x.App, x => x.UseTranscriptionProvider, Guid.Empty, this.Id, settingsTable, dryRun); diff --git a/app/MindWork AI Studio/Tools/PluginSystem/PluginFactory.Loading.cs b/app/MindWork AI Studio/Tools/PluginSystem/PluginFactory.Loading.cs index b46409a5..0437fb4f 100644 --- a/app/MindWork AI Studio/Tools/PluginSystem/PluginFactory.Loading.cs +++ b/app/MindWork AI Studio/Tools/PluginSystem/PluginFactory.Loading.cs @@ -229,6 +229,12 @@ public static partial class PluginFactory if(ManagedConfiguration.IsConfigurationLeftOver(x => x.Chat, x => x.SendToChatDataSourceBehavior, AVAILABLE_PLUGINS)) wasConfigurationChanged = true; + + if(ManagedConfiguration.IsConfigurationLeftOver(x => x.Chat, x => x.EnablePromptInjectionProtection, AVAILABLE_PLUGINS)) + wasConfigurationChanged = true; + + if(ManagedConfiguration.IsConfigurationLeftOver(x => x.Chat, x => x.ShowPromptInjectionAlert, AVAILABLE_PLUGINS)) + wasConfigurationChanged = true; // Check for the update interval: if(ManagedConfiguration.IsConfigurationLeftOver(x => x.App, x => x.UpdateInterval, AVAILABLE_PLUGINS)) diff --git a/app/MindWork AI Studio/Tools/RAG/IRetrievalContextExtensions.cs b/app/MindWork AI Studio/Tools/RAG/IRetrievalContextExtensions.cs index 24b1d24e..ddef2011 100644 --- a/app/MindWork AI Studio/Tools/RAG/IRetrievalContextExtensions.cs +++ b/app/MindWork AI Studio/Tools/RAG/IRetrievalContextExtensions.cs @@ -1,6 +1,7 @@ using System.Text; using AIStudio.Chat; +using AIStudio.Tools.Security; namespace AIStudio.Tools.RAG; @@ -16,7 +17,18 @@ public static class IRetrievalContextExtensions foreach(var retrievalContext in retrievalContexts) { index++; - await retrievalContext.AsMarkdown(sb, index, retrievalContexts.Count, token); + try + { + await retrievalContext.AsMarkdown(sb, index, retrievalContexts.Count, token); + } + catch (PromptInjectionBlockedException exception) + { + LOGGER.LogWarning( + exception, + "Skipping retrieval context '{DataSourceName}' at '{Path}' because it was blocked by prompt-injection protection.", + retrievalContext.DataSourceName, + retrievalContext.Path); + } } return sb.ToString(); @@ -25,74 +37,80 @@ public static class IRetrievalContextExtensions public static async Task AsMarkdown(this IRetrievalContext retrievalContext, StringBuilder? sb = null, int index = -1, int numTotalRetrievalContexts = -1, CancellationToken token = default) { sb ??= new StringBuilder(); + var contextBuilder = new StringBuilder(); switch (index) { case > 0 when numTotalRetrievalContexts is -1: - sb.AppendLine($"# Retrieval context {index}"); + contextBuilder.AppendLine($"# Retrieval context {index}"); break; case > 0 when numTotalRetrievalContexts > 0: - sb.AppendLine($"# Retrieval context {index} of {numTotalRetrievalContexts}"); + contextBuilder.AppendLine($"# Retrieval context {index} of {numTotalRetrievalContexts}"); break; default: - sb.AppendLine("# Retrieval context"); + contextBuilder.AppendLine("# Retrieval context"); break; } - sb.AppendLine($"Data source name: {retrievalContext.DataSourceName}"); - sb.AppendLine($"Content category: {retrievalContext.Category}"); - sb.AppendLine($"Content type: {retrievalContext.Type}"); - sb.AppendLine($"Content path: {retrievalContext.Path}"); + contextBuilder.AppendLine($"Data source name: {retrievalContext.DataSourceName}"); + contextBuilder.AppendLine($"Content category: {retrievalContext.Category}"); + contextBuilder.AppendLine($"Content type: {retrievalContext.Type}"); + contextBuilder.AppendLine($"Content path: {retrievalContext.Path}"); if(retrievalContext.Links.Count > 0) { - sb.AppendLine("Additional links:"); + contextBuilder.AppendLine("Additional links:"); foreach(var link in retrievalContext.Links) - sb.AppendLine($"- {link}"); + contextBuilder.AppendLine($"- {link}"); } - + + var guardService = Program.SERVICE_PROVIDER.GetRequiredService(); + var source = PromptInjectionSource.RetrievalContext(retrievalContext.DataSourceName, retrievalContext.Path); switch(retrievalContext) { case RetrievalTextContext textContext: - sb.AppendLine(); - sb.AppendLine("Matched text content:"); - sb.AppendLine("````"); - sb.AppendLine(textContext.MatchedText); - sb.AppendLine("````"); + contextBuilder.AppendLine(); + contextBuilder.AppendLine("Matched text content:"); + contextBuilder.AppendLine("````"); + contextBuilder.AppendLine(textContext.MatchedText); + contextBuilder.AppendLine("````"); if(textContext.SurroundingContent.Count > 0) { - sb.AppendLine(); - sb.AppendLine("Surrounding text content:"); + contextBuilder.AppendLine(); + contextBuilder.AppendLine("Surrounding text content:"); foreach(var surrounding in textContext.SurroundingContent) { - sb.AppendLine(); - sb.AppendLine("````"); - sb.AppendLine(surrounding); - sb.AppendLine("````"); + contextBuilder.AppendLine(); + contextBuilder.AppendLine("````"); + contextBuilder.AppendLine(surrounding); + contextBuilder.AppendLine("````"); } } - - + + await guardService.EnsureSafeForLlmAsync(contextBuilder.ToString(), source); break; case RetrievalImageContext imageContext: - sb.AppendLine(); - sb.AppendLine("Matched image content as base64-encoded data:"); - sb.AppendLine("````"); - sb.AppendLine(await imageContext.TryAsBase64(token) is (success: true, { } base64Image) + await guardService.EnsureSafeForLlmAsync(contextBuilder.ToString(), source); + contextBuilder.AppendLine(); + contextBuilder.AppendLine("Matched image content as base64-encoded data:"); + contextBuilder.AppendLine("````"); + contextBuilder.AppendLine(await imageContext.TryAsBase64(token) is (success: true, { } base64Image) ? base64Image : string.Empty); - sb.AppendLine("````"); + contextBuilder.AppendLine("````"); break; default: + await guardService.EnsureSafeForLlmAsync(contextBuilder.ToString(), source); LOGGER.LogWarning($"The retrieval content type '{retrievalContext.Type}' of data source '{retrievalContext.DataSourceName}' at location '{retrievalContext.Path}' is not supported yet."); break; } - sb.AppendLine(); + contextBuilder.AppendLine(); + sb.Append(contextBuilder); return sb.ToString(); } } \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/Security/PromptInjectionAlertMessage.cs b/app/MindWork AI Studio/Tools/Security/PromptInjectionAlertMessage.cs new file mode 100644 index 00000000..0d789a7e --- /dev/null +++ b/app/MindWork AI Studio/Tools/Security/PromptInjectionAlertMessage.cs @@ -0,0 +1,3 @@ +namespace AIStudio.Tools.Security; + +public sealed record PromptInjectionAlertMessage(PromptInjectionScanResult Result); \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/Security/PromptInjectionBlockedException.cs b/app/MindWork AI Studio/Tools/Security/PromptInjectionBlockedException.cs new file mode 100644 index 00000000..2d79ea90 --- /dev/null +++ b/app/MindWork AI Studio/Tools/Security/PromptInjectionBlockedException.cs @@ -0,0 +1,6 @@ +namespace AIStudio.Tools.Security; + +public sealed class PromptInjectionBlockedException(PromptInjectionScanResult result, string message) : Exception(message) +{ + public PromptInjectionScanResult Result { get; } = result; +} \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/Security/PromptInjectionFinding.cs b/app/MindWork AI Studio/Tools/Security/PromptInjectionFinding.cs new file mode 100644 index 00000000..738ccb18 --- /dev/null +++ b/app/MindWork AI Studio/Tools/Security/PromptInjectionFinding.cs @@ -0,0 +1,3 @@ +namespace AIStudio.Tools.Security; + +public sealed record PromptInjectionFinding(string RuleId, string Category, string DetectionStage, string Snippet); \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/Security/PromptInjectionGuardService.cs b/app/MindWork AI Studio/Tools/Security/PromptInjectionGuardService.cs new file mode 100644 index 00000000..22bc90f8 --- /dev/null +++ b/app/MindWork AI Studio/Tools/Security/PromptInjectionGuardService.cs @@ -0,0 +1,44 @@ +using AIStudio.Settings; +using AIStudio.Tools.PluginSystem; + +namespace AIStudio.Tools.Security; + +public sealed class PromptInjectionGuardService( + PromptInjectionScanner scanner, + SettingsManager settingsManager, + ILogger logger) +{ + public const string WIKI_URL = "https://de.wikipedia.org/wiki/Prompt-Engineering#Prompt_Injection"; + + private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(PromptInjectionGuardService).Namespace, nameof(PromptInjectionGuardService)); + + public async Task EnsureSafeForLlmAsync(string text, PromptInjectionSource source) + { + if (!settingsManager.ConfigurationData.Chat.EnablePromptInjectionProtection || string.IsNullOrWhiteSpace(text)) + return text; + + var result = scanner.Scan(text, source); + if (!result.IsBlocked) + return text; + + var message = string.Format(TB("AI Studio blocked content from '{0}' because it looks like a prompt-injection attempt."), source.Label); + await this.HandleDetectionAsync(result, message); + throw new PromptInjectionBlockedException(result, message); + } + + private async Task HandleDetectionAsync(PromptInjectionScanResult result, string message) + { + logger.LogWarning( + "Blocked suspected prompt injection in {SourceKind} '{SourceLabel}'. RuleIds={RuleIds}", + result.Source.Kind, + result.Source.Label, + string.Join(", ", result.RuleIds)); + + await MessageBus.INSTANCE.SendWarning(new( + Icons.Material.Filled.GppMaybe, + message)); + + if (settingsManager.ConfigurationData.Chat.ShowPromptInjectionAlert) + await MessageBus.INSTANCE.SendMessage(null, Event.SHOW_PROMPT_INJECTION_ALERT, new(result)); + } +} \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/Security/PromptInjectionPatterns.cs b/app/MindWork AI Studio/Tools/Security/PromptInjectionPatterns.cs new file mode 100644 index 00000000..97ae69c8 --- /dev/null +++ b/app/MindWork AI Studio/Tools/Security/PromptInjectionPatterns.cs @@ -0,0 +1,141 @@ +using System.Text.RegularExpressions; + +namespace AIStudio.Tools.Security; + +internal readonly record struct PromptInjectionRegexRule(string Id, string Category, Regex Regex); + +internal static partial class PromptInjectionPatterns +{ + internal static readonly IReadOnlyList RULES = + [ + new("instruction_override", "override", InstructionOverrideRegex()), + new("instruction_priority_override", "override", InstructionPriorityOverrideRegex()), + new("system_prompt_spoofing", "role_override", SystemPromptSpoofingRegex()), + new("system_prompt_exfiltration", "exfiltration", SystemPromptExfiltrationRegex()), + new("prompt_echo_exfiltration", "exfiltration", PromptEchoExfiltrationRegex()), + new("policy_bypass", "override", PolicyBypassRegex()), + new("role_reassignment", "role_override", RoleReassignmentRegex()), + new("privileged_persona_activation", "jailbreak", PrivilegedPersonaActivationRegex()), + new("tool_or_secret_exfiltration", "exfiltration", ToolOrSecretExfiltrationRegex()), + new("conversation_memory_exfiltration", "exfiltration", ConversationMemoryExfiltrationRegex()), + new("tool_call_manipulation", "agent_manipulation", ToolCallManipulationRegex()), + new("agent_thought_injection", "agent_manipulation", AgentThoughtInjectionRegex()), + new("delimiter_wrapped_attack", "delimiter_evasion", DelimiterWrappedAttackRegex()), + new("hidden_markup_injection", "markup_evasion", HiddenMarkupInjectionRegex()), + new("latex_invisible_text", "markup_evasion", LatexInvisibleTextRegex()), + new("unicode_smuggling", "encoding_evasion", UnicodeSmugglingRegex()), + new("ignore_safety_after_data", "override", IgnoreSafetyAfterDataRegex()), + new("persistent_or_delayed_trigger", "persistence", PersistentOrDelayedTriggerRegex()), + new("jailbreak_marker", "jailbreak", JailbreakMarkerRegex()), + ]; + + private const RegexOptions RULE_OPTIONS = RegexOptions.IgnoreCase | RegexOptions.CultureInvariant; + private const int MATCH_TIMEOUT_MILLISECONDS = 100; + + private const string INSTRUCTION_OVERRIDE_PATTERN = """(?:ignore|disregard|forget|bypass|override|replace|drop)\s+(?:all\s+)?(?:previous|prior|above|earlier)\s+(?:instructions?|prompts?|messages?|rules?)"""; + private const string INSTRUCTION_PRIORITY_OVERRIDE_PATTERN = """(?:(?:new|following|these)\s+(?:instructions?|rules?|prompts?)\s+(?:are|is)\s+(?:now\s+)?(?:the\s+)?(?:highest|top|only)\s+priority|(?:take|takes|treat)\s+(?:the\s+)?(?:following|these|this)\s+as\s+(?:the\s+)?(?:new\s+)?(?:system|developer)\s+(?:prompt|message|instructions?)|(?:supersede|replace|override)\s+(?:the\s+)?(?:system|developer|previous|prior|earlier)\s+(?:prompt|message|instructions?|rules?))"""; + private const string SYSTEM_PROMPT_SPOOFING_PATTERN = """(?:(?:this|the\s+following)\s+is\s+(?:a\s+)?(?:system|developer)\s+(?:prompt|message|instruction)|(?:prepend|insert|write)\s+(?:a\s+)?(?:system|developer)\s+(?:prompt|message|instruction)|(?:system|developer|assistant)\s*[:>#-]\s*(?:ignore|bypass|override|reveal|you\s+are\s+now))"""; + private const string SYSTEM_PROMPT_EXFILTRATION_PATTERN = """(?:reveal|show|print|display|dump|expose|leak|tell\s+me|return|quote|repeat\s+back)\s+(?:the\s+)?(?:hidden\s+|full\s+|exact\s+|verbatim\s+|initial\s+|original\s+)?(?:system|developer|assistant)\s+(?:prompt|message|instructions?)"""; + private const string PROMPT_ECHO_EXFILTRATION_PATTERN = """(?:(?:what\s+(?:were|are))\s+your\s+(?:exact|full|hidden|original)\s+(?:instructions?|prompt)|(?:repeat|quote|print|output|display)\s+(?:the\s+)?text\s+(?:above|before|from\s+the\s+top)\s+(?:verbatim|exactly)?|starting\s+with\s+["']?you\s+are)"""; + private const string POLICY_BYPASS_PATTERN = """(?:do\s+not|don't|stop\s+to|never)\s+(?:follow|obey|respect|apply|enforce)\s+(?:the\s+)?(?:system|developer|safety|security|content|usage)\s+(?:prompt|message|instructions?|policy|policies|guardrails?|restrictions?)"""; + private const string ROLE_REASSIGNMENT_PATTERN = """(?:you\s+are\s+now|you\s+are\s+no\s+longer|act\s+as|pretend\s+to\s+be|simulate\s+being|assume\s+the\s+role\s+of|from\s+now\s+on\s+you\s+are)\s+(?:an\s+)?(?:unfiltered|unrestricted|developer|system|root|admin|jailbroken|evil|dan|do\s+anything\s+now)"""; + private const string PRIVILEGED_PERSONA_ACTIVATION_PATTERN = """\b(?:developer\s+mode|debug\s+mode|admin\s+mode|root\s+mode|god\s+mode|maintenance\s+mode|dan\s*(?:mode)?|do\s+anything\s+now|grandmother\s+trick)\b"""; + private const string TOOL_OR_SECRET_EXFILTRATION_PATTERN = """(?:export|send|return|reveal|show|print|list|dump|exfiltrate)\s+(?:all\s+)?(?:tools?|functions?|plugins?|api\s*keys?|keys?|tokens?|credentials?|secrets?|passwords?|hidden\s+instructions?|environment\s+variables?|system\s+information)"""; + private const string CONVERSATION_MEMORY_EXFILTRATION_PATTERN = """(?:(?:show|print|reveal|return|dump|list)\s+(?:the\s+)?(?:conversation\s+history|chat\s+history|memory|scratchpad|chain\s+of\s+thought|reasoning|previous\s+user\s+messages?|prior\s+messages?)|(?:what\s+did\s+(?:the\s+)?previous\s+user\s+say))"""; + private const string TOOL_CALL_MANIPULATION_PATTERN = """(?:(?:call|invoke|execute|run|use|trigger)\s+(?:the\s+)?(?:tool|function|plugin|api|browser|web|shell|terminal|command)[^\n]{0,120}(?:with|using|to)\s+(?:these\s+)?(?:arguments|params?|parameters)|(?:do\s+not|don't)\s+ask\s+for\s+(?:confirmation|approval|permission)|(?:silently|secretly|without\s+asking)\s+(?:call|invoke|execute|run|use))"""; + private const string AGENT_THOUGHT_INJECTION_PATTERN = """(?:(?:thought|observation|reasoning|scratchpad|tool\s+output|assistant|system|developer)\s*[:=]\s*(?:ignore|bypass|override|reveal|call|execute)|forge\s+(?:an\s+)?(?:observation|tool\s+output|assistant\s+message)|pretend\s+(?:the\s+)?tool\s+(?:returned|said))"""; + private const string DELIMITER_WRAPPED_ATTACK_PATTERN = """(?:^|\n)\s*(?:<{2,}|>{2,}|`{3,}|#{1,6}\s*)(?:\s*(?:system|developer|assistant|instructions?|prompt)\b)"""; + private const string HIDDEN_MARKUP_INJECTION_PATTERN = """(?:|<(?:span|div|p|font|section)[^>]{0,200}(?:display\s*:\s*none|visibility\s*:\s*hidden|opacity\s*:\s*0|font-size\s*:\s*0|color\s*:\s*(?:white|#fff(?:fff)?|rgb\(\s*255\s*,\s*255\s*,\s*255\s*\)))[^>]{0,200}>)"""; + private const string LATEX_INVISIBLE_TEXT_PATTERN = """(?:\\(?:color|textcolor)\s*\{\s*white\s*\}\s*\{[^}]{0,300}\}|\\(?:fontsize|tiny|scriptsize)\b[^\r\n]{0,120}(?:ignore|bypass|override|reveal))"""; + private const string UNICODE_SMUGGLING_PATTERN = """[\u200B-\u200F\u2060-\u2064\u2066-\u2069\uFEFF]"""; + private const string IGNORE_SAFETY_AFTER_DATA_PATTERN = """(?:after\s+reading|once\s+you\s+read|when\s+you\s+see)\s+.*?(?:ignore|bypass|override)\s+.*?(?:instructions?|safety|rules?)"""; + private const string PERSISTENT_OR_DELAYED_TRIGGER_PATTERN = """(?:(?:remember|store|save|persist|memorize)\s+(?:this|these|the\s+following)\s+(?:instructions?|rules?|message)|(?:later|in\s+the\s+next\s+message|when\s+you\s+see|whenever\s+you\s+read|if\s+you\s+encounter)\s+.{0,120}(?:ignore|bypass|override|reveal|exfiltrate))"""; + private const string JAILBREAK_MARKER_PATTERN = """\b(?:jailbreak|prompt\s+injection|ignore\s+your\s+guardrails?|bypass\s+(?:your\s+)?(?:guardrails?|safety)|unfiltered\s+mode|do\s+anything\s+now|developer\s+mode|admin\s+mode|root\s+mode)\b"""; + + private const string ANY_RULE_PATTERN = + "(?:" + INSTRUCTION_OVERRIDE_PATTERN + ")|(?:" + + INSTRUCTION_PRIORITY_OVERRIDE_PATTERN + ")|(?:" + SYSTEM_PROMPT_SPOOFING_PATTERN + ")|(?:" + + SYSTEM_PROMPT_EXFILTRATION_PATTERN + ")|(?:" + PROMPT_ECHO_EXFILTRATION_PATTERN + ")|(?:" + + POLICY_BYPASS_PATTERN + ")|(?:" + ROLE_REASSIGNMENT_PATTERN + ")|(?:" + + PRIVILEGED_PERSONA_ACTIVATION_PATTERN + ")|(?:" + TOOL_OR_SECRET_EXFILTRATION_PATTERN + ")|(?:" + + CONVERSATION_MEMORY_EXFILTRATION_PATTERN + ")|(?:" + TOOL_CALL_MANIPULATION_PATTERN + ")|(?:" + + AGENT_THOUGHT_INJECTION_PATTERN + ")|(?:" + DELIMITER_WRAPPED_ATTACK_PATTERN + ")|(?:" + + HIDDEN_MARKUP_INJECTION_PATTERN + ")|(?:" + LATEX_INVISIBLE_TEXT_PATTERN + ")|(?:" + + UNICODE_SMUGGLING_PATTERN + ")|(?:" + IGNORE_SAFETY_AFTER_DATA_PATTERN + ")|(?:" + + PERSISTENT_OR_DELAYED_TRIGGER_PATTERN + ")|(?:" + JAILBREAK_MARKER_PATTERN + ")"; + + [GeneratedRegex(ANY_RULE_PATTERN, RULE_OPTIONS, MATCH_TIMEOUT_MILLISECONDS)] + internal static partial Regex AnyRuleRegex(); + + [GeneratedRegex(@"\b[a-zA-Z](?:[\s._:/\\|-]+[a-zA-Z]){2,}\b", RegexOptions.CultureInvariant)] + internal static partial Regex SpacedLetterSequenceRegex(); + + [GeneratedRegex(@"\b[a-zA-Z]{5,12}\b", RegexOptions.CultureInvariant)] + internal static partial Regex WordRegex(); + + [GeneratedRegex(@"(? findings) +{ + public PromptInjectionSource Source { get; } = source; + + public IReadOnlyList Findings { get; } = findings; + + public bool IsBlocked => this.Findings.Count > 0; + + public IReadOnlyList RuleIds => this.Findings.Select(finding => finding.RuleId).Distinct(StringComparer.Ordinal).ToList(); +} \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/Security/PromptInjectionScanner.cs b/app/MindWork AI Studio/Tools/Security/PromptInjectionScanner.cs new file mode 100644 index 00000000..ae425b24 --- /dev/null +++ b/app/MindWork AI Studio/Tools/Security/PromptInjectionScanner.cs @@ -0,0 +1,293 @@ +using System.Buffers; +using System.Text; +using System.Text.RegularExpressions; + +namespace AIStudio.Tools.Security; + +public sealed class PromptInjectionScanner(ILogger logger) +{ + private const int MAX_DECODED_CANDIDATES_PER_ENCODING = 12; + private const int MAX_DECODED_TEXT_LENGTH = 12_000; + private const int MAX_FINDINGS = 8; + private const int SNIPPET_RADIUS = 80; + + private static readonly IReadOnlyDictionary<(int Length, char First, char Last), string[]> TYPOGLYCEMIA_KEYWORDS = + CreateTypoglycemiaKeywordIndex(); + + public PromptInjectionScanResult Scan(string text, PromptInjectionSource source) + { + if (string.IsNullOrWhiteSpace(text)) + return new(source, []); + + var findings = new List(); + var findingKeys = new HashSet(StringComparer.Ordinal); + + this.ScanVariant(text, "raw", findings, findingKeys); + if (findings.Count >= MAX_FINDINGS) + return new(source, findings); + + var collapsed = CollapseCharacterSpacedContent(text); + if (!string.Equals(text, collapsed, StringComparison.Ordinal)) + this.ScanVariant(collapsed, "character_spacing", findings, findingKeys); + + if (findings.Count < MAX_FINDINGS) + this.ScanDecodedCandidates(text, findings, findingKeys); + + if (findings.Count < MAX_FINDINGS) + this.ScanTypoglycemia(text, findings, findingKeys); + + return new(source, findings); + } + + private void ScanVariant(string text, string stage, List findings, HashSet findingKeys) + { + try + { + if (!PromptInjectionPatterns.AnyRuleRegex().IsMatch(text)) + return; + } + catch (RegexMatchTimeoutException exception) + { + logger.LogWarning(exception, "Prompt-injection regex prefilter timed out during stage '{Stage}'. Falling back to individual rules.", stage); + } + + foreach (var rule in PromptInjectionPatterns.RULES) + { + if (findings.Count >= MAX_FINDINGS) + return; + + Match match; + try + { + match = rule.Regex.Match(text); + } + catch (RegexMatchTimeoutException exception) + { + logger.LogWarning(exception, "Prompt-injection regex '{RuleId}' timed out during stage '{Stage}'.", rule.Id, stage); + continue; + } + + if (!match.Success) + continue; + + var snippet = ExtractSnippet(text, match.Index, match.Length, SNIPPET_RADIUS); + AddFinding(findings, findingKeys, new(rule.Id, rule.Category, stage, snippet)); + } + } + + private void ScanDecodedCandidates(string text, List findings, HashSet findingKeys) + { + var processed = 0; + var seenDecodedTexts = new HashSet(StringComparer.Ordinal); + foreach (var match in PromptInjectionPatterns.Base64Regex().EnumerateMatches(text)) + { + if (processed++ >= MAX_DECODED_CANDIDATES_PER_ENCODING || findings.Count >= MAX_FINDINGS) + break; + + var decoded = TryDecodeBase64(text.AsSpan(match.Index, match.Length)); + if (decoded is not null && seenDecodedTexts.Add(decoded)) + this.ScanVariant(decoded, "decoded_base64", findings, findingKeys); + } + + processed = 0; + seenDecodedTexts.Clear(); + foreach (var match in PromptInjectionPatterns.HexPairRegex().EnumerateMatches(text)) + { + if (processed++ >= MAX_DECODED_CANDIDATES_PER_ENCODING || findings.Count >= MAX_FINDINGS) + break; + + var decoded = TryDecodeHex(text.AsSpan(match.Index, match.Length)); + if (decoded is not null && seenDecodedTexts.Add(decoded)) + this.ScanVariant(decoded, "decoded_hex_pairs", findings, findingKeys); + } + + processed = 0; + seenDecodedTexts.Clear(); + foreach (var match in PromptInjectionPatterns.HexCompactRegex().EnumerateMatches(text)) + { + if (processed++ >= MAX_DECODED_CANDIDATES_PER_ENCODING || findings.Count >= MAX_FINDINGS) + break; + + var decoded = TryDecodeHex(text.AsSpan(match.Index, match.Length)); + if (decoded is not null && seenDecodedTexts.Add(decoded)) + this.ScanVariant(decoded, "decoded_hex", findings, findingKeys); + } + } + + private void ScanTypoglycemia(string text, List findings, HashSet findingKeys) + { + foreach (var match in PromptInjectionPatterns.WordRegex().EnumerateMatches(text)) + { + if (findings.Count >= MAX_FINDINGS) + return; + + var token = text.AsSpan(match.Index, match.Length); + var key = (token.Length, char.ToLowerInvariant(token[0]), char.ToLowerInvariant(token[^1])); + if (!TYPOGLYCEMIA_KEYWORDS.TryGetValue(key, out var keywords)) + continue; + + foreach (var keyword in keywords) + { + if (!IsTypoglycemiaVariant(token, keyword)) + continue; + + var snippet = ExtractSnippet(text, match.Index, match.Length, SNIPPET_RADIUS); + AddFinding(findings, findingKeys, new($"typoglycemia:{keyword}", "evasion", "typoglycemia", snippet)); + break; + } + } + } + + private static bool IsTypoglycemiaVariant(ReadOnlySpan token, string keyword) + { + if (token.Equals(keyword, StringComparison.OrdinalIgnoreCase)) + return false; + + Span characterCounts = stackalloc int[26]; + for (var index = 1; index < token.Length - 1; index++) + { + characterCounts[char.ToLowerInvariant(token[index]) - 'a']++; + characterCounts[keyword[index] - 'a']--; + } + + foreach (var count in characterCounts) + { + if (count != 0) + return false; + } + + return true; + } + + private static string CollapseCharacterSpacedContent(string text) + { + return PromptInjectionPatterns.SpacedLetterSequenceRegex().Replace(text, static match => + { + var builder = new StringBuilder(match.Value.Length); + foreach (var character in match.Value) + { + if (char.IsLetter(character)) + builder.Append(character); + } + + return builder.ToString(); + }); + } + + private static string? TryDecodeBase64(ReadOnlySpan candidate) + { + var encodedLength = Math.Min(candidate.Length, ((MAX_DECODED_TEXT_LENGTH + 2) / 3) * 4); + encodedLength -= encodedLength % 4; + if (encodedLength == 0) + return null; + + var bytes = ArrayPool.Shared.Rent(MAX_DECODED_TEXT_LENGTH); + try + { + if (!Convert.TryFromBase64Chars(candidate[..encodedLength], bytes, out var bytesWritten)) + return null; + + return ConvertDecodedBytesToText(bytes.AsSpan(0, bytesWritten)); + } + finally + { + ArrayPool.Shared.Return(bytes); + } + } + + private static string? TryDecodeHex(ReadOnlySpan candidate) + { + var bytes = ArrayPool.Shared.Rent(MAX_DECODED_TEXT_LENGTH); + try + { + var bytesWritten = 0; + var highNibble = -1; + foreach (var character in candidate) + { + var nibble = HexValue(character); + if (nibble < 0) + continue; + + if (highNibble < 0) + { + highNibble = nibble; + continue; + } + + bytes[bytesWritten++] = (byte)((highNibble << 4) | nibble); + highNibble = -1; + if (bytesWritten >= MAX_DECODED_TEXT_LENGTH) + break; + } + + return bytesWritten == 0 ? null : ConvertDecodedBytesToText(bytes.AsSpan(0, bytesWritten)); + } + finally + { + ArrayPool.Shared.Return(bytes); + } + } + + private static int HexValue(char character) + { + if (character is >= '0' and <= '9') + return character - '0'; + if (character is >= 'A' and <= 'F') + return character - 'A' + 10; + if (character is >= 'a' and <= 'f') + return character - 'a' + 10; + return -1; + } + + private static string? ConvertDecodedBytesToText(ReadOnlySpan bytes) + { + if (bytes.IsEmpty) + return null; + + var text = Encoding.UTF8.GetString(bytes); + return LooksTextLike(text) ? text : null; + } + + private static bool LooksTextLike(string text) + { + if (string.IsNullOrWhiteSpace(text)) + return false; + + var printableCount = 0; + foreach (var character in text) + { + if (!char.IsControl(character) || character is '\r' or '\n' or '\t') + printableCount++; + } + + return printableCount >= text.Length * 0.85; + } + + private static void AddFinding(List findings, HashSet findingKeys, PromptInjectionFinding finding) + { + var key = $"{finding.RuleId}|{finding.DetectionStage}|{finding.Snippet}"; + if (findingKeys.Add(key)) + findings.Add(finding); + } + + private static string ExtractSnippet(string text, int index, int length, int radius) + { + var start = Math.Max(0, index - radius); + var end = Math.Min(text.Length, index + length + radius); + var snippet = text[start..end].Replace('\r', ' ').Replace('\n', ' ').Trim(); + return snippet.Length > radius * 2 ? $"{snippet[..(radius * 2)]}..." : snippet; + } + + private static IReadOnlyDictionary<(int Length, char First, char Last), string[]> CreateTypoglycemiaKeywordIndex() + { + string[] keywords = + [ + "ignore", "bypass", "override", "reveal", "forget", "disregard", "delete", "reset", "expose", + "system", "prompt", "policy", "safety", "developer", "instructions", "admin", "secret", "token", "credential", + ]; + + return keywords + .GroupBy(keyword => (keyword.Length, keyword[0], keyword[^1])) + .ToDictionary(group => group.Key, group => group.ToArray()); + } +} \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/Security/PromptInjectionSource.cs b/app/MindWork AI Studio/Tools/Security/PromptInjectionSource.cs new file mode 100644 index 00000000..bdb283dd --- /dev/null +++ b/app/MindWork AI Studio/Tools/Security/PromptInjectionSource.cs @@ -0,0 +1,12 @@ +namespace AIStudio.Tools.Security; + +public readonly record struct PromptInjectionSource(string Kind, string Label) +{ + public static PromptInjectionSource WebContent(string url) => new("Web content", url); + + public static PromptInjectionSource FileContent(string filePath) => new("File content", filePath); + + public static PromptInjectionSource ChatAttachment(string filePath) => new("Chat attachment", filePath); + + public static PromptInjectionSource RetrievalContext(string dataSourceName, string path) => new("Retrieval context", $"{dataSourceName}: {path}"); +} \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/UserFile.cs b/app/MindWork AI Studio/Tools/UserFile.cs index 14fc0fb4..18544e0c 100644 --- a/app/MindWork AI Studio/Tools/UserFile.cs +++ b/app/MindWork AI Studio/Tools/UserFile.cs @@ -1,4 +1,5 @@ using AIStudio.Dialogs; +using AIStudio.Tools.Security; using AIStudio.Tools.PluginSystem; using AIStudio.Tools.Services; using DialogOptions = AIStudio.Dialogs.DialogOptions; @@ -46,6 +47,7 @@ public static class UserFile } var fileContent = await rustService.ReadArbitraryFileData(filePath, int.MaxValue); - return fileContent; + var guardService = Program.SERVICE_PROVIDER.GetRequiredService(); + return await guardService.EnsureSafeForLlmAsync(fileContent, PromptInjectionSource.FileContent(filePath)); } } \ No newline at end of file diff --git a/app/MindWork AI Studio/wwwroot/changelog/v26.7.3.md b/app/MindWork AI Studio/wwwroot/changelog/v26.7.3.md index 24e7b9e3..ebb57789 100644 --- a/app/MindWork AI Studio/wwwroot/changelog/v26.7.3.md +++ b/app/MindWork AI Studio/wwwroot/changelog/v26.7.3.md @@ -1,5 +1,6 @@ # v26.7.3, build 245 (2026-07-xx xx:xx UTC) - Improved the "My Tasks Assistant": you can now provide one or more documents in addition to text or use documents alone when asking to identify tasks. +- Added checks for prompt injections in file attachments, images and web pages to process large documents faster and with lower memory usage. - Fixed an issue that could leave AI Studio unresponsive after waking the computer from sleep. - Fixed enterprise configuration plugins from Windows-created ZIP files not loading correctly on Linux when the ZIP contained plugin files inside a folder. - Upgraded Rust to v1.97.0.