feat: add `withSetCookies` to request options

dist/lib/request.d.ts

@@ -3,5 +3,6 @@ export declare const axios: import("axios").AxiosInstance;
 export interface RequestConfig extends AxiosRequestConfig {
     timeout?: number;
     maxRetries?: number;
+    withSetCookies?: boolean;
 }
 export declare const request: (config: RequestConfig, retries?: number) => Promise<AxiosResponse<any>>;

dist/lib/request.js

@@ -6,8 +6,11 @@ Object.defineProperty(exports, "__esModule", { value: true });
 exports.request = exports.axios = void 0;
 const chalk_1 = __importDefault(require("chalk"));
 const axios_1 = __importDefault(require("axios"));
+const psl_1 = require("psl");
+const set_cookie_parser_1 = require("set-cookie-parser");
 const DEF_TIMEOUT = 10000;
-const DEF_RETRIES = 2;
+const DEF_MAX_RETRIES = 2;
+const DEF_WITH_SET_COOKIES = false;
 const DEF_UA = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/86.0.4240.80 Safari/537.36';
 exports.axios = axios_1.default.create({
     headers: {
@@ -15,13 +18,63 @@ exports.axios = axios_1.default.create({
     },
     withCredentials: true
 });
+class CookieController {
+    static parseHostToDomains(host) {
+        const domains = [host, `.${host}`];
+        const domain = psl_1.get(host);
+        if (domain && domain !== host) {
+            domains.push(domain, `.${domain}`);
+        }
+        return domains;
+    }
+    static sync(setCookies, host) {
+        const cookies = set_cookie_parser_1.parse(setCookies);
+        for (const item of cookies) {
+            console.debug(chalk_1.default.white(`└─ set-cookie`, JSON.stringify(item)));
+            const { name, domain = host } = item;
+            const cached = this.cacheMap.get(domain) || [];
+            const existIndex = cached.findIndex(c => c.name === name);
+            if (existIndex > -1) {
+                cached.splice(existIndex, 1, item);
+            }
+            else {
+                cached.push(item);
+            }
+            this.cacheMap.set(domain, cached);
+        }
+    }
+    static inject(url, headers) {
+        const domains = this.parseHostToDomains(new URL(url).host);
+        const cookies = [];
+        const now = Date.now();
+        for (const domain of domains) {
+            const cache = this.cacheMap.get(domain);
+            if (cache) {
+                const cookieStrs = cache
+                    .filter(c => !c.expires || c.expires.valueOf() > now)
+                    .map(c => `${c.name}=${c.value}`);
+                cookies.push(...cookieStrs);
+            }
+        }
+        const cookie = (headers.cookie ? `${headers.cookie}; ` : '') + cookies.join('; ');
+        headers.cookie = cookie;
+        cookie && console.debug(chalk_1.default.white(`└─ cookie:`, cookie));
+    }
+}
+CookieController.cacheMap = new Map();
 const request = (config, retries = 0) => {
     config.timeout = config.timeout || DEF_TIMEOUT;
-    config.maxRetries = config.maxRetries || DEF_RETRIES;
+    config.maxRetries = config.maxRetries || DEF_MAX_RETRIES;
+    config.withSetCookies = config.withSetCookies != null ? config.withSetCookies : DEF_WITH_SET_COOKIES;
     const _config = JSON.stringify(config);
+    const fullUrl = `${config.baseURL || exports.axios.defaults.baseURL || ''}${config.url}`;
     const cancelTokenSource = axios_1.default.CancelToken.source();
     config.cancelToken = cancelTokenSource.token;
-    console.debug(chalk_1.default.white(`[request] ${config.method || 'GET'} ${config.baseURL || exports.axios.defaults.baseURL || ''}${config.url}`, config.params ? JSON.stringify(config.params) : ''));
+    config.headers = config.headers || {};
+    console.debug(chalk_1.default.white(`[request] ${config.method || 'GET'} ${fullUrl}`, config.params ? JSON.stringify(config.params) : ''));
+    if (config.withSetCookies) {
+        CookieController.inject(fullUrl, config.headers);
+    }
     return new Promise((resolve, reject) => {
         let done = false;
         const timer = setTimeout(() => {
@@ -42,6 +95,11 @@ const request = (config, retries = 0) => {
         exports.axios(config)
             .then((res) => {
             console.debug(chalk_1.default.white(`└─ ${res.status} ${Date.now() - startAt}ms`));
+            const setCookies = res.headers['set-cookie'];
+            if (config.withSetCookies && setCookies) {
+                const { host } = new URL(res.request.res.responseUrl);
+                CookieController.sync(setCookies, host);
+            }
             resolve(res);
         })
             .catch((e) => {

package.json

@@ -1,6 +1,6 @@
 {
   "name": "crawler-lib",
-  "version": "1.0.0",
+  "version": "1.1.0",
   "main": "dist/lib/index.js",
   "module": "dist/lib/index.js",
   "author": "Acathur",
@@ -9,11 +9,15 @@
     "build": "rm -rf dist && tsc"
   },
   "devDependencies": {
+    "@types/node": "^14.14.16",
+    "@types/psl": "^1.1.0",
+    "@types/set-cookie-parser": "^2.4.0",
     "typescript": "^4.1.3"
   },
   "dependencies": {
-    "@types/node": "^14.14.16",
     "axios": "^0.21.1",
-    "chalk": "^4.1.0"
+    "chalk": "^4.1.0",
+    "psl": "^1.8.0",
+    "set-cookie-parser": "^2.4.6"
   }
 }

src/lib/request.ts

@@ -1,8 +1,11 @@
 import chalk from 'chalk'
 import Axios, { AxiosRequestConfig, AxiosResponse } from 'axios'
+import { get as pslGet } from 'psl'
+import { Cookie, parse as parseSetCookie } from 'set-cookie-parser'
 
 const DEF_TIMEOUT = 10000
-const DEF_RETRIES = 2
+const DEF_MAX_RETRIES = 2
+const DEF_WITH_SET_COOKIES = false
 const DEF_UA = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/86.0.4240.80 Safari/537.36'
 
 export const axios = Axios.create({
@@ -15,16 +18,77 @@ export const axios = Axios.create({
 export interface RequestConfig extends AxiosRequestConfig {
   timeout?: number
   maxRetries?: number
+  withSetCookies?: boolean
+}
+
+class CookieController {
+  static cacheMap: Map<string, Cookie[]> = new Map()
+
+  private static parseHostToDomains(host: string) {
+    const domains = [host, `.${host}`]
+    const domain = pslGet(host)
+
+    if (domain && domain !== host) {
+      domains.push(domain, `.${domain}`)
+    }
+    return domains
+  }
+
+  static sync(setCookies: string[], host: string) {
+    const cookies = parseSetCookie(setCookies)
+
+    for (const item of cookies) {
+      console.debug(chalk.white(`└─ set-cookie`, JSON.stringify(item)))
+      const { name, domain = host } = item
+      const cached = this.cacheMap.get(domain) || []
+      const existIndex = cached.findIndex(c => c.name === name)
+
+      if (existIndex > -1) {
+        cached.splice(existIndex, 1, item)
+      } else {
+        cached.push(item)
+      }
+      this.cacheMap.set(domain, cached)
+    }
+  }
+
+  static inject(url: string, headers: any) {
+    const domains = this.parseHostToDomains(new URL(url).host)
+    const cookies: string[] = []
+    const now = Date.now()
+
+    for (const domain of domains) {
+      const cache = this.cacheMap.get(domain)
+
+      if (cache) {
+        const cookieStrs = cache
+          .filter(c => !c.expires || c.expires.valueOf() > now)
+          .map(c => `${c.name}=${c.value}`)
+        cookies.push(...cookieStrs)
+      }
+    }
+
+    const cookie = (headers.cookie ? `${headers.cookie}; ` : '') + cookies.join('; ')
+    headers.cookie = cookie
+    cookie && console.debug(chalk.white(`└─ cookie:`, cookie))
+  }
 }
 
 export const request = (config: RequestConfig, retries = 0): Promise<AxiosResponse<any>> => {
   config.timeout = config.timeout || DEF_TIMEOUT
-  config.maxRetries = config.maxRetries || DEF_RETRIES
+  config.maxRetries = config.maxRetries || DEF_MAX_RETRIES
+  config.withSetCookies = config.withSetCookies != null ? config.withSetCookies : DEF_WITH_SET_COOKIES
   const _config = JSON.stringify(config)
+  const fullUrl = `${config.baseURL || axios.defaults.baseURL || ''}${config.url}`
   const cancelTokenSource = Axios.CancelToken.source()
   config.cancelToken = cancelTokenSource.token
+  config.headers = config.headers || {}
 
-  console.debug(chalk.white(`[request] ${config.method || 'GET'} ${config.baseURL || axios.defaults.baseURL || ''}${config.url}`, config.params ? JSON.stringify(config.params) : ''))
+  console.debug(chalk.white(`[request] ${config.method || 'GET'} ${fullUrl}`, config.params ? JSON.stringify(config.params) : ''))
+
+  if (config.withSetCookies) {
+    CookieController.inject(fullUrl, config.headers)
+  }
 
   return new Promise((resolve, reject) => {
     let done = false
@@ -49,6 +113,13 @@ export const request = (config: RequestConfig, retries = 0): Promise<AxiosRespon
     axios(config)
       .then((res) => {
         console.debug(chalk.white(`└─ ${res.status} ${Date.now() - startAt}ms`))
+        const setCookies = res.headers['set-cookie']
+
+        if (config.withSetCookies && setCookies) {
+          const { host } = new URL(res.request.res.responseUrl)
+          CookieController.sync(setCookies, host)
+        }
+
         resolve(res)
       })
       .catch((e) => {

src/test/request.ts

@@ -1,8 +1,25 @@
 import { request } from '../lib'
 
 (async () => {
+  await request({
+    url: 'https://www.baidu.com',
+    withSetCookies: true
+  })
+
+  await request({
+    url: 'https://www.baidu.com',
+    withSetCookies: true
+  })
+
+  await request({
+    url: 'https://cn.bing.com',
+    withSetCookies: true
+  })
+
   const res = await request({
-    url: 'https://www.google.com'
+    url: 'https://cn.bing.com',
+    withSetCookies: true
   })
+
   console.info('@res', res)
 })()

yarn.lock

@@ -2,11 +2,23 @@
 # yarn lockfile v1
 
 
-"@types/node@^14.14.16":
+"@types/node@*", "@types/node@^14.14.16":
   version "14.14.16"
   resolved "https://registry.yarnpkg.com/@types/node/-/node-14.14.16.tgz#3cc351f8d48101deadfed4c9e4f116048d437b4b"
   integrity sha512-naXYePhweTi+BMv11TgioE2/FXU4fSl29HAH1ffxVciNsH3rYXjNP2yM8wqmSm7jS20gM8TIklKiTen+1iVncw==
 
+"@types/psl@^1.1.0":
+  version "1.1.0"
+  resolved "https://registry.yarnpkg.com/@types/psl/-/psl-1.1.0.tgz#390c5df1613b166ce3c3eb9fda4d93dc3eeec7b5"
+  integrity sha512-HhZnoLAvI2koev3czVPzBNRYvdrzJGLjQbWZhqFmS9Q6a0yumc5qtfSahBGb5g+6qWvA8iiQktqGkwoIXa/BNQ==
+
+"@types/set-cookie-parser@^2.4.0":
+  version "2.4.0"
+  resolved "https://registry.yarnpkg.com/@types/set-cookie-parser/-/set-cookie-parser-2.4.0.tgz#10cc0446bad372827671a5195fbd14ebce4a9baf"
+  integrity sha512-w7BFUq81sy7H/0jN0K5cax8MwRN6NOSURpY4YuO4+mOgoicxCZ33BUYz+gyF/sUf7uDl2We2yGJfppxzEXoAXQ==
+  dependencies:
+    "@types/node" "*"
+
 ansi-styles@^4.1.0:
   version "4.3.0"
   resolved "https://registry.yarnpkg.com/ansi-styles/-/ansi-styles-4.3.0.tgz#edd803628ae71c04c85ae7a0906edad34b648937"
@@ -51,6 +63,16 @@ has-flag@^4.0.0:
   resolved "https://registry.yarnpkg.com/has-flag/-/has-flag-4.0.0.tgz#944771fd9c81c81265c4d6941860da06bb59479b"
   integrity sha512-EykJT/Q1KjTWctppgIAgfSO0tKVuZUjhgMr17kqTumMl6Afv3EISleU7qZUzoXDFTAHTDC4NOoG/ZxU3EvlMPQ==
 
+psl@^1.8.0:
+  version "1.8.0"
+  resolved "https://registry.yarnpkg.com/psl/-/psl-1.8.0.tgz#9326f8bcfb013adcc005fdff056acce020e51c24"
+  integrity sha512-RIdOzyoavK+hA18OGGWDqUTsCLhtA7IcZ/6NCs4fFJaHBDab+pDDmDIByWFRQJq2Cd7r1OoQxBGKOaztq+hjIQ==
+
+set-cookie-parser@^2.4.6:
+  version "2.4.6"
+  resolved "https://registry.yarnpkg.com/set-cookie-parser/-/set-cookie-parser-2.4.6.tgz#43bdea028b9e6f176474ee5298e758b4a44799c3"
+  integrity sha512-mNCnTUF0OYPwYzSHbdRdCfNNHqrne+HS5tS5xNb6yJbdP9wInV0q5xPLE0EyfV/Q3tImo3y/OXpD8Jn0Jtnjrg==
+
 supports-color@^7.1.0:
   version "7.2.0"
   resolved "https://registry.yarnpkg.com/supports-color/-/supports-color-7.2.0.tgz#1b7dcdcb32b8138801b3e478ba6a51caa89648da"