From 9adf50c71bed204dea9152d2f1e16b87b8791786 Mon Sep 17 00:00:00 2001 From: hart_s3 Date: Fri, 7 Aug 2026 14:54:34 +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 | 6 +- app/MindWork AI Studio/Chat/ContentText.cs | 7 +- .../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 | 5 + .../Layout/MainLayout.razor.cs | 33 +- .../Plugins/configuration/plugin.lua | 8 + app/MindWork AI Studio/Program.cs | 5 +- .../Settings/DataModel/DataChat.cs | 9 + .../Tools/AIJobs/AIJobService.cs | 7 + app/MindWork AI Studio/Tools/Event.cs | 5 + .../Tools/PluginSystem/PluginConfiguration.cs | 2 + .../Tools/RAG/IRetrievalContextExtensions.cs | 76 +++-- .../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.8.1.md | 1 + 31 files changed, 827 insertions(+), 36 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 395f8055..6e911a90 100644 --- a/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs +++ b/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs @@ -5,6 +5,7 @@ using AIStudio.Dialogs.Settings; using AIStudio.Tools.AIJobs; using AIStudio.Tools.AssistantSessions; using AIStudio.Tools.Media; +using AIStudio.Tools.Security; using AIStudio.Tools.Services; using Microsoft.AspNetCore.Components; @@ -273,6 +274,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; @@ -461,6 +468,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 0fed4451..5b674b1f 100644 --- a/app/MindWork AI Studio/Assistants/DocumentAnalysis/DocumentAnalysisAssistant.razor.cs +++ b/app/MindWork AI Studio/Assistants/DocumentAnalysis/DocumentAnalysisAssistant.razor.cs @@ -716,7 +716,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 +386,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 cf23d97c..ca65f0bd 100644 --- a/app/MindWork AI Studio/Components/ReadFileContent.razor.cs +++ b/app/MindWork AI Studio/Components/ReadFileContent.razor.cs @@ -3,6 +3,7 @@ using AIStudio.Tools.Media; using AIStudio.Tools.Rust; using AIStudio.Tools.Services; using AIStudio.Tools.Validation; +using AIStudio.Tools.Security; using Microsoft.AspNetCore.Components; @@ -329,6 +330,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..263fb800 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; @@ -7,6 +8,9 @@ namespace AIStudio.Components; public partial class ReadWebContent : MSGComponentBase { + [Inject] + private PromptInjectionGuardService PromptInjectionGuardService { get; init; } = null!; + [Inject] private HTMLParser HTMLParser { get; init; } = null!; @@ -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..46d82126 100644 --- a/app/MindWork AI Studio/Dialogs/DocumentCheckDialog.razor.cs +++ b/app/MindWork AI Studio/Dialogs/DocumentCheckDialog.razor.cs @@ -2,6 +2,7 @@ using AIStudio.Components; using AIStudio.Tools.Services; using Microsoft.AspNetCore.Components; +using AIStudio.Tools.Security; namespace AIStudio.Dialogs; @@ -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..e51b7a63 --- /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..4e87fe2e --- /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..64ccb576 100644 --- a/app/MindWork AI Studio/Dialogs/Settings/SettingsDialogChat.razor +++ b/app/MindWork AI Studio/Dialogs/Settings/SettingsDialogChat.razor @@ -15,6 +15,11 @@ + + + + + diff --git a/app/MindWork AI Studio/Layout/MainLayout.razor.cs b/app/MindWork AI Studio/Layout/MainLayout.razor.cs index a28f6a5c..aa58004b 100644 --- a/app/MindWork AI Studio/Layout/MainLayout.razor.cs +++ b/app/MindWork AI Studio/Layout/MainLayout.razor.cs @@ -7,6 +7,7 @@ using AIStudio.Tools.Media; using AIStudio.Tools.PluginSystem; using AIStudio.Tools.Rust; using AIStudio.Tools.Services; +using AIStudio.Tools.Security; using Microsoft.AspNetCore.Components; using Microsoft.AspNetCore.Components.Routing; @@ -71,6 +72,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 = []; @@ -112,7 +114,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.SHOW_INFO, 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, ]); @@ -253,6 +255,12 @@ public partial class MainLayout : LayoutComponentBase, IMessageBusReceiver, ILan success.Show(this.Snackbar); 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) @@ -347,6 +355,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) { diff --git a/app/MindWork AI Studio/Plugins/configuration/plugin.lua b/app/MindWork AI Studio/Plugins/configuration/plugin.lua index f0b86a74..c6d74dab 100644 --- a/app/MindWork AI Studio/Plugins/configuration/plugin.lua +++ b/app/MindWork AI Studio/Plugins/configuration/plugin.lua @@ -259,6 +259,14 @@ CONFIG["SETTINGS"] = {} -- Examples are PRE_WRITER_MODE_2024 and PRE_RAG_2024. -- CONFIG["SETTINGS"]["DataApp.EnabledPreviewFeatures"] = { "PRE_RAG_2024" } +-- 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. -- It must be one of the provider IDs defined in CONFIG["LLM_PROVIDERS"]. -- Please note: using an empty string ("") will lock the preselected provider selection, even though no valid preselected provider is found. diff --git a/app/MindWork AI Studio/Program.cs b/app/MindWork AI Studio/Program.cs index 5a47a261..d4a129f8 100644 --- a/app/MindWork AI Studio/Program.cs +++ b/app/MindWork AI Studio/Program.cs @@ -10,6 +10,7 @@ using AIStudio.Tools.PluginSystem; using AIStudio.Tools.PluginSystem.Assistants; using AIStudio.Tools.Rust; using AIStudio.Tools.Services; +using AIStudio.Tools.Security; using Microsoft.AspNetCore.DataProtection; using Microsoft.AspNetCore.Server.Kestrel.Core; @@ -193,6 +194,8 @@ internal sealed class Program builder.Services.AddHostedService(); builder.Services.AddScoped(); builder.Services.AddScoped(); + builder.Services.AddSingleton(); + builder.Services.AddSingleton(); // ReSharper disable AccessToDisposedClosure builder.Services.AddHostedService(_ => rust); @@ -300,4 +303,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..7843c5e0 100644 --- a/app/MindWork AI Studio/Settings/DataModel/DataChat.cs +++ b/app/MindWork AI Studio/Settings/DataModel/DataChat.cs @@ -92,7 +92,16 @@ public sealed class DataChat(Expression>? configSelection = this.PreselectedDataSourceIds = [..value.PreselectedDataSourceIds]; } } + /// + /// 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 96354087..4e3be458 100644 --- a/app/MindWork AI Studio/Tools/Event.cs +++ b/app/MindWork AI Studio/Tools/Event.cs @@ -72,6 +72,11 @@ public enum Event /// Requests display of a success notification. /// SHOW_SUCCESS, + + /// + /// Requests display of a prompt-injection alert dialog. + /// + SHOW_PROMPT_INJECTION_ALERT, /// /// Requests display of an informational notification. diff --git a/app/MindWork AI Studio/Tools/PluginSystem/PluginConfiguration.cs b/app/MindWork AI Studio/Tools/PluginSystem/PluginConfiguration.cs index 87acb0ea..295a7b7f 100644 --- a/app/MindWork AI Studio/Tools/PluginSystem/PluginConfiguration.cs +++ b/app/MindWork AI Studio/Tools/PluginSystem/PluginConfiguration.cs @@ -281,6 +281,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/RAG/IRetrievalContextExtensions.cs b/app/MindWork AI Studio/Tools/RAG/IRetrievalContextExtensions.cs index 24b1d24e..ca8d14c2 100644 --- a/app/MindWork AI Studio/Tools/RAG/IRetrievalContextExtensions.cs +++ b/app/MindWork AI Studio/Tools/RAG/IRetrievalContextExtensions.cs @@ -1,5 +1,5 @@ using System.Text; - +using AIStudio.Tools.Security; using AIStudio.Chat; namespace AIStudio.Tools.RAG; @@ -16,7 +16,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 +36,81 @@ 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..17008025 --- /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..ef19096c --- /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..20c95167 --- /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..5f612247 --- /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..ad56775c --- /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..99e0c419 --- /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..a33534b0 --- /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..5603fe5a 100644 --- a/app/MindWork AI Studio/Tools/UserFile.cs +++ b/app/MindWork AI Studio/Tools/UserFile.cs @@ -2,6 +2,7 @@ using AIStudio.Tools.PluginSystem; using AIStudio.Tools.Services; using DialogOptions = AIStudio.Dialogs.DialogOptions; +using AIStudio.Tools.Security; namespace AIStudio.Tools; @@ -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.8.1.md b/app/MindWork AI Studio/wwwroot/changelog/v26.8.1.md index 1ea9e7bc..846c55ab 100644 --- a/app/MindWork AI Studio/wwwroot/changelog/v26.8.1.md +++ b/app/MindWork AI Studio/wwwroot/changelog/v26.8.1.md @@ -6,6 +6,7 @@ - Added the option to import plugins by dropping a plugin archive onto the plugin page. - Added the dedicated file extension `.mwplugin` for plugin archives. - Added an option for organizations to disable importing, sharing, and exporting plugins. +- Added checks for prompt injections in file attachments, images and web pages to process large documents faster and with lower memory usage. - Improved how your organization's configuration behaves when a configuration plugin is present but cannot be loaded, e.g. because of an error in the plugin. Such a plugin still manages your app, so its settings, providers, data sources, profiles, and chat templates now stay in place instead of being removed. - Fixed reset buttons in assistants. As you may have noticed in the Document Analysis Assistant, resetting it could leave content from the previous analysis visible. Reset buttons now clear previous results completely. - Fixed dropping files after you closed a dialog that accepts files itself. Such a dialog takes over dropped files while it is open, but never handed that role back when you closed it. Afterwards, the chat and the assistants silently ignored dropped files until you switched to another page. Each time you opened such a dialog again, the problem got worse.