From d51da99b8f24ae9c57cb832d791ef3a5c10d8e9c Mon Sep 17 00:00:00 2001 From: SuYao Date: Thu, 27 Feb 2025 16:46:05 +0800 Subject: [PATCH] fix: some proxy errors in main process (#2294) * formatter proxy config type * fix: some proxy errors in main process * chore: Remove debug logging in ProxyManager --- package.json | 1 + src/main/ipc.ts | 12 +-- src/main/services/GeminiService.ts | 7 +- src/main/services/KnowledgeService.ts | 4 +- src/main/services/ProxyManager.ts | 115 ++++++++++++++++++++++++++ src/main/services/WebDav.ts | 11 ++- yarn.lock | 8 ++ 7 files changed, 149 insertions(+), 9 deletions(-) create mode 100644 src/main/services/ProxyManager.ts diff --git a/package.json b/package.json index de79431a..a5902bb5 100644 --- a/package.json +++ b/package.json @@ -77,6 +77,7 @@ "markdown-it": "^14.1.0", "officeparser": "^4.1.1", "tokenx": "^0.4.1", + "undici": "^7.3.0", "webdav": "4.11.4" }, "devDependencies": { diff --git a/src/main/ipc.ts b/src/main/ipc.ts index 04e7287f..f3118d52 100644 --- a/src/main/ipc.ts +++ b/src/main/ipc.ts @@ -2,7 +2,7 @@ import fs from 'node:fs' import path from 'node:path' import { 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,12 +14,12 @@ import FileService from './services/FileService' import FileStorage from './services/FileStorage' import { GeminiService } from './services/GeminiService' import KnowledgeService from './services/KnowledgeService' +import { ProxyConfig, proxyManager } from './services/ProxyManager' import { registerShortcuts, unregisterAllShortcuts } from './services/ShortcutService' import { TrayService } from './services/TrayService' import { windowService } from './services/WindowService' import { getResourcePath } from './utils' -import { decrypt } from './utils/aes' -import { encrypt } from './utils/aes' +import { decrypt, encrypt } from './utils/aes' import { compress, decompress } from './utils/zip' const fileManager = new FileStorage() @@ -40,9 +40,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..da6c21e0 100644 --- a/src/main/services/GeminiService.ts +++ b/src/main/services/GeminiService.ts @@ -1,14 +1,16 @@ import { FileMetadataResponse, FileState, GoogleAIFileManager } from '@google/generative-ai/server' import { FileType } from '@types' import fs from 'fs' +import { ProxyAgent, setGlobalDispatcher } from 'undici' 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) { + setGlobalDispatcher(new ProxyAgent(proxyManager.getProxyUrl() || '')) 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 { + setGlobalDispatcher(new ProxyAgent(proxyManager.getProxyUrl() || '')) 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) { + setGlobalDispatcher(new ProxyAgent(proxyManager.getProxyUrl() || '')) const fileManager = new GoogleAIFileManager(apiKey) return await fileManager.listFiles() } static async deleteFile(_: Electron.IpcMainInvokeEvent, apiKey: string, fileId: string) { + setGlobalDispatcher(new ProxyAgent(proxyManager.getProxyUrl() || '')) const fileManager = new GoogleAIFileManager(apiKey) await fileManager.deleteFile(fileId) } diff --git a/src/main/services/KnowledgeService.ts b/src/main/services/KnowledgeService.ts index ae116bb4..89d991b5 100644 --- a/src/main/services/KnowledgeService.ts +++ b/src/main/services/KnowledgeService.ts @@ -15,6 +15,7 @@ import { FileType, KnowledgeBaseParams, KnowledgeItem } from '@types' import { app } from 'electron' import { v4 as uuidv4 } from 'uuid' +import { proxyManager } from './ProxyManager' import { windowService } from './WindowService' class KnowledgeService { @@ -48,13 +49,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 }) diff --git a/src/main/services/ProxyManager.ts b/src/main/services/ProxyManager.ts new file mode 100644 index 00000000..6a4b0fc2 --- /dev/null +++ b/src/main/services/ProxyManager.ts @@ -0,0 +1,115 @@ +import { ProxyConfig as _ProxyConfig, session } from 'electron' +import { HttpsProxyAgent } from 'https-proxy-agent' + +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: '' + } + } + + private async setSessionsProxy(config: _ProxyConfig): Promise { + const sessions = [session.defaultSession, session.fromPartition('persist:webview')] + await Promise.all(sessions.map((session) => session.setProxy(config))) + } + + 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 { + this.proxyUrl = await this.resolveSystemProxy() + if (this.proxyUrl) { + this.proxyAgent = new HttpsProxyAgent(this.proxyUrl) + this.setEnvironment(this.proxyUrl) + await this.setSessionsProxy({ mode: 'system' }) + } + } catch (error) { + console.error('Failed to set system proxy:', error) + throw error + } + } + + private async setCustomProxy(): Promise { + try { + if (this.config.url) { + this.proxyAgent = new HttpsProxyAgent(this.config.url) + this.setEnvironment(this.config.url) + await this.setSessionsProxy({ proxyRules: this.config.url }) + } + } 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' } + } + + 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 + } +} + +export const proxyManager = new ProxyManager() diff --git a/src/main/services/WebDav.ts b/src/main/services/WebDav.ts index b3bc4173..6d9465ca 100644 --- a/src/main/services/WebDav.ts +++ b/src/main/services/WebDav.ts @@ -1,8 +1,12 @@ import { WebDavConfig } from '@types' import Logger from 'electron-log' +import { HttpProxyAgent } from 'http-proxy-agent' +import { HttpsProxyAgent } from 'https-proxy-agent' import Stream from 'stream' import { BufferLike, createClient, GetFileContentsOptions, PutFileContentsOptions, WebDAVClient } from 'webdav' +import { proxyManager } from './ProxyManager' + export default class WebDav { public instance: WebDAVClient | undefined private webdavPath: string @@ -10,11 +14,16 @@ export default class WebDav { constructor(params: WebDavConfig) { this.webdavPath = params.webdavPath + const httpAgent = new HttpProxyAgent(proxyManager.getProxyUrl() || '') + const httpsAgent = new HttpsProxyAgent(proxyManager.getProxyUrl() || '') + this.instance = createClient(params.webdavHost, { username: params.webdavUser, password: params.webdavPass, maxBodyLength: Infinity, - maxContentLength: Infinity + maxContentLength: Infinity, + httpAgent: httpAgent, + httpsAgent: httpsAgent }) this.putFileContents = this.putFileContents.bind(this) diff --git a/yarn.lock b/yarn.lock index 5e87d52e..d44bf1d8 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3094,6 +3094,7 @@ __metadata: tinycolor2: "npm:^1.6.0" tokenx: "npm:^0.4.1" typescript: "npm:^5.6.2" + undici: "npm:^7.3.0" uuid: "npm:^10.0.0" vite: "npm:^5.0.12" webdav: "npm:4.11.4" @@ -13943,6 +13944,13 @@ __metadata: languageName: node linkType: hard +"undici@npm:^7.3.0": + version: 7.3.0 + resolution: "undici@npm:7.3.0" + checksum: 10c0/62c5e335725cadb02e19950932c7823fc330cbfd80106e6836daa6db1379aa727510b77de0a4e6f912087b288ded93f7daf4b8c154ad36fd5c9c4b96b26888b8 + languageName: node + linkType: hard + "unified@npm:^11.0.0": version: 11.0.5 resolution: "unified@npm:11.0.5"