From 1c3ce56b48923772d1e4156561739f29662502ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89mile=20R=C3=A9?= <emile@probo.com>
Date: Fri, 22 May 2026 18:04:17 +0200
Subject: [PATCH] Mark page-world extension writes with EXTENSION source
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The previous cleanup deleted every isExtensionCaller() site, including
the one in cookie/storage detectors that did fire reliably for the
residual case: page-world extensions (MV3 main world, userscripts with
@grant none) whose stack contains a chrome-/moz-/safari-web-extension
frame at the synchronous write. Recover that signal for free by
returning fromExtension from getInitiatorURL (it already walks the
stack and discards extension frames via continue), and have the cookie
and storage detectors report source: "extension" instead of "script"
when the flag is set.

End-to-end plumbing reuses the existing source column: extend the
cookie_source Postgres enum with EXTENSION, add the CookieSourceExtension
constant with a doc block describing each bucket's actual semantics,
add the handler.go switch cases, expose EXTENSION on the GraphQL and
MCP CookieSource enums, and add the Extension option to the console
source filter.

Update bestSource in the pattern analysis worker so a glob merging
only extension-attributed exact patterns is no longer silently rolled
up to PRE_EXISTING. New precedence is SCRIPT > EXTENSION > PRE_EXISTING,
matching the upsert SQL's "page-script wins" rule and the asymmetric
signal strength of each bucket.

Out of scope: any behavioural use of EXTENSION (auto-exclusion,
denylist classification, dashboard surfacing) -- that belongs in the
follow-up backend denylist plan.

Signed-off-by: Émile Ré <emile@probo.com>
---
 .../trackers/CookieBannerTrackersPage.tsx     |  1 +
 .../src/detectors/cookie-detector.ts          |  4 +-
 .../cookie-banner/src/detectors/initiator.ts  | 48 +++++++++++++------
 .../src/detectors/storage-detector.ts         | 11 +++--
 packages/cookie-banner/src/detectors/types.ts |  4 +-
 pkg/cookiebanner/pattern_analysis_worker.go   | 27 ++++++++++-
 pkg/coredata/cookie_source.go                 | 37 +++++++++++++-
 pkg/coredata/migrations/20260522T155558Z.sql  | 15 ++++++
 .../console/v1/graphql/cookie_banner.graphql  |  4 ++
 pkg/server/api/cookiebanner/v1/handler.go     |  4 ++
 pkg/server/api/mcp/v1/specification.yaml      |  2 +-
 11 files changed, 130 insertions(+), 27 deletions(-)
 create mode 100644 pkg/coredata/migrations/20260522T155558Z.sql

diff --git a/apps/console/src/pages/organizations/cookie-banners/configuration/trackers/CookieBannerTrackersPage.tsx b/apps/console/src/pages/organizations/cookie-banners/configuration/trackers/CookieBannerTrackersPage.tsx
index c2eb96608..8fe92694b 100644
--- a/apps/console/src/pages/organizations/cookie-banners/configuration/trackers/CookieBannerTrackersPage.tsx
+++ b/apps/console/src/pages/organizations/cookie-banners/configuration/trackers/CookieBannerTrackersPage.tsx
@@ -209,6 +209,7 @@ export default function CookieBannerTrackersPage({
           <Option value="ALL">{__("All sources")}</Option>
           <Option value="SCRIPT">{__("Script")}</Option>
           <Option value="PRE_EXISTING">{__("Pre-existing")}</Option>
+          <Option value="EXTENSION">{__("Extension")}</Option>
         </Select>
         <Select
           value={categoryFilter ?? "ALL"}
diff --git a/packages/cookie-banner/src/detectors/cookie-detector.ts b/packages/cookie-banner/src/detectors/cookie-detector.ts
index d30f87ce1..17247b212 100644
--- a/packages/cookie-banner/src/detectors/cookie-detector.ts
+++ b/packages/cookie-banner/src/detectors/cookie-detector.ts
@@ -83,12 +83,12 @@ export class CookieDetector implements Detector {
     if (!name || this.knownNames.has(name)) return;
 
     const maxAgeSeconds = parseMaxAgeSeconds(raw);
-    const initiatorUrl = getInitiatorURL(this.apiOrigin);
+    const { url: initiatorUrl, fromExtension } = getInitiatorURL(this.apiOrigin);
 
     const entry: DetectedCookieEntry = {
       name,
       max_age_seconds: maxAgeSeconds,
-      source: "script",
+      source: fromExtension ? "extension" : "script",
     };
     if (initiatorUrl) entry.initiator_url = initiatorUrl;
     this.queue.reportCookie(entry);
diff --git a/packages/cookie-banner/src/detectors/initiator.ts b/packages/cookie-banner/src/detectors/initiator.ts
index bc0f9bc39..7401cb399 100644
--- a/packages/cookie-banner/src/detectors/initiator.ts
+++ b/packages/cookie-banner/src/detectors/initiator.ts
@@ -17,23 +17,41 @@ const STACK_URL_RE = /https?:\/\/[^\s)'"`]+/g;
 const LINE_COL_SUFFIX_RE = /:\d+(?::\d+)?$/;
 const MAX_INITIATOR_URL_LENGTH = 1024;
 
-// getInitiatorURL walks the current call stack and returns the first
-// third-party script URL (as origin+pathname). It deliberately skips:
-//   - browser extension URLs (chrome/moz/safari-web-extension://)
-//   - the SDK's API origin (instrumentation frames)
-//   - the page's own origin (we want the third-party loader, not first-party code)
+export interface InitiatorContext {
+  url: string | null;
+  fromExtension: boolean;
+}
+
+// getInitiatorURL walks the current call stack once and returns:
+//   - `url`: the first third-party HTTP(S) script URL it finds
+//     (as origin+pathname), skipping the SDK's API origin and the
+//     page's own origin. Returns null when no such frame exists.
+//   - `fromExtension`: true if any frame on the stack was a
+//     browser-extension URL (chrome/moz/safari-web-extension://).
+//     Page-world extensions (MV3 main world, userscripts with
+//     @grant none) reliably leave such a frame; isolated-world
+//     content scripts use a different realm and never reach this
+//     function in the first place.
 //
-// Returns null when no third-party frame is found (anonymous/eval/inline
-// scripts, or writes originating from the page itself).
-export function getInitiatorURL(apiOrigin: string): string | null {
+// Both signals come from the same single stack walk, so callers
+// that need either or both pay only one `new Error().stack` cost.
+export function getInitiatorURL(apiOrigin: string): InitiatorContext {
   const stack = new Error().stack;
-  if (!stack) return null;
+  if (!stack) return { url: null, fromExtension: false };
+
+  let fromExtension = false;
+  let url: string | null = null;
 
   STACK_URL_RE.lastIndex = 0;
   let m: RegExpExecArray | null;
   while ((m = STACK_URL_RE.exec(stack)) !== null) {
     const raw = m[0];
-    if (EXTENSION_URL_RE.test(raw)) continue;
+    if (EXTENSION_URL_RE.test(raw)) {
+      fromExtension = true;
+      continue;
+    }
+
+    if (url !== null) continue;
 
     const cleaned = raw.replace(LINE_COL_SUFFIX_RE, "");
 
@@ -48,10 +66,10 @@ export function getInitiatorURL(apiOrigin: string): string | null {
     if (parsed.origin === location.origin) continue;
 
     const result = parsed.origin + parsed.pathname;
-    if (result.length > MAX_INITIATOR_URL_LENGTH) {
-      return result.slice(0, MAX_INITIATOR_URL_LENGTH);
-    }
-    return result;
+    url = result.length > MAX_INITIATOR_URL_LENGTH
+      ? result.slice(0, MAX_INITIATOR_URL_LENGTH)
+      : result;
   }
-  return null;
+
+  return { url, fromExtension };
 }
diff --git a/packages/cookie-banner/src/detectors/storage-detector.ts b/packages/cookie-banner/src/detectors/storage-detector.ts
index e01799948..afc44f6ff 100644
--- a/packages/cookie-banner/src/detectors/storage-detector.ts
+++ b/packages/cookie-banner/src/detectors/storage-detector.ts
@@ -100,7 +100,8 @@ export class StorageDetector implements Detector {
     const self = this;
 
     caches.open = function (name: string): Promise<Cache> {
-      self.onCacheStorageOpen(name, "script");
+      const { fromExtension } = getInitiatorURL(self.apiOrigin);
+      self.onCacheStorageOpen(name, fromExtension ? "extension" : "script");
       return originalOpen(name);
     };
   }
@@ -112,24 +113,26 @@ export class StorageDetector implements Detector {
   ): void {
     if (key.startsWith(OWN_KEY_PREFIX)) return;
 
-    const initiatorUrl = getInitiatorURL(this.apiOrigin);
+    const { url: initiatorUrl, fromExtension } = getInitiatorURL(this.apiOrigin);
 
     const entry: DetectedStorageEntry = {
       key,
       storage_type: storageType,
       value_size: value.length * 2,
-      source: "script",
+      source: fromExtension ? "extension" : "script",
     };
     if (initiatorUrl) entry.initiator_url = initiatorUrl;
     this.queue.reportStorage(entry);
   }
 
   private onIndexedDBOpen(name: string): void {
+    const { fromExtension } = getInitiatorURL(this.apiOrigin);
+
     this.queue.reportStorage({
       key: name,
       storage_type: "indexed_db",
       value_size: null,
-      source: "script",
+      source: fromExtension ? "extension" : "script",
     });
   }
 
diff --git a/packages/cookie-banner/src/detectors/types.ts b/packages/cookie-banner/src/detectors/types.ts
index 3e9a9b413..7cc548eeb 100644
--- a/packages/cookie-banner/src/detectors/types.ts
+++ b/packages/cookie-banner/src/detectors/types.ts
@@ -17,9 +17,9 @@
 // pkg/server/api/cookiebanner/v1/handler.go. Any change here must be
 // matched server-side or the request will be rejected.
 
-export type CookieSource = "script" | "pre-existing" | "http";
+export type CookieSource = "script" | "pre-existing" | "http" | "extension";
 
-export type StorageSource = "script" | "pre-existing";
+export type StorageSource = "script" | "pre-existing" | "extension";
 
 export type StorageType =
   | "local_storage"
diff --git a/pkg/cookiebanner/pattern_analysis_worker.go b/pkg/cookiebanner/pattern_analysis_worker.go
index 4d55cbaeb..f46e4148a 100644
--- a/pkg/cookiebanner/pattern_analysis_worker.go
+++ b/pkg/cookiebanner/pattern_analysis_worker.go
@@ -633,11 +633,34 @@ func globMatch(pattern, name string) bool {
 	return true
 }
 
+// bestSource rolls up the source values of a group of exact patterns
+// being merged into a single glob. Precedence is SCRIPT > EXTENSION
+// > PRE_EXISTING, mirroring both the upsert SQL's "page-script wins"
+// rule and the asymmetric signal strength of each bucket: SCRIPT is
+// high-confidence page evidence (a real page tracker), EXTENSION is
+// high-confidence extension evidence, and PRE_EXISTING is the
+// catch-all that may include extension state injected before SDK
+// load. HTTP and nil collapse into PRE_EXISTING here, preserving
+// the original two-value rollup behaviour for non-script values.
 func bestSource(patterns []*coredata.TrackerPattern) *coredata.CookieSource {
+	var hasExtension bool
+
 	for _, p := range patterns {
-		if p.Source != nil && *p.Source == coredata.CookieSourceScript {
-			return p.Source
+		if p.Source == nil {
+			continue
 		}
+
+		switch *p.Source {
+		case coredata.CookieSourceScript:
+			return p.Source
+		case coredata.CookieSourceExtension:
+			hasExtension = true
+		}
+	}
+
+	if hasExtension {
+		src := coredata.CookieSourceExtension
+		return &src
 	}
 
 	src := coredata.CookieSourcePreExisting
diff --git a/pkg/coredata/cookie_source.go b/pkg/coredata/cookie_source.go
index 6291114cf..6320c4ea0 100644
--- a/pkg/coredata/cookie_source.go
+++ b/pkg/coredata/cookie_source.go
@@ -21,10 +21,43 @@ import (
 
 type CookieSource string
 
+// CookieSourceScript: JS write observed via a detector hook on the
+// page realm's prototypes. In practice this is a page-script write
+// with high confidence -- isolated-world content scripts use their
+// own realm's prototypes and never trip the hook (they don't land
+// in this bucket at all), and page-world extensions (MV3 main world,
+// userscripts with @grant none) reliably leave a browser-extension
+// frame on the stack and classify as CookieSourceExtension. The
+// only residual contamination is rare cases where a page-world
+// extension's frame gets stripped from the stack (deep async,
+// page-side trampolines).
+//
+// CookieSourceExtension: synchronous JS write whose stack at the
+// hook contained at least one chrome-/moz-/safari-web-extension
+// frame. A page-world extension write is confirmed.
+//
+// CookieSourcePreExisting: enumerated from the storage at SDK init
+// rather than observed at write time. This is the catch-all bucket:
+// it bundles real pre-existing cookies/storage from prior sessions,
+// HTTP-set cookies that landed before our SDK ran, and -- crucially
+// -- writes from any extension realm (including isolated-world
+// content scripts) that happened before SDK init. Many extensions
+// inject at document_start specifically to set state before page
+// scripts run, so a meaningful share of PRE_EXISTING rows can be
+// extension-origin even though we cannot prove it. Treat this value
+// as low-signal for "is this a real page tracker" decisions.
+//
+// CookieSourceHTTP: cookie change observed via the CookieStore API
+// change event. Set by the server (Set-Cookie response header).
+//
+// Rows persisted before CookieSourceExtension was introduced cannot
+// be backfilled -- the stack at write time is gone -- so historic
+// SCRIPT rows retain the (mild) ambiguity above for that period.
 const (
 	CookieSourceScript      CookieSource = "SCRIPT"
 	CookieSourcePreExisting CookieSource = "PRE_EXISTING"
 	CookieSourceHTTP        CookieSource = "HTTP"
+	CookieSourceExtension   CookieSource = "EXTENSION"
 )
 
 var (
@@ -38,6 +71,7 @@ func CookieSources() []CookieSource {
 		CookieSourceScript,
 		CookieSourcePreExisting,
 		CookieSourceHTTP,
+		CookieSourceExtension,
 	}
 }
 
@@ -46,7 +80,8 @@ func (v CookieSource) IsValid() bool {
 	case
 		CookieSourceScript,
 		CookieSourcePreExisting,
-		CookieSourceHTTP:
+		CookieSourceHTTP,
+		CookieSourceExtension:
 		return true
 	}
 
diff --git a/pkg/coredata/migrations/20260522T155558Z.sql b/pkg/coredata/migrations/20260522T155558Z.sql
new file mode 100644
index 000000000..78fae1ffc
--- /dev/null
+++ b/pkg/coredata/migrations/20260522T155558Z.sql
@@ -0,0 +1,15 @@
+-- Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
+--
+-- Permission to use, copy, modify, and/or distribute this software for any
+-- purpose with or without fee is hereby granted, provided that the above
+-- copyright notice and this permission notice appear in all copies.
+--
+-- THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+-- REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+-- AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+-- INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+-- LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+-- OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+-- PERFORMANCE OF THIS SOFTWARE.
+
+ALTER TYPE cookie_source ADD VALUE IF NOT EXISTS 'EXTENSION';
diff --git a/pkg/server/api/console/v1/graphql/cookie_banner.graphql b/pkg/server/api/console/v1/graphql/cookie_banner.graphql
index fa172b436..9eab52b06 100644
--- a/pkg/server/api/console/v1/graphql/cookie_banner.graphql
+++ b/pkg/server/api/console/v1/graphql/cookie_banner.graphql
@@ -32,6 +32,10 @@ enum CookieSource
         @goEnum(
             value: "go.probo.inc/probo/pkg/coredata.CookieSourcePreExisting"
         )
+    EXTENSION
+        @goEnum(
+            value: "go.probo.inc/probo/pkg/coredata.CookieSourceExtension"
+        )
 }
 
 enum TrackerType
diff --git a/pkg/server/api/cookiebanner/v1/handler.go b/pkg/server/api/cookiebanner/v1/handler.go
index 9d84b8911..89cd9e9fe 100644
--- a/pkg/server/api/cookiebanner/v1/handler.go
+++ b/pkg/server/api/cookiebanner/v1/handler.go
@@ -439,6 +439,8 @@ func (h *Handler) handleReportDetectedTrackers(w http.ResponseWriter, r *http.Re
 			source = coredata.CookieSourcePreExisting
 		case "http":
 			source = coredata.CookieSourceHTTP
+		case "extension":
+			source = coredata.CookieSourceExtension
 		default:
 			source = coredata.CookieSourceScript
 		}
@@ -480,6 +482,8 @@ func (h *Handler) handleReportDetectedTrackers(w http.ResponseWriter, r *http.Re
 		switch strings.TrimSpace(s.Source) {
 		case "pre-existing":
 			source = coredata.CookieSourcePreExisting
+		case "extension":
+			source = coredata.CookieSourceExtension
 		default:
 			source = coredata.CookieSourceScript
 		}
diff --git a/pkg/server/api/mcp/v1/specification.yaml b/pkg/server/api/mcp/v1/specification.yaml
index f938a5b4f..a58c6b61d 100644
--- a/pkg/server/api/mcp/v1/specification.yaml
+++ b/pkg/server/api/mcp/v1/specification.yaml
@@ -9470,7 +9470,7 @@ components:
           description: Pattern description
         source:
           type: string
-          enum: [SCRIPT, PRE_EXISTING]
+          enum: [SCRIPT, PRE_EXISTING, EXTENSION]
           description: How the pattern was discovered
         excluded:
           type: boolean