深度思考模式修复，跟进官方的插冷鸡策略

LLM-Red-Team · Dec 10, 2024 · 78e0ef8 · 78e0ef8
1 parent 74e5bf4
commit 78e0ef8
Show file tree

Hide file tree

Showing 4 changed files with 155 additions and 1 deletion.
diff --git a/Dockerfile b/Dockerfile
@@ -12,6 +12,7 @@ COPY --from=BUILD_IMAGE /app/configs /app/configs
 COPY --from=BUILD_IMAGE /app/package.json /app/package.json
 COPY --from=BUILD_IMAGE /app/dist /app/dist
 COPY --from=BUILD_IMAGE /app/public /app/public
+COPY --from=BUILD_IMAGE /app/challenge-worker.js /app/challenge-worker.js
 COPY --from=BUILD_IMAGE /app/node_modules /app/node_modules
 
 WORKDIR /app

diff --git a/challenge-worker.js b/challenge-worker.js
@@ -0,0 +1,31 @@
+import crypto from 'crypto';
+import { parentPort } from 'worker_threads';
+
+parentPort?.on('message', (data) => {
+  try {
+    const { algorithm, challenge, salt, difficulty, expire_at, signature } = data;
+    let answer = 0;
+    let i = difficulty - 1;
+
+    for (let r = 0; r <= i; r++) {
+      const str = "".concat(salt, "_").concat(expire_at, "_").concat(r.toString());
+      const hash = crypto.createHash('sha256').update(str).digest('hex');
+      if (hash === challenge) {
+        answer = r;
+        break;
+      }
+    }
+
+    if(answer === 0) throw new Error('No solution found');
+
+    parentPort?.postMessage({
+      algorithm,
+      challenge, 
+      salt,
+      answer,
+      signature
+    });
+  } catch (error) {
+    parentPort?.postMessage({ error: error.message });
+  }
+}); 
diff --git a/src/api/controllers/chat.ts b/src/api/controllers/chat.ts
@@ -1,6 +1,8 @@
 import { PassThrough } from "stream";
 import _ from "lodash";
 import axios, { AxiosResponse } from "axios";
+import { Worker } from 'worker_threads';
+import path from 'path';
 
 import APIException from "@/lib/exceptions/APIException.ts";
 import EX from "@/api/consts/exceptions.ts";
@@ -44,6 +46,27 @@ const accessTokenMap = new Map();
 // access_token请求队列映射
 const accessTokenRequestQueueMap: Record<string, Function[]> = {};
 
+// 添加 worker 池
+const workerPool: (Worker & { inUse?: boolean })[] = [];
+const MAX_WORKERS = 4; // 可以根据需要调整
+
+function getWorker() {
+  // 从池中获取空闲worker或创建新worker
+  let worker = workerPool.find(w => !w.inUse);
+  if (!worker && workerPool.length < MAX_WORKERS) {
+    worker = new Worker(path.join(path.resolve(), 'challenge-worker.js'));
+    workerPool.push(worker);
+  }
+  if (worker) {
+    worker.inUse = true;
+  }
+  return worker;
+}
+
+function releaseWorker(worker: (Worker & { inUse?: boolean })) {
+  worker.inUse = false;
+}
+
 async function getIPAddress() {
   if (ipAddress) return ipAddress;
   const result = await axios.get('https://chat.deepseek.com/', {
@@ -169,6 +192,59 @@ async function createSession(model: string, refreshToken: string): Promise<strin
   return biz_data.id;
 }
 
+/**
+ * 碰撞challenge答案
+ * 
+ * 厂商这个反逆向的策略不错哦
+ * 相当于把计算量放在浏览器侧的话，用户分摊了这个计算量
+ * 但是如果逆向在服务器上算，那这个成本都在服务器集中，并发一高就GG
+ */
+function answerChallenge(response: any): Promise<any> {
+  return new Promise((resolve, reject) => {
+    const worker = getWorker();
+    if (!worker) {
+      reject(new Error('No available workers'));
+      return;
+    }
+
+    worker.once('message', (result) => {
+      releaseWorker(worker);
+      if (result.error) {
+        reject(new Error(result.error));
+      } else {
+        resolve(result);
+      }
+    });
+
+    worker.once('error', (error) => {
+      releaseWorker(worker);
+      reject(error);
+    });
+
+    worker.postMessage(response);
+  });
+}
+
+/**
+ * 获取challenge响应
+ *
+ * @param refreshToken 用于刷新access_token的refresh_token
+ */
+async function getChallengeResponse(refreshToken: string) {
+  const token = await acquireToken(refreshToken);
+  const result = await axios.get('https://chat.deepseek.com/api/v0/chat/challenge', {
+    headers: {
+      Authorization: `Bearer ${token}`,
+      ...FAKE_HEADERS,
+        Cookie: generateCookie()
+    },
+    timeout: 15000,
+    validateStatus: () => true,
+  });
+  const { biz_data: { challenge } } = checkResult(result, refreshToken);
+  return challenge;
+}
+
 /**
  * 同步对话补全
  *
@@ -205,19 +281,24 @@ async function createCompletion(
 
     const isSearchModel = model.includes('search') || prompt.includes('联网搜索');
     const isThinkingModel = model.includes('think') || model.includes('r1') || prompt.includes('深度思考');
-
+
+    let challenge;
     if (isThinkingModel) {
       const thinkingQuota = await getThinkingQuota(refreshToken);
       if (thinkingQuota <= 0) {
         throw new APIException(EX.API_REQUEST_FAILED, '深度思考配额不足');
       }
+      const challengeResponse = await getChallengeResponse(refreshToken);
+      challenge = await answerChallenge(challengeResponse);
+      logger.info(`插冷鸡: ${JSON.stringify(challenge)}`);
     }
 
     const result = await axios.post(
       "https://chat.deepseek.com/api/v0/chat/completion",
       {
         chat_session_id: sessionId,
         parent_message_id: refParentMsgId || null,
+        challenge_response: challenge,
         prompt,
         ref_file_ids: [],
         search_enabled: isSearchModel,
@@ -306,11 +387,15 @@ async function createCompletionStream(
     const isSearchModel = model.includes('search') || prompt.includes('联网搜索');
     const isThinkingModel = model.includes('think') || model.includes('r1') || prompt.includes('深度思考');
 
+    let challenge;
     if (isThinkingModel) {
       const thinkingQuota = await getThinkingQuota(refreshToken);
       if (thinkingQuota <= 0) {
         throw new APIException(EX.API_REQUEST_FAILED, '深度思考配额不足');
       }
+      const challengeResponse = await getChallengeResponse(refreshToken);
+      challenge = await answerChallenge(challengeResponse);
+      logger.info(`插冷鸡: ${JSON.stringify(challenge)}`);
     }
 
     // 创建会话
@@ -324,6 +409,7 @@ async function createCompletionStream(
         chat_session_id: sessionId,
         parent_message_id: refParentMsgId || null,
         prompt,
+        challenge_response: challenge,
         ref_file_ids: [],
         search_enabled: isSearchModel,
         thinking_enabled: isThinkingModel
@@ -1162,6 +1248,11 @@ getIPAddress().then(() => {
   logger.error('获取 IP 地址失败:', err);
 });
 
+// 在程序退出时清理workers
+process.on('exit', () => {
+  workerPool.forEach(worker => worker.terminate());
+});
+
 export default {
   createCompletion,
   createCompletionStream,

diff --git a/src/workers/challengeWorker.ts b/src/workers/challengeWorker.ts
@@ -0,0 +1,31 @@
+import crypto from 'crypto';
+import { parentPort } from 'worker_threads';
+
+parentPort?.on('message', (data) => {
+  try {
+    const { algorithm, challenge, salt, difficulty, expire_at, signature } = data;
+    let answer = 0;
+    let i = difficulty - 1;
+
+    for (let r = 0; r <= i; r++) {
+      const str = "".concat(salt, "_").concat(expire_at, "_").concat(r.toString());
+      const hash = crypto.createHash('sha256').update(str).digest('hex');
+      if (hash === challenge) {
+        answer = r;
+        break;
+      }
+    }
+
+    if(answer === 0) throw new Error('No solution found');
+
+    parentPort?.postMessage({
+      algorithm,
+      challenge, 
+      salt,
+      answer,
+      signature
+    });
+  } catch (error) {
+    parentPort?.postMessage({ error: error.message });
+  }
+});