diff --git a/PRIVACY.md b/PRIVACY.md index 3df511a..48f7fe0 100644 --- a/PRIVACY.md +++ b/PRIVACY.md @@ -79,7 +79,9 @@ When you interact with these AI providers through Panelize, you are subject to t ### Cookie Access -Panelize does not access cookie data directly and does not use the cookies API. Provider websites may use your existing browser session cookies when they load inside panels, just like in a normal tab. +Provider websites use your existing browser session cookies when they load inside panels. Because Chromium omits `SameSite=Lax` and `SameSite=Strict` cookies on WebSocket handshakes initiated from embedded extension panels, Panelize uses the `cookies` API locally in your browser to forward existing `grok.com` session cookies to Grok's own `wss://grok.com/*` WebSocket requests made from Panelize's Multi-Panel tabs. The forwarding applies only while the Multi-Panel page is the top-level page of a tab; it does not apply when another website embeds a Panelize page. The forwarding rule lasts only for the browser session; cookie values are not saved in Panelize's settings or sent to Panelize-operated servers or unrelated third parties. + +Panelize also preserves your existing Grok cookie consent choices by setting the `OptanonAlertBoxClosed` and `OptanonConsent` cookies to `SameSite=None; Secure`, allowing Grok to read and save those choices inside its embedded panel. This happens only when you open Grok in the Multi-Panel page. This preserves the choice you make in Grok's cookie controls and does not automatically accept cookies on your behalf. ### Update Checks @@ -155,6 +157,7 @@ Panelize requests the following Chrome extension permissions: | Permission | Purpose | |------------|---------| | `storage` | Store your prompts and settings locally | +| `cookies` | Attach existing `grok.com` session cookies to Grok's embedded WebSocket connections and preserve your Grok cookie consent choices inside Panelize | | `contextMenus` | Add "Send to Panelize" option when right-clicking | | `declarativeNetRequest` | Allow AI provider websites to load in embedded panels (bypass X-Frame-Options) | | `declarativeNetRequestWithHostAccess` | Apply header modifications for specific AI provider domains | diff --git a/background/service-worker.js b/background/service-worker.js index d3dd37d..821b599 100644 --- a/background/service-worker.js +++ b/background/service-worker.js @@ -5,6 +5,17 @@ import { removeDeprecatedClaudeModelMode } from '../modules/settings.js'; import { reconcileOptionalProviderAccess } from '../modules/optional-provider-access.js'; +import { + GROK_WEBSOCKET_COOKIE_SYNC_ACTION, + createGrokWebSocketCookieReconcileScheduler, + handleGrokTabNavigationComplete, + handleGrokTabNavigationStart, + handleGrokTabRemoved, + handleMultiPanelSyncRequest, + isGrokWebSocketCookieChange, + isMultiPanelTopFrameSender, + reconcileGrokWebSocketCookies +} from '../modules/grok-websocket-cookies.js'; // Install event - setup context menus const DEFAULT_SHORTCUT_SETTING = { keyboardShortcutEnabled: true }; @@ -126,12 +137,20 @@ function scheduleOptionalProviderReconcile() { return optionalProviderReconcileQueue; } +const scheduleGrokWebSocketCookieReconcile = createGrokWebSocketCookieReconcileScheduler( + reconcileGrokWebSocketCookies, + (error) => { + console.error('[Background] Failed to synchronize Grok WebSocket cookies:', error); + } +); + chrome.runtime.onInstalled.addListener(async (details) => { await migrateProviderSettingsForUpdate(details); if (details?.reason === 'update') { await removeDeprecatedClaudeModelMode(); } await scheduleOptionalProviderReconcile(); + await scheduleGrokWebSocketCookieReconcile(); await createContextMenus(); await loadShortcutSetting(); await loadOpenModeSetting(); @@ -139,6 +158,7 @@ chrome.runtime.onInstalled.addListener(async (details) => { chrome.runtime.onStartup.addListener(async () => { await scheduleOptionalProviderReconcile(); + await scheduleGrokWebSocketCookieReconcile(); await loadShortcutSetting(); await loadOpenModeSetting(); }); @@ -172,9 +192,45 @@ chrome.storage.onChanged.addListener((changes, namespace) => { if (namespace === 'sync' && changes.enabledProviders) { scheduleOptionalProviderReconcile(); + scheduleGrokWebSocketCookieReconcile(); } }); +if (chrome.cookies?.onChanged) { + chrome.cookies.onChanged.addListener((changeInfo) => { + if (isGrokWebSocketCookieChange(changeInfo)) { + scheduleGrokWebSocketCookieReconcile(); + } + }); +} + +if (chrome.tabs?.onUpdated) { + chrome.tabs.onUpdated.addListener((tabId, changeInfo) => { + if (changeInfo?.status === 'loading') { + handleGrokTabNavigationStart(tabId).catch((error) => { + console.error('[Background] Failed to drop navigating tab from Grok cookie rule:', error); + scheduleGrokWebSocketCookieReconcile(); + }); + } else if (changeInfo?.status === 'complete' && handleGrokTabNavigationComplete(tabId)) { + scheduleGrokWebSocketCookieReconcile(); + } + }); +} + +if (chrome.tabs?.onRemoved) { + chrome.tabs.onRemoved.addListener((tabId) => { + handleGrokTabRemoved(tabId) + .then((inRule) => { + if (inRule) { + scheduleGrokWebSocketCookieReconcile(); + } + }) + .catch((error) => { + console.error('[Background] Failed to handle closed Grok tab:', error); + }); + }); +} + if (chrome.permissions?.onAdded) { chrome.permissions.onAdded.addListener(() => { scheduleOptionalProviderReconcile(); @@ -295,7 +351,19 @@ chrome.runtime.onMessage.addListener((message, sender, sendResponse) => { handleFetchLatestCommit().then(sendResponse); return true; // Keep channel open for async response } - return true; + + if (message.action === GROK_WEBSOCKET_COOKIE_SYNC_ACTION) { + if (!isMultiPanelTopFrameSender(sender)) { + sendResponse({ success: false, error: 'Grok cookie sync is only available to the Multi-Panel page' }); + return false; + } + + handleMultiPanelSyncRequest(sender.tab.id); + scheduleGrokWebSocketCookieReconcile({ upgradeConsentCookies: true }).then(sendResponse); + return true; + } + + return false; }); // Handle version check by fetching latest commit from GitHub API diff --git a/content-scripts/grok-cookie-consent.js b/content-scripts/grok-cookie-consent.js new file mode 100644 index 0000000..51ac0dc --- /dev/null +++ b/content-scripts/grok-cookie-consent.js @@ -0,0 +1,60 @@ +// Ensure OneTrust cookie consent choices (OptanonAlertBoxClosed, OptanonConsent) +// written via document.cookie inside an embedded Grok iframe use SameSite=None; Secure +// so Chromium's cross-site script cookie check accepts and persists them. +(() => { + if (window.top === window.self) { + return; + } + + // Only rewrite inside extension pages such as the Multi-Panel page. Grok + // frames embedded by ordinary websites (or by a website framing the + // web-accessible Multi-Panel page) keep the browser's default behavior. + const ancestorOrigins = location.ancestorOrigins; + const topOrigin = ancestorOrigins?.length ? ancestorOrigins[ancestorOrigins.length - 1] : ''; + if (!topOrigin.startsWith('chrome-extension://')) { + return; + } + + const OPTANON_COOKIE_PATTERN = /^\s*Optanon(?:AlertBoxClosed|Consent)=/; + + function rewriteGrokConsentCookieString(cookieString) { + if (typeof cookieString !== 'string' || !OPTANON_COOKIE_PATTERN.test(cookieString)) { + return cookieString; + } + + if (typeof location !== 'undefined' && location.protocol !== 'https:') { + return cookieString; + } + + let rewritten = cookieString.replace(/;\s*SameSite=(?:Lax|Strict)/gi, '; SameSite=None'); + if (!/;\s*SameSite=None\b/i.test(rewritten)) { + rewritten += '; SameSite=None'; + } + if (!/;\s*Secure\b/i.test(rewritten)) { + rewritten += '; Secure'; + } + return rewritten; + } + + // Patch the prototype that actually owns the accessor (Document.prototype in + // Chromium) so no own property further down the chain shadows the override. + let owner = Object.getPrototypeOf(document); + while (owner && !Object.prototype.hasOwnProperty.call(owner, 'cookie')) { + owner = Object.getPrototypeOf(owner); + } + const descriptor = owner && Object.getOwnPropertyDescriptor(owner, 'cookie'); + if (!descriptor || typeof descriptor.get !== 'function' || typeof descriptor.set !== 'function') { + return; + } + + Object.defineProperty(owner, 'cookie', { + configurable: true, + enumerable: descriptor.enumerable ?? true, + get() { + return descriptor.get.call(this); + }, + set(value) { + return descriptor.set.call(this, rewriteGrokConsentCookieString(value)); + } + }); +})(); diff --git a/manifest.json b/manifest.json index 92cfcc8..5f166f2 100644 --- a/manifest.json +++ b/manifest.json @@ -6,6 +6,7 @@ "default_locale": "en", "permissions": [ "storage", + "cookies", "contextMenus", "activeTab", "scripting", @@ -18,6 +19,7 @@ "*://claude.ai/*", "*://gemini.google.com/*", "*://grok.com/*", + "wss://grok.com/*", "*://chat.deepseek.com/*", "*://www.kimi.com/*", "*://kimi.com/*", @@ -114,6 +116,17 @@ "run_at": "document_start", "all_frames": true }, + { + "matches": [ + "https://grok.com/*" + ], + "js": [ + "content-scripts/grok-cookie-consent.js" + ], + "run_at": "document_start", + "all_frames": true, + "world": "MAIN" + }, { "matches": [ "https://grok.com/*" diff --git a/modules/grok-websocket-cookies.js b/modules/grok-websocket-cookies.js new file mode 100644 index 0000000..8b29817 --- /dev/null +++ b/modules/grok-websocket-cookies.js @@ -0,0 +1,478 @@ +import { DEFAULT_PROVIDER_IDS } from './provider-defaults.js'; + +export const GROK_PROVIDER_ID = 'grok'; +export const GROK_COOKIE_URL = 'https://grok.com/'; +export const GROK_DOMAIN = 'grok.com'; +export const GROK_WEBSOCKET_COOKIE_RULE_ID = 2001; +export const GROK_WEBSOCKET_COOKIE_SYNC_ACTION = 'syncGrokWebSocketCookies'; +export const MULTI_PANEL_PAGE_PATH = 'multi-panel/multi-panel.html'; +export const GROK_CONSENT_COOKIE_NAMES = Object.freeze([ + 'OptanonAlertBoxClosed', + 'OptanonConsent' +]); + +// Tabs whose top-level document started navigating. They stay out of the +// forwarding rule until the navigation completes or a Multi-Panel page in that +// tab asks for a sync, so a Multi-Panel tab that navigates elsewhere never +// keeps receiving forwarded cookies while its old document is still listed. +const navigatingTabIds = new Set(); +// Navigating tabs that were left out of the rule only because of the mark +// above; they need a reconcile once their navigation settles. +const suppressedTabIds = new Set(); +// Tab IDs of the installed rule, or null when unknown (e.g. after a service +// worker restart) and the session rules must be read again. +let appliedRuleTabIds = null; + +function isValidTabId(tabId) { + return Number.isInteger(tabId) && tabId >= 0; +} + +function normalizeTabIds(tabIds) { + return [...new Set( + (Array.isArray(tabIds) ? tabIds : []).filter(isValidTabId) + )].sort((a, b) => a - b); +} + +export function resetGrokWebSocketCookieState() { + navigatingTabIds.clear(); + suppressedTabIds.clear(); + appliedRuleTabIds = null; +} + +export function isGrokCookieDomain(domain) { + if (typeof domain !== 'string' || domain.length === 0) { + return false; + } + + const normalized = domain.startsWith('.') + ? domain.slice(1).toLowerCase() + : domain.toLowerCase(); + return normalized === GROK_DOMAIN || normalized.endsWith(`.${GROK_DOMAIN}`); +} + +function isForwardedCookie(cookie) { + return Boolean( + cookie && + typeof cookie.name === 'string' && + cookie.name.length > 0 && + typeof cookie.value === 'string' && + cookie.sameSite !== 'no_restriction' && + !cookie.partitionKey + ); +} + +/** + * Formats Grok cookies that Chromium omits from cross-site WebSocket handshakes + * inside extension iframes. SameSite=None cookies (including partitioned CHIPS + * cookies such as cf_clearance) are already attached by Chromium, so only + * SameSite=Lax, SameSite=Strict, and unspecified cookies are appended. + */ +export function formatGrokWebSocketCookieHeader(cookies) { + if (!Array.isArray(cookies)) { + return ''; + } + + return cookies + .filter(isForwardedCookie) + .map((cookie) => `${cookie.name}=${cookie.value}`) + .join('; '); +} + +/** + * Whether a chrome.cookies.onChanged event can change the forwarded header. + * Overwrites report the old cookie as removed and the new one as added, so + * a cookie moving out of the header (e.g. to SameSite=None) is still caught + * by its removal event. + */ +export function isGrokWebSocketCookieChange(changeInfo) { + const cookie = changeInfo?.cookie; + return isGrokCookieDomain(cookie?.domain) && isForwardedCookie(cookie); +} + +/** + * Top-level domains the rule is restricted to: this extension's own pages. + * RuleCondition.topDomains exists since Chrome 145, and older versions reject + * unknown condition keys, so it is only used when the browser advertises it. + * There the rule cannot apply to a Multi-Panel tab that already navigated to + * another site, or to a third-party page framing multi-panel.html, even + * before the tab list is updated. + */ +export function getExtensionTopDomains() { + if ( + typeof chrome === 'undefined' || + chrome.declarativeNetRequest?.RuleConditionKeys?.TOP_DOMAINS !== 'topDomains' || + typeof chrome.runtime?.id !== 'string' || + chrome.runtime.id.length === 0 + ) { + return null; + } + return [chrome.runtime.id]; +} + +export function createGrokWebSocketCookieRule(cookieHeaderValue, tabIds, topDomains = getExtensionTopDomains()) { + return { + id: GROK_WEBSOCKET_COOKIE_RULE_ID, + priority: 1, + action: { + type: 'modifyHeaders', + requestHeaders: [ + { + header: 'Cookie', + operation: 'append', + value: cookieHeaderValue + } + ] + }, + condition: { + requestDomains: [GROK_DOMAIN], + initiatorDomains: [GROK_DOMAIN], + ...(Array.isArray(topDomains) && topDomains.length > 0 ? { topDomains: [...topDomains] } : {}), + resourceTypes: ['websocket'], + tabIds: normalizeTabIds(tabIds) + } + }; +} + +export function isMultiPanelPageUrl(url) { + if (typeof url !== 'string' || typeof chrome === 'undefined' || !chrome.runtime?.getURL) { + return false; + } + + try { + const actual = new URL(url); + const expected = new URL(chrome.runtime.getURL(MULTI_PANEL_PAGE_PATH)); + // Compare scheme and host explicitly: URL.origin is "null" for + // chrome-extension: URLs outside the browser. + return ( + actual.protocol === expected.protocol && + actual.host === expected.host && + actual.pathname === expected.pathname + ); + } catch { + return false; + } +} + +/** + * Only a Multi-Panel page that is the top-level document of its tab may ask + * for cookie forwarding. multi-panel.html is web-accessible, so third-party + * pages can frame it; those frames have a non-zero frameId. + */ +export function isMultiPanelTopFrameSender(sender) { + return Boolean( + sender && + typeof chrome !== 'undefined' && + sender.id === chrome.runtime?.id && + sender.frameId === 0 && + isValidTabId(sender.tab?.id) && + isMultiPanelPageUrl(sender.url) + ); +} + +async function getMultiPanelTopFrameTabIds() { + if (typeof chrome === 'undefined' || !chrome.runtime?.getContexts) { + // Without getContexts there is no trustworthy way to tell which tabs show + // the Multi-Panel page, so forward nothing. + return []; + } + + const contexts = await chrome.runtime.getContexts({ contextTypes: ['TAB'] }); + return normalizeTabIds((Array.isArray(contexts) ? contexts : []) + .filter((context) => context?.frameId === 0 && isMultiPanelPageUrl(context.documentUrl)) + .map((context) => context.tabId)); +} + +export async function getMultiPanelTabIds() { + const tabIds = await getMultiPanelTopFrameTabIds(); + return tabIds.filter((tabId) => { + if (navigatingTabIds.has(tabId)) { + suppressedTabIds.add(tabId); + return false; + } + return true; + }); +} + +async function getAppliedRuleTabIds() { + if (appliedRuleTabIds) { + return appliedRuleTabIds; + } + + if (typeof chrome === 'undefined' || !chrome.declarativeNetRequest?.getSessionRules) { + return []; + } + + const sessionRules = await chrome.declarativeNetRequest.getSessionRules(); + const rule = sessionRules.find(({ id }) => id === GROK_WEBSOCKET_COOKIE_RULE_ID); + appliedRuleTabIds = normalizeTabIds(rule?.condition?.tabIds); + return appliedRuleTabIds; +} + +/** + * Handles a top-level navigation start in any tab. A tab covered by the + * forwarding rule is removed from it right away, without a full reconcile, so + * the next document in that tab (e.g. grok.com top-level) does not receive + * duplicate cookies. Returns true when the rule changed. + */ +export async function handleGrokTabNavigationStart(tabId) { + if (!isValidTabId(tabId)) { + return false; + } + + // Mark synchronously so a reconcile racing with this handler already skips + // the tab. + navigatingTabIds.add(tabId); + const ruleTabIds = await getAppliedRuleTabIds(); + if (!ruleTabIds.includes(tabId)) { + return false; + } + + const sessionRules = await chrome.declarativeNetRequest.getSessionRules(); + const rule = sessionRules.find(({ id }) => id === GROK_WEBSOCKET_COOKIE_RULE_ID); + const remainingTabIds = normalizeTabIds(rule?.condition?.tabIds).filter((id) => id !== tabId); + await chrome.declarativeNetRequest.updateSessionRules({ + removeRuleIds: [GROK_WEBSOCKET_COOKIE_RULE_ID], + addRules: rule && remainingTabIds.length > 0 + ? [{ ...rule, condition: { ...rule.condition, tabIds: remainingTabIds } }] + : [] + }); + appliedRuleTabIds = remainingTabIds; + suppressedTabIds.add(tabId); + return true; +} + +/** + * Handles a finished navigation. Returns true when the tab was left out of the + * rule while navigating and must be re-evaluated against its new document. + */ +export function handleGrokTabNavigationComplete(tabId) { + navigatingTabIds.delete(tabId); + return suppressedTabIds.delete(tabId); +} + +/** + * Handles a sync request from a trusted Multi-Panel top frame: its document is + * live, so the tab no longer counts as navigating. + */ +export function handleMultiPanelSyncRequest(tabId) { + navigatingTabIds.delete(tabId); + suppressedTabIds.delete(tabId); +} + +/** + * Handles a closed tab. Returns true when the forwarding rule still lists it. + */ +export async function handleGrokTabRemoved(tabId) { + navigatingTabIds.delete(tabId); + suppressedTabIds.delete(tabId); + const ruleTabIds = await getAppliedRuleTabIds(); + return ruleTabIds.includes(tabId); +} + +function isSameList(left = [], right = []) { + return left.length === right.length && left.every((value, index) => value === right[index]); +} + +function isSameRule(leftRule, rightRule) { + if (!leftRule || !rightRule) { + return false; + } + + const leftHeader = leftRule.action?.requestHeaders?.[0]; + const rightHeader = rightRule.action?.requestHeaders?.[0]; + + return ( + leftRule.id === rightRule.id && + leftRule.priority === rightRule.priority && + leftRule.action?.type === rightRule.action?.type && + leftHeader?.header === rightHeader?.header && + leftHeader?.operation === rightHeader?.operation && + leftHeader?.value === rightHeader?.value && + isSameList(leftRule.condition?.tabIds, rightRule.condition?.tabIds) && + isSameList(leftRule.condition?.topDomains, rightRule.condition?.topDomains) + ); +} + +export async function syncGrokWebSocketCookies({ + enabledProviderIds = DEFAULT_PROVIDER_IDS, + tabIds = [], + cookies = [] +} = {}) { + if ( + typeof chrome === 'undefined' || + !chrome.declarativeNetRequest?.getSessionRules || + !chrome.declarativeNetRequest?.updateSessionRules + ) { + return; + } + + const grokEnabled = Array.isArray(enabledProviderIds) && enabledProviderIds.includes(GROK_PROVIDER_ID); + const normalizedTabIds = normalizeTabIds(tabIds); + const cookieHeaderValue = grokEnabled && normalizedTabIds.length > 0 + ? formatGrokWebSocketCookieHeader(cookies) + : ''; + const desiredRule = cookieHeaderValue + ? createGrokWebSocketCookieRule(cookieHeaderValue, normalizedTabIds) + : null; + + const sessionRules = await chrome.declarativeNetRequest.getSessionRules(); + const existingRule = sessionRules.find(({ id }) => id === GROK_WEBSOCKET_COOKIE_RULE_ID) || null; + + if (!desiredRule) { + if (existingRule) { + await chrome.declarativeNetRequest.updateSessionRules({ + removeRuleIds: [GROK_WEBSOCKET_COOKIE_RULE_ID], + addRules: [] + }); + } + appliedRuleTabIds = []; + return; + } + + if (!isSameRule(existingRule, desiredRule)) { + // Always remove the ID first so the update stays valid even if another + // context installed the rule after getSessionRules() returned. + await chrome.declarativeNetRequest.updateSessionRules({ + removeRuleIds: [GROK_WEBSOCKET_COOKIE_RULE_ID], + addRules: [desiredRule] + }); + } + appliedRuleTabIds = desiredRule.condition.tabIds; +} + +/** + * Rewrites existing OneTrust consent cookies to SameSite=None; Secure so Grok + * can read a choice the user already made top-level from inside its embedded + * panel; Chromium hides SameSite=Lax cookies from document.cookie there. + * Value, path, domain and expiry are kept. New choices made inside the panel + * are handled by content-scripts/grok-cookie-consent.js instead. + */ +export async function upgradeGrokConsentCookies(cookies) { + if (!Array.isArray(cookies) || typeof chrome === 'undefined' || !chrome.cookies?.set) { + return cookies; + } + + const updatedCookies = []; + for (const cookie of cookies) { + if ( + cookie && + GROK_CONSENT_COOKIE_NAMES.includes(cookie.name) && + !cookie.partitionKey && + (cookie.sameSite !== 'no_restriction' || !cookie.secure) + ) { + const details = { + url: GROK_COOKIE_URL, + name: cookie.name, + value: cookie.value, + path: cookie.path || '/', + secure: true, + httpOnly: Boolean(cookie.httpOnly), + sameSite: 'no_restriction' + }; + if (!cookie.hostOnly && cookie.domain) { + details.domain = cookie.domain; + } + if (typeof cookie.expirationDate === 'number') { + details.expirationDate = cookie.expirationDate; + } + if (typeof cookie.storeId === 'string') { + details.storeId = cookie.storeId; + } + + try { + await chrome.cookies.set(details); + updatedCookies.push({ + ...cookie, + secure: true, + sameSite: 'no_restriction' + }); + continue; + } catch (error) { + console.warn(`[Grok] Failed to update consent cookie ${cookie.name}:`, error); + } + } + + updatedCookies.push(cookie); + } + + return updatedCookies; +} + +/** + * Recomputes the forwarding rule from the live Multi-Panel tabs and cookies. + * @param {object} [options] + * @param {boolean} [options.upgradeConsentCookies] - Also rewrite existing + * consent cookies; only set when a Multi-Panel page is about to load Grok. + */ +export async function reconcileGrokWebSocketCookies({ upgradeConsentCookies = false } = {}) { + if ( + typeof chrome === 'undefined' || + !chrome.declarativeNetRequest?.getSessionRules || + !chrome.declarativeNetRequest?.updateSessionRules + ) { + return; + } + + const settings = await chrome.storage.sync.get({ + enabledProviders: DEFAULT_PROVIDER_IDS + }); + const enabledProviderIds = Array.isArray(settings.enabledProviders) + ? settings.enabledProviders + : DEFAULT_PROVIDER_IDS; + const grokEnabled = enabledProviderIds.includes(GROK_PROVIDER_ID); + const tabIds = grokEnabled ? await getMultiPanelTabIds() : []; + const active = grokEnabled && tabIds.length > 0; + let cookies = active && chrome.cookies?.getAll + ? await chrome.cookies.getAll({ url: GROK_COOKIE_URL }) + : []; + if (active && upgradeConsentCookies) { + cookies = await upgradeGrokConsentCookies(cookies); + } + + await syncGrokWebSocketCookies({ + enabledProviderIds, + tabIds, + cookies + }); +} + +/** + * Serializes reconciles and coalesces bursts: requests that arrive while a run + * is queued share it, so a flurry of cookie or tab events costs at most one + * run in progress plus one queued run. The returned promise never rejects. + * @param {(options: {upgradeConsentCookies?: boolean}) => Promise} reconcile + * @param {(error: unknown) => void} [onError] + * @returns {(options?: {upgradeConsentCookies?: boolean}) => + * Promise<{success: true} | {success: false, error: string}>} + */ +export function createGrokWebSocketCookieReconcileScheduler(reconcile, onError = () => {}) { + let tail = Promise.resolve(); + let pendingRun = null; + let pendingOptions = {}; + + return function scheduleReconcile(options = {}) { + if (options.upgradeConsentCookies) { + pendingOptions = { ...pendingOptions, upgradeConsentCookies: true }; + } + if (pendingRun) { + return pendingRun; + } + + pendingRun = tail + .then(() => { + const runOptions = pendingOptions; + pendingRun = null; + pendingOptions = {}; + return reconcile(runOptions); + }) + .then( + () => ({ success: true }), + (error) => { + onError(error); + return { success: false, error: error?.message || String(error) }; + } + ); + tail = pendingRun; + return pendingRun; + }; +} diff --git a/multi-panel/multi-panel.js b/multi-panel/multi-panel.js index 6ec59fe..64c0267 100644 --- a/multi-panel/multi-panel.js +++ b/multi-panel/multi-panel.js @@ -7,6 +7,7 @@ import { PROVIDERS, getProviderById, getEnabledProviders, getProviderIcon } from '../modules/providers.js'; import { DEFAULT_PROVIDER_IDS } from '../modules/provider-defaults.js'; +import { GROK_WEBSOCKET_COOKIE_SYNC_ACTION } from '../modules/grok-websocket-cookies.js'; import { DEFAULT_GOOGLE_PROVIDER_MODE, GOOGLE_PROVIDER_MODE_AI, @@ -244,6 +245,19 @@ function isChatgptProvider(providerId) { return providerId === 'chatgpt'; } +// The service worker owns the forwarding rule and serializes updates to it; +// a failed sync must not block the panel from loading. +async function requestGrokWebSocketCookieSync() { + try { + const response = await chrome.runtime.sendMessage({ action: GROK_WEBSOCKET_COOKIE_SYNC_ACTION }); + if (!response?.success) { + console.warn('[Multi-Panel] Grok WebSocket cookie sync failed:', response?.error); + } + } catch (error) { + console.warn('[Multi-Panel] Grok WebSocket cookie sync failed:', error); + } +} + function isTemporaryChatSupportedProvider(providerId) { return TEMP_CHAT_SUPPORTED_PROVIDERS.has(providerId); } @@ -1164,11 +1178,25 @@ async function initializePanels() { // ===== Panel Management ===== async function addPanel(providerId) { + // Sync before the capacity check so no await separates the check from the + // panel being added. + if (providerId === 'grok') { + await requestGrokWebSocketCookieSync(); + } + if (panels.length >= MAX_PANELS) { showToast(`Maximum number of panels reached (${MAX_PANELS})`); return; } + const provider = getProviderById(providerId); + if (!provider) { + console.error('Provider not found:', providerId); + return; + } + + // Everything from the capacity check to panels.push() below runs without + // awaiting, so concurrent calls cannot both pass the check. // Auto layout adjustment: upgrade from 1xN to 1x(N+1) when adding panel exceeds capacity const newPanelCount = panels.length + 1; const adjustedLayout = getAutoAdjustedLayout(currentLayout, newPanelCount); @@ -1183,15 +1211,7 @@ async function addPanel(providerId) { document.querySelectorAll('.layout-option').forEach(btn => { btn.classList.toggle('active', btn.dataset.layout === adjustedLayout); }); - - // Save configuration - await saveProviderConfiguration(); - } - - const provider = getProviderById(providerId); - if (!provider) { - console.error('Provider not found:', providerId); - return; + // The adjusted layout is saved together with the new panel below. } const panelId = `panel-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`; @@ -1322,6 +1342,10 @@ async function switchPanelProvider(panelId, newProviderId) { syncGoogleModeControls(); } + if (newProviderId === 'grok') { + await requestGrokWebSocketCookieSync(); + } + // Update panel header const headerIcon = panelEl.querySelector('.panel-header-left img'); const headerName = panelEl.querySelector('.panel-header-left span'); diff --git a/tests/e2e/grok-websocket.e2e.test.js b/tests/e2e/grok-websocket.e2e.test.js new file mode 100644 index 0000000..d9eab8b --- /dev/null +++ b/tests/e2e/grok-websocket.e2e.test.js @@ -0,0 +1,402 @@ +import { test, expect } from '@playwright/test'; +import { createHash } from 'node:crypto'; +import http from 'node:http'; +import { + createPatchedExtensionCopy, + launchExtension, + openMultiPanel, +} from './extension-harness.js'; +import { PROVIDERS } from '../../modules/providers.js'; + +const WEBSOCKET_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'; +const GROK_WEBSOCKET_COOKIE_RULE_ID = 2001; + +function buildWebSocketTextFrame(text) { + const payload = Buffer.from(text, 'utf8'); + return Buffer.concat([Buffer.from([0x81, payload.length]), payload]); +} + +function renderGrokPage() { + return ` +

Connecting

+ + `; +} + +function renderThirdPartyPage(frameUrl) { + return ` + + `; +} + +function startGrokWebSocketFixtureServer() { + const wsHandshakes = []; + const sockets = new Set(); + const server = http.createServer((request, response) => { + const url = new URL(request.url, `http://${request.headers.host}`); + response.writeHead(200, { 'Content-Type': 'text/html' }); + if (url.hostname === 'attacker.test') { + response.end(renderThirdPartyPage(url.searchParams.get('frame'))); + return; + } + response.end(renderGrokPage()); + }); + + server.on('connection', (socket) => { + sockets.add(socket); + socket.on('close', () => sockets.delete(socket)); + }); + + server.on('upgrade', (request, socket) => { + sockets.add(socket); + socket.on('close', () => sockets.delete(socket)); + const tag = new URL(request.url, 'http://grok.com').searchParams.get('tag'); + const cookieHeader = request.headers.cookie || ''; + const ssoMatches = cookieHeader.match(/(?:^|;\s*)sso=([^;]+)/g) || []; + const ssoRwMatches = cookieHeader.match(/(?:^|;\s*)sso-rw=([^;]+)/g) || []; + const deviceMatches = cookieHeader.match(/(?:^|;\s*)grok_device_id=([^;]+)/g) || []; + const ssoValue = cookieHeader.match(/(?:^|;\s*)sso=([^;]+)/)?.[1] || null; + const authenticated = Boolean( + ssoValue && + ssoMatches.length === 1 && + ssoRwMatches.length === 1 && + deviceMatches.length === 1 + ); + + wsHandshakes.push({ + tag, + authenticated, + ssoValue, + ssoCount: ssoMatches.length, + ssoRwCount: ssoRwMatches.length, + }); + + if (!authenticated) { + socket.write('HTTP/1.1 401 Unauthorized\r\nConnection: close\r\n\r\n'); + socket.destroy(); + return; + } + + const acceptKey = createHash('sha1') + .update(`${request.headers['sec-websocket-key']}${WEBSOCKET_GUID}`) + .digest('base64'); + + socket.write([ + 'HTTP/1.1 101 Switching Protocols', + 'Upgrade: websocket', + 'Connection: Upgrade', + `Sec-WebSocket-Accept: ${acceptKey}`, + '\r\n', + ].join('\r\n')); + socket.write(buildWebSocketTextFrame(`OK:${ssoValue}`)); + }); + + return new Promise((resolve) => { + server.listen(0, '127.0.0.1', () => { + resolve({ + port: server.address().port, + wsHandshakes, + close: () => new Promise((closeResolve, closeReject) => { + for (const socket of sockets) { + socket.destroy(); + } + server.close((error) => (error ? closeReject(error) : closeResolve())); + }), + }); + }); + }); +} + +// Grok's real session and device cookies use SameSite=Lax and SameSite=Strict +// with HttpOnly, which Chromium strips on cross-site WebSocket handshakes. +function grokSessionCookies(ssoValue) { + return [ + { + name: 'grok_device_id', + value: 'fixture-device', + domain: '.grok.com', + path: '/', + httpOnly: true, + sameSite: 'Lax', + }, + { + name: 'sso', + value: ssoValue, + domain: '.grok.com', + path: '/', + httpOnly: true, + sameSite: 'Lax', + }, + { + name: 'sso-rw', + value: 'fixture-session-rw', + domain: '.grok.com', + path: '/', + httpOnly: true, + sameSite: 'Strict', + }, + ]; +} + +async function launchGrokFixture({ grantOptionalProviders = false } = {}) { + const server = await startGrokWebSocketFixtureServer(); + const extensionCopy = await createPatchedExtensionCopy((manifest) => { + if (grantOptionalProviders) { + // Optional providers need a user gesture to grant; grant them up front. + manifest.host_permissions.push(...manifest.optional_host_permissions); + } + // Map the wss:// permission to ws:// for the loopback fixture server. + manifest.host_permissions.push('ws://grok.com/*'); + // Run the Grok content scripts on the plain-http fixture as well. + for (const contentScript of manifest.content_scripts) { + if (contentScript.matches.includes('https://grok.com/*')) { + contentScript.matches.push('http://grok.com/*'); + } + } + }); + const extension = await launchExtension({ + extensionPath: extensionCopy.extensionPath, + args: [ + '--disable-features=HttpsUpgrades', + '--proxy-bypass-list=*', + '--host-resolver-rules=MAP grok.com 127.0.0.1, MAP attacker.test 127.0.0.1', + ], + }); + + const pageErrors = []; + extension.context.on('page', (page) => { + page.on('pageerror', (error) => pageErrors.push(error.message)); + }); + + await extension.context.addCookies(grokSessionCookies('fixture-session-1')); + // Other provider sites are irrelevant here; serve blank pages instead of the network. + await extension.context.route((url) => url.protocol === 'https:' && url.hostname !== 'grok.com', (route) => ( + route.fulfill({ contentType: 'text/html', body: 'provider' }) + )); + await extension.context.route('https://grok.com/**', (route) => { + const target = new URL(route.request().url()); + return route.fulfill({ + status: 302, + headers: { Location: `http://grok.com:${server.port}${target.pathname}${target.search}` }, + }); + }); + + return { + server, + extension, + pageErrors, + grokUrl: (tag, delay = 0) => `http://grok.com:${server.port}/?tag=${tag}&delay=${delay}`, + // Simulates Chrome < 145, which has no RuleCondition.topDomains. + async disableTopDomains() { + await extension.serviceWorker.evaluate(() => { + Object.defineProperty(chrome.declarativeNetRequest, 'RuleConditionKeys', { + configurable: true, + value: undefined, + }); + }); + }, + async getRule() { + return extension.serviceWorker.evaluate(async (ruleId) => { + const rules = await chrome.declarativeNetRequest.getSessionRules(); + return rules.find(({ id }) => id === ruleId) ?? null; + }, GROK_WEBSOCKET_COOKIE_RULE_ID); + }, + thirdPartyUrl: (frameUrl) => ( + `http://attacker.test:${server.port}/?frame=${encodeURIComponent(frameUrl)}` + ), + async getRuleTabIds() { + return extension.serviceWorker.evaluate(async (ruleId) => { + const rules = await chrome.declarativeNetRequest.getSessionRules(); + return rules.find(({ id }) => id === ruleId)?.condition?.tabIds ?? null; + }, GROK_WEBSOCKET_COOKIE_RULE_ID); + }, + async getTabId(page) { + const url = page.url(); + return extension.serviceWorker.evaluate(async (pageUrl) => { + const contexts = await chrome.runtime.getContexts({ contextTypes: ['TAB'] }); + const match = contexts.find((context) => context.frameId === 0 && context.documentUrl === pageUrl); + return match?.tabId ?? null; + }, url); + }, + async close() { + await extension.close(); + await server.close(); + await extensionCopy.cleanup(); + }, + }; +} + +test('connects Grok WebSockets in an embedded panel with SameSite session cookies', async () => { + test.setTimeout(60000); + const fixture = await launchGrokFixture(); + const { extension, server } = fixture; + + try { + // Top-level Grok works natively and must not receive duplicate cookies. + const topLevel = await extension.context.newPage(); + await topLevel.goto(fixture.grokUrl('top')); + await expect(topLevel.locator('#ws-state')).toHaveText('OK:fixture-session-1'); + + // Embedded Grok panel inside Panelize must also establish its WebSocket. + const page = await openMultiPanel(extension, { providers: ['grok'] }); + const panelState = page.frameLocator('iframe').locator('#ws-state'); + await expect(panelState).toHaveText('OK:fixture-session-1'); + + // Rotate the session cookie and verify the updated cookie is used on reload. + await extension.context.addCookies(grokSessionCookies('fixture-session-2')); + + await page.reload(); + await expect(page.frameLocator('iframe').locator('#ws-state')) + .toHaveText('OK:fixture-session-2'); + + // Re-check top-level tab while the Multi-Panel session rule is active. + await topLevel.reload(); + await expect(topLevel.locator('#ws-state')).toHaveText('OK:fixture-session-2'); + + expect(server.wsHandshakes.map(({ authenticated, ssoValue, ssoCount, ssoRwCount }) => ({ + authenticated, ssoValue, ssoCount, ssoRwCount, + }))).toEqual([ + { authenticated: true, ssoValue: 'fixture-session-1', ssoCount: 1, ssoRwCount: 1 }, + { authenticated: true, ssoValue: 'fixture-session-1', ssoCount: 1, ssoRwCount: 1 }, + { authenticated: true, ssoValue: 'fixture-session-2', ssoCount: 1, ssoRwCount: 1 }, + { authenticated: true, ssoValue: 'fixture-session-2', ssoCount: 1, ssoRwCount: 1 }, + ]); + expect(await fixture.getRuleTabIds()).toEqual([await fixture.getTabId(page)]); + expect(fixture.pageErrors).toEqual([]); + } finally { + await fixture.close(); + } +}); + +for (const legacy of [false, true]) { + const variant = legacy ? ' (without topDomains)' : ''; + + test(`never forwards Grok session cookies to Grok frames embedded by third-party pages${variant}`, async () => { + test.setTimeout(60000); + const fixture = await launchGrokFixture(); + const { extension, server } = fixture; + + try { + if (legacy) { + await fixture.disableTopDomains(); + } + const multiPanel = await openMultiPanel(extension, { providers: ['grok'] }); + await expect(multiPanel.frameLocator('iframe').locator('#ws-state')) + .toHaveText('OK:fixture-session-1'); + const multiPanelTabId = await fixture.getTabId(multiPanel); + expect((await fixture.getRule()).condition.topDomains) + .toEqual(legacy ? undefined : [extension.extensionId]); + + // A third-party page frames grok.com directly. + const thirdParty = await extension.context.newPage(); + await thirdParty.goto(fixture.thirdPartyUrl(fixture.grokUrl('third-party'))); + await expect(thirdParty.frameLocator('#embedded').locator('#ws-state')) + .toHaveText('Unable to start conversation'); + + // A third-party page frames the web-accessible Multi-Panel page, which in + // turn frames Grok. + const framedMultiPanel = await extension.context.newPage(); + await framedMultiPanel.goto(fixture.thirdPartyUrl( + extension.extensionUrl('multi-panel/multi-panel.html') + )); + await expect(framedMultiPanel.frameLocator('#embedded').locator('#panel-grid iframe')) + .toHaveCount(1); + await expect(framedMultiPanel.frameLocator('#embedded').frameLocator('#panel-grid iframe') + .locator('#ws-state')).toHaveText('Unable to start conversation'); + + expect(await fixture.getRuleTabIds()).toEqual([multiPanelTabId]); + expect(server.wsHandshakes.filter(({ tag }) => tag === 'third-party')) + .toEqual([{ tag: 'third-party', authenticated: false, ssoValue: null, ssoCount: 0, ssoRwCount: 0 }]); + // Panel frames load https://grok.com/ without a tag: first the real + // Multi-Panel tab, then the one framed by the third-party page. + expect(server.wsHandshakes.filter(({ tag }) => tag === 'default')).toEqual([ + { tag: 'default', authenticated: true, ssoValue: 'fixture-session-1', ssoCount: 1, ssoRwCount: 1 }, + { tag: 'default', authenticated: false, ssoValue: null, ssoCount: 0, ssoRwCount: 0 }, + ]); + } finally { + await fixture.close(); + } + }); + + test(`stops forwarding Grok cookies once a Multi-Panel tab navigates to grok.com${variant}`, async () => { + test.setTimeout(60000); + const fixture = await launchGrokFixture(); + const { extension, server } = fixture; + + try { + if (legacy) { + await fixture.disableTopDomains(); + } + const page = await openMultiPanel(extension, { providers: ['grok'] }); + await expect(page.frameLocator('iframe').locator('#ws-state')) + .toHaveText('OK:fixture-session-1'); + expect(await fixture.getRuleTabIds()).toHaveLength(1); + + // The same tab now shows Grok top-level, where Chromium already sends + // SameSite cookies natively; appending them again would duplicate them. + // With topDomains the rule never matches there, even if the WebSocket + // opens immediately. Without it the tab is dropped from the rule when + // the navigation starts, ahead of Grok loading its app bundle. + await page.goto(fixture.grokUrl('navigated', legacy ? 300 : 0)); + await expect(page.locator('#ws-state')).toHaveText('OK:fixture-session-1'); + + expect(server.wsHandshakes.filter(({ tag }) => tag === 'navigated')) + .toEqual([{ tag: 'navigated', authenticated: true, ssoValue: 'fixture-session-1', ssoCount: 1, ssoRwCount: 1 }]); + await expect.poll(() => fixture.getRuleTabIds()).toBeNull(); + } finally { + await fixture.close(); + } + }); +} + +test('adding Grok twice at once never exceeds the panel limit', async () => { + test.setTimeout(90000); + const fixture = await launchGrokFixture({ grantOptionalProviders: true }); + const { extension } = fixture; + const MAX_PANELS = 12; + const otherProviderIds = PROVIDERS.map(({ id }) => id).filter((id) => id !== 'grok').slice(0, MAX_PANELS - 1); + expect(otherProviderIds).toHaveLength(MAX_PANELS - 1); + + try { + // Multi-Panel opens min(enabled providers, layout capacity) panels, so + // start with exactly 11 enabled and enable Grok only for the add menu. + const page = await openMultiPanel(extension, { + providers: otherProviderIds, + layout: '1x12', + }); + await extension.serviceWorker.evaluate(async (enabledProviders) => { + await chrome.storage.sync.set({ enabledProviders, providerOrder: enabledProviders }); + }, [...otherProviderIds, 'grok']); + + await page.click('#add-panel-btn'); + const grokItem = page.locator('.add-panel-item[data-provider-id="grok"]'); + await expect(grokItem).toHaveCount(1); + // Both clicks start addPanel() before either Grok cookie sync resolves. + await grokItem.evaluate((item) => { + item.click(); + item.click(); + }); + + await expect(page.locator('#panel-grid iframe')).toHaveCount(MAX_PANELS); + await expect.poll(() => extension.serviceWorker.evaluate(async () => ( + (await chrome.storage.sync.get('multiPanelProviders')).multiPanelProviders.length + ))).toBe(MAX_PANELS); + await page.waitForTimeout(500); + await expect(page.locator('#panel-grid iframe')).toHaveCount(MAX_PANELS); + } finally { + await fixture.close(); + } +}); diff --git a/tests/grok-websocket-cookies.test.js b/tests/grok-websocket-cookies.test.js new file mode 100644 index 0000000..527c201 --- /dev/null +++ b/tests/grok-websocket-cookies.test.js @@ -0,0 +1,494 @@ +import { readFileSync } from 'node:fs'; +import { resolve } from 'node:path'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { + GROK_COOKIE_URL, + GROK_WEBSOCKET_COOKIE_RULE_ID, + createGrokWebSocketCookieReconcileScheduler, + createGrokWebSocketCookieRule, + formatGrokWebSocketCookieHeader, + getMultiPanelTabIds, + handleGrokTabNavigationComplete, + handleGrokTabNavigationStart, + handleGrokTabRemoved, + handleMultiPanelSyncRequest, + isGrokCookieDomain, + isGrokWebSocketCookieChange, + isMultiPanelTopFrameSender, + reconcileGrokWebSocketCookies, + resetGrokWebSocketCookieState, + syncGrokWebSocketCookies, + upgradeGrokConsentCookies, +} from '../modules/grok-websocket-cookies.js'; + +const manifest = JSON.parse( + readFileSync(resolve(process.cwd(), 'manifest.json'), 'utf8') +); +const grokCookieConsentScript = readFileSync( + resolve(process.cwd(), 'content-scripts/grok-cookie-consent.js'), + 'utf8' +); + +const MULTI_PANEL_URL = 'chrome-extension://test-ext-id/multi-panel/multi-panel.html'; + +function mockMultiPanelTabs(...tabIds) { + chrome.runtime.getContexts.mockResolvedValue(tabIds.map((tabId) => ({ + contextType: 'TAB', + documentUrl: MULTI_PANEL_URL, + frameId: 0, + tabId, + }))); +} + +describe('Grok WebSocket cookie synchronization', () => { + beforeEach(() => { + resetGrokWebSocketCookieState(); + chrome.runtime.id = 'test-ext-id'; + chrome.runtime.getURL = vi.fn((path) => `chrome-extension://test-ext-id/${path}`); + chrome.runtime.getContexts = vi.fn(() => Promise.resolve([])); + chrome.cookies = { + getAll: vi.fn(() => Promise.resolve([])), + set: vi.fn(() => Promise.resolve({})), + }; + let sessionRules = []; + chrome.declarativeNetRequest = { + getSessionRules: vi.fn(() => Promise.resolve(sessionRules)), + updateSessionRules: vi.fn(({ removeRuleIds = [], addRules = [] }) => { + sessionRules = [ + ...sessionRules.filter(({ id }) => !removeRuleIds.includes(id)), + ...addRules, + ]; + return Promise.resolve(); + }), + }; + chrome.storage.sync.get.mockImplementation((defaults) => Promise.resolve(defaults)); + }); + + it('declares cookies and wss://grok.com/* permissions in manifest.json', () => { + expect(manifest.permissions).toContain('cookies'); + expect(manifest.host_permissions).toEqual(expect.arrayContaining([ + '*://grok.com/*', + 'wss://grok.com/*', + ])); + expect(manifest.content_scripts).toContainEqual({ + matches: ['https://grok.com/*'], + js: ['content-scripts/grok-cookie-consent.js'], + run_at: 'document_start', + all_frames: true, + world: 'MAIN', + }); + }); + + it('matches Grok cookie domains and rejects unrelated domains', () => { + expect(isGrokCookieDomain('grok.com')).toBe(true); + expect(isGrokCookieDomain('.grok.com')).toBe(true); + expect(isGrokCookieDomain('assets.grok.com')).toBe(true); + expect(isGrokCookieDomain('chatgpt.com')).toBe(false); + expect(isGrokCookieDomain('notgrok.com')).toBe(false); + expect(isGrokCookieDomain('')).toBe(false); + expect(isGrokCookieDomain(null)).toBe(false); + }); + + it('formats SameSite=Lax, Strict, and unspecified cookies while omitting SameSite=None and partitioned cookies', () => { + const header = formatGrokWebSocketCookieHeader([ + { name: 'grok_device_id', value: 'device-1', sameSite: 'lax' }, + { name: 'sso', value: 'sso-token', sameSite: 'lax', httpOnly: true }, + { name: 'sso-rw', value: 'sso-rw-token', sameSite: 'strict', httpOnly: true }, + { name: '__cuid', value: 'cuid-1', sameSite: 'unspecified' }, + { name: '__cf_bm', value: 'cf-bm', sameSite: 'no_restriction' }, + { + name: 'cf_clearance', + value: 'cf-clearance', + sameSite: 'no_restriction', + partitionKey: { topLevelSite: 'chrome-extension://test-ext-id' }, + }, + ]); + + expect(header).toBe( + 'grok_device_id=device-1; sso=sso-token; sso-rw=sso-rw-token; __cuid=cuid-1' + ); + }); + + it('creates a session rule scoped to Grok WebSockets in Multi-Panel tabs', () => { + expect(createGrokWebSocketCookieRule('sso=token-1; sso-rw=token-2', [42, 7, 42])).toEqual({ + id: GROK_WEBSOCKET_COOKIE_RULE_ID, + priority: 1, + action: { + type: 'modifyHeaders', + requestHeaders: [ + { + header: 'Cookie', + operation: 'append', + value: 'sso=token-1; sso-rw=token-2', + }, + ], + }, + condition: { + requestDomains: ['grok.com'], + initiatorDomains: ['grok.com'], + resourceTypes: ['websocket'], + tabIds: [7, 42], + }, + }); + }); + + it('restricts the rule to this extension as top-level site when topDomains is supported', async () => { + chrome.declarativeNetRequest.RuleConditionKeys = { TOP_DOMAINS: 'topDomains' }; + + expect(createGrokWebSocketCookieRule('sso=a', [3]).condition.topDomains).toEqual(['test-ext-id']); + + // A rule installed before topDomains was available gets replaced. + chrome.declarativeNetRequest.getSessionRules.mockResolvedValue([ + createGrokWebSocketCookieRule('sso=a', [3], null), + ]); + await syncGrokWebSocketCookies({ + enabledProviderIds: ['grok'], + tabIds: [3], + cookies: [{ name: 'sso', value: 'a', sameSite: 'lax' }], + }); + expect(chrome.declarativeNetRequest.updateSessionRules).toHaveBeenCalledWith({ + removeRuleIds: [GROK_WEBSOCKET_COOKIE_RULE_ID], + addRules: [expect.objectContaining({ + condition: expect.objectContaining({ topDomains: ['test-ext-id'], tabIds: [3] }), + })], + }); + }); + + it('omits topDomains on browsers that reject unknown rule condition keys', () => { + expect(createGrokWebSocketCookieRule('sso=a', [3]).condition).not.toHaveProperty('topDomains'); + }); + + it('discovers only Multi-Panel pages that are the top-level document of their tab', async () => { + chrome.runtime.getContexts.mockResolvedValue([ + { contextType: 'TAB', documentUrl: `${MULTI_PANEL_URL}?mode=popup`, frameId: 0, tabId: 8 }, + { contextType: 'TAB', documentUrl: 'chrome-extension://test-ext-id/options/options.html', frameId: 0, tabId: 99 }, + // multi-panel.html is web-accessible, so third-party pages can frame it. + { contextType: 'TAB', documentUrl: MULTI_PANEL_URL, frameId: 5, tabId: 77 }, + { contextType: 'TAB', documentUrl: 'chrome-extension://other-ext/multi-panel/multi-panel.html', frameId: 0, tabId: 66 }, + ]); + + await expect(getMultiPanelTabIds()).resolves.toEqual([8]); + }); + + it('forwards nothing when runtime.getContexts is unavailable', async () => { + delete chrome.runtime.getContexts; + + await expect(getMultiPanelTabIds()).resolves.toEqual([]); + }); + + it('accepts sync requests only from the top-level Multi-Panel page of this extension', () => { + const trusted = { id: 'test-ext-id', frameId: 0, tab: { id: 3 }, url: MULTI_PANEL_URL }; + + expect(isMultiPanelTopFrameSender(trusted)).toBe(true); + expect(isMultiPanelTopFrameSender({ ...trusted, frameId: 4 })).toBe(false); + expect(isMultiPanelTopFrameSender({ ...trusted, url: 'https://grok.com/' })).toBe(false); + expect(isMultiPanelTopFrameSender({ ...trusted, url: 'chrome-extension://test-ext-id/options/options.html' })).toBe(false); + expect(isMultiPanelTopFrameSender({ ...trusted, id: 'other-ext' })).toBe(false); + expect(isMultiPanelTopFrameSender({ ...trusted, tab: undefined })).toBe(false); + expect(isMultiPanelTopFrameSender(undefined)).toBe(false); + }); + + it('installs and updates the Grok WebSocket session rule when Grok is enabled', async () => { + mockMultiPanelTabs(12); + chrome.cookies.getAll.mockResolvedValue([ + { name: 'sso', value: 'session-a', sameSite: 'lax', httpOnly: true }, + { name: 'sso-rw', value: 'session-rw-a', sameSite: 'strict', httpOnly: true }, + ]); + + await reconcileGrokWebSocketCookies(); + + expect(chrome.cookies.getAll).toHaveBeenCalledWith({ url: GROK_COOKIE_URL }); + expect(chrome.declarativeNetRequest.updateSessionRules).toHaveBeenCalledWith({ + removeRuleIds: [GROK_WEBSOCKET_COOKIE_RULE_ID], + addRules: [ + createGrokWebSocketCookieRule('sso=session-a; sso-rw=session-rw-a', [12]), + ], + }); + }); + + it('propagates failures instead of swallowing them', async () => { + mockMultiPanelTabs(12); + chrome.cookies.getAll.mockResolvedValue([{ name: 'sso', value: 'session-a', sameSite: 'lax' }]); + chrome.declarativeNetRequest.updateSessionRules.mockRejectedValue(new Error('quota')); + + await expect(reconcileGrokWebSocketCookies()).rejects.toThrow('quota'); + }); + + it('skips redundant session rule updates when the installed rule already matches', async () => { + const rule = createGrokWebSocketCookieRule('sso=session-a', [12]); + chrome.declarativeNetRequest.getSessionRules.mockResolvedValue([rule]); + + await syncGrokWebSocketCookies({ + enabledProviderIds: ['grok'], + tabIds: [12], + cookies: [{ name: 'sso', value: 'session-a', sameSite: 'lax' }], + }); + + expect(chrome.declarativeNetRequest.updateSessionRules).not.toHaveBeenCalled(); + }); + + it('removes the session rule when Grok is disabled or Multi-Panel tabs close', async () => { + const existingRule = createGrokWebSocketCookieRule('sso=session-a', [12]); + mockMultiPanelTabs(12); + chrome.cookies.getAll.mockResolvedValue([{ name: 'sso', value: 'session-a', sameSite: 'lax' }]); + chrome.declarativeNetRequest.getSessionRules.mockResolvedValue([existingRule]); + chrome.storage.sync.get.mockResolvedValue({ enabledProviders: ['chatgpt'] }); + + await reconcileGrokWebSocketCookies(); + + expect(chrome.runtime.getContexts).not.toHaveBeenCalled(); + expect(chrome.declarativeNetRequest.updateSessionRules).toHaveBeenCalledWith({ + removeRuleIds: [GROK_WEBSOCKET_COOKIE_RULE_ID], + addRules: [], + }); + + chrome.declarativeNetRequest.updateSessionRules.mockClear(); + chrome.storage.sync.get.mockResolvedValue({ enabledProviders: ['grok'] }); + mockMultiPanelTabs(); + + await reconcileGrokWebSocketCookies(); + + expect(chrome.declarativeNetRequest.updateSessionRules).toHaveBeenCalledWith({ + removeRuleIds: [GROK_WEBSOCKET_COOKIE_RULE_ID], + addRules: [], + }); + }); + + it('drops a Multi-Panel tab from the rule while its top-level document navigates', async () => { + mockMultiPanelTabs(12, 20); + chrome.cookies.getAll.mockResolvedValue([{ name: 'sso', value: 'session-a', sameSite: 'lax' }]); + await reconcileGrokWebSocketCookies(); + chrome.declarativeNetRequest.updateSessionRules.mockClear(); + + // Unrelated tab navigations do not require a reconcile. + await expect(handleGrokTabNavigationStart(30)).resolves.toBe(false); + expect(handleGrokTabNavigationComplete(30)).toBe(false); + + // Tab 12 starts navigating: it leaves the rule immediately. + await expect(handleGrokTabNavigationStart(12)).resolves.toBe(true); + expect(chrome.declarativeNetRequest.updateSessionRules).toHaveBeenLastCalledWith({ + removeRuleIds: [GROK_WEBSOCKET_COOKIE_RULE_ID], + addRules: [createGrokWebSocketCookieRule('sso=session-a', [20])], + }); + + // A reconcile while getContexts still lists the old document keeps it out. + chrome.declarativeNetRequest.updateSessionRules.mockClear(); + await reconcileGrokWebSocketCookies(); + expect(chrome.declarativeNetRequest.updateSessionRules).not.toHaveBeenCalled(); + + // Navigation finished on a page that is no longer the Multi-Panel page. + mockMultiPanelTabs(20); + expect(handleGrokTabNavigationComplete(12)).toBe(true); + chrome.declarativeNetRequest.updateSessionRules.mockClear(); + await reconcileGrokWebSocketCookies(); + expect(chrome.declarativeNetRequest.updateSessionRules).not.toHaveBeenCalled(); + }); + + it('re-admits a reloading Multi-Panel tab once its new page requests a sync', async () => { + mockMultiPanelTabs(12); + chrome.cookies.getAll.mockResolvedValue([{ name: 'sso', value: 'session-a', sameSite: 'lax' }]); + await reconcileGrokWebSocketCookies(); + + await expect(handleGrokTabNavigationStart(12)).resolves.toBe(true); + expect(chrome.declarativeNetRequest.updateSessionRules).toHaveBeenLastCalledWith({ + removeRuleIds: [GROK_WEBSOCKET_COOKIE_RULE_ID], + addRules: [], + }); + + handleMultiPanelSyncRequest(12); + await reconcileGrokWebSocketCookies(); + expect(chrome.declarativeNetRequest.updateSessionRules).toHaveBeenLastCalledWith({ + removeRuleIds: [GROK_WEBSOCKET_COOKIE_RULE_ID], + addRules: [createGrokWebSocketCookieRule('sso=session-a', [12])], + }); + expect(handleGrokTabNavigationComplete(12)).toBe(false); + }); + + it('reconciles on tab close only when the rule covers the closed tab', async () => { + chrome.declarativeNetRequest.getSessionRules.mockResolvedValue([ + createGrokWebSocketCookieRule('sso=session-a', [12]), + ]); + + await expect(handleGrokTabRemoved(99)).resolves.toBe(false); + await expect(handleGrokTabRemoved(12)).resolves.toBe(true); + // The installed tab set is cached after the first read. + expect(chrome.declarativeNetRequest.getSessionRules).toHaveBeenCalledTimes(1); + }); + + it('only treats cookie changes that alter the forwarded header as relevant', () => { + expect(isGrokWebSocketCookieChange({ cookie: { domain: '.grok.com', name: 'sso', value: 'a', sameSite: 'lax' } })).toBe(true); + expect(isGrokWebSocketCookieChange({ removed: true, cause: 'overwrite', cookie: { domain: '.grok.com', name: 'OptanonConsent', value: 'a', sameSite: 'lax' } })).toBe(true); + expect(isGrokWebSocketCookieChange({ cookie: { domain: '.grok.com', name: '__cf_bm', value: 'a', sameSite: 'no_restriction' } })).toBe(false); + expect(isGrokWebSocketCookieChange({ cookie: { domain: 'grok.com', name: 'cf', value: 'a', sameSite: 'lax', partitionKey: { topLevelSite: 'chrome-extension://x' } } })).toBe(false); + expect(isGrokWebSocketCookieChange({ cookie: { domain: '.x.com', name: 'sso', value: 'a', sameSite: 'lax' } })).toBe(false); + }); + + it('coalesces queued reconcile requests and reports failures without rejecting', async () => { + const calls = []; + let releaseFirst; + const reconcile = vi.fn((options) => { + calls.push(options); + if (calls.length === 1) { + return new Promise((resolve) => { releaseFirst = resolve; }); + } + if (calls.length === 3) { + return Promise.reject(new Error('boom')); + } + return Promise.resolve(); + }); + const onError = vi.fn(); + const schedule = createGrokWebSocketCookieReconcileScheduler(reconcile, onError); + + const first = schedule(); + await Promise.resolve(); + await Promise.resolve(); + // Three requests while the first run is in progress share one queued run. + const second = schedule(); + const third = schedule({ upgradeConsentCookies: true }); + const fourth = schedule(); + expect(third).toBe(second); + expect(fourth).toBe(second); + + releaseFirst(); + await expect(first).resolves.toEqual({ success: true }); + await expect(second).resolves.toEqual({ success: true }); + expect(reconcile).toHaveBeenCalledTimes(2); + expect(calls[1]).toEqual({ upgradeConsentCookies: true }); + + await expect(schedule()).resolves.toEqual({ success: false, error: 'boom' }); + expect(onError).toHaveBeenCalledWith(expect.objectContaining({ message: 'boom' })); + await expect(schedule()).resolves.toEqual({ success: true }); + expect(calls[2]).toEqual({}); + }); + + it('upgrades consent cookies only for explicit Multi-Panel sync requests', async () => { + mockMultiPanelTabs(12); + const consentCookie = { + name: 'OptanonConsent', + value: 'groups=C0001:1', + domain: '.grok.com', + path: '/', + secure: false, + sameSite: 'lax', + }; + chrome.cookies.getAll.mockResolvedValue([ + consentCookie, + { name: 'sso', value: 'session-a', sameSite: 'lax' }, + ]); + + await reconcileGrokWebSocketCookies(); + expect(chrome.cookies.set).not.toHaveBeenCalled(); + + await reconcileGrokWebSocketCookies({ upgradeConsentCookies: true }); + expect(chrome.cookies.set).toHaveBeenCalledTimes(1); + // The upgraded cookie is now attached natively, so it leaves the header. + expect(chrome.declarativeNetRequest.updateSessionRules).toHaveBeenLastCalledWith({ + removeRuleIds: [GROK_WEBSOCKET_COOKIE_RULE_ID], + addRules: [createGrokWebSocketCookieRule('sso=session-a', [12])], + }); + }); + + it('upgrades existing OneTrust consent cookies to SameSite=None; Secure for embedded frames', async () => { + const cookies = [ + { + name: 'OptanonAlertBoxClosed', + value: '2026-10-16T12:00:00.000Z', + domain: '.grok.com', + path: '/', + secure: false, + httpOnly: false, + sameSite: 'lax', + expirationDate: 1821756479, + storeId: '0', + }, + { + name: 'sso', + value: 'session-token', + domain: '.grok.com', + path: '/', + secure: true, + httpOnly: true, + sameSite: 'lax', + }, + ]; + + const upgraded = await upgradeGrokConsentCookies(cookies); + + expect(chrome.cookies.set).toHaveBeenCalledTimes(1); + expect(chrome.cookies.set).toHaveBeenCalledWith({ + url: GROK_COOKIE_URL, + name: 'OptanonAlertBoxClosed', + value: '2026-10-16T12:00:00.000Z', + domain: '.grok.com', + path: '/', + secure: true, + httpOnly: false, + sameSite: 'no_restriction', + expirationDate: 1821756479, + storeId: '0', + }); + expect(upgraded[0]).toMatchObject({ + name: 'OptanonAlertBoxClosed', + secure: true, + sameSite: 'no_restriction', + }); + expect(upgraded[1]).toMatchObject({ + name: 'sso', + sameSite: 'lax', + }); + }); + + function runConsentScript({ cookieOwner = 'document', ancestorOrigins = ['chrome-extension://test-ext-id'] } = {}) { + const writtenValues = []; + const documentProto = {}; + const htmlDocumentProto = Object.create(documentProto); + const owner = cookieOwner === 'htmlDocument' ? htmlDocumentProto : documentProto; + Object.defineProperty(owner, 'cookie', { + configurable: true, + enumerable: true, + get() { + return writtenValues.join('; '); + }, + set(value) { + writtenValues.push(value); + }, + }); + const fakeDoc = Object.create(htmlDocumentProto); + const fakeLocation = { protocol: 'https:', ancestorOrigins }; + const fakeWindow = { self: {}, top: {}, location: fakeLocation }; + + new Function('window', 'document', 'location', grokCookieConsentScript)( + fakeWindow, + fakeDoc, + fakeLocation + ); + + fakeDoc.cookie = 'OptanonAlertBoxClosed=2026-10-16T13:00:00.000Z; path=/; domain=.grok.com; Samesite=Lax'; + fakeDoc.cookie = 'other_cookie=1; path=/; SameSite=Lax'; + return writtenValues; + } + + it('rewrites Optanon consent cookies written via document.cookie in Grok frames embedded by an extension page', () => { + expect(runConsentScript()).toEqual([ + 'OptanonAlertBoxClosed=2026-10-16T13:00:00.000Z; path=/; domain=.grok.com; SameSite=None; Secure', + 'other_cookie=1; path=/; SameSite=Lax', + ]); + }); + + it('patches the cookie accessor on whichever prototype owns it', () => { + expect(runConsentScript({ cookieOwner: 'htmlDocument' })[0]).toBe( + 'OptanonAlertBoxClosed=2026-10-16T13:00:00.000Z; path=/; domain=.grok.com; SameSite=None; Secure' + ); + }); + + it('leaves consent cookies untouched in Grok frames embedded by websites', () => { + const untouched = 'OptanonAlertBoxClosed=2026-10-16T13:00:00.000Z; path=/; domain=.grok.com; Samesite=Lax'; + + expect(runConsentScript({ ancestorOrigins: ['https://attacker.test'] })[0]).toBe(untouched); + // A website framing the web-accessible Multi-Panel page. + expect(runConsentScript({ + ancestorOrigins: ['chrome-extension://test-ext-id', 'https://attacker.test'], + })[0]).toBe(untouched); + }); +});