Update index.js

index.js

@@ -2,16 +2,20 @@ import fetch from 'node-fetch';
 import express from 'express';
 import cors from 'cors';
 import dotenv from 'dotenv';
+import puppeteer from 'puppeteer';
 
 dotenv.config();
 
-const Tokens =[];
+const Tokens = [];
 let tokenManager;
+let redisClient;
 let currentIndex = 0;
 const CONFIG = {
   API: {
     BASE_URL: "https://partyrock.aws",
-    API_KEY: process.env.API_KEY || "sk-123456"//自定义你自己的认证密钥，记得修改
+    API_KEY: process.env.API_KEY || "sk-123456",//自定义你自己的认证密钥，记得修改
+    RedisUrl: process.env.RedisUrl,
+    RedisToken: process.env.RedisToken
   },
   RETRY: {
     MAX_ATTEMPTS: 1,
@@ -49,33 +53,127 @@ const CONFIG = {
     "Cookie": "",
     "accept-language": "zh-CN,zh;q=0.9",
     "priority": "u=1, i"
+  },
+  CHROME_PATH: process.env.CHROME_PATH || "C:\\Program Files\\Google\\Chrome\\Application\\chrome.exe"
+};
+var RedisClient = class {
+  constructor() {
+    this.url = CONFIG.API.RedisUrl;
+    this.token = CONFIG.API.RedisToken;
+    console.log(this.url, this.token);
+  }
+  async get(key) {
+    const response = await fetch(`${this.url}/get/${key}`, {
+      headers: {
+        Authorization: `Bearer ${this.token}`
+      }
+    });
+    if (!response.ok) {
+      console.log("redis获取内容失败", response.status);
+    }
+    const data = await response.json();
+    return data.result;
+  }
+  async set(key, value) {
+    const url = `${this.url}/set/${key}`;
+    const response = await fetch(url, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${this.token}`
+      },
+      body: `${value}`
+    });
+    if (!response.ok) {
+      console.log("redis设置内容失败", response.status);
+    }
   }
 };
-
 class TokenManager {
+  async updateRedisTokens() {
+    await redisClient.set(`tokens_${currentIndex}`, JSON.stringify(Tokens[currentIndex]));
+  }
+  async getRedisTokens() {
+    var checkRedis = JSON.parse(await redisClient.get(`tokens_${currentIndex}`));
+    return checkRedis;
+  }
+
   async updateCacheTokens() {
     CONFIG.DEFAULT_HEADERS["anti-csrftoken-a2z"] = Tokens[currentIndex]["anti_csrftoken_a2z"];
     CONFIG.DEFAULT_HEADERS.Cookie = `idToken=${Tokens[currentIndex]["idToken"]}; pr_refresh_token=${Tokens[currentIndex]["pr_refresh_token"]};aws-waf-token=${Tokens[currentIndex]["aws_waf_token"]}`;
     CONFIG.DEFAULT_HEADERS.referer = Tokens[currentIndex]["refreshUrl"];
   }
 
-  async updateTokens(response) {
-    const newCsrfToken = response.headers.get('anti-csrftoken-a2z');
-    if (newCsrfToken) {
-      Tokens[currentIndex]["anti_csrftoken_a2z"] = newCsrfToken;
+  async updateTokens(response, isWaf = false) {
+    if (isWaf) {
+      console.log("更新waf");
+      if (await Utils.extractWaf()) {
+        Tokens[currentIndex]["aws_waf_token"] = await Utils.extractWaf();
+        return;
+      } else {
+        console.log("提取aws-waf-token失败");
+        currentIndex = (currentIndex + 1) % Tokens.length;
+      }
     }
-
+    const newCsrfToken = response.headers.get('anti-csrftoken-a2z');
     const cookies = response.headers.get('set-cookie');
-    if (cookies) {
+    if (newCsrfToken && cookies) {
+      console.log("更新缓存");
+      Tokens[currentIndex]["anti_csrftoken_a2z"] = newCsrfToken;
       const idTokenMatch = cookies.match(/idToken=([^;]+)/);
       if (idTokenMatch && idTokenMatch[1]) {
         Tokens[currentIndex]["idToken"] = idTokenMatch[1];
       }
+      await this.updateRedisTokens();//最后更新redis数据库缓存
+      console.log("更新缓存完毕");
     }
     currentIndex = (currentIndex + 1) % Tokens.length;
   }
 }
+
+
 class Utils {
+  static async extractWaf() {
+    const browser = await puppeteer.launch({
+      headless: true,
+      args: [
+        '--no-sandbox',
+        '--disable-setuid-sandbox',
+        '--disable-dev-shm-usage',
+        '--disable-gpu'
+      ],
+      executablePath: CONFIG.CHROME_PATH
+    });
+    try {
+      const page = await browser.newPage();
+      await page.setExtraHTTPHeaders({
+        cookie: `pr_refresh_token=${Tokens[currentIndex]["pr_refresh_token"]};aws-waf-token=${Tokens[currentIndex]["aws_waf_token"]};idToken=${Tokens[currentIndex]["idToken"]}`
+      });
+
+      await page.goto(Tokens[currentIndex]["refreshUrl"], { waitUntil: 'networkidle0' });
+
+      // 直接从页面 cookies 中提取 aws-waf-token
+      const awsWafToken = (await page.cookies()).find(
+        cookie => cookie.name.toLowerCase() === 'aws-waf-token'
+      )?.value;
+
+      if (awsWafToken) {
+        Tokens[currentIndex]["aws_waf_token"] = awsWafToken;
+        console.log("成功提取 aws-waf-token");
+        await browser.close();
+        return awsWafToken;
+      } else {
+        console.log("提取aws-waf-token失败");
+        await browser.close();
+        return null;
+      }
+
+    } catch (error) {
+      console.error('获取 aws-waf-token 出错:', error);
+      await browser.close();
+      return null;
+    }
+  }
+
   static async extractTokens(cookieString) {
     const tokens = {};
     const cookiePairs = cookieString.split(';').map(pair => pair.trim());
@@ -116,30 +214,55 @@ class Utils {
 }
 async function initializeService() {
   console.log('服务初始化中...');
-  // 遍历所有可能的 token 组
+  tokenManager = new TokenManager();
+  redisClient = new RedisClient();
+
   let index = 0;
   while (true) {
-    const refreshUrl = process.env[`AUTH_TOKENS_${index}_REFRESH_URL`];
-    const anti_csrftoken_a2z = process.env[`AUTH_TOKENS_${index}_ANTI_CSRF_TOKEN`];
-    const cookie = process.env[`AUTH_TOKENS_${index}_COOKIE`];
-    if (!refreshUrl && !anti_csrftoken_a2z && !cookie) {
-      break;
-    }
-    const cookies = await Utils.extractTokens(cookie);
-    // 只有当所有属性都存在时才添加 token
-    if (refreshUrl && anti_csrftoken_a2z && cookie) {
-      Tokens.push({
-        refreshUrl,
-        anti_csrftoken_a2z,
-        pr_refresh_token: cookies["pr_refresh_token"],
-        aws_waf_token: cookies["aws-waf-token"],
-        idToken: cookies["idToken"]
-      });
+    console.log(index, '开始检测是否有缓存');
+    // 使用 JSON.parse 确保正确解析
+    var checkRedis = await redisClient.get(`tokens_${index}`);
+    if (checkRedis) {
+      // 尝试解析 JSON 字符串
+      try {
+        const parsedRedis = typeof checkRedis === 'string'
+          ? JSON.parse(checkRedis)
+          : checkRedis;
+        Tokens.push({
+          refreshUrl: parsedRedis.refreshUrl,
+          anti_csrftoken_a2z: parsedRedis.anti_csrftoken_a2z,
+          pr_refresh_token: parsedRedis.pr_refresh_token,
+          aws_waf_token: parsedRedis.aws_waf_token,
+          idToken: parsedRedis.idToken
+        });
+        console.log(`成功添加第 ${index} 组 Token`);
+      } catch (error) {
+        console.error(`解析第 ${index} 组 Token 时出错:`, error);
+      }
+    } else {
+      console.log(index, '没有缓存，开始提取环境变量');
+      const refreshUrl = process.env[`AUTH_TOKENS_${index}_REFRESH_URL`];
+      const anti_csrftoken_a2z = process.env[`AUTH_TOKENS_${index}_ANTI_CSRF_TOKEN`];
+      const cookie = process.env[`AUTH_TOKENS_${index}_COOKIE`];
+
+      if (!refreshUrl && !anti_csrftoken_a2z && !cookie) {
+        break;
+      }
+      const cookies = await Utils.extractTokens(cookie);
+
+      if (refreshUrl && anti_csrftoken_a2z && cookie) {
+        Tokens.push({
+          refreshUrl,
+          anti_csrftoken_a2z,
+          pr_refresh_token: cookies["pr_refresh_token"],
+          aws_waf_token: cookies["aws-waf-token"],
+          idToken: cookies["idToken"]
+        });
+      }
     }
     index++;
   }
-  
-  tokenManager = new TokenManager();
+  console.log('服务初始化完毕');
 }
 
 await initializeService();
@@ -285,6 +408,10 @@ class ResponseHandler {
           if (trimmedLine && trimmedLine.startsWith('data: ')) {
             const data = trimmedLine.substring(6);
             if (!data) continue;
+            if (data == "[DONE]") {
+              res.write('data: [DONE]\n\n');
+              return res.end();
+            }
             try {
               const json = JSON.parse(data);
               if (json?.text) {
@@ -366,6 +493,7 @@ app.get('/hf/v1/models', (req, res) => {
 });
 
 app.post('/hf/v1/chat/completions', async (req, res) => {
+  var reqStatus = 500;
   try {
     const authToken = req.headers.authorization?.replace('Bearer ', '');
     if (authToken !== CONFIG.API.API_KEY) {
@@ -389,16 +517,20 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
           },
           body: JSON.stringify(requestPayload)
         });
-        await tokenManager.updateTokens(response);
         if (response.status == 200) {
           console.log("请求成功");
           break; // 如果请求成功，跳出重试循环
         }
         if (response.status != 200) {
-          console.error(JSON.stringify(await response.text(), null, 2));
+          reqStatus = response.status;
+          if (response.status == 202) {
+            console.log("请求失败，已达到速率，重新获取waf然后更新token库");
+            await tokenManager.updateTokens(response, true);
+          }
         }
         retryCount++;
         if (retryCount >= CONFIG.RETRY.MAX_ATTEMPTS) {
+          reqStatus = response.status;
           throw new Error(`上游服务请求失败! status: ${response.status}`);
         }
         // 等待一段时间后重试
@@ -418,9 +550,10 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
     } else {
       await ResponseHandler.handleNormalResponse(response, req.body.model, res);
     }
-
+    //请求完毕后再尝试更新redis数据库缓存，避免阻塞
+    await tokenManager.updateTokens(response);
   } catch (error) {
-    res.status(500).json({
+    res.status(reqStatus).json({
       error: {
         message: error.message,
         type: 'server_error',