Merge remote-tracking branch 'origin/main' into feat/account-level-proxy
# Conflicts: # src/server/routes/proxy/completions.ts # src/server/services/siteProxy.ts
This commit is contained in:
@@ -60,6 +60,41 @@ export function isMaskedTokenValue(token: string | null | undefined): boolean {
|
||||
return value.includes('*') || value.includes('•');
|
||||
}
|
||||
|
||||
function normalizeMaskedTokenForCompare(token: string | null | undefined): string {
|
||||
return normalizeTokenForDisplay(token).replace(/•/g, '*');
|
||||
}
|
||||
|
||||
function matchesMaskedTokenValue(
|
||||
fullToken: string | null | undefined,
|
||||
maskedToken: string | null | undefined,
|
||||
): boolean {
|
||||
const normalizedFull = normalizeTokenForDisplay(fullToken);
|
||||
const normalizedMasked = normalizeMaskedTokenForCompare(maskedToken);
|
||||
if (!normalizedFull || !normalizedMasked) return false;
|
||||
|
||||
if (!isMaskedTokenValue(normalizedMasked)) {
|
||||
return normalizedFull === normalizedMasked;
|
||||
}
|
||||
|
||||
const firstMaskIndex = normalizedMasked.search(/[\*]/);
|
||||
const lastMaskIndex = Math.max(
|
||||
normalizedMasked.lastIndexOf('*'),
|
||||
normalizedMasked.lastIndexOf('•'),
|
||||
);
|
||||
if (firstMaskIndex < 0 || lastMaskIndex < firstMaskIndex) {
|
||||
return normalizedFull === normalizedMasked;
|
||||
}
|
||||
|
||||
const prefix = normalizedMasked.slice(0, firstMaskIndex);
|
||||
const suffix = normalizedMasked.slice(lastMaskIndex + 1);
|
||||
const visiblePrefix = prefix.replace(/^sk-/i, '');
|
||||
if (!visiblePrefix && !suffix) return false;
|
||||
if (normalizedFull.length < prefix.length + suffix.length) return false;
|
||||
if (prefix && !normalizedFull.startsWith(prefix)) return false;
|
||||
if (suffix && !normalizedFull.endsWith(suffix)) return false;
|
||||
return true;
|
||||
}
|
||||
|
||||
function normalizeTokenValueStatus(value: string | null | undefined): AccountTokenValueStatus {
|
||||
const normalized = (value || '').trim().toLowerCase();
|
||||
return normalized === ACCOUNT_TOKEN_VALUE_STATUS_MASKED_PENDING
|
||||
@@ -306,6 +341,62 @@ export async function syncTokensFromUpstream(accountId: number, upstreamTokens:
|
||||
continue;
|
||||
}
|
||||
|
||||
const matchingReadyByMaskedValue = nextValueStatus === ACCOUNT_TOKEN_VALUE_STATUS_MASKED_PENDING
|
||||
? existing.filter((row) => (
|
||||
resolveAccountTokenValueStatus(row) === ACCOUNT_TOKEN_VALUE_STATUS_READY
|
||||
&& matchesMaskedTokenValue(row.token, tokenValue)
|
||||
&& row.name === tokenName
|
||||
&& sameTokenGroup(row.tokenGroup, row.name, tokenGroup, tokenName)
|
||||
))
|
||||
: [];
|
||||
const readyMaskedMatch = matchingReadyByMaskedValue.length === 1
|
||||
? matchingReadyByMaskedValue[0]
|
||||
: null;
|
||||
if (readyMaskedMatch) {
|
||||
const staleMaskedPlaceholders = existing.filter((row) => (
|
||||
row.id !== readyMaskedMatch.id
|
||||
&& resolveAccountTokenValueStatus(row) === ACCOUNT_TOKEN_VALUE_STATUS_MASKED_PENDING
|
||||
&& matchesMaskedTokenValue(row.token, tokenValue)
|
||||
&& row.name === tokenName
|
||||
&& sameTokenGroup(row.tokenGroup, row.name, tokenGroup, tokenName)
|
||||
));
|
||||
|
||||
await db.update(schema.accountTokens)
|
||||
.set({
|
||||
name: tokenName,
|
||||
tokenGroup,
|
||||
valueStatus: ACCOUNT_TOKEN_VALUE_STATUS_READY,
|
||||
source: 'sync',
|
||||
enabled,
|
||||
updatedAt: now,
|
||||
})
|
||||
.where(eq(schema.accountTokens.id, readyMaskedMatch.id))
|
||||
.run();
|
||||
readyMaskedMatch.name = tokenName;
|
||||
readyMaskedMatch.tokenGroup = tokenGroup;
|
||||
readyMaskedMatch.valueStatus = ACCOUNT_TOKEN_VALUE_STATUS_READY;
|
||||
readyMaskedMatch.enabled = enabled;
|
||||
readyMaskedMatch.source = 'sync';
|
||||
readyMaskedMatch.updatedAt = now;
|
||||
|
||||
if (staleMaskedPlaceholders.length > 0) {
|
||||
for (const placeholder of staleMaskedPlaceholders) {
|
||||
await db.delete(schema.accountTokens)
|
||||
.where(eq(schema.accountTokens.id, placeholder.id))
|
||||
.run();
|
||||
}
|
||||
for (const placeholder of staleMaskedPlaceholders) {
|
||||
const placeholderIndex = existing.findIndex((row) => row.id === placeholder.id);
|
||||
if (placeholderIndex >= 0) {
|
||||
existing.splice(placeholderIndex, 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
updated++;
|
||||
continue;
|
||||
}
|
||||
|
||||
const matchingPlaceholder = existing.find((row) => (
|
||||
isMaskedPendingAccountToken(row)
|
||||
&& row.name === tokenName
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest';
|
||||
import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
import { once } from 'node:events';
|
||||
import { mkdtempSync } from 'node:fs';
|
||||
import { createServer } from 'node:http';
|
||||
import { connect as connectSocket } from 'node:net';
|
||||
import { tmpdir } from 'node:os';
|
||||
import { join } from 'node:path';
|
||||
import { sql } from 'drizzle-orm';
|
||||
import { Headers } from 'undici';
|
||||
import { SocksClient } from 'socks';
|
||||
import { Headers, fetch } from 'undici';
|
||||
|
||||
type DbModule = typeof import('../db/index.js');
|
||||
|
||||
@@ -71,6 +75,72 @@ describe('siteProxy', () => {
|
||||
expect('dispatcher' in requestInit).toBe(true);
|
||||
});
|
||||
|
||||
it('injects a working dispatcher for socks5 system proxies', async () => {
|
||||
const upstreamServer = createServer((_request, response) => {
|
||||
response.writeHead(200, { 'content-type': 'application/json' });
|
||||
response.end(JSON.stringify({ ok: true }));
|
||||
});
|
||||
upstreamServer.listen(0, '127.0.0.1');
|
||||
await once(upstreamServer, 'listening');
|
||||
const upstreamAddress = upstreamServer.address();
|
||||
if (!upstreamAddress || typeof upstreamAddress === 'string') {
|
||||
throw new Error('Failed to determine upstream server address');
|
||||
}
|
||||
const requestUrl = `http://proxy-site.example.com:${upstreamAddress.port}/v1/chat/completions`;
|
||||
|
||||
await db.insert(schema.settings).values({
|
||||
key: 'system_proxy_url',
|
||||
value: JSON.stringify('socks5h://127.0.0.1:1080'),
|
||||
}).run();
|
||||
await db.run(sql`
|
||||
INSERT INTO sites (name, url, platform, use_system_proxy)
|
||||
VALUES ('proxy-site', ${`http://proxy-site.example.com:${upstreamAddress.port}`}, 'new-api', 1)
|
||||
`);
|
||||
|
||||
const createConnectionSpy = vi.spyOn(SocksClient, 'createConnection').mockImplementation(async () => {
|
||||
const socket = connectSocket(upstreamAddress.port, '127.0.0.1');
|
||||
await once(socket, 'connect');
|
||||
return { socket } as Awaited<ReturnType<typeof SocksClient.createConnection>>;
|
||||
});
|
||||
|
||||
try {
|
||||
const { withSiteProxyRequestInit } = await import('./siteProxy.js');
|
||||
const requestInit = await withSiteProxyRequestInit(requestUrl, {
|
||||
method: 'GET',
|
||||
});
|
||||
|
||||
expect('dispatcher' in requestInit).toBe(true);
|
||||
|
||||
const response = await fetch(requestUrl, requestInit);
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(createConnectionSpy).toHaveBeenCalledTimes(1);
|
||||
expect(createConnectionSpy).toHaveBeenCalledWith(expect.objectContaining({
|
||||
command: 'connect',
|
||||
proxy: expect.objectContaining({
|
||||
host: '127.0.0.1',
|
||||
port: 1080,
|
||||
type: 5,
|
||||
}),
|
||||
destination: expect.objectContaining({
|
||||
host: 'proxy-site.example.com',
|
||||
port: upstreamAddress.port,
|
||||
}),
|
||||
}));
|
||||
} finally {
|
||||
createConnectionSpy.mockRestore();
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
upstreamServer.close((error) => {
|
||||
if (error) {
|
||||
reject(error);
|
||||
return;
|
||||
}
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
it('merges site custom headers by matched request url and keeps explicit headers authoritative', async () => {
|
||||
await db.insert(schema.sites).values({
|
||||
name: 'headers-site',
|
||||
|
||||
@@ -2,8 +2,12 @@ import { AsyncLocalStorage } from 'node:async_hooks';
|
||||
import { db, schema } from '../db/index.js';
|
||||
import { eq } from 'drizzle-orm';
|
||||
import { config } from '../config.js';
|
||||
import { lookup as dnsLookup } from 'node:dns';
|
||||
import { isIP, type Socket } from 'node:net';
|
||||
import { connect as tlsConnect, type TLSSocket } from 'node:tls';
|
||||
import { SocksClient } from 'socks';
|
||||
import type { Dispatcher, RequestInit as UndiciRequestInit } from 'undici';
|
||||
import { ProxyAgent } from 'undici';
|
||||
import { Agent as UndiciAgent, ProxyAgent } from 'undici';
|
||||
import { mergeHeadersWithSiteCustomHeaders } from './siteCustomHeaders.js';
|
||||
import { getProxyUrlFromExtraConfig } from './accountExtraConfig.js';
|
||||
|
||||
@@ -17,6 +21,15 @@ const SUPPORTED_PROXY_PROTOCOLS = new Set([
|
||||
'socks5:',
|
||||
'socks5h:',
|
||||
]);
|
||||
const SOCKS_PROXY_PROTOCOLS = new Set([
|
||||
'socks:',
|
||||
'socks4:',
|
||||
'socks4a:',
|
||||
'socks5:',
|
||||
'socks5h:',
|
||||
]);
|
||||
const DEFAULT_PROXY_CONNECT_TIMEOUT_MS = 10_000;
|
||||
const DEFAULT_PROXY_KEEPALIVE_INITIAL_DELAY_MS = 60_000;
|
||||
|
||||
type SiteProxyRow = {
|
||||
siteUrl: string;
|
||||
@@ -58,6 +71,27 @@ export function withAccountProxyOverride<T>(
|
||||
return accountProxyOverride.run(normalized, fn);
|
||||
}
|
||||
|
||||
type ParsedSocksProxyConfig = {
|
||||
shouldLookup: boolean;
|
||||
proxy: {
|
||||
host: string;
|
||||
port: number;
|
||||
type: 4 | 5;
|
||||
userId?: string;
|
||||
password?: string;
|
||||
};
|
||||
};
|
||||
|
||||
type UndiciConnectOptions = {
|
||||
hostname: string;
|
||||
host?: string;
|
||||
protocol: string;
|
||||
port: string;
|
||||
servername?: string;
|
||||
localAddress?: string | null;
|
||||
httpSocket?: Socket;
|
||||
};
|
||||
|
||||
function normalizeSiteUrl(value: string): string {
|
||||
const trimmed = (value || '').trim();
|
||||
if (!trimmed) return '';
|
||||
@@ -128,7 +162,10 @@ function getDispatcherByProxyUrl(proxyUrl: string): Dispatcher | undefined {
|
||||
if (cached) return cached;
|
||||
|
||||
try {
|
||||
const dispatcher = new ProxyAgent(normalized);
|
||||
const parsedProxyUrl = new URL(normalized);
|
||||
const dispatcher = SOCKS_PROXY_PROTOCOLS.has(parsedProxyUrl.protocol.toLowerCase())
|
||||
? createSocksDispatcher(parsedProxyUrl)
|
||||
: new ProxyAgent(normalized);
|
||||
dispatcherCache.set(normalized, dispatcher);
|
||||
return dispatcher;
|
||||
} catch {
|
||||
@@ -136,6 +173,131 @@ function getDispatcherByProxyUrl(proxyUrl: string): Dispatcher | undefined {
|
||||
}
|
||||
}
|
||||
|
||||
function parseSocksProxyUrl(proxyUrl: URL): ParsedSocksProxyConfig {
|
||||
let shouldLookup = false;
|
||||
let type: 4 | 5 = 5;
|
||||
|
||||
switch (proxyUrl.protocol.toLowerCase()) {
|
||||
case 'socks4:':
|
||||
shouldLookup = true;
|
||||
type = 4;
|
||||
break;
|
||||
case 'socks4a:':
|
||||
type = 4;
|
||||
break;
|
||||
case 'socks5:':
|
||||
shouldLookup = true;
|
||||
type = 5;
|
||||
break;
|
||||
case 'socks:':
|
||||
case 'socks5h:':
|
||||
type = 5;
|
||||
break;
|
||||
default:
|
||||
throw new TypeError(`Unsupported SOCKS proxy protocol: ${proxyUrl.protocol}`);
|
||||
}
|
||||
|
||||
const proxy: ParsedSocksProxyConfig['proxy'] = {
|
||||
host: proxyUrl.hostname,
|
||||
port: Number.parseInt(proxyUrl.port, 10) || 1080,
|
||||
type,
|
||||
};
|
||||
|
||||
if (proxyUrl.username) {
|
||||
proxy.userId = decodeURIComponent(proxyUrl.username);
|
||||
}
|
||||
if (proxyUrl.password) {
|
||||
proxy.password = decodeURIComponent(proxyUrl.password);
|
||||
}
|
||||
|
||||
return { shouldLookup, proxy };
|
||||
}
|
||||
|
||||
function applySocketDefaults(socket: Socket | TLSSocket) {
|
||||
socket.setNoDelay(true);
|
||||
socket.setKeepAlive(true, DEFAULT_PROXY_KEEPALIVE_INITIAL_DELAY_MS);
|
||||
}
|
||||
|
||||
async function resolveSocksDestinationHost(hostname: string): Promise<string> {
|
||||
return new Promise((resolve, reject) => {
|
||||
dnsLookup(hostname, {}, (error, address) => {
|
||||
if (error) {
|
||||
reject(error);
|
||||
return;
|
||||
}
|
||||
resolve(address);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async function createSocksSocket(
|
||||
connectOptions: UndiciConnectOptions,
|
||||
socksProxy: ParsedSocksProxyConfig,
|
||||
): Promise<Socket | TLSSocket> {
|
||||
if (!connectOptions.hostname) {
|
||||
throw new Error('Missing hostname for SOCKS proxy request');
|
||||
}
|
||||
|
||||
const destinationHost = socksProxy.shouldLookup
|
||||
? await resolveSocksDestinationHost(connectOptions.hostname)
|
||||
: connectOptions.hostname;
|
||||
const destinationPort = Number.parseInt(connectOptions.port, 10)
|
||||
|| (connectOptions.protocol === 'https:' ? 443 : 80);
|
||||
|
||||
const { socket } = await SocksClient.createConnection({
|
||||
proxy: socksProxy.proxy,
|
||||
destination: {
|
||||
host: destinationHost,
|
||||
port: destinationPort,
|
||||
},
|
||||
command: 'connect',
|
||||
timeout: DEFAULT_PROXY_CONNECT_TIMEOUT_MS,
|
||||
socket_options: connectOptions.localAddress
|
||||
? { localAddress: connectOptions.localAddress } as any
|
||||
: undefined,
|
||||
});
|
||||
applySocketDefaults(socket);
|
||||
|
||||
if (connectOptions.protocol !== 'https:') {
|
||||
return socket;
|
||||
}
|
||||
|
||||
return await new Promise<TLSSocket>((resolve, reject) => {
|
||||
const tlsSocket = tlsConnect({
|
||||
socket,
|
||||
host: connectOptions.hostname,
|
||||
servername: connectOptions.servername || (!isIP(connectOptions.hostname) ? connectOptions.hostname : undefined),
|
||||
ALPNProtocols: ['http/1.1'],
|
||||
});
|
||||
|
||||
const cleanup = (error: Error) => {
|
||||
socket.destroy();
|
||||
tlsSocket.destroy();
|
||||
reject(error);
|
||||
};
|
||||
|
||||
tlsSocket.once('secureConnect', () => {
|
||||
tlsSocket.off('error', cleanup);
|
||||
applySocketDefaults(tlsSocket);
|
||||
resolve(tlsSocket);
|
||||
});
|
||||
tlsSocket.once('error', cleanup);
|
||||
});
|
||||
}
|
||||
|
||||
function createSocksDispatcher(proxyUrl: URL): Dispatcher {
|
||||
const socksProxy = parseSocksProxyUrl(proxyUrl);
|
||||
return new UndiciAgent({
|
||||
connect: (connectOptions, callback) => {
|
||||
void createSocksSocket(connectOptions, socksProxy)
|
||||
.then((socket) => callback(null, socket))
|
||||
.catch((error) => {
|
||||
callback(error instanceof Error ? error : new Error(String(error)), null as any);
|
||||
});
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
export function normalizeSiteProxyUrl(input: unknown): string | null {
|
||||
if (typeof input !== 'string') return null;
|
||||
const trimmed = input.trim();
|
||||
|
||||
Reference in New Issue
Block a user