disables keys shared between too many IP addresses

.env.example

@@ -16,6 +16,8 @@
 
 # Optional settings for user management. See docs/user-management.md.
 # GATEKEEPER=none
+# GATEKEEPER_STORE=memory
+# MAX_IPS_PER_USER=20
 
 # Optional settings for prompt logging. See docs/logging-sheets.md.
 # PROMPT_LOGGING=false

src/config.ts

@@ -49,6 +49,8 @@ type Config = {
   firebaseRtdbUrl?: string;
   /** Base64-encoded Firebase service account key if using the Firebase RTDB store. */
   firebaseKey?: string;
+  /** Maximum number of IPs per user, after which their token is disabled. */
+  maxIpsPerUser: number;
   /** Per-IP limit for requests per minute to OpenAI's completions endpoint. */
   modelRateLimit: number;
   /** Max number of tokens to generate. Requests which specify a higher value will be rewritten to use this value. */
@@ -100,6 +102,7 @@ export const config: Config = {
   adminKey: getEnvWithDefault("ADMIN_KEY", ""),
   gatekeeper: getEnvWithDefault("GATEKEEPER", "none"),
   gatekeeperStore: getEnvWithDefault("GATEKEEPER_STORE", "memory"),
+  maxIpsPerUser: getEnvWithDefault("MAX_IPS_PER_USER", 20),
   firebaseRtdbUrl: getEnvWithDefault("FIREBASE_RTDB_URL", undefined),
   firebaseKey: getEnvWithDefault("FIREBASE_KEY", undefined),
   modelRateLimit: getEnvWithDefault("MODEL_RATE_LIMIT", 4),

src/proxy/auth/user-store.ts

@@ -43,6 +43,8 @@ export type UserType = "normal" | "special";
 
 type UserUpdate = Partial<User> & Pick<User, "token">;
 
+const MAX_IPS_PER_USER = config.maxIpsPerUser;
+
 const users: Map<string, User> = new Map();
 const usersToFlush = new Set<string>();
 
@@ -98,12 +100,12 @@ export function upsertUser(user: UserUpdate) {
     ...user,
   });
   usersToFlush.add(user.token);
-  
+
   // Immediately schedule a flush to the database if we're using Firebase.
   if (config.gatekeeperStore === "firebase_rtdb") {
     setImmediate(flushUsers);
   }
-  
+
   return users.get(user.token);
 }
 
@@ -132,6 +134,13 @@ export function authenticate(token: string, ip: string) {
   const user = users.get(token);
   if (!user || user.disabledAt) return;
   if (!user.ip.includes(ip)) user.ip.push(ip);
+
+  // If too many IPs are associated with the user, disable the account.
+  if (user.ip.length > MAX_IPS_PER_USER) {
+    disableUser(token, "Too many IP addresses associated with this token.");
+    return;
+  }
+
   user.lastUsedAt = Date.now();
   usersToFlush.add(token);
   return user;