Merge pull request #72 from kevinconan/main

feat: 增加APIKey 客户端限制功能
2026-01-23 00:53:33 +00:00 · 2025-07-26 11:03:06 +08:00
parent ff612685f6 441361f3b4
commit 383f33fa77
7 changed files with 302 additions and 17 deletions
--- a/src/middleware/auth.js
+++ b/src/middleware/auth.js
@@ -2,6 +2,7 @@ const apiKeyService = require('../services/apiKeyService');
 const logger = require('../utils/logger');
 const redis = require('../models/redis');
 const { RateLimiterRedis } = require('rate-limiter-flexible');
+const config = require('../../config/config');

 // 🔑 API Key验证中间件（优化版）
 const authenticateApiKey = async (req, res, next) => {
@@ -42,6 +43,52 @@ const authenticateApiKey = async (req, res, next) => {
      });
    }

+    // 🔒 检查客户端限制
+    if (validation.keyData.enableClientRestriction && validation.keyData.allowedClients?.length > 0) {
+      const userAgent = req.headers['user-agent'] || '';
+      const clientIP = req.ip || req.connection?.remoteAddress || 'unknown';
+      
+      // 记录客户端限制检查开始
+      logger.api(`🔍 Checking client restriction for key: ${validation.keyData.id} (${validation.keyData.name})`);
+      logger.api(`   User-Agent: "${userAgent}"`);
+      logger.api(`   Allowed clients: ${validation.keyData.allowedClients.join(', ')}`);
+      
+      let clientAllowed = false;
+      let matchedClient = null;
+      
+      // 遍历允许的客户端列表
+      for (const allowedClientId of validation.keyData.allowedClients) {
+        // 在预定义客户端列表中查找
+        const predefinedClient = config.clientRestrictions.predefinedClients.find(
+          client => client.id === allowedClientId
+        );
+        
+        if (predefinedClient) {
+          // 使用预定义的正则表达式匹配 User-Agent
+          if (predefinedClient.userAgentPattern.test(userAgent)) {
+            clientAllowed = true;
+            matchedClient = predefinedClient.name;
+            break;
+          }
+        } else if (config.clientRestrictions.allowCustomClients) {
+          // 如果允许自定义客户端，这里可以添加自定义客户端的验证逻辑
+          // 目前暂时跳过自定义客户端
+          continue;
+        }
+      }
+      
+      if (!clientAllowed) {
+        logger.security(`🚫 Client restriction failed for key: ${validation.keyData.id} (${validation.keyData.name}) from ${clientIP}, User-Agent: ${userAgent}`);
+        return res.status(403).json({
+          error: 'Client not allowed',
+          message: 'Your client is not authorized to use this API key',
+          allowedClients: validation.keyData.allowedClients
+        });
+      }
+      
+      logger.api(`✅ Client validated: ${matchedClient} for key: ${validation.keyData.id} (${validation.keyData.name})`);
+      logger.api(`   Matched client: ${matchedClient} with User-Agent: "${userAgent}"`);
+    }

    // 检查并发限制
    const concurrencyLimit = validation.keyData.concurrencyLimit || 0;
@@ -205,12 +252,16 @@ const authenticateApiKey = async (req, res, next) => {
      rateLimitRequests: validation.keyData.rateLimitRequests,
      enableModelRestriction: validation.keyData.enableModelRestriction,
      restrictedModels: validation.keyData.restrictedModels,
+      enableClientRestriction: validation.keyData.enableClientRestriction,
+      allowedClients: validation.keyData.allowedClients,
      usage: validation.keyData.usage
    };
    req.usage = validation.keyData.usage;
    
    const authDuration = Date.now() - startTime;
+    const userAgent = req.headers['user-agent'] || 'No User-Agent';
    logger.api(`🔓 Authenticated request from key: ${validation.keyData.name} (${validation.keyData.id}) in ${authDuration}ms`);
+    logger.api(`   User-Agent: "${userAgent}"`);
    
    next();
  } catch (error) {
--- a/src/routes/admin.js
+++ b/src/routes/admin.js
@@ -12,6 +12,7 @@ const claudeCodeHeadersService = require('../services/claudeCodeHeadersService')
 const axios = require('axios');
 const fs = require('fs');
 const path = require('path');
+const config = require('../../config/config');

 const router = express.Router();

@@ -236,6 +237,21 @@ router.get('/api-keys', authenticateAdmin, async (req, res) => {
  }
 });

+// 获取支持的客户端列表
+router.get('/supported-clients', authenticateAdmin, async (req, res) => {
+  try {
+    const clients = config.clientRestrictions.predefinedClients.map(client => ({
+      id: client.id,
+      name: client.name,
+      description: client.description
+    }));
+    res.json({ success: true, data: clients });
+  } catch (error) {
+    logger.error('❌ Failed to get supported clients:', error);
+    res.status(500).json({ error: 'Failed to get supported clients', message: error.message });
+  }
+});
+
 // 创建新的API Key
 router.post('/api-keys', authenticateAdmin, async (req, res) => {
  try {
@@ -251,7 +267,9 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => {
      rateLimitWindow,
      rateLimitRequests,
      enableModelRestriction,
-      restrictedModels
+      restrictedModels,
+      enableClientRestriction,
+      allowedClients
    } = req.body;

    // 输入验证
@@ -293,6 +311,15 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => {
      return res.status(400).json({ error: 'Restricted models must be an array' });
    }

+    // 验证客户端限制字段
+    if (enableClientRestriction !== undefined && typeof enableClientRestriction !== 'boolean') {
+      return res.status(400).json({ error: 'Enable client restriction must be a boolean' });
+    }
+
+    if (allowedClients !== undefined && !Array.isArray(allowedClients)) {
+      return res.status(400).json({ error: 'Allowed clients must be an array' });
+    }
+
    const newKey = await apiKeyService.generateApiKey({
      name,
      description,
@@ -305,7 +332,9 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => {
      rateLimitWindow,
      rateLimitRequests,
      enableModelRestriction,
-      restrictedModels
+      restrictedModels,
+      enableClientRestriction,
+      allowedClients
    });

    logger.success(`🔑 Admin created new API key: ${name}`);
@@ -320,7 +349,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, rateLimitWindow, rateLimitRequests, claudeAccountId, geminiAccountId, permissions, enableModelRestriction, restrictedModels, expiresAt } = req.body;
+    const { tokenLimit, concurrencyLimit, rateLimitWindow, rateLimitRequests, claudeAccountId, geminiAccountId, permissions, enableModelRestriction, restrictedModels, enableClientRestriction, allowedClients, expiresAt } = req.body;

    // 只允许更新指定字段
    const updates = {};
@@ -386,6 +415,21 @@ router.put('/api-keys/:keyId', authenticateAdmin, async (req, res) => {
      updates.restrictedModels = restrictedModels;
    }

+    // 处理客户端限制字段
+    if (enableClientRestriction !== undefined) {
+      if (typeof enableClientRestriction !== 'boolean') {
+        return res.status(400).json({ error: 'Enable client restriction must be a boolean' });
+      }
+      updates.enableClientRestriction = enableClientRestriction;
+    }
+
+    if (allowedClients !== undefined) {
+      if (!Array.isArray(allowedClients)) {
+        return res.status(400).json({ error: 'Allowed clients must be an array' });
+      }
+      updates.allowedClients = allowedClients;
+    }
+
    // 处理过期时间字段
    if (expiresAt !== undefined) {
      if (expiresAt === null) {
--- a/src/services/apiKeyService.js
+++ b/src/services/apiKeyService.js
@@ -24,7 +24,9 @@ class ApiKeyService {
      rateLimitWindow = null,
      rateLimitRequests = null,
      enableModelRestriction = false,
-      restrictedModels = []
+      restrictedModels = [],
+      enableClientRestriction = false,
+      allowedClients = []
    } = options;

    // 生成简单的API Key (64字符十六进制)
@@ -47,6 +49,8 @@ class ApiKeyService {
      permissions: permissions || 'all',
      enableModelRestriction: String(enableModelRestriction),
      restrictedModels: JSON.stringify(restrictedModels || []),
+      enableClientRestriction: String(enableClientRestriction || false),
+      allowedClients: JSON.stringify(allowedClients || []),
      createdAt: new Date().toISOString(),
      lastUsedAt: '',
      expiresAt: expiresAt || '',
@@ -73,6 +77,8 @@ class ApiKeyService {
      permissions: keyData.permissions,
      enableModelRestriction: keyData.enableModelRestriction === 'true',
      restrictedModels: JSON.parse(keyData.restrictedModels),
+      enableClientRestriction: keyData.enableClientRestriction === 'true',
+      allowedClients: JSON.parse(keyData.allowedClients || '[]'),
      createdAt: keyData.createdAt,
      expiresAt: keyData.expiresAt,
      createdBy: keyData.createdBy
@@ -122,6 +128,14 @@ class ApiKeyService {
        restrictedModels = [];
      }

+      // 解析允许的客户端
+      let allowedClients = [];
+      try {
+        allowedClients = keyData.allowedClients ? JSON.parse(keyData.allowedClients) : [];
+      } catch (e) {
+        allowedClients = [];
+      }
+
      return {
        valid: true,
        keyData: {
@@ -136,6 +150,8 @@ class ApiKeyService {
          rateLimitRequests: parseInt(keyData.rateLimitRequests || 0),
          enableModelRestriction: keyData.enableModelRestriction === 'true',
          restrictedModels: restrictedModels,
+          enableClientRestriction: keyData.enableClientRestriction === 'true',
+          allowedClients: allowedClients,
          usage
        }
      };
@@ -160,12 +176,18 @@ class ApiKeyService {
        key.currentConcurrency = await redis.getConcurrency(key.id);
        key.isActive = key.isActive === 'true';
        key.enableModelRestriction = key.enableModelRestriction === 'true';
+        key.enableClientRestriction = key.enableClientRestriction === 'true';
        key.permissions = key.permissions || 'all'; // 兼容旧数据
        try {
          key.restrictedModels = key.restrictedModels ? JSON.parse(key.restrictedModels) : [];
        } catch (e) {
          key.restrictedModels = [];
        }
+        try {
+          key.allowedClients = key.allowedClients ? JSON.parse(key.allowedClients) : [];
+        } catch (e) {
+          key.allowedClients = [];
+        }
        delete key.apiKey; // 不返回哈希后的key
      }

@@ -185,15 +207,15 @@ class ApiKeyService {
      }

      // 允许更新的字段
-      const allowedUpdates = ['name', 'description', 'tokenLimit', 'concurrencyLimit', 'rateLimitWindow', 'rateLimitRequests', 'isActive', 'claudeAccountId', 'geminiAccountId', 'permissions', 'expiresAt', 'enableModelRestriction', 'restrictedModels'];
+      const allowedUpdates = ['name', 'description', 'tokenLimit', 'concurrencyLimit', 'rateLimitWindow', 'rateLimitRequests', 'isActive', 'claudeAccountId', 'geminiAccountId', 'permissions', 'expiresAt', 'enableModelRestriction', 'restrictedModels', 'enableClientRestriction', 'allowedClients'];
      const updatedData = { ...keyData };

      for (const [field, value] of Object.entries(updates)) {
        if (allowedUpdates.includes(field)) {
-          if (field === 'restrictedModels') {
-            // 特殊处理 restrictedModels 数组
+          if (field === 'restrictedModels' || field === 'allowedClients') {
+            // 特殊处理数组字段
            updatedData[field] = JSON.stringify(value || []);
-          } else if (field === 'enableModelRestriction') {
+          } else if (field === 'enableModelRestriction' || field === 'enableClientRestriction') {
            // 布尔值转字符串
            updatedData[field] = String(value);
          } else {