From 4823bcf4ff5209fa182bdda60fc63c68a8e5bd24 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Wed, 24 Jan 2024 15:47:54 +0200
Subject: [PATCH] Add option to forbid external images

---
 public/index.html            |  4 ++++
 public/script.js             | 19 +++++++++++++++++++
 public/scripts/power-user.js |  8 ++++++++
 3 files changed, 31 insertions(+)

diff --git a/public/index.html b/public/index.html
index d494f5c4a..8c28f9b41 100644
--- a/public/index.html
+++ b/public/index.html
@@ -3420,6 +3420,10 @@
                                         <span class="fa-solid fa-circle-question note-link-span"></span>
                                     </a>
                                 </label>
+                                <label class="checkbox_label" for="forbid_external_images" title="Disalow embedded images from other domains in chat messages.">
+                                    <input id="forbid_external_images" type="checkbox" />
+                                    <span data-i18n="Forbid External Images">Forbid External Images</span>
+                                </label>
                                 <label data-newbie-hidden class="checkbox_label" for="allow_name2_display">
                                     <input id="allow_name2_display" type="checkbox" />
                                     <span data-i18n="Allow {{char}}: in bot messages">Show {{char}}: in responses</span>
diff --git a/public/script.js b/public/script.js
index 13c9cbb5e..673af9d75 100644
--- a/public/script.js
+++ b/public/script.js
@@ -296,6 +296,25 @@ DOMPurify.addHook('uponSanitizeAttribute', (_, data, config) => {
     }
 });
 
+DOMPurify.addHook('uponSanitizeElement', (node, _, config) => {
+    if (!config['MESSAGE_SANITIZE']) {
+        return;
+    }
+
+    switch (node.tagName) {
+        case 'IMG': {
+            const isExternalUrl = (url) => (url.indexOf('://') > 0 || url.indexOf('//') === 0) && !url.startsWith(window.location.origin);
+            const src = node.getAttribute('src');
+
+            if (power_user.forbid_external_images && isExternalUrl(src)) {
+                console.warn('External image blocked', src);
+                node.remove();
+            }
+        }
+        break;
+    }
+});
+
 // API OBJECT FOR EXTERNAL WIRING
 window['SillyTavern'] = {};
 
diff --git a/public/scripts/power-user.js b/public/scripts/power-user.js
index 1f40c72b2..996ca4820 100644
--- a/public/scripts/power-user.js
+++ b/public/scripts/power-user.js
@@ -237,6 +237,7 @@ let power_user = {
     compact_input_area: true,
     auto_connect: false,
     auto_load_chat: false,
+    forbid_external_images: false,
 };
 
 let themes = [];
@@ -1529,6 +1530,7 @@ function loadPowerUserSettings(settings, data) {
     $('#reduced_motion').prop('checked', power_user.reduced_motion);
     $('#auto-connect-checkbox').prop('checked', power_user.auto_connect);
     $('#auto-load-chat-checkbox').prop('checked', power_user.auto_load_chat);
+    $('#forbid_external_images').prop('checked', power_user.forbid_external_images);
 
     for (const theme of themes) {
         const option = document.createElement('option');
@@ -3234,6 +3236,12 @@ $(document).ready(() => {
         saveSettingsDebounced();
     });
 
+    $('#forbid_external_images').on('input', function () {
+        power_user.forbid_external_images = !!$(this).prop('checked');
+        saveSettingsDebounced();
+        reloadCurrentChat();
+    });
+
     $(document).on('click', '#debug_table [data-debug-function]', function () {
         const functionId = $(this).data('debug-function');
         const functionRecord = debug_functions.find(f => f.functionId === functionId);