fix: fix cookie rules update strategy

Author: gera2ld

src/handler/cookie.ts

@@ -0,0 +1,106 @@
+import { textTester, urlTester } from '@/common/util';
+import { CookieData, CookieMatchResult } from '@/types';
+import { debounce, pick } from 'es-toolkit';
+import browser from 'webextension-polyfill';
+import { dataLoaded } from './data';
+import { getUrl } from './util';
+
+let cookieRules: CookieData[] = [];
+let processing = false;
+const updates = new Map<string, browser.Cookies.SetDetailsType>();
+let dataReady: Promise<void>;
+
+const debouncedReload = debounce(reloadRules, 500);
+const updateCookiesLater = debounce(updateCookies, 100);
+
+export const cookieActions = {
+  reload: debouncedReload,
+};
+
+browser.cookies.onChanged.addListener(handleCookieChange);
+
+async function reloadRules() {
+  const lists = await dataLoaded;
+  cookieRules = lists.cookie
+    .filter((list) => list.enabled)
+    .flatMap((list) => list.rules.filter((rule) => rule.enabled));
+  console.info(`[cookie] rules reloaded (${cookieRules.length})`);
+}
+
+async function handleCookieChange(
+  changeInfo: browser.Cookies.OnChangedChangeInfoType,
+) {
+  // if (['expired', 'evicted'].includes(changeInfo.cause)) return;
+  if (changeInfo.cause !== 'explicit') return;
+  if (processing) return;
+  await (dataReady ||= reloadRules());
+  let update: CookieMatchResult | undefined;
+  for (const rule of cookieRules) {
+    const url = getUrl(changeInfo.cookie);
+    const matches = url.match(urlTester(rule.url));
+    if (!matches) continue;
+    if (rule.name && !changeInfo.cookie.name.match(textTester(rule.name)))
+      continue;
+    const { ttl } = rule;
+    if (changeInfo.removed && !(ttl && ttl > 0)) {
+      // If cookie is removed and no positive ttl, ignore since change will not persist
+      continue;
+    }
+    update = pick(rule, ['sameSite', 'httpOnly', 'secure']);
+    if (ttl != null) {
+      // If ttl is 0, set to undefined to mark the cookie as a session cookie
+      update.expirationDate = ttl
+        ? Math.floor(Date.now() / 1000 + ttl)
+        : undefined;
+    }
+    if (update.sameSite === 'no_restriction') update.secure = true;
+    break;
+  }
+  if (update) {
+    const { cookie } = changeInfo;
+    const hasUpdate = Object.entries(update).some(([key, value]) => {
+      return cookie[key as keyof browser.Cookies.Cookie] !== value;
+    });
+    if (!hasUpdate) {
+      console.info(`[cookie] no update: ${cookie.name} ${getUrl(cookie)}`);
+      return;
+    }
+    const details: browser.Cookies.SetDetailsType = {
+      url: getUrl(pick(cookie, ['domain', 'path', 'secure'])),
+      domain: cookie.hostOnly ? undefined : cookie.domain,
+      expirationDate: cookie.session ? undefined : cookie.expirationDate,
+      ...pick(cookie, [
+        'name',
+        'path',
+        'httpOnly',
+        'sameSite',
+        'secure',
+        'storeId',
+        'value',
+      ]),
+      ...update,
+    };
+    console.info(`[cookie] matched: ${details.name} ${details.url}`, details);
+    updates.set(
+      [details.storeId, details.url, details.name].join('\n'),
+      details,
+    );
+    updateCookiesLater();
+  }
+}
+
+async function updateCookies() {
+  if (processing) return;
+  processing = true;
+  const items = Array.from(updates.values());
+  updates.clear();
+  for (const item of items) {
+    console.info(`[cookie] set: ${item.name} ${item.url}`, item);
+    try {
+      await browser.cookies.set(item);
+    } catch (err) {
+      console.error(err);
+    }
+  }
+  processing = false;
+}

src/handler/data.ts

@@ -0,0 +1,84 @@
+import { normalizeCookieRule, normalizeRequestRule } from '@/common/list';
+import type { ListData, ListGroups } from '@/types';
+import { flatMap, groupBy } from 'es-toolkit';
+import browser from 'webextension-polyfill';
+import { dumpExactData, getExactData } from './util';
+
+const LIST_PREFIX = 'list:';
+const KEY_LISTS = 'lists';
+
+let lastId = -1;
+
+export const dataLoaded = loadData();
+
+export async function loadData() {
+  let ids = await getExactData<number[]>(KEY_LISTS);
+  const lists: ListGroups = { request: [], cookie: [] };
+  if (Array.isArray(ids)) {
+    const allData = await browser.storage.local.get(
+      ids.map((id) => `${LIST_PREFIX}${id}`),
+    );
+    const allLists = ids
+      .map((id) => allData[`${LIST_PREFIX}${id}`])
+      .filter(Boolean) as ListData[];
+    const groups = groupBy(allLists, (item) => item.type);
+    Object.assign(lists, groups);
+  } else {
+    const allData = await browser.storage.local.get();
+    const allLists = Object.keys(allData)
+      .filter((key) => key.startsWith(LIST_PREFIX))
+      .map((key) => allData[key]) as ListData[];
+    const groups = groupBy(allLists, (item) => item.type);
+    Object.assign(lists, groups);
+  }
+  if (import.meta.env.DEV) console.log('loadData:raw', lists);
+  ids = Object.values(lists).flatMap((group: ListData[]) =>
+    group.map((item) => item.id),
+  );
+  lastId = Math.max(0, ...ids);
+  lists.request.forEach((list) => {
+    list.enabled ??= true;
+    list.rules = flatMap(list.rules, normalizeRequestRule);
+  });
+  lists.cookie.forEach((list) => {
+    list.enabled ??= true;
+    list.rules = flatMap(list.rules, normalizeCookieRule);
+  });
+  if (import.meta.env.DEV) console.log('loadData', lists);
+  return lists;
+}
+
+export function getKey(id: number) {
+  return `${LIST_PREFIX}${id}`;
+}
+
+export async function dumpLists(lists: ListGroups) {
+  await dumpExactData(
+    KEY_LISTS,
+    Object.values(lists).flatMap((group: ListData[]) =>
+      group.map((group) => group.id),
+    ),
+  );
+}
+
+export async function saveList(data: Partial<ListData>) {
+  const list: ListData = {
+    id: 0,
+    name: 'No name',
+    subscribeUrl: '',
+    lastUpdated: 0,
+    enabled: true,
+    type: 'request',
+    rules: [],
+    ...data,
+  };
+  if (!list.rules) throw new Error('Invalid list data');
+  list.name ||= 'No name';
+  if (!list.id) {
+    if (lastId < 0) throw new Error('Data is not loaded yet');
+    list.id = ++lastId;
+  }
+  if (import.meta.env.DEV) console.log('saveList', list);
+  await dumpExactData(getKey(list.id), list);
+  return list;
+}