Implement rate limit for risuai.xyz (khanon/oai-reverse-proxy!31)

src/proxy/rate-limit.ts

@@ -1,12 +1,15 @@
+import axios from "axios";
 import { Request, Response, NextFunction } from "express";
 import { config } from "../config";
 
 export const AGNAI_DOT_CHAT_IP = "157.230.249.32";
+const RISUAI_TOKEN_CHECKER_URL = "https://sv.risuai.xyz/public/api/checktoken";
 const RATE_LIMIT_ENABLED = Boolean(config.modelRateLimit);
 const RATE_LIMIT = Math.max(1, config.modelRateLimit);
 const ONE_MINUTE_MS = 60 * 1000;
 
 const lastAttempts = new Map<string, number[]>();
+const validRisuTokens = new Set<string>();
 
 const expireOldAttempts = (now: number) => (attempt: number) =>
   attempt > now - ONE_MINUTE_MS;
@@ -52,7 +55,11 @@ export const getUniqueIps = () => {
   return lastAttempts.size;
 };
 
-export const ipLimiter = (req: Request, res: Response, next: NextFunction) => {
+export const ipLimiter = async (
+  req: Request,
+  res: Response,
+  next: NextFunction
+) => {
   if (!RATE_LIMIT_ENABLED) {
     next();
     return;
@@ -66,9 +73,46 @@ export const ipLimiter = (req: Request, res: Response, next: NextFunction) => {
     return;
   }
 
+  // makes risuai.xyz rate limiting by x-risu-tk header since it's shared between a lot of users.
+  let risuToken = req.header("x-risu-tk") || null;
+  if (risuToken) {
+    try {
+      // checks the token only when it is not in freshRisuTokens or bitFreshRisuTokens
+      if (!validRisuTokens.has(risuToken)) {
+        req.log.info(
+          { token: `${risuToken.slice(0, 4)}...` },
+          "Authenticating new RisuAI token"
+        );
+        // checks the token is vaild (fresh) to prevend abuse
+        const validCheck = await axios.post<{ vaild: boolean }>(
+          RISUAI_TOKEN_CHECKER_URL,
+          { token: risuToken },
+          { headers: { "Content-Type": "application/json" } }
+        );
+
+        if (!validCheck.data.vaild) {
+          //if its invaild, uses ip instead
+          req.log.warn("Invalid RisuAI token; rate limiting by IP instead");
+          risuToken = null;
+        } else {
+          req.log.info("RisuAI token authenticated; adding to known tokens");
+          validRisuTokens.add(risuToken);
+        }
+      }
+    } catch (e: any) {
+      //if request throws error, uses ip
+      // TODO: probably need a backoff here to avoid spamming RisuAI
+      req.log.warn(
+        { error: e.message },
+        "Error authenticating RisuAI token; rate limiting by IP instead"
+      );
+      risuToken = null;
+    }
+  }
+
   // If user is authenticated, key rate limiting by their token. Otherwise, key
   // rate limiting by their IP address. Mitigates key sharing.
-  const rateLimitKey = req.user?.token || req.ip;
+  const rateLimitKey = req.user?.token || risuToken || req.ip;
 
   const { remaining, reset } = getStatus(rateLimitKey);
   res.set("X-RateLimit-Limit", config.modelRateLimit.toString());