diff --git a/packages/chrome-extension-mock/declarativ_net_request.ts b/packages/chrome-extension-mock/declarativ_net_request.ts index abbf15c3e..af50dae54 100644 --- a/packages/chrome-extension-mock/declarativ_net_request.ts +++ b/packages/chrome-extension-mock/declarativ_net_request.ts @@ -1,4 +1,8 @@ export default class DeclarativeNetRequest { + MAX_NUMBER_OF_SESSION_RULES = 5000; + + private _sessionRules: chrome.declarativeNetRequest.Rule[] = []; + HeaderOperation = { APPEND: "append", SET: "set", @@ -30,9 +34,44 @@ export default class DeclarativeNetRequest { OTHER: "other", }; - updateSessionRules() { + updateSessionRules(arg1: any, arg2: any): Promise { + let options: { + addRules?: chrome.declarativeNetRequest.Rule[]; + removeRuleIds?: number[]; + } = {}; + let callback: undefined | ((...args: any) => any) = undefined; + + if (typeof arg1 === "function") { + callback = arg1; + } else if (typeof arg2 === "function") { + callback = arg2; + } + if (typeof arg1 === "object" && arg1) options = arg1; + return new Promise((resolve) => { + const { addRules = [], removeRuleIds = [] } = options; + + // Remove rules by ID + if (removeRuleIds.length > 0) { + this._sessionRules = this._sessionRules.filter((rule) => !removeRuleIds.includes(rule.id)); + } + + // Add or update rules (upsert by ID) + for (const newRule of addRules) { + const existingIndex = this._sessionRules.findIndex((rule) => rule.id === newRule.id); + if (existingIndex !== -1) { + this._sessionRules[existingIndex] = newRule; // update + } else { + this._sessionRules.push(newRule); // add + } + } + resolve(); + callback?.(); }); } + + getSessionRules(): Promise { + return Promise.resolve([...this._sessionRules]); + } } diff --git a/src/app/service/service_worker/gm_api/dnr_id_controller.test.ts b/src/app/service/service_worker/gm_api/dnr_id_controller.test.ts new file mode 100644 index 000000000..1b5f218ca --- /dev/null +++ b/src/app/service/service_worker/gm_api/dnr_id_controller.test.ts @@ -0,0 +1,200 @@ +import { sleep } from "@App/pkg/utils/utils"; +import { describe, expect, it } from "vitest"; +import { + getSessionRuleIds, + LIMIT_SESSION_RULES, + nextSessionRuleId, + removeSessionRuleIdEntry, +} from "./dnr_id_controller"; + +describe("getSessionRuleIds", () => { + it("initializes from existing chrome session rules", async () => { + const ids = await getSessionRuleIds(); + expect(ids.size).lessThan(100); + await nextSessionRuleId(); + expect(ids.size).greaterThanOrEqual(1); + await nextSessionRuleId(); + expect(ids.size).greaterThanOrEqual(2); + }); +}); + +describe("nextSessionRuleId", () => { + it("returns unique incrementing IDs on each call", async () => { + const id1 = await nextSessionRuleId(); + const id2 = await nextSessionRuleId(); + const id3 = await nextSessionRuleId(); + + expect(id2).toBeGreaterThan(id1); + expect(id3).toBeGreaterThan(id2); + }); + + it("skips IDs that already exist in session rules", async () => { + const ids = await getSessionRuleIds(); + const next = await nextSessionRuleId(); + + ids.add(next + 1); + const skipped = await nextSessionRuleId(); + + expect(skipped).toBeGreaterThan(next + 1); + }); +}); + +describe("removeSessionRuleIdEntry", () => { + it("removes the given ID from the tracked set", async () => { + const ids = await getSessionRuleIds(); + const id = await nextSessionRuleId(); + + ids.add(id); + removeSessionRuleIdEntry(id); + + expect(ids.has(id)).toBe(false); + }); + + it("is a no-op when called before sessionRuleIds is initialized", () => { + expect(() => removeSessionRuleIdEntry(10404)).not.toThrow(); + }); + + it("throws when ruleId is <= 10000", () => { + expect(() => removeSessionRuleIdEntry(10000)).toThrow(); + expect(() => removeSessionRuleIdEntry(1)).toThrow(); + }); + + it("rewinds SESSION_RULE_ID_BEGIN so the removed id gets reused next", async () => { + const ids = await getSessionRuleIds(); + + const id1 = await nextSessionRuleId(); + const id2 = await nextSessionRuleId(); + const id3 = await nextSessionRuleId(); + + ids.add(id1); + ids.add(id2); + ids.add(id3); + + // Remove the latest — counter should rewind and reuse it + removeSessionRuleIdEntry(id3); + const reused = await nextSessionRuleId(); + expect(reused).toBe(id3); + }); + + it("rewinds SESSION_RULE_ID_BEGIN when removed id is behind the counter", async () => { + const ids = await getSessionRuleIds(); + + const id1 = await nextSessionRuleId(); + const id2 = await nextSessionRuleId(); + const id3 = await nextSessionRuleId(); + + ids.add(id1); + ids.add(id2); + ids.add(id3); + + // Remove an older id — counter rewinds to id1 - 1, so id1 gets reused + removeSessionRuleIdEntry(id1); + const reused = await nextSessionRuleId(); + expect(reused).toBe(id1); + }); + + it("does not rewind SESSION_RULE_ID_BEGIN when removed id was pre-existing (ahead of counter)", async () => { + const id1 = await nextSessionRuleId(); + const id2 = await nextSessionRuleId(); + const id3 = await nextSessionRuleId(); + const _id4 = await nextSessionRuleId(); + const id5 = await nextSessionRuleId(); + const id6 = await nextSessionRuleId(); + + removeSessionRuleIdEntry(id1); + removeSessionRuleIdEntry(id5); + removeSessionRuleIdEntry(id3); + // removeSessionRuleIdEntry(id4); + removeSessionRuleIdEntry(id2); + const nextBefore = await nextSessionRuleId(); // e.g. 10001 + removeSessionRuleIdEntry(id6); + + expect(await nextSessionRuleId()).toBe(nextBefore + 1); // counter unchanged, just increments normally + expect(await nextSessionRuleId()).toBe(nextBefore + 2); // counter unchanged, just increments normally + // expect(await nextSessionRuleId()).toBe(nextBefore + 3); // counter unchanged, just increments normally + expect(await nextSessionRuleId()).toBe(nextBefore + 4); // counter unchanged, just increments normally + expect(await nextSessionRuleId()).toBe(nextBefore + 5); // counter unchanged, just increments normally + }); +}); + +describe("nextSessionRuleId limit control", () => { + it("locks when session rules reach the limit and unlocks when an entry is removed", async () => { + const ids = await getSessionRuleIds(); + expect(ids.size).toBeLessThan(100); + + const added = []; + for (let w = ids.size; w < LIMIT_SESSION_RULES; w++) { + const j = await nextSessionRuleId(); + added.push(j); + ids.add(j); + } + expect(ids.size).toBeGreaterThan(1000); + + const lockedPromise = nextSessionRuleId(); + + const raceResult1 = await Promise.race([lockedPromise.then(() => "resolved"), sleep(5).then(() => "pending")]); + expect(raceResult1).toBe("pending"); + + const m1 = Math.floor(Math.random() * (added.length - 9)); + const p1 = added[m1]; + const p2 = added[m1 + 6]; + removeSessionRuleIdEntry(p1); + removeSessionRuleIdEntry(p2); + + const raceResult2 = await Promise.race([lockedPromise.then(() => "resolved"), sleep(5).then(() => "pending")]); + expect(raceResult2).toBe("resolved"); + + const id1 = await lockedPromise; + const id2 = await nextSessionRuleId(); + expect(id1).toBe(p1); + expect(id2).toBe(p2); + + for (const k of added) { + removeSessionRuleIdEntry(k); + } + + const res = await getSessionRuleIds(); + expect(res).toBe(ids); + expect(res.size).toBeLessThan(100); + }); + + it("only one lock is created even with concurrent nextSessionRuleId calls", async () => { + const ids = await getSessionRuleIds(); + expect(ids.size).toBeLessThan(100); + + const added = []; + for (let w = ids.size; w < LIMIT_SESSION_RULES; w++) { + const j = await nextSessionRuleId(); + added.push(j); + ids.add(j); + } + expect(ids.size).toBeGreaterThan(1000); + + // Fire multiple concurrent calls while locked + const p1 = nextSessionRuleId(); + const p2 = nextSessionRuleId(); + const p3 = nextSessionRuleId(); + + const raceResult = await Promise.race([ + Promise.all([p1, p2, p3]).then(() => "resolved"), + sleep(5).then(() => "pending"), + ]); + expect(raceResult).toBe("pending"); + + // Single removal should unlock all waiters sequentially + const toRemove = [...ids].find((id) => id > 10000)!; + removeSessionRuleIdEntry(toRemove); + + // All three should eventually resolve + const results = await Promise.race([Promise.all([p1, p2, p3]).then((ids) => ids), sleep(50).then(() => null)]); + expect(results).not.toBeNull(); + + for (const k of added) { + removeSessionRuleIdEntry(k); + } + + const res = await getSessionRuleIds(); + expect(res).toBe(ids); + expect(res.size).toBeLessThan(100); + }); +}); diff --git a/src/app/service/service_worker/gm_api/dnr_id_controller.ts b/src/app/service/service_worker/gm_api/dnr_id_controller.ts new file mode 100644 index 000000000..83db1ce6c --- /dev/null +++ b/src/app/service/service_worker/gm_api/dnr_id_controller.ts @@ -0,0 +1,57 @@ +import type { Deferred } from "@App/pkg/utils/utils"; +import { deferred } from "@App/pkg/utils/utils"; + +let sessionRuleIdsPromise: Promise> | null = null; +let sessionRuleIds: Set | null = null; + +let SESSION_RULE_ID_BEGIN = 10000; +export const LIMIT_SESSION_RULES = + process.env.VI_TESTING === "true" ? 1234 : chrome.declarativeNetRequest.MAX_NUMBER_OF_SESSION_RULES - 300; +let lockSessionRuleCreation: Deferred | null = null; + +export const getSessionRuleIds = async (): Promise> => { + if (!sessionRuleIdsPromise) { + sessionRuleIdsPromise = chrome.declarativeNetRequest + .getSessionRules() + .then((rules) => { + sessionRuleIds = new Set(rules.map((rule) => rule.id).filter(Boolean)); + return sessionRuleIds; + }) + .catch((e) => { + console.warn(e); + sessionRuleIds = new Set([]); + return sessionRuleIds; + }); + } + const ruleIds = sessionRuleIds || (await sessionRuleIdsPromise); + return ruleIds; +}; + +export const removeSessionRuleIdEntry = (ruleId: number) => { + if (ruleId <= 10000) { + throw new Error("removeSessionRuleIdEntry cannot remove ids not created by nextSessionRuleId"); + } + if (sessionRuleIds) { + if (sessionRuleIds.delete(ruleId) === true) { + if (ruleId <= SESSION_RULE_ID_BEGIN + 1) { + SESSION_RULE_ID_BEGIN = ruleId - 1; + } + if (sessionRuleIds.size < LIMIT_SESSION_RULES) { + lockSessionRuleCreation?.resolve(); + lockSessionRuleCreation = null; + } + } + } +}; + +export const nextSessionRuleId = async () => { + const ruleIds = await getSessionRuleIds(); + if (!lockSessionRuleCreation && ruleIds.size + 1 > LIMIT_SESSION_RULES) lockSessionRuleCreation = deferred(); + if (lockSessionRuleCreation) await lockSessionRuleCreation.promise; + let id; + do { + id = ++SESSION_RULE_ID_BEGIN; + } while (ruleIds.has(id)); + ruleIds.add(id); + return id; +}; diff --git a/src/app/service/service_worker/gm_api/gm_api.ts b/src/app/service/service_worker/gm_api/gm_api.ts index 1ecc56f94..d087a58cb 100644 --- a/src/app/service/service_worker/gm_api/gm_api.ts +++ b/src/app/service/service_worker/gm_api/gm_api.ts @@ -53,6 +53,7 @@ import { headerModifierMap, headersReceivedMap } from "./gm_xhr"; import { BgGMXhr } from "@App/pkg/utils/xhr/bg_gm_xhr"; import { mightPrepareSetClipboard, setClipboard } from "../clipboard"; import { nativePageWindowOpen } from "../../offscreen/gm_api"; +import { nextSessionRuleId, removeSessionRuleIdEntry } from "./dnr_id_controller"; let generatedUniqueMarkerIDs = ""; let generatedUniqueMarkerIDWhen = ""; @@ -685,8 +686,7 @@ export default class GMApi { } const redirectNotManual = params.redirect !== "manual"; - // 使用 cacheInstance 避免SW重启造成重复 DNR Rule ID - const ruleId = 10000 + (await cacheInstance.incr("gmXhrRequestId", 1)); + const ruleId = await nextSessionRuleId(); const rule = { id: ruleId, action: { @@ -1609,6 +1609,7 @@ export default class GMApi { if (lastError) { console.error("chrome.declarativeNetRequest.updateSessionRules:", lastError); } + removeSessionRuleIdEntry(rule.id); } ); } diff --git a/tests/vitest.setup.ts b/tests/vitest.setup.ts index 8dcc95fc4..178d5a5e7 100644 --- a/tests/vitest.setup.ts +++ b/tests/vitest.setup.ts @@ -9,6 +9,8 @@ import { mockFetch } from "./mocks/fetch"; vi.stubGlobal("chrome", chromeMock); chromeMock.init(); +// vi.spyOn(chromeMock.declarativeNetRequest, "getSessionRules"); +// vi.spyOn(chromeMock.declarativeNetRequest, "updateSessionRules"); initTestEnv(); chromeMock.runtime.getURL = vi.fn().mockImplementation((path: string) => {