From 088ce266ba460c2c62f6f036eee45a0e3e100430 Mon Sep 17 00:00:00 2001 From: shaw Date: Sun, 20 Jul 2025 15:58:00 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0API=20Key=E6=97=B6?= =?UTF-8?q?=E9=97=B4=E7=AA=97=E5=8F=A3=E9=99=90=E6=B5=81=E5=8A=9F=E8=83=BD?= =?UTF-8?q?=E5=B9=B6=E7=A7=BB=E9=99=A4=E7=B4=AF=E8=AE=A1=E6=80=BB=E9=87=8F?= =?UTF-8?q?=E9=99=90=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增时间窗口限流功能，支持按分钟设置时间窗口 - 支持在时间窗口内限制请求次数和Token使用量 - 移除原有的累计总量限制，只保留时间窗口限制 - Token统计包含所有4种类型：输入、输出、缓存创建、缓存读取 - 前端UI优化，明确显示限流参数的作用范围 - 限流触发时提供友好的错误提示和重置时间 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- src/middleware/auth.js | 91 ++++++++++++++++++++++++++++++ src/routes/admin.js | 28 ++++++++- src/routes/api.js | 17 ++++++ src/services/apiKeyService.js | 20 ++++--- src/services/claudeRelayService.js | 4 +- web/admin/app.js | 14 ++++- web/admin/index.html | 66 ++++++++++++++++++++-- 7 files changed, 224 insertions(+), 16 deletions(-) diff --git a/src/middleware/auth.js b/src/middleware/auth.js index fdb448d6..ee39577e 100644 --- a/src/middleware/auth.js +++ b/src/middleware/auth.js @@ -103,6 +103,95 @@ const authenticateApiKey = async (req, res, next) => { }; } + // 检查时间窗口限流 + const rateLimitWindow = validation.keyData.rateLimitWindow || 0; + const rateLimitRequests = validation.keyData.rateLimitRequests || 0; + + if (rateLimitWindow > 0 && (rateLimitRequests > 0 || validation.keyData.tokenLimit > 0)) { + const windowStartKey = `rate_limit:window_start:${validation.keyData.id}`; + const requestCountKey = `rate_limit:requests:${validation.keyData.id}`; + const tokenCountKey = `rate_limit:tokens:${validation.keyData.id}`; + + const now = Date.now(); + const windowDuration = rateLimitWindow * 60 * 1000; // 转换为毫秒 + + // 获取窗口开始时间 + let windowStart = await redis.getClient().get(windowStartKey); + + if (!windowStart) { + // 第一次请求，设置窗口开始时间 + await redis.getClient().set(windowStartKey, now, 'PX', windowDuration); + await redis.getClient().set(requestCountKey, 0, 'PX', windowDuration); + await redis.getClient().set(tokenCountKey, 0, 'PX', windowDuration); + windowStart = now; + } else { + windowStart = parseInt(windowStart); + + // 检查窗口是否已过期 + if (now - windowStart >= windowDuration) { + // 窗口已过期，重置 + await redis.getClient().set(windowStartKey, now, 'PX', windowDuration); + await redis.getClient().set(requestCountKey, 0, 'PX', windowDuration); + await redis.getClient().set(tokenCountKey, 0, 'PX', windowDuration); + windowStart = now; + } + } + + // 获取当前计数 + const currentRequests = parseInt(await redis.getClient().get(requestCountKey) || '0'); + const currentTokens = parseInt(await redis.getClient().get(tokenCountKey) || '0'); + + // 检查请求次数限制 + if (rateLimitRequests > 0 && currentRequests >= rateLimitRequests) { + const resetTime = new Date(windowStart + windowDuration); + const remainingMinutes = Math.ceil((resetTime - now) / 60000); + + logger.security(`🚦 Rate limit exceeded (requests) for key: ${validation.keyData.id} (${validation.keyData.name}), requests: ${currentRequests}/${rateLimitRequests}`); + + return res.status(429).json({ + error: 'Rate limit exceeded', + message: `已达到请求次数限制 (${rateLimitRequests} 次)，将在 ${remainingMinutes} 分钟后重置`, + currentRequests, + requestLimit: rateLimitRequests, + resetAt: resetTime.toISOString(), + remainingMinutes + }); + } + + // 检查Token使用量限制 + const tokenLimit = parseInt(validation.keyData.tokenLimit); + if (tokenLimit > 0 && currentTokens >= tokenLimit) { + const resetTime = new Date(windowStart + windowDuration); + const remainingMinutes = Math.ceil((resetTime - now) / 60000); + + logger.security(`🚦 Rate limit exceeded (tokens) for key: ${validation.keyData.id} (${validation.keyData.name}), tokens: ${currentTokens}/${tokenLimit}`); + + return res.status(429).json({ + error: 'Rate limit exceeded', + message: `已达到 Token 使用限制 (${tokenLimit} tokens)，将在 ${remainingMinutes} 分钟后重置`, + currentTokens, + tokenLimit, + resetAt: resetTime.toISOString(), + remainingMinutes + }); + } + + // 增加请求计数 + await redis.getClient().incr(requestCountKey); + + // 存储限流信息到请求对象 + req.rateLimitInfo = { + windowStart, + windowDuration, + requestCountKey, + tokenCountKey, + currentRequests: currentRequests + 1, + currentTokens, + rateLimitRequests, + tokenLimit + }; + } + // 将验证信息添加到请求对象（只包含必要信息） req.apiKey = { id: validation.keyData.id, @@ -110,6 +199,8 @@ const authenticateApiKey = async (req, res, next) => { tokenLimit: validation.keyData.tokenLimit, claudeAccountId: validation.keyData.claudeAccountId, concurrencyLimit: validation.keyData.concurrencyLimit, + rateLimitWindow: validation.keyData.rateLimitWindow, + rateLimitRequests: validation.keyData.rateLimitRequests, enableModelRestriction: validation.keyData.enableModelRestriction, restrictedModels: validation.keyData.restrictedModels }; diff --git a/src/routes/admin.js b/src/routes/admin.js index 2978b28b..c2f993b5 100644 --- a/src/routes/admin.js +++ b/src/routes/admin.js @@ -33,6 +33,8 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => { expiresAt, claudeAccountId, concurrencyLimit, + rateLimitWindow, + rateLimitRequests, enableModelRestriction, restrictedModels } = req.body; @@ -58,6 +60,14 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => { if (concurrencyLimit !== undefined && concurrencyLimit !== null && concurrencyLimit !== '' && (!Number.isInteger(Number(concurrencyLimit)) || Number(concurrencyLimit) < 0)) { return res.status(400).json({ error: 'Concurrency limit must be a non-negative integer' }); } + + if (rateLimitWindow !== undefined && rateLimitWindow !== null && rateLimitWindow !== '' && (!Number.isInteger(Number(rateLimitWindow)) || Number(rateLimitWindow) < 1)) { + return res.status(400).json({ error: 'Rate limit window must be a positive integer (minutes)' }); + } + + if (rateLimitRequests !== undefined && rateLimitRequests !== null && rateLimitRequests !== '' && (!Number.isInteger(Number(rateLimitRequests)) || Number(rateLimitRequests) < 1)) { + return res.status(400).json({ error: 'Rate limit requests must be a positive integer' }); + } // 验证模型限制字段 if (enableModelRestriction !== undefined && typeof enableModelRestriction !== 'boolean') { @@ -75,6 +85,8 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => { expiresAt, claudeAccountId, concurrencyLimit, + rateLimitWindow, + rateLimitRequests, enableModelRestriction, restrictedModels }); @@ -91,7 +103,7 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => { router.put('/api-keys/:keyId', authenticateAdmin, async (req, res) => { try { const { keyId } = req.params; - const { tokenLimit, concurrencyLimit, claudeAccountId, enableModelRestriction, restrictedModels } = req.body; + const { tokenLimit, concurrencyLimit, rateLimitWindow, rateLimitRequests, claudeAccountId, enableModelRestriction, restrictedModels } = req.body; // 只允许更新指定字段 const updates = {}; @@ -109,6 +121,20 @@ router.put('/api-keys/:keyId', authenticateAdmin, async (req, res) => { } updates.concurrencyLimit = Number(concurrencyLimit); } + + if (rateLimitWindow !== undefined && rateLimitWindow !== null && rateLimitWindow !== '') { + if (!Number.isInteger(Number(rateLimitWindow)) || Number(rateLimitWindow) < 0) { + return res.status(400).json({ error: 'Rate limit window must be a non-negative integer (minutes)' }); + } + updates.rateLimitWindow = Number(rateLimitWindow); + } + + if (rateLimitRequests !== undefined && rateLimitRequests !== null && rateLimitRequests !== '') { + if (!Number.isInteger(Number(rateLimitRequests)) || Number(rateLimitRequests) < 0) { + return res.status(400).json({ error: 'Rate limit requests must be a non-negative integer' }); + } + updates.rateLimitRequests = Number(rateLimitRequests); + } if (claudeAccountId !== undefined) { // 空字符串表示解绑，null或空字符串都设置为空字符串 diff --git a/src/routes/api.js b/src/routes/api.js index 140a7d12..ed70ed2a 100644 --- a/src/routes/api.js +++ b/src/routes/api.js @@ -3,6 +3,7 @@ const claudeRelayService = require('../services/claudeRelayService'); const apiKeyService = require('../services/apiKeyService'); const { authenticateApiKey } = require('../middleware/auth'); const logger = require('../utils/logger'); +const redis = require('../models/redis'); const router = express.Router(); @@ -66,6 +67,15 @@ router.post('/v1/messages', authenticateApiKey, async (req, res) => { logger.error('❌ Failed to record stream usage:', error); }); + // 更新时间窗口内的token计数 + if (req.rateLimitInfo) { + const totalTokens = inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens; + redis.getClient().incrby(req.rateLimitInfo.tokenCountKey, totalTokens).catch(error => { + logger.error('❌ Failed to update rate limit token count:', error); + }); + logger.api(`📊 Updated rate limit token count: +${totalTokens} tokens`); + } + usageDataCaptured = true; logger.api(`📊 Stream usage recorded (real) - Model: ${model}, Input: ${inputTokens}, Output: ${outputTokens}, Cache Create: ${cacheCreateTokens}, Cache Read: ${cacheReadTokens}, Total: ${inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens} tokens`); } else { @@ -122,6 +132,13 @@ router.post('/v1/messages', authenticateApiKey, async (req, res) => { // 记录真实的token使用量（包含模型信息和所有4种token） await apiKeyService.recordUsage(req.apiKey.id, inputTokens, outputTokens, cacheCreateTokens, cacheReadTokens, model); + // 更新时间窗口内的token计数 + if (req.rateLimitInfo) { + const totalTokens = inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens; + await redis.getClient().incrby(req.rateLimitInfo.tokenCountKey, totalTokens); + logger.api(`📊 Updated rate limit token count: +${totalTokens} tokens`); + } + usageRecorded = true; logger.api(`📊 Non-stream usage recorded (real) - Model: ${model}, Input: ${inputTokens}, Output: ${outputTokens}, Cache Create: ${cacheCreateTokens}, Cache Read: ${cacheReadTokens}, Total: ${inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens} tokens`); } else { diff --git a/src/services/apiKeyService.js b/src/services/apiKeyService.js index d696a623..ca38d193 100644 --- a/src/services/apiKeyService.js +++ b/src/services/apiKeyService.js @@ -19,6 +19,8 @@ class ApiKeyService { claudeAccountId = null, isActive = true, concurrencyLimit = 0, + rateLimitWindow = null, + rateLimitRequests = null, enableModelRestriction = false, restrictedModels = [] } = options; @@ -35,6 +37,8 @@ class ApiKeyService { apiKey: hashedKey, tokenLimit: String(tokenLimit ?? 0), concurrencyLimit: String(concurrencyLimit ?? 0), + rateLimitWindow: String(rateLimitWindow ?? 0), + rateLimitRequests: String(rateLimitRequests ?? 0), isActive: String(isActive), claudeAccountId: claudeAccountId || '', enableModelRestriction: String(enableModelRestriction), @@ -57,6 +61,8 @@ class ApiKeyService { description: keyData.description, tokenLimit: parseInt(keyData.tokenLimit), concurrencyLimit: parseInt(keyData.concurrencyLimit), + rateLimitWindow: parseInt(keyData.rateLimitWindow || 0), + rateLimitRequests: parseInt(keyData.rateLimitRequests || 0), isActive: keyData.isActive === 'true', claudeAccountId: keyData.claudeAccountId, enableModelRestriction: keyData.enableModelRestriction === 'true', @@ -94,14 +100,8 @@ class ApiKeyService { return { valid: false, error: 'API key has expired' }; } - // 检查使用限制 + // 获取使用统计（供返回数据使用） const usage = await redis.getUsageStats(keyData.id); - const tokenLimit = parseInt(keyData.tokenLimit); - - if (tokenLimit > 0 && usage.total.tokens >= tokenLimit) { - return { valid: false, error: 'Token limit exceeded' }; - } - // 更新最后使用时间（优化：只在实际API调用时更新，而不是验证时） // 注意：lastUsedAt的更新已移至recordUsage方法中 @@ -124,6 +124,8 @@ class ApiKeyService { claudeAccountId: keyData.claudeAccountId, tokenLimit: parseInt(keyData.tokenLimit), concurrencyLimit: parseInt(keyData.concurrencyLimit || 0), + rateLimitWindow: parseInt(keyData.rateLimitWindow || 0), + rateLimitRequests: parseInt(keyData.rateLimitRequests || 0), enableModelRestriction: keyData.enableModelRestriction === 'true', restrictedModels: restrictedModels, usage @@ -145,6 +147,8 @@ class ApiKeyService { key.usage = await redis.getUsageStats(key.id); key.tokenLimit = parseInt(key.tokenLimit); key.concurrencyLimit = parseInt(key.concurrencyLimit || 0); + key.rateLimitWindow = parseInt(key.rateLimitWindow || 0); + key.rateLimitRequests = parseInt(key.rateLimitRequests || 0); key.currentConcurrency = await redis.getConcurrency(key.id); key.isActive = key.isActive === 'true'; key.enableModelRestriction = key.enableModelRestriction === 'true'; @@ -172,7 +176,7 @@ class ApiKeyService { } // 允许更新的字段 - const allowedUpdates = ['name', 'description', 'tokenLimit', 'concurrencyLimit', 'isActive', 'claudeAccountId', 'expiresAt', 'enableModelRestriction', 'restrictedModels']; + const allowedUpdates = ['name', 'description', 'tokenLimit', 'concurrencyLimit', 'rateLimitWindow', 'rateLimitRequests', 'isActive', 'claudeAccountId', 'expiresAt', 'enableModelRestriction', 'restrictedModels']; const updatedData = { ...keyData }; for (const [field, value] of Object.entries(updates)) { diff --git a/src/services/claudeRelayService.js b/src/services/claudeRelayService.js index 4e5221ec..8d62a96b 100644 --- a/src/services/claudeRelayService.js +++ b/src/services/claudeRelayService.js @@ -23,7 +23,7 @@ class ClaudeRelayService { try { // 调试日志：查看API Key数据 - logger.info(`🔍 API Key data received:`, { + logger.info('🔍 API Key data received:', { apiKeyName: apiKeyData.name, enableModelRestriction: apiKeyData.enableModelRestriction, restrictedModels: apiKeyData.restrictedModels, @@ -448,7 +448,7 @@ class ClaudeRelayService { async relayStreamRequestWithUsageCapture(requestBody, apiKeyData, responseStream, clientHeaders, usageCallback) { try { // 调试日志：查看API Key数据（流式请求） - logger.info(`🔍 [Stream] API Key data received:`, { + logger.info('🔍 [Stream] API Key data received:', { apiKeyName: apiKeyData.name, enableModelRestriction: apiKeyData.enableModelRestriction, restrictedModels: apiKeyData.restrictedModels, diff --git a/web/admin/app.js b/web/admin/app.js index 39de13f9..014703a2 100644 --- a/web/admin/app.js +++ b/web/admin/app.js @@ -117,6 +117,8 @@ const app = createApp({ tokenLimit: '', description: '', concurrencyLimit: '', + rateLimitWindow: '', + rateLimitRequests: '', claudeAccountId: '', enableModelRestriction: false, restrictedModels: [], @@ -158,6 +160,8 @@ const app = createApp({ name: '', tokenLimit: '', concurrencyLimit: '', + rateLimitWindow: '', + rateLimitRequests: '', claudeAccountId: '', enableModelRestriction: false, restrictedModels: [], @@ -1211,6 +1215,8 @@ const app = createApp({ tokenLimit: this.apiKeyForm.tokenLimit && this.apiKeyForm.tokenLimit.trim() ? parseInt(this.apiKeyForm.tokenLimit) : null, description: this.apiKeyForm.description || '', concurrencyLimit: this.apiKeyForm.concurrencyLimit && this.apiKeyForm.concurrencyLimit.trim() ? parseInt(this.apiKeyForm.concurrencyLimit) : 0, + rateLimitWindow: this.apiKeyForm.rateLimitWindow && this.apiKeyForm.rateLimitWindow.trim() ? parseInt(this.apiKeyForm.rateLimitWindow) : null, + rateLimitRequests: this.apiKeyForm.rateLimitRequests && this.apiKeyForm.rateLimitRequests.trim() ? parseInt(this.apiKeyForm.rateLimitRequests) : null, claudeAccountId: this.apiKeyForm.claudeAccountId || null, enableModelRestriction: this.apiKeyForm.enableModelRestriction, restrictedModels: this.apiKeyForm.restrictedModels @@ -1231,7 +1237,7 @@ const app = createApp({ // 关闭创建弹窗并清理表单 this.showCreateApiKeyModal = false; - this.apiKeyForm = { name: '', tokenLimit: '', description: '', concurrencyLimit: '', claudeAccountId: '', enableModelRestriction: false, restrictedModels: [], modelInput: '' }; + this.apiKeyForm = { name: '', tokenLimit: '', description: '', concurrencyLimit: '', rateLimitWindow: '', rateLimitRequests: '', claudeAccountId: '', enableModelRestriction: false, restrictedModels: [], modelInput: '' }; // 重新加载API Keys列表 await this.loadApiKeys(); @@ -1275,6 +1281,8 @@ const app = createApp({ name: key.name, tokenLimit: key.tokenLimit || '', concurrencyLimit: key.concurrencyLimit || '', + rateLimitWindow: key.rateLimitWindow || '', + rateLimitRequests: key.rateLimitRequests || '', claudeAccountId: key.claudeAccountId || '', enableModelRestriction: key.enableModelRestriction || false, restrictedModels: key.restrictedModels ? [...key.restrictedModels] : [], @@ -1290,6 +1298,8 @@ const app = createApp({ name: '', tokenLimit: '', concurrencyLimit: '', + rateLimitWindow: '', + rateLimitRequests: '', claudeAccountId: '', enableModelRestriction: false, restrictedModels: [], @@ -1309,6 +1319,8 @@ const app = createApp({ body: JSON.stringify({ tokenLimit: this.editApiKeyForm.tokenLimit && this.editApiKeyForm.tokenLimit.toString().trim() !== '' ? parseInt(this.editApiKeyForm.tokenLimit) : 0, concurrencyLimit: this.editApiKeyForm.concurrencyLimit && this.editApiKeyForm.concurrencyLimit.toString().trim() !== '' ? parseInt(this.editApiKeyForm.concurrencyLimit) : 0, + rateLimitWindow: this.editApiKeyForm.rateLimitWindow && this.editApiKeyForm.rateLimitWindow.toString().trim() !== '' ? parseInt(this.editApiKeyForm.rateLimitWindow) : 0, + rateLimitRequests: this.editApiKeyForm.rateLimitRequests && this.editApiKeyForm.rateLimitRequests.toString().trim() !== '' ? parseInt(this.editApiKeyForm.rateLimitRequests) : 0, claudeAccountId: this.editApiKeyForm.claudeAccountId || null, enableModelRestriction: this.editApiKeyForm.enableModelRestriction, restrictedModels: this.editApiKeyForm.restrictedModels diff --git a/web/admin/index.html b/web/admin/index.html index 51a6c17c..6b1c7d43 100644 --- a/web/admin/index.html +++ b/web/admin/index.html @@ -534,6 +534,16 @@ / {{ key.concurrencyLimit }} + +
+ 时间窗口: + {{ key.rateLimitWindow }} 分钟 +
+ +
+ 请求限制: + {{ key.rateLimitRequests }} 次/窗口 +
输入: {{ formatNumber((key.usage && key.usage.total && key.usage.total.inputTokens) || 0) }} @@ -1792,14 +1802,38 @@
- + + +

设置时间窗口（分钟），在此时间内限制请求次数或Token使用量

+
+ +
+ + +

在时间窗口内允许的最大请求次数

+
+ +
+ -

设置此 API Key 的最大 token 使用量

+

设置在时间窗口内的最大 Token 使用量（需同时设置时间窗口）

@@ -1951,7 +1985,31 @@
- + + +

设置时间窗口（分钟），在此时间内限制请求次数或Token使用量

+
+ +
+ + +

在时间窗口内允许的最大请求次数

+
+ +
+ -

设置此 API Key 的最大 token 使用量，0 或留空表示无限制

+

设置在时间窗口内的最大 Token 使用量（需同时设置时间窗口），0 或留空表示无限制