From b9d97e8a355ff60ee642dd789c088be410e6fc6a Mon Sep 17 00:00:00 2001 From: suyao Date: Tue, 11 Mar 2025 00:01:07 +0800 Subject: [PATCH] feat(Proxy): Implement proxy management system - Add ProxyManager service to handle system, custom, and no proxy configurations - Integrate proxy support for Gemini, Knowledge, and WebDav services - Add fetch-socks and undici for advanced proxy handling - Enhance proxy configuration with environment variable and session management --- package.json | 2 + src/main/ipc.ts | 9 +- src/main/services/GeminiService.ts | 5 + src/main/services/KnowledgeService.ts | 5 +- src/main/services/ProxyManager.ts | 145 ++++++++++++++++++++++++++ src/main/services/WebDav.ts | 8 +- yarn.lock | 21 +++- 7 files changed, 187 insertions(+), 8 deletions(-) create mode 100644 src/main/services/ProxyManager.ts diff --git a/package.json b/package.json index 5136195f..5a70af6f 100644 --- a/package.json +++ b/package.json @@ -78,11 +78,13 @@ "electron-updater": "^6.3.9", "electron-window-state": "^5.0.3", "epub": "^1.3.0", + "fetch-socks": "^1.3.2", "fs-extra": "^11.2.0", "markdown-it": "^14.1.0", "officeparser": "^4.1.1", "p-queue": "^8.1.0", "tokenx": "^0.4.1", + "undici": "^7.4.0", "webdav": "4.11.4" }, "devDependencies": { diff --git a/src/main/ipc.ts b/src/main/ipc.ts index fe603091..d3140d26 100644 --- a/src/main/ipc.ts +++ b/src/main/ipc.ts @@ -1,7 +1,7 @@ import fs from 'node:fs' import { MCPServer, Shortcut, ThemeMode } from '@types' -import { BrowserWindow, ipcMain, ProxyConfig, session, shell } from 'electron' +import { BrowserWindow, ipcMain, session, shell } from 'electron' import log from 'electron-log' import { titleBarOverlayDark, titleBarOverlayLight } from './config' @@ -14,6 +14,7 @@ import FileStorage from './services/FileStorage' import { GeminiService } from './services/GeminiService' import KnowledgeService from './services/KnowledgeService' import MCPService from './services/mcp' +import { ProxyConfig, proxyManager } from './services/ProxyManager' import { registerShortcuts, unregisterAllShortcuts } from './services/ShortcutService' import { TrayService } from './services/TrayService' import { windowService } from './services/WindowService' @@ -41,9 +42,9 @@ export function registerIpc(mainWindow: BrowserWindow, app: Electron.App) { })) ipcMain.handle('app:proxy', async (_, proxy: string) => { - const sessions = [session.defaultSession, session.fromPartition('persist:webview')] - const proxyConfig: ProxyConfig = proxy === 'system' ? { mode: 'system' } : proxy ? { proxyRules: proxy } : {} - await Promise.all(sessions.map((session) => session.setProxy(proxyConfig))) + const proxyConfig: ProxyConfig = + proxy === 'system' ? { mode: 'system' } : proxy ? { mode: 'custom', url: proxy } : { mode: 'none' } + await proxyManager.configureProxy(proxyConfig) }) ipcMain.handle('app:reload', () => mainWindow.reload()) diff --git a/src/main/services/GeminiService.ts b/src/main/services/GeminiService.ts index b79193ff..5b9510c0 100644 --- a/src/main/services/GeminiService.ts +++ b/src/main/services/GeminiService.ts @@ -3,12 +3,14 @@ import { FileType } from '@types' import fs from 'fs' import { CacheService } from './CacheService' +import { proxyManager } from './ProxyManager' export class GeminiService { private static readonly FILE_LIST_CACHE_KEY = 'gemini_file_list' private static readonly CACHE_DURATION = 3000 static async uploadFile(_: Electron.IpcMainInvokeEvent, file: FileType, apiKey: string) { + proxyManager.setGlobalProxy() const fileManager = new GoogleAIFileManager(apiKey) const uploadResult = await fileManager.uploadFile(file.path, { mimeType: 'application/pdf', @@ -29,6 +31,7 @@ export class GeminiService { file: FileType, apiKey: string ): Promise { + proxyManager.setGlobalProxy() const fileManager = new GoogleAIFileManager(apiKey) const cachedResponse = CacheService.get(GeminiService.FILE_LIST_CACHE_KEY) @@ -52,11 +55,13 @@ export class GeminiService { } static async listFiles(_: Electron.IpcMainInvokeEvent, apiKey: string) { + proxyManager.setGlobalProxy() const fileManager = new GoogleAIFileManager(apiKey) return await fileManager.listFiles() } static async deleteFile(_: Electron.IpcMainInvokeEvent, apiKey: string, fileId: string) { + proxyManager.setGlobalProxy() const fileManager = new GoogleAIFileManager(apiKey) await fileManager.deleteFile(fileId) } diff --git a/src/main/services/KnowledgeService.ts b/src/main/services/KnowledgeService.ts index f36d7deb..8272c039 100644 --- a/src/main/services/KnowledgeService.ts +++ b/src/main/services/KnowledgeService.ts @@ -23,6 +23,7 @@ import { SitemapLoader } from '@llm-tools/embedjs-loader-sitemap' import { WebLoader } from '@llm-tools/embedjs-loader-web' import { AzureOpenAiEmbeddings, OpenAiEmbeddings } from '@llm-tools/embedjs-openai' import { addFileLoader } from '@main/loader' +import { proxyManager } from '@main/services/ProxyManager' import { windowService } from '@main/services/WindowService' import { getInstanceName } from '@main/utils' import { getAllFiles } from '@main/utils/file' @@ -123,13 +124,14 @@ class KnowledgeService { azureOpenAIApiVersion: apiVersion, azureOpenAIApiDeploymentName: model, azureOpenAIApiInstanceName: getInstanceName(baseURL), + configuration: { httpAgent: proxyManager.getProxyAgent() }, dimensions, batchSize }) : new OpenAiEmbeddings({ model, apiKey, - configuration: { baseURL }, + configuration: { baseURL, httpAgent: proxyManager.getProxyAgent() }, dimensions, batchSize }) @@ -424,6 +426,7 @@ class KnowledgeService { } public add = (_: Electron.IpcMainInvokeEvent, options: KnowledgeBaseAddItemOptions): Promise => { + proxyManager.setGlobalProxy() return new Promise((resolve) => { const { base, item, forceReload = false } = options const optionsNonNullableAttribute = { base, item, forceReload } diff --git a/src/main/services/ProxyManager.ts b/src/main/services/ProxyManager.ts new file mode 100644 index 00000000..f1b455ac --- /dev/null +++ b/src/main/services/ProxyManager.ts @@ -0,0 +1,145 @@ +import { ProxyConfig as _ProxyConfig, session } from 'electron' +import { socksDispatcher } from 'fetch-socks' +import { HttpsProxyAgent } from 'https-proxy-agent' +import { ProxyAgent, setGlobalDispatcher } from 'undici' + +type ProxyMode = 'system' | 'custom' | 'none' + +export interface ProxyConfig { + mode: ProxyMode + url?: string | null +} + +export class ProxyManager { + private config: ProxyConfig + private proxyAgent: HttpsProxyAgent | null = null + private proxyUrl: string | null = null + + constructor() { + this.config = { + mode: 'system', + url: '' + } + this.monitorSystemProxy() + } + + private async setSessionsProxy(config: _ProxyConfig): Promise { + const sessions = [session.defaultSession, session.fromPartition('persist:webview')] + await Promise.all(sessions.map((session) => session.setProxy(config))) + } + + private async monitorSystemProxy(): Promise { + setInterval(async () => { + await this.setSystemProxy() + }, 10000) + } + + async configureProxy(config: ProxyConfig): Promise { + try { + this.config = config + if (this.config.mode === 'system') { + await this.setSystemProxy() + } else if (this.config.mode == 'custom') { + await this.setCustomProxy() + } else { + await this.clearProxy() + } + } catch (error) { + console.error('Failed to config proxy:', error) + throw error + } + } + + private setEnvironment(url: string): void { + process.env.grpc_proxy = url + process.env.HTTP_PROXY = url + process.env.HTTPS_PROXY = url + process.env.http_proxy = url + process.env.https_proxy = url + } + + private async setSystemProxy(): Promise { + try { + await this.setSessionsProxy({ mode: 'system' }) + const url = await this.resolveSystemProxy() + if (url && url !== this.proxyUrl) { + this.proxyUrl = url.toLowerCase() + this.proxyAgent = new HttpsProxyAgent(this.proxyUrl) + this.setEnvironment(this.proxyUrl) + } + } catch (error) { + console.error('Failed to set system proxy:', error) + throw error + } + } + + private async setCustomProxy(): Promise { + try { + if (this.config.url) { + this.proxyUrl = this.config.url.toLowerCase() + this.proxyAgent = new HttpsProxyAgent(this.proxyUrl) + this.setEnvironment(this.proxyUrl) + await this.setSessionsProxy({ proxyRules: this.proxyUrl }) + } + } catch (error) { + console.error('Failed to set custom proxy:', error) + throw error + } + } + + private async clearProxy(): Promise { + delete process.env.HTTP_PROXY + delete process.env.HTTPS_PROXY + await this.setSessionsProxy({}) + this.config = { mode: 'none' } + this.proxyAgent = null + this.proxyUrl = null + } + + private async resolveSystemProxy(): Promise { + try { + return await this.resolveElectronProxy() + } catch (error) { + console.error('Failed to resolve system proxy:', error) + return null + } + } + + private async resolveElectronProxy(): Promise { + try { + const proxyString = await session.defaultSession.resolveProxy('https://dummy.com') + const [protocol, address] = proxyString.split(';')[0].split(' ') + return protocol === 'PROXY' ? `http://${address}` : null + } catch (error) { + console.error('Failed to resolve electron proxy:', error) + return null + } + } + + getProxyAgent(): HttpsProxyAgent | null { + return this.proxyAgent + } + + getProxyUrl(): string | null { + return this.proxyUrl + } + + setGlobalProxy() { + const proxyUrl = this.proxyUrl + if (proxyUrl) { + const [protocol, host, port] = proxyUrl.split(':') + if (!protocol.includes('socks')) { + setGlobalDispatcher(new ProxyAgent(proxyUrl)) + } else { + const dispatcher = socksDispatcher({ + port: parseInt(port), + type: protocol === 'socks5' ? 5 : 4, + host: host + }) + global[Symbol.for('undici.globalDispatcher.1')] = dispatcher + } + } + } +} + +export const proxyManager = new ProxyManager() diff --git a/src/main/services/WebDav.ts b/src/main/services/WebDav.ts index b3bc4173..f9e9549f 100644 --- a/src/main/services/WebDav.ts +++ b/src/main/services/WebDav.ts @@ -1,20 +1,24 @@ +import { proxyManager } from '@main/services/ProxyManager' import { WebDavConfig } from '@types' import Logger from 'electron-log' +import { HttpProxyAgent } from 'http-proxy-agent' import Stream from 'stream' import { BufferLike, createClient, GetFileContentsOptions, PutFileContentsOptions, WebDAVClient } from 'webdav' - export default class WebDav { public instance: WebDAVClient | undefined private webdavPath: string constructor(params: WebDavConfig) { this.webdavPath = params.webdavPath + const url = proxyManager.getProxyUrl() this.instance = createClient(params.webdavHost, { username: params.webdavUser, password: params.webdavPass, maxBodyLength: Infinity, - maxContentLength: Infinity + maxContentLength: Infinity, + httpAgent: url ? new HttpProxyAgent(url) : undefined, + httpsAgent: proxyManager.getProxyAgent() }) this.putFileContents = this.putFileContents.bind(this) diff --git a/yarn.lock b/yarn.lock index b1c24a84..fafeec42 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3176,6 +3176,7 @@ __metadata: eslint-plugin-react-hooks: "npm:^4.6.2" eslint-plugin-simple-import-sort: "npm:^12.1.1" eslint-plugin-unused-imports: "npm:^4.0.0" + fetch-socks: "npm:^1.3.2" fs-extra: "npm:^11.2.0" html-to-image: "npm:^1.11.13" i18next: "npm:^23.11.5" @@ -3211,6 +3212,7 @@ __metadata: tinycolor2: "npm:^1.6.0" tokenx: "npm:^0.4.1" typescript: "npm:^5.6.2" + undici: "npm:^7.4.0" uuid: "npm:^10.0.0" vite: "npm:^5.0.12" webdav: "npm:4.11.4" @@ -6500,6 +6502,16 @@ __metadata: languageName: node linkType: hard +"fetch-socks@npm:^1.3.2": + version: 1.3.2 + resolution: "fetch-socks@npm:1.3.2" + dependencies: + socks: "npm:^2.8.2" + undici: "npm:>=6" + checksum: 10c0/6a3f20142c82d3eaef0bfe6b53a0af61381ffbe8bfeb1fdfe5c285c863f9648159ba5ab9b771fac6d3c726e0b894ba52e1069947de0ec97dc287645b40e5d24c + languageName: node + linkType: hard + "fflate@npm:0.8.1": version: 0.8.1 resolution: "fflate@npm:0.8.1" @@ -13541,7 +13553,7 @@ __metadata: languageName: node linkType: hard -"socks@npm:^2.6.2, socks@npm:^2.8.3": +"socks@npm:^2.6.2, socks@npm:^2.8.2, socks@npm:^2.8.3": version: 2.8.4 resolution: "socks@npm:2.8.4" dependencies: @@ -14621,6 +14633,13 @@ __metadata: languageName: node linkType: hard +"undici@npm:>=6, undici@npm:^7.4.0": + version: 7.4.0 + resolution: "undici@npm:7.4.0" + checksum: 10c0/0d8d8d627c87e72cf58148d257a79d019ce058b6761363ee5752103aa0ab57d132448fce4ef15171671ee138ef156a695ec1daeb72cd09ae408afa74dee070b5 + languageName: node + linkType: hard + "unified@npm:^11.0.0": version: 11.0.5 resolution: "unified@npm:11.0.5"