diff --git a/Dockerfile b/Dockerfile index 77f6e7d..ab606c3 100644 --- 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 new file mode 100644 index 0000000..1431cb9 --- /dev/null +++ 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 }); + } +}); \ No newline at end of file diff --git a/src/api/controllers/chat.ts b/src/api/controllers/chat.ts index f6f623e..ad8eede 100644 --- 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 = {}; +// 添加 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 { + 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,12 +281,16 @@ 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( @@ -218,6 +298,7 @@ async function createCompletion( { 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 new file mode 100644 index 0000000..1431cb9 --- /dev/null +++ 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 }); + } +}); \ No newline at end of file