fix: allow querying with both url and hostname when caching headers and cookies (#123)

src/entries/Background/db.ts

@@ -1,7 +1,8 @@
 import { Level } from 'level';
 import type { RequestHistory } from './rpc';
-import { PluginConfig, PluginMetadata, sha256 } from '../../utils/misc';
+import { PluginConfig, PluginMetadata, sha256, urlify } from '../../utils/misc';
 import mutex from './mutex';
+import { minimatch } from 'minimatch';
 const charwise = require('charwise');
 
 export const db = new Level('./ext-db', {
@@ -325,18 +326,32 @@ export async function clearCookies(host: string) {
   });
 }
 
-export async function getCookies(host: string, name: string) {
+export async function getCookies(link: string, name: string) {
   try {
-    const existing = await cookiesDb.sublevel(host).get(name);
+    const existing = await cookiesDb.sublevel(link).get(name);
     return existing;
   } catch (e) {
     return null;
   }
 }
 
-export async function getCookiesByHost(host: string) {
+export async function getCookiesByHost(link: string) {
   const ret: { [key: string]: string } = {};
-  for await (const [key, value] of cookiesDb.sublevel(host).iterator()) {
+  const links: { [k: string]: boolean } = {};
+  const url = urlify(link);
+
+  for await (const sublevel of cookiesDb.keys({ keyEncoding: 'utf8' })) {
+    const l = sublevel.split('!')[1];
+    links[l] = true;
+  }
+
+  const cookieLink = url
+    ? Object.keys(links).filter((l) => minimatch(l, link))[0]
+    : Object.keys(links).filter((l) => urlify(l)?.host === link)[0];
+
+  if (!cookieLink) return ret;
+
+  for await (const [key, value] of cookiesDb.sublevel(cookieLink).iterator()) {
     ret[key] = value;
   }
   return ret;
@@ -359,10 +374,10 @@ export async function getConnection(origin: string) {
   }
 }
 
-export async function setHeaders(host: string, name: string, value?: string) {
+export async function setHeaders(link: string, name: string, value?: string) {
   if (!value) return null;
   return mutex.runExclusive(async () => {
-    await headersDb.sublevel(host).put(name, value);
+    await headersDb.sublevel(link).put(name, value);
     return true;
   });
 }
@@ -382,9 +397,23 @@ export async function getHeaders(host: string, name: string) {
     return null;
   }
 }
-export async function getHeadersByHost(host: string) {
+export async function getHeadersByHost(link: string) {
   const ret: { [key: string]: string } = {};
-  for await (const [key, value] of headersDb.sublevel(host).iterator()) {
+  const url = urlify(link);
+
+  const links: { [k: string]: boolean } = {};
+  for await (const sublevel of headersDb.keys({ keyEncoding: 'utf8' })) {
+    const l = sublevel.split('!')[1];
+    links[l] = true;
+  }
+
+  const headerLink = url
+    ? Object.keys(links).filter((l) => minimatch(l, link))[0]
+    : Object.keys(links).filter((l) => urlify(l)?.host === link)[0];
+
+  if (!headerLink) return ret;
+
+  for await (const [key, value] of headersDb.sublevel(headerLink).iterator()) {
     ret[key] = value;
   }
   return ret;

src/entries/Background/handlers.ts

@@ -4,7 +4,7 @@ import mutex from './mutex';
 import browser from 'webextension-polyfill';
 import { addRequest } from '../../reducers/requests';
 import { urlify } from '../../utils/misc';
-import { setCookies, setHeaders } from './db';
+import { getHeadersByHost, setCookies, setHeaders } from './db';
 export const onSendHeaders = (
   details: browser.WebRequest.OnSendHeadersDetailsType,
 ) => {
@@ -14,9 +14,11 @@ export const onSendHeaders = (
     if (method !== 'OPTIONS') {
       const cache = getCacheByTabId(tabId);
       const existing = cache.get<RequestLog>(requestId);
-      const { hostname } = urlify(details.url) || {};
+      const { origin, pathname } = urlify(details.url) || {};
 
-      if (hostname && details.requestHeaders) {
+      const link = [origin, pathname].join('');
+
+      if (link && details.requestHeaders) {
         details.requestHeaders.forEach((header) => {
           const { name, value } = header;
           if (/^cookie$/i.test(name) && value) {
@@ -24,10 +26,10 @@ export const onSendHeaders = (
               .split(';')
               .map((v) => v.split('='))
               .forEach((cookie) => {
-                setCookies(hostname, cookie[0].trim(), cookie[1]);
+                setCookies(link, cookie[0].trim(), cookie[1]);
               });
           } else {
-            setHeaders(hostname, name, value);
+            setHeaders(link, name, value);
           }
         });
       }

src/utils/misc.ts

@@ -303,20 +303,20 @@ export const makePlugin = async (
   }
 
   if (config?.cookies) {
-    const cookies: { [hostname: string]: { [key: string]: string } } = {};
-    for (const host of config.cookies) {
-      const cache = await getCookiesByHost(host);
-      cookies[host] = cache;
+    const cookies: { [link: string]: { [key: string]: string } } = {};
+    for (const link of config.cookies) {
+      const cache = await getCookiesByHost(link);
+      cookies[link] = cache;
     }
     // @ts-ignore
     injectedConfig.cookies = JSON.stringify(cookies);
   }
 
   if (config?.headers) {
-    const headers: { [hostname: string]: { [key: string]: string } } = {};
-    for (const host of config.headers) {
-      const cache = await getHeadersByHost(host);
-      headers[host] = cache;
+    const headers: { [link: string]: { [key: string]: string } } = {};
+    for (const link of config.headers) {
+      const cache = await getHeadersByHost(link);
+      headers[link] = cache;
     }
     // @ts-ignore
     injectedConfig.headers = JSON.stringify(headers);