From 5cc40435c4c29a4afaeec53bfa7ac4f27f248396 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daphn=C3=A9=20Popin?= <daphne.popin@gmail.com>
Date: Fri, 3 Jan 2025 15:45:08 +0100
Subject: [PATCH] Ext: Disable context menu items if domain is blacklisted

---
 extension/app/background.ts                   | 54 ++++++++++++++++++-
 .../input_bar/InputBarCitations.tsx           | 44 ++++++++-------
 2 files changed, 74 insertions(+), 24 deletions(-)

diff --git a/extension/app/background.ts b/extension/app/background.ts
index 0319dfb9bd20b..02d79c49198f2 100644
--- a/extension/app/background.ts
+++ b/extension/app/background.ts
@@ -1,5 +1,5 @@
 import type { PendingUpdate } from "@extension/lib/storage";
-import { savePendingUpdate } from "@extension/lib/storage";
+import { getStoredUser, savePendingUpdate } from "@extension/lib/storage";
 
 import {
   AUTH0_CLIENT_DOMAIN,
@@ -73,6 +73,58 @@ chrome.runtime.onInstalled.addListener(() => {
   });
 });
 
+/**
+ * Util & listeners to disable context menu items based on the domain.
+ */
+const shouldDisableContextMenuForDomain = async (
+  url: string
+): Promise<boolean> => {
+  if (url.startsWith("chrome://")) {
+    return true;
+  }
+
+  const user = await getStoredUser();
+  if (!user || !user.selectedWorkspace) {
+    return false;
+  }
+
+  const blacklistedDomains =
+    user.workspaces.find((w) => w.sId === user.selectedWorkspace)
+      ?.blacklistedDomains || [];
+
+  return blacklistedDomains.some((d) => url.includes(d));
+};
+
+const toggleContextMenus = (isDisabled: boolean) => {
+  [
+    "ask_dust",
+    "add_tab_content",
+    "add_tab_screenshot",
+    "add_selection",
+  ].forEach((menuId) => {
+    chrome.contextMenus.update(menuId, { enabled: !isDisabled });
+  });
+};
+
+// Add URL change listener to update context menu state.
+chrome.tabs.onUpdated.addListener(async (tabId, changeInfo, tab) => {
+  console.log(tab.url);
+  if (changeInfo.status === "complete" && tab.url) {
+    const isDisabled = await shouldDisableContextMenuForDomain(tab.url);
+    toggleContextMenus(isDisabled);
+  }
+});
+
+// Also add URL change listener for active tab changes.
+chrome.tabs.onActivated.addListener(async (activeInfo) => {
+  const tab = await chrome.tabs.get(activeInfo.tabId);
+  console.log(tab.url);
+  if (tab.url) {
+    const isDisabled = await shouldDisableContextMenuForDomain(tab.url);
+    toggleContextMenus(isDisabled);
+  }
+});
+
 chrome.runtime.onConnect.addListener((port) => {
   if (port.name === "sidepanel-connection") {
     console.log("Sidepanel is there");
diff --git a/extension/app/src/components/input_bar/InputBarCitations.tsx b/extension/app/src/components/input_bar/InputBarCitations.tsx
index db8d44e902c1f..b2ca599d53e04 100644
--- a/extension/app/src/components/input_bar/InputBarCitations.tsx
+++ b/extension/app/src/components/input_bar/InputBarCitations.tsx
@@ -26,31 +26,29 @@ export function InputBarCitations({
       const isImage = Boolean(blob.preview);
 
       nodes.push(
-        <>
-          <Citation
-            disabled={disabled}
-            key={`cf-${blob.id}`}
-            className="w-48"
-            isLoading={blob.isUploading}
-          >
-            {isImage ? (
-              <>
-                <CitationImage imgSrc={blob.preview ?? ""} />
-                <CitationIcons>
-                  <Icon visual={ImageIcon} />
-                </CitationIcons>
-              </>
-            ) : (
+        <Citation
+          disabled={disabled}
+          key={`cf-${blob.id}`}
+          className="w-48"
+          isLoading={blob.isUploading}
+        >
+          {isImage ? (
+            <>
+              <CitationImage imgSrc={blob.preview ?? ""} />
               <CitationIcons>
-                <Icon visual={DocumentIcon} />
+                <Icon visual={ImageIcon} />
               </CitationIcons>
-            )}
-            <CitationTitle>{blob.id}</CitationTitle>
-            <CitationClose
-              onClick={() => fileUploaderService.removeFile(blob.id)}
-            />
-          </Citation>
-        </>
+            </>
+          ) : (
+            <CitationIcons>
+              <Icon visual={DocumentIcon} />
+            </CitationIcons>
+          )}
+          <CitationTitle>{blob.id}</CitationTitle>
+          <CitationClose
+            onClick={() => fileUploaderService.removeFile(blob.id)}
+          />
+        </Citation>
       );
     }