From 759c03cd23547573eb68088bad422d600c17caee Mon Sep 17 00:00:00 2001
From: DIYgod <i@diygod.me>
Date: Thu, 28 Mar 2024 02:45:08 +0800
Subject: [PATCH] feat: intercept undici and use fetch dispatcher

---
 lib/app.tsx                      |  2 +-
 lib/utils/proxy/index.ts         |  7 +++
 lib/utils/request-interceptor.ts | 80 ++++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+), 1 deletion(-)
 create mode 100644 lib/utils/request-interceptor.ts

diff --git a/lib/app.tsx b/lib/app.tsx
index c3b27ab41f..f2a549eef0 100644
--- a/lib/app.tsx
+++ b/lib/app.tsx
@@ -1,4 +1,4 @@
-import '@/utils/request-wrapper';
+import '@/utils/request-interceptor';
 
 import { Hono } from 'hono';
 
diff --git a/lib/utils/proxy/index.ts b/lib/utils/proxy/index.ts
index b49730693b..016501b2ac 100644
--- a/lib/utils/proxy/index.ts
+++ b/lib/utils/proxy/index.ts
@@ -2,6 +2,7 @@ import { config } from '@/config';
 import { PacProxyAgent } from 'pac-proxy-agent';
 import { HttpsProxyAgent } from 'https-proxy-agent';
 import { SocksProxyAgent } from 'socks-proxy-agent';
+import { ProxyAgent } from 'undici';
 
 const proxyIsPAC = config.pacUri || config.pacScript;
 
@@ -24,11 +25,16 @@ if (proxyIsPAC) {
 }
 
 let agent: PacProxyAgent<string> | HttpsProxyAgent<string> | SocksProxyAgent | null = null;
+let dispatcher: ProxyAgent | null = null;
 if (proxyIsPAC) {
     agent = new PacProxyAgent(`pac+${proxyUri}`);
 } else if (proxyUri) {
     if (proxyUri.startsWith('http')) {
         agent = new HttpsProxyAgent(proxyUri);
+        dispatcher = new ProxyAgent({
+            uri: proxyUri,
+            token: proxyObj?.auth ? `Basic ${proxyObj.auth}` : undefined,
+        });
     } else if (proxyUri.startsWith('socks')) {
         agent = new SocksProxyAgent(proxyUri);
     }
@@ -36,6 +42,7 @@ if (proxyIsPAC) {
 
 export default {
     agent,
+    dispatcher,
     proxyUri,
     proxyObj,
     proxyUrlHandler,
diff --git a/lib/utils/request-interceptor.ts b/lib/utils/request-interceptor.ts
new file mode 100644
index 0000000000..c85ba01cd7
--- /dev/null
+++ b/lib/utils/request-interceptor.ts
@@ -0,0 +1,80 @@
+import { setupServer } from 'msw/node';
+import { http, passthrough } from 'msw';
+import logger from '@/utils/logger';
+import { config } from '@/config';
+import { fetch, Headers, FormData, ProxyAgent, Request, RequestInfo, RequestInit, Response } from 'undici';
+import proxy from '@/utils/proxy';
+import type nodehttp from 'node:http';
+
+class ExtendedRequest extends Request {
+    public dispatcher: ProxyAgent | undefined;
+    public agent: nodehttp.ClientRequestArgs['agent'];
+
+    constructor(input: RequestInfo, init?: RequestInit & nodehttp.ClientRequestArgs) {
+        super(input, init);
+        this.dispatcher = init?.dispatcher; // fetch
+        this.agent = init?.agent; // http
+    }
+}
+
+Object.defineProperties(globalThis, {
+    fetch: { value: fetch, writable: true },
+    Headers: { value: Headers },
+    FormData: { value: FormData },
+    Request: { value: ExtendedRequest },
+    Response: { value: Response },
+});
+
+const handler = (request: ExtendedRequest) => {
+    request.headers.set('debug', '1');
+    // ua
+    if (!request.headers.get('user-agent')) {
+        request.headers.set('user-agent', config.ua);
+    }
+
+    // accept
+    if (!request.headers.get('accept')) {
+        request.headers.set('accept', '*/*');
+    }
+
+    // referer
+    if (!request.headers.get('referer')) {
+        try {
+            const urlHandler = new URL(request.url);
+            request.headers.set('referer', urlHandler.origin);
+        } catch {
+            // ignore
+        }
+    }
+
+    // proxy
+    if (!request.dispatcher && !request.agent && proxy.dispatcher) {
+        const proxyRegex = new RegExp(proxy.proxyObj.url_regex);
+        let urlHandler;
+        try {
+            urlHandler = new URL(request.url);
+        } catch {
+            // ignore
+        }
+
+        if (proxyRegex.test(request.url) && request.url.startsWith('http') && !(urlHandler && urlHandler.host === proxy.proxyUrlHandler?.host)) {
+            // fetch
+            request.dispatcher = proxy.dispatcher;
+
+            // http
+            request.agent = proxy.agent || undefined;
+            if (proxy.proxyObj.auth) {
+                request.headers.set('Proxy-Authorization', `Basic ${proxy.proxyObj.auth}`);
+            }
+        }
+    }
+};
+
+const server = setupServer(
+    http.all('*', ({ request }) => {
+        logger.debug(`Outgoing request: ${request.method} ${request.url}`);
+        handler(request);
+        return passthrough();
+    })
+);
+server.listen();