diff --git a/src/proxy/check-risu-token.ts b/src/proxy/check-risu-token.ts index 115c8f9..3e8e329 100644 --- a/src/proxy/check-risu-token.ts +++ b/src/proxy/check-risu-token.ts @@ -21,7 +21,7 @@ kYuIJbnAGw5Oq0L6dXFW2DFwlcLz51kPVOmDc159FsQjyuPnta7NiZAANS8KM1CJ pwIDAQAB`; let IMPORTED_RISU_KEY: CryptoKey | null = null; -type RisuToken = { id: Uint8Array; expiresIn: number }; +type RisuToken = { id: string; expiresIn: number }; type SignedToken = { data: RisuToken; sig: string }; (async () => { @@ -54,14 +54,14 @@ export async function checkRisuToken( try { const { valid, data } = await validCheck(header); - if (!valid) { + if (!valid || !data) { req.log.warn( { token: header, data }, "Invalid RisuAI token; using IP instead" ); } else { req.log.info("RisuAI token validated"); - req.risuToken = header; + req.risuToken = String(data.id); } } catch (err) { req.log.warn( @@ -81,12 +81,13 @@ async function validCheck(header: string) { ); } catch (err) { log.warn({ error: err.message }, "Provided unparseable RisuAI token"); - return { valid: false, data: "[unparseable]" }; + return { valid: false }; } const data: RisuToken = tk.data; const sig = Buffer.from(tk.sig, "base64"); if (data.expiresIn < Math.floor(Date.now() / 1000)) { + log.warn({ token: header }, "Provided expired RisuAI token"); return { valid: false }; } @@ -97,5 +98,9 @@ async function validCheck(header: string) { Buffer.from(JSON.stringify(data)) ); + if (!valid) { + log.warn({ token: header }, "RisuAI token failed signature check"); + } + return { valid, data }; } diff --git a/src/proxy/queue.ts b/src/proxy/queue.ts index 0e66413..4354bee 100644 --- a/src/proxy/queue.ts +++ b/src/proxy/queue.ts @@ -32,7 +32,7 @@ const queue: Request[] = []; const log = logger.child({ module: "request-queue" }); /** Maximum number of queue slots for Agnai.chat requests. */ -const AGNAI_CONCURRENCY_LIMIT = 15; +const AGNAI_CONCURRENCY_LIMIT = 5; /** Maximum number of queue slots for individual users. */ const USER_CONCURRENCY_LIMIT = 1; @@ -68,12 +68,11 @@ export function enqueue(req: Request) { // more spots in the queue. Can't make it unlimited because people will // intentionally abuse it. // Authenticated users always get a single spot in the queue. + const isAgnai = AGNAI_DOT_CHAT_IP.includes(req.ip); const maxConcurrentQueuedRequests = - isGuest && req.ip === AGNAI_DOT_CHAT_IP - ? AGNAI_CONCURRENCY_LIMIT - : USER_CONCURRENCY_LIMIT; + isGuest && isAgnai ? AGNAI_CONCURRENCY_LIMIT : USER_CONCURRENCY_LIMIT; if (enqueuedRequestCount >= maxConcurrentQueuedRequests) { - if (req.ip === AGNAI_DOT_CHAT_IP) { + if (isAgnai) { // Re-enqueued requests are not counted towards the limit since they // already made it through the queue once. if (req.retryCount === 0) { @@ -336,11 +335,7 @@ function killQueuedRequest(req: Request) { try { const message = `Your request has been terminated by the proxy because it has been in the queue for more than 5 minutes. The queue is currently ${queue.length} requests long.`; if (res.headersSent) { - const fakeErrorEvent = buildFakeSse( - "proxy queue error", - message, - req - ); + const fakeErrorEvent = buildFakeSse("proxy queue error", message, req); res.write(fakeErrorEvent); res.end(); } else { diff --git a/src/proxy/rate-limit.ts b/src/proxy/rate-limit.ts index e4269d4..2d98f22 100644 --- a/src/proxy/rate-limit.ts +++ b/src/proxy/rate-limit.ts @@ -1,7 +1,12 @@ import { Request, Response, NextFunction } from "express"; import { config } from "../config"; -export const AGNAI_DOT_CHAT_IP = "157.230.249.32"; +export const AGNAI_DOT_CHAT_IP = [ + "157.230.249.32", // old + "157.245.148.56", + "174.138.29.50", + "209.97.162.44", +]; const RATE_LIMIT_ENABLED = Boolean(config.modelRateLimit); const RATE_LIMIT = Math.max(1, config.modelRateLimit); @@ -66,7 +71,8 @@ export const ipLimiter = async ( // Exempt Agnai.chat from rate limiting since it's shared between a lot of // users. Dunno how to prevent this from being abused without some sort of // identifier sent from Agnaistic to identify specific users. - if (req.ip === AGNAI_DOT_CHAT_IP) { + if (AGNAI_DOT_CHAT_IP.includes(req.ip)) { + req.log.info("Exempting Agnai request from rate limiting."); next(); return; }