From daf6a123d58e486d425f606e0e7a63bc7c8098d9 Mon Sep 17 00:00:00 2001
From: nai-degen <khoners@protonmail.com>
Date: Wed, 4 Oct 2023 09:39:59 -0500
Subject: [PATCH] adjusts Agnai.chat and RisuAI rate limiting

---
 src/proxy/check-risu-token.ts | 13 +++++++++----
 src/proxy/queue.ts            | 15 +++++----------
 src/proxy/rate-limit.ts       | 10 ++++++++--
 3 files changed, 22 insertions(+), 16 deletions(-)

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;
   }