Skip to content

Commit

Permalink
🐛 fix: support request headers for chat (#4934)
Browse files Browse the repository at this point in the history
  • Loading branch information
arvinxx authored Dec 8, 2024
1 parent a1b5c26 commit 8cdc062
Show file tree
Hide file tree
Showing 4 changed files with 96 additions and 1 deletion.
7 changes: 7 additions & 0 deletions src/libs/agent-runtime/types/chat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,14 @@ export interface ChatStreamPayload {

export interface ChatCompetitionOptions {
callback?: ChatStreamCallbacks;
/**
* response headers
*/
headers?: Record<string, any>;
/**
* send the request to the ai api endpoint
*/
requestHeaders?: Record<string, any>;
signal?: AbortSignal;
/**
* userId for the chat completion
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -211,7 +211,7 @@ export const LobeOpenAICompatibleFactory = <T extends Record<string, any> = any>
},
{
// https://github.com/lobehub/lobe-chat/pull/318
headers: { Accept: '*/*' },
headers: { Accept: '*/*', ...options?.requestHeaders },
signal: options?.signal,
},
);
Expand Down
54 changes: 54 additions & 0 deletions src/utils/clientIP.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
import { describe, expect, it } from 'vitest';

import { getClientIP } from './clientIP';

describe('getClientIP', () => {
// Helper function to create Headers object
const createHeaders = (entries: [string, string][]) => {
return new Headers(entries);
};

it('should return null when no IP headers are present', () => {
const headers = createHeaders([]);
expect(getClientIP(headers)).toBe('');
});

it('should handle Cloudflare IP header', () => {
const headers = createHeaders([['cf-connecting-ip', '1.2.3.4']]);
expect(getClientIP(headers)).toBe('1.2.3.4');
});

it('should handle x-forwarded-for with single IP', () => {
const headers = createHeaders([['x-forwarded-for', '5.6.7.8']]);
expect(getClientIP(headers)).toBe('5.6.7.8');
});

it('should handle x-forwarded-for with multiple IPs and return the first one', () => {
const headers = createHeaders([['x-forwarded-for', '9.10.11.12, 13.14.15.16, 17.18.19.20']]);
expect(getClientIP(headers)).toBe('9.10.11.12');
});

it('should handle x-real-ip header', () => {
const headers = createHeaders([['x-real-ip', '21.22.23.24']]);
expect(getClientIP(headers)).toBe('21.22.23.24');
});

it('should trim whitespace from IP addresses', () => {
const headers = createHeaders([['x-client-ip', ' 25.26.27.28 ']]);
expect(getClientIP(headers)).toBe('25.26.27.28');
});

it('should respect header priority order', () => {
const headers = createHeaders([
['x-forwarded-for', '1.1.1.1'],
['cf-connecting-ip', '2.2.2.2'], // Should take precedence
['x-real-ip', '3.3.3.3'],
]);
expect(getClientIP(headers)).toBe('2.2.2.2');
});

it('should handle empty x-forwarded-for value', () => {
const headers = createHeaders([['x-forwarded-for', '']]);
expect(getClientIP(headers)).toBe('');
});
});
34 changes: 34 additions & 0 deletions src/utils/clientIP.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
/**
* 获取客户端 IP
* @param headers HTTP 请求头
*/
export const getClientIP = (headers: Headers): string => {
// 按优先级顺序检查各种 IP 头
const ipHeaders = [
'cf-connecting-ip', // Cloudflare
'x-real-ip', // Nginx proxy
'x-forwarded-for', // 标准代理头
'x-client-ip', // Apache
'true-client-ip', // Akamai and Cloudflare
'x-cluster-client-ip', // 负载均衡
'forwarded', // RFC 7239
'fastly-client-ip', // Fastly CDN
'x-forwarded', // General forward
'x-original-forwarded-for', // Original forwarded
];

for (const header of ipHeaders) {
const value = headers.get(header);
if (!value) continue;

// 处理可能包含多个 IP 的情况(比如 x-forwarded-for)
if (header.toLowerCase() === 'x-forwarded-for') {
const firstIP = value.split(',')[0].trim();
if (firstIP) return firstIP;
}

return value.trim();
}

return '';
};

0 comments on commit 8cdc062

Please sign in to comment.