| 1 | import { beforeEach, describe, expect, it, vi } from 'vitest' |
| 2 | |
| 3 | type MockDispatcher = { |
| 4 | kind: string |
| 5 | url: string |
| 6 | close: ReturnType<typeof vi.fn> |
| 7 | destroy: ReturnType<typeof vi.fn> |
| 8 | dispatch: ReturnType<typeof vi.fn> |
| 9 | } |
| 10 | |
| 11 | const proxyTestState = vi.hoisted(() => { |
| 12 | const createDispatcher = (kind: string, url = ''): MockDispatcher => ({ |
| 13 | kind, |
| 14 | url, |
| 15 | close: vi.fn().mockResolvedValue(undefined), |
| 16 | destroy: vi.fn().mockResolvedValue(undefined), |
| 17 | dispatch: vi.fn(() => true) |
| 18 | }) |
| 19 | |
| 20 | const originalDispatcher = createDispatcher('original') |
| 21 | const createdDispatchers: MockDispatcher[] = [] |
| 22 | const logMock = { |
| 23 | info: vi.fn() |
| 24 | } |
| 25 | const getGlobalDispatcherMock = vi.fn(() => originalDispatcher) |
| 26 | const setGlobalDispatcherMock = vi.fn() |
| 27 | const ProxyAgentMock = vi.fn(function (options: unknown) { |
| 28 | const url = typeof options === 'string' ? options : String((options as { uri: string }).uri) |
| 29 | const dispatcher = createDispatcher('proxy', url) |
| 30 | createdDispatchers.push(dispatcher) |
| 31 | return dispatcher |
| 32 | }) |
| 33 | const Socks5ProxyAgentMock = vi.fn(function (url: string) { |
| 34 | const dispatcher = createDispatcher('socks', url) |
| 35 | createdDispatchers.push(dispatcher) |
| 36 | return dispatcher |
| 37 | }) |
| 38 | const AgentMock = vi.fn(function () { |
| 39 | const dispatcher = createDispatcher('direct') |
| 40 | createdDispatchers.push(dispatcher) |
| 41 | return dispatcher |
| 42 | }) |
| 43 | |
| 44 | return { |
| 45 | AgentMock, |
| 46 | createdDispatchers, |
| 47 | getGlobalDispatcherMock, |
| 48 | logMock, |
| 49 | originalDispatcher, |
| 50 | ProxyAgentMock, |
| 51 | setGlobalDispatcherMock, |
| 52 | Socks5ProxyAgentMock |
| 53 | } |
| 54 | }) |
| 55 | |
| 56 | vi.mock('electron-log/main.js', () => ({ |
| 57 | default: proxyTestState.logMock |
| 58 | })) |
| 59 | |
| 60 | vi.mock('undici', () => { |
| 61 | class DispatcherMock { |
| 62 | dispatch(): boolean { |
| 63 | return false |
| 64 | } |
| 65 | |
| 66 | async close(): Promise<void> {} |
| 67 | |
| 68 | async destroy(): Promise<void> {} |
| 69 | } |
| 70 | |
| 71 | return { |
| 72 | Agent: proxyTestState.AgentMock, |
| 73 | Dispatcher: DispatcherMock, |
| 74 | ProxyAgent: proxyTestState.ProxyAgentMock, |
| 75 | Socks5ProxyAgent: proxyTestState.Socks5ProxyAgentMock, |
| 76 | getGlobalDispatcher: proxyTestState.getGlobalDispatcherMock, |
| 77 | setGlobalDispatcher: proxyTestState.setGlobalDispatcherMock |
| 78 | } |
| 79 | }) |
| 80 | |
| 81 | async function loadProxyModule() { |
| 82 | vi.resetModules() |
| 83 | return import('../../../src/main/utils/proxy') |
| 84 | } |
| 85 | |
| 86 | describe('proxy dispatcher management', () => { |
| 87 | beforeEach(() => { |
| 88 | proxyTestState.createdDispatchers.length = 0 |
| 89 | proxyTestState.getGlobalDispatcherMock.mockClear() |
| 90 | proxyTestState.getGlobalDispatcherMock.mockReturnValue(proxyTestState.originalDispatcher) |
| 91 | proxyTestState.logMock.info.mockClear() |
| 92 | proxyTestState.originalDispatcher.close.mockClear() |
| 93 | proxyTestState.AgentMock.mockClear() |
| 94 | proxyTestState.ProxyAgentMock.mockClear() |
| 95 | proxyTestState.setGlobalDispatcherMock.mockClear() |
| 96 | proxyTestState.Socks5ProxyAgentMock.mockClear() |
| 97 | }) |
| 98 | |
| 99 | it('applies HTTP proxy via ProxyAgent and trims the URL', async () => { |
| 100 | const { applyProxy } = await loadProxyModule() |
| 101 | |
| 102 | applyProxy({ url: ' http://127.0.0.1:7890 ' }) |
| 103 | |
| 104 | expect(proxyTestState.ProxyAgentMock).toHaveBeenCalledWith('http://127.0.0.1:7890') |
| 105 | expect(proxyTestState.Socks5ProxyAgentMock).not.toHaveBeenCalled() |
| 106 | expect(proxyTestState.setGlobalDispatcherMock).toHaveBeenCalledWith( |
| 107 | proxyTestState.createdDispatchers[0] |
| 108 | ) |
| 109 | expect(proxyTestState.logMock.info).toHaveBeenCalledWith( |
| 110 | '[proxy] applied', |
| 111 | expect.objectContaining({ url: 'http://127.0.0.1:7890', hasAuth: false }) |
| 112 | ) |
| 113 | }) |
| 114 | |
| 115 | it('sends proxy auth as a Basic token for HTTP proxies', async () => { |
| 116 | const { applyProxy } = await loadProxyModule() |
| 117 | |
| 118 | applyProxy({ |
| 119 | url: 'http://127.0.0.1:7890', |
| 120 | username: 'corp\\user', |
| 121 | password: 'secret' |
| 122 | }) |
| 123 | |
| 124 | expect(proxyTestState.ProxyAgentMock).toHaveBeenCalledWith({ |
| 125 | uri: 'http://127.0.0.1:7890', |
| 126 | token: `Basic ${Buffer.from('corp\\user:secret').toString('base64')}` |
| 127 | }) |
| 128 | }) |
| 129 | |
| 130 | it('applies SOCKS proxy with username and password options', async () => { |
| 131 | const { applyProxy } = await loadProxyModule() |
| 132 | |
| 133 | applyProxy({ |
| 134 | url: 'socks5://127.0.0.1:1080', |
| 135 | username: 'user', |
| 136 | password: 'pass' |
| 137 | }) |
| 138 | |
| 139 | expect(proxyTestState.Socks5ProxyAgentMock).toHaveBeenCalledWith('socks5://127.0.0.1:1080', { |
| 140 | username: 'user', |
| 141 | password: 'pass' |
| 142 | }) |
| 143 | expect(proxyTestState.ProxyAgentMock).not.toHaveBeenCalled() |
| 144 | }) |
| 145 | |
| 146 | it('routes no_proxy hosts direct and everything else through the proxy', async () => { |
| 147 | const { applyProxy } = await loadProxyModule() |
| 148 | |
| 149 | applyProxy({ |
| 150 | url: 'http://127.0.0.1:7890', |
| 151 | noProxy: 'localhost,10.0.0.0/8,*.corp.internal' |
| 152 | }) |
| 153 | |
| 154 | // 直连 dispatcher（Agent）和代理 dispatcher 都被创建，全局 dispatcher 是分流包装器 |
| 155 | const directDispatcher = proxyTestState.createdDispatchers.find((d) => d.kind === 'direct') |
| 156 | const proxyDispatcher = proxyTestState.createdDispatchers.find((d) => d.kind === 'proxy') |
| 157 | expect(directDispatcher).toBeDefined() |
| 158 | expect(proxyDispatcher).toBeDefined() |
| 159 | expect(proxyTestState.setGlobalDispatcherMock).toHaveBeenCalledWith( |
| 160 | expect.objectContaining({ constructor: expect.any(Function) }) |
| 161 | ) |
| 162 | const activeDispatcher = proxyTestState.setGlobalDispatcherMock.mock.calls[0][0] as { |
| 163 | dispatch: (options: { origin: string }, handler: unknown) => boolean |
| 164 | } |
| 165 | |
| 166 | const handler = {} |
| 167 | activeDispatcher.dispatch({ origin: 'http://10.1.2.3:8000/v1/models' }, handler) |
| 168 | expect(directDispatcher?.dispatch).toHaveBeenCalled() |
| 169 | expect(proxyDispatcher?.dispatch).not.toHaveBeenCalled() |
| 170 | |
| 171 | directDispatcher?.dispatch.mockClear() |
| 172 | activeDispatcher.dispatch({ origin: 'https://api.openai.com/v1/models' }, handler) |
| 173 | expect(proxyDispatcher?.dispatch).toHaveBeenCalled() |
| 174 | expect(directDispatcher?.dispatch).not.toHaveBeenCalled() |
| 175 | }) |
| 176 | |
| 177 | it('bypasses the proxy entirely for the * wildcard', async () => { |
| 178 | const { applyProxy } = await loadProxyModule() |
| 179 | |
| 180 | applyProxy({ url: 'http://127.0.0.1:7890', noProxy: '*' }) |
| 181 | |
| 182 | expect(proxyTestState.ProxyAgentMock).not.toHaveBeenCalled() |
| 183 | expect(proxyTestState.setGlobalDispatcherMock).toHaveBeenLastCalledWith( |
| 184 | proxyTestState.originalDispatcher |
| 185 | ) |
| 186 | }) |
| 187 | |
| 188 | it('closes the previous proxy dispatcher before switching to a new one', async () => { |
| 189 | const { applyProxy } = await loadProxyModule() |
| 190 | |
| 191 | applyProxy({ url: 'http://127.0.0.1:7890' }) |
| 192 | const firstDispatcher = proxyTestState.createdDispatchers[0] |
| 193 | |
| 194 | applyProxy({ url: 'http://127.0.0.1:7891' }) |
| 195 | |
| 196 | expect(firstDispatcher.close).toHaveBeenCalledTimes(1) |
| 197 | expect(proxyTestState.setGlobalDispatcherMock).toHaveBeenNthCalledWith( |
| 198 | 2, |
| 199 | proxyTestState.createdDispatchers[1] |
| 200 | ) |
| 201 | }) |
| 202 | |
| 203 | it('restores the original dispatcher when the proxy URL is empty', async () => { |
| 204 | const { applyProxy } = await loadProxyModule() |
| 205 | |
| 206 | applyProxy({ url: 'http://127.0.0.1:7890' }) |
| 207 | const activeProxyDispatcher = proxyTestState.createdDispatchers[0] |
| 208 | |
| 209 | applyProxy({ url: ' ' }) |
| 210 | |
| 211 | expect(activeProxyDispatcher.close).toHaveBeenCalledTimes(1) |
| 212 | expect(proxyTestState.setGlobalDispatcherMock).toHaveBeenLastCalledWith( |
| 213 | proxyTestState.originalDispatcher |
| 214 | ) |
| 215 | expect(proxyTestState.logMock.info).toHaveBeenCalledWith( |
| 216 | '[proxy] cleared, restored default dispatcher' |
| 217 | ) |
| 218 | }) |
| 219 | |
| 220 | it('ignores async close rejections from the previous proxy dispatcher', async () => { |
| 221 | const { applyProxy, clearProxy } = await loadProxyModule() |
| 222 | |
| 223 | applyProxy({ url: 'http://127.0.0.1:7890' }) |
| 224 | const activeProxyDispatcher = proxyTestState.createdDispatchers[0] |
| 225 | activeProxyDispatcher.close.mockRejectedValueOnce(new Error('close failed')) |
| 226 | |
| 227 | expect(() => clearProxy()).not.toThrow() |
| 228 | await Promise.resolve() |
| 229 | |
| 230 | expect(activeProxyDispatcher.close).toHaveBeenCalledTimes(1) |
| 231 | expect(proxyTestState.setGlobalDispatcherMock).toHaveBeenLastCalledWith( |
| 232 | proxyTestState.originalDispatcher |
| 233 | ) |
| 234 | }) |
| 235 | }) |
| 236 |