User authentication (khanon/oai-reverse-proxy!7)

package-lock.json

@@ -18,12 +18,15 @@
         "pino": "^8.11.0",
         "pino-http": "^8.3.3",
         "showdown": "^2.1.0",
-        "zlib": "^1.0.5"
+        "uuid": "^9.0.0",
+        "zlib": "^1.0.5",
+        "zod": "^3.21.4"
       },
       "devDependencies": {
         "@types/cors": "^2.8.13",
         "@types/express": "^4.17.17",
         "@types/showdown": "^2.0.0",
+        "@types/uuid": "^9.0.1",
         "concurrently": "^8.0.1",
         "esbuild": "^0.17.16",
         "esbuild-register": "^3.4.2",
@@ -547,6 +550,12 @@
       "integrity": "sha512-70xBJoLv+oXjB5PhtA8vo7erjLDp9/qqI63SRHm4REKrwuPOLs8HhXwlZJBJaB4kC18cCZ1UUZ6Fb/PLFW4TCA==",
       "dev": true
     },
+    "node_modules/@types/uuid": {
+      "version": "9.0.1",
+      "resolved": "https://registry.npmjs.org/@types/uuid/-/uuid-9.0.1.tgz",
+      "integrity": "sha512-rFT3ak0/2trgvp4yYZo5iKFEPsET7vKydKF+VRCxlQ9bpheehyAJH89dAkaLEq/j/RZXJIqcgsmPJKUP1Z28HA==",
+      "dev": true
+    },
     "node_modules/abbrev": {
       "version": "1.1.1",
       "resolved": "https://registry.npmjs.org/abbrev/-/abbrev-1.1.1.tgz",
@@ -2776,6 +2785,14 @@
       "engines": {
         "node": ">=0.2.0"
       }
+    },
+    "node_modules/zod": {
+      "version": "3.21.4",
+      "resolved": "https://registry.npmjs.org/zod/-/zod-3.21.4.tgz",
+      "integrity": "sha512-m46AKbrzKVzOzs/DZgVnG5H55N1sv1M8qZU3A8RIKbs3mrACDNeIOeilDymVb2HdmP8uwshOCF4uJ8uM9rCqJw==",
+      "funding": {
+        "url": "https://github.com/sponsors/colinhacks"
+      }
     }
   }
 }

package.json

@@ -27,12 +27,15 @@
     "pino": "^8.11.0",
     "pino-http": "^8.3.3",
     "showdown": "^2.1.0",
-    "zlib": "^1.0.5"
+    "uuid": "^9.0.0",
+    "zlib": "^1.0.5",
+    "zod": "^3.21.4"
   },
   "devDependencies": {
     "@types/cors": "^2.8.13",
     "@types/express": "^4.17.17",
     "@types/showdown": "^2.0.0",
+    "@types/uuid": "^9.0.1",
     "concurrently": "^8.0.1",
     "esbuild": "^0.17.16",
     "esbuild-register": "^3.4.2",

src/admin/admin-users.yaml

@@ -0,0 +1,204 @@
+# Shat out by GPT-4, I did not check for correctness beyond a cursory glance
+openapi: 3.0.0
+info:
+  version: 1.0.0
+  title: User Management API
+paths:
+  /admin/users:
+    get:
+      summary: List all users
+      operationId: getUsers
+      responses:
+        "200":
+          description: A list of users
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  users:
+                    type: array
+                    items:
+                      $ref: "#/components/schemas/User"
+                  count:
+                    type: integer
+                    format: int32
+    post:
+      summary: Create a new user
+      operationId: createUser
+      responses:
+        "200":
+          description: The created user's token
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  token:
+                    type: string
+    put:
+      summary: Bulk upsert users
+      operationId: bulkUpsertUsers
+      requestBody:
+        content:
+          application/json:
+            schema:
+              type: object
+              properties:
+                users:
+                  type: array
+                  items:
+                    $ref: "#/components/schemas/User"
+      responses:
+        "200":
+          description: The upserted users
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  upserted_users:
+                    type: array
+                    items:
+                      $ref: "#/components/schemas/User"
+                  count:
+                    type: integer
+                    format: int32
+        "400":
+          description: Bad request
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  error:
+                    type: string
+
+  /admin/users/{token}:
+    get:
+      summary: Get a user by token
+      operationId: getUser
+      parameters:
+        - name: token
+          in: path
+          required: true
+          schema:
+            type: string
+      responses:
+        "200":
+          description: A user
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/User"
+        "404":
+          description: Not found
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  error:
+                    type: string
+    put:
+      summary: Update a user by token
+      operationId: upsertUser
+      parameters:
+        - name: token
+          in: path
+          required: true
+          schema:
+            type: string
+      requestBody:
+        content:
+          application/json:
+            schema:
+              $ref: "#/components/schemas/User"
+      responses:
+        "200":
+          description: The updated user
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/User"
+        "400":
+          description: Bad request
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  error:
+                    type: string
+    delete:
+      summary: Disables the user with the given token
+      description: Optionally accepts a `disabledReason` query parameter. Returns the disabled user.
+      parameters:
+        - in: path
+          name: token
+          required: true
+          schema:
+            type: string
+          description: The token of the user to disable
+        - in: query
+          name: disabledReason
+          required: false
+          schema:
+            type: string
+          description: The reason for disabling the user
+      responses:
+        '200':
+          description: The disabled user
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/User'
+        '400':
+          description: Bad request
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  error:
+                    type: string
+        '404':
+          description: Not found
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  error:
+                    type: string
+components:
+  schemas:
+    User:
+      type: object
+      properties:
+        token:
+          type: string
+        ip:
+          type: array
+          items:
+            type: string
+        type:
+          type: string
+          enum: ["normal", "special"]
+        promptCount:
+          type: integer
+          format: int32
+        tokenCount:
+          type: integer
+          format: int32
+        createdAt:
+          type: integer
+          format: int64
+        lastUsedAt:
+          type: integer
+          format: int64
+        disabledAt:
+          type: integer
+          format: int64
+        disabledReason:
+          type: string

src/admin/routes.ts

@@ -0,0 +1,36 @@
+import { RequestHandler, Router } from "express";
+import { config } from "../config";
+import { usersRouter } from "./users";
+
+const ADMIN_KEY = config.adminKey;
+const failedAttempts = new Map<string, number>();
+
+const adminRouter = Router();
+
+const auth: RequestHandler = (req, res, next) => {
+  const token = req.headers.authorization?.slice("Bearer ".length);
+  const attempts = failedAttempts.get(req.ip) ?? 0;
+  if (attempts > 5) {
+    req.log.warn(
+      { ip: req.ip, token },
+      `Blocked request to admin API due to too many failed attempts`
+    );
+    return res.status(401).json({ error: "Too many attempts" });
+  }
+
+  if (token !== ADMIN_KEY) {
+    const newAttempts = attempts + 1;
+    failedAttempts.set(req.ip, newAttempts);
+    req.log.warn(
+      { ip: req.ip, attempts: newAttempts, token },
+      `Attempted admin API request with invalid token`
+    );
+    return res.status(401).json({ error: "Unauthorized" });
+  }
+
+  next();
+};
+
+adminRouter.use(auth);
+adminRouter.use("/users", usersRouter);
+export { adminRouter };

src/admin/users.ts

@@ -0,0 +1,114 @@
+import { Router } from "express";
+import { z } from "zod";
+import * as userStore from "../proxy/auth/user-store";
+
+const usersRouter = Router();
+
+const UserSchema = z
+  .object({
+    ip: z.array(z.string()).optional(),
+    type: z.enum(["normal", "special"]).optional(),
+    promptCount: z.number().optional(),
+    tokenCount: z.number().optional(),
+    createdAt: z.number().optional(),
+    lastUsedAt: z.number().optional(),
+    disabledAt: z.number().optional(),
+    disabledReason: z.string().optional(),
+  })
+  .strict();
+
+const UserSchemaWithToken = UserSchema.extend({
+  token: z.string(),
+}).strict();
+
+/**
+ * Returns a list of all users, sorted by prompt count and then last used time.
+ * GET /admin/users
+ */
+usersRouter.get("/", (_req, res) => {
+  const users = userStore.getUsers().sort((a, b) => {
+    if (a.promptCount !== b.promptCount) {
+      return b.promptCount - a.promptCount;
+    }
+    return (b.lastUsedAt ?? 0) - (a.lastUsedAt ?? 0);
+  });
+  res.json({ users, count: users.length });
+});
+
+/**
+ * Returns the user with the given token.
+ * GET /admin/users/:token
+ */
+usersRouter.get("/:token", (req, res) => {
+  const user = userStore.getUser(req.params.token);
+  if (!user) {
+    return res.status(404).json({ error: "Not found" });
+  }
+  res.json(user);
+});
+
+/**
+ * Creates a new user.
+ * Returns the created user's token.
+ * POST /admin/users
+ */
+usersRouter.post("/", (_req, res) => {
+  res.json({ token: userStore.createUser() });
+});
+
+/**
+ * Updates the user with the given token, creating them if they don't exist.
+ * Accepts a JSON body containing at least one field on the User type.
+ * Returns the upserted user.
+ * PUT /admin/users/:token
+ */
+usersRouter.put("/:token", (req, res) => {
+  const result = UserSchema.safeParse(req.body);
+  if (!result.success) {
+    return res.status(400).json({ error: result.error });
+  }
+  userStore.upsertUser({ ...result.data, token: req.params.token });
+  res.json(userStore.getUser(req.params.token));
+});
+
+/**
+ * Bulk-upserts users given a list of User updates.
+ * Accepts a JSON body with the field `users` containing an array of updates.
+ * Returns an object containing the upserted users and the number of upserts.
+ * PUT /admin/users
+ */
+usersRouter.put("/", (req, res) => {
+  const result = z.array(UserSchemaWithToken).safeParse(req.body.users);
+  if (!result.success) {
+    return res.status(400).json({ error: result.error });
+  }
+  const upserts = result.data.map((user) => userStore.upsertUser(user));
+  res.json({
+    upserted_users: upserts,
+    count: upserts.length,
+  });
+});
+
+/**
+ * Disables the user with the given token. Optionally accepts a `disabledReason`
+ * query parameter.
+ * Returns the disabled user.
+ * DELETE /admin/users/:token
+ */
+usersRouter.delete("/:token", (req, res) => {
+  const user = userStore.getUser(req.params.token);
+  const disabledReason = z
+    .string()
+    .optional()
+    .safeParse(req.query.disabledReason);
+  if (!disabledReason.success) {
+    return res.status(400).json({ error: disabledReason.error });
+  }
+  if (!user) {
+    return res.status(404).json({ error: "Not found" });
+  }
+  userStore.disableUser(req.params.token, disabledReason.data);
+  res.json(userStore.getUser(req.params.token));
+});
+
+export { usersRouter };

src/config.ts

@@ -11,8 +11,30 @@ type Config = {
   port: number;
   /** OpenAI API key, either a single key or a comma-delimeted list of keys. */
   openaiKey?: string;
-  /** Proxy key. If set, requests must provide this key in the Authorization header to use the proxy. */
+  /**
+   * The proxy key to require for requests. Only applicable if the user
+   * management mode is set to 'proxy_key', and required if so.
+   **/
   proxyKey?: string;
+  /**
+   * The admin key to used for accessing the /admin API. Required if the user
+   * management mode is set to 'user_token'.
+   **/
+  adminKey?: string;
+  /**
+   * Which user management mode to use.
+   *
+   * `none`: No user management. Proxy is open to all requests with basic
+   *  abuse protection.
+   *
+   * `proxy_key`: A specific proxy key must be provided in the Authorization
+   *  header to use the proxy.
+   *
+   * `user_token`: Users must be created via the /admin REST API and provide
+   *  their personal access token in the Authorization header to use the proxy.
+   *  Configure this function and add users via the /admin API.
+   */
+  gatekeeper: "none" | "proxy_key" | "user_token";
   /** 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. */
@@ -35,16 +57,22 @@ type Config = {
   checkKeys?: boolean;
   /**
    * How to display quota information on the info page.
-   * 'none' - Hide quota information
-   * 'partial' - Display quota information only as a percentage
-   * 'full' - Display quota information as usage against total capacity
+   *
+   * `none` - Hide quota information
+   *
+   * `partial` - Display quota information only as a percentage
+   *
+   * `full` - Display quota information as usage against total capacity
    */
   quotaDisplayMode: "none" | "partial" | "full";
   /**
    * Which request queueing strategy to use when keys are over their rate limit.
-   * 'fair' - Requests are serviced in the order they were received (default)
-   * 'random' - Requests are serviced randomly
-   * 'none' - Requests are not queued and users have to retry manually
+   *
+   * `fair` - Requests are serviced in the order they were received (default)
+   *
+   * `random` - Requests are serviced randomly
+   *
+   * `none` - Requests are not queued and users have to retry manually
    */
   queueMode: DequeueMode;
 };
@@ -55,6 +83,8 @@ export const config: Config = {
   port: getEnvWithDefault("PORT", 7860),
   openaiKey: getEnvWithDefault("OPENAI_KEY", ""),
   proxyKey: getEnvWithDefault("PROXY_KEY", ""),
+  adminKey: getEnvWithDefault("ADMIN_KEY", ""),
+  gatekeeper: getEnvWithDefault("GATEKEEPER", "none"),
   modelRateLimit: getEnvWithDefault("MODEL_RATE_LIMIT", 4),
   maxOutputTokens: getEnvWithDefault("MAX_OUTPUT_TOKENS", 300),
   rejectDisallowed: getEnvWithDefault("REJECT_DISALLOWED", false),
@@ -75,22 +105,62 @@ export const config: Config = {
   queueMode: getEnvWithDefault("QUEUE_MODE", "fair"),
 } as const;
 
+/** Prevents the server from starting if config state is invalid. */
+export function assertConfigIsValid(): void {
+  // Ensure gatekeeper mode is valid.
+  if (!["none", "proxy_key", "user_token"].includes(config.gatekeeper)) {
+    throw new Error(
+      `Invalid gatekeeper mode: ${config.gatekeeper}. Must be one of: none, proxy_key, user_token.`
+    );
+  }
+
+  // Don't allow `user_token` mode without `ADMIN_KEY`.
+  if (config.gatekeeper === "user_token" && !config.adminKey) {
+    throw new Error(
+      "`user_token` gatekeeper mode requires an `ADMIN_KEY` to be set."
+    );
+  }
+
+  // Don't allow `proxy_key` mode without `PROXY_KEY`.
+  if (config.gatekeeper === "proxy_key" && !config.proxyKey) {
+    throw new Error(
+      "`proxy_key` gatekeeper mode requires a `PROXY_KEY` to be set."
+    );
+  }
+
+  // Don't allow `PROXY_KEY` to be set for other modes.
+  if (config.gatekeeper !== "proxy_key" && config.proxyKey) {
+    throw new Error(
+      "`PROXY_KEY` is set, but gatekeeper mode is not `proxy_key`. Make sure to set `GATEKEEPER=proxy_key`."
+    );
+  }
+}
+
+/** Masked, but not omitted as users may wish to see if they're set. */
 export const SENSITIVE_KEYS: (keyof Config)[] = [
-  "proxyKey",
-  "openaiKey",
   "googleSheetsKey",
   "googleSheetsSpreadsheetId",
 ];
+
+/** Omitted as they're not useful to display, masked or not. */
+export const OMITTED_KEYS: (keyof Config)[] = [
+  "port",
+  "logLevel",
+  "openaiKey",
+  "proxyKey",
+  "adminKey",
+];
+
 const getKeys = Object.keys as <T extends object>(obj: T) => Array<keyof T>;
 export function listConfig(): Record<string, string> {
   const result: Record<string, string> = {};
   for (const key of getKeys(config)) {
     const value = config[key]?.toString() || "";
-    
-    if (value === "" || value === "undefined") {
+
+    if (value === "" || value === "undefined" || OMITTED_KEYS.includes(key)) {
       continue;
     }
-    
+
     if (value && SENSITIVE_KEYS.includes(key)) {
       result[key] = "********";
     } else {

src/key-management/key-pool.ts

@@ -112,8 +112,8 @@ export class KeyPool {
       };
       this.keys.push(newKey);
 
-      this.log.info({ key: newKey.hash }, "Key added");
     }
+    this.log.info({ keyCount: this.keys.length }, "Loaded keys");
   }
 
   public init() {

src/proxy/auth.ts

@@ -1,17 +0,0 @@
-import type { Request, Response, NextFunction } from "express";
-import { config } from "../config";
-
-const PROXY_KEY = config.proxyKey;
-
-export const auth = (req: Request, res: Response, next: NextFunction) => {
-  if (!PROXY_KEY) {
-    next();
-    return;
-  }
-  if (req.headers.authorization === `Bearer ${PROXY_KEY}`) {
-    delete req.headers.authorization;
-    next();
-  } else {
-    res.status(401).json({ error: "Unauthorized" });
-  }
-};

src/proxy/auth/gatekeeper.ts

@@ -0,0 +1,46 @@
+import type { RequestHandler } from "express";
+import { config } from "../../config";
+import { authenticate, getUser } from "./user-store";
+
+const GATEKEEPER = config.gatekeeper;
+const PROXY_KEY = config.proxyKey;
+const ADMIN_KEY = config.adminKey;
+
+export const gatekeeper: RequestHandler = (req, res, next) => {
+  const token = req.headers.authorization?.slice("Bearer ".length);
+  delete req.headers.authorization;
+
+  // TODO: Generate anonymous users based on IP address for public or proxy_key
+  // modes so that all middleware can assume a user of some sort is present.
+
+  if (token === ADMIN_KEY) {
+    return next();
+  }
+
+  if (GATEKEEPER === "none") {
+    return next();
+  }
+
+  if (GATEKEEPER === "proxy_key" && token === PROXY_KEY) {
+    return next();
+  }
+
+  if (GATEKEEPER === "user_token" && token) {
+    const user = authenticate(token, req.ip);
+    if (user) {
+      req.user = user;
+      return next();
+    } else {
+      const maybeBannedUser = getUser(token);
+      if (maybeBannedUser?.disabledAt) {
+        return res.status(403).json({
+          error: `Forbidden: ${
+            maybeBannedUser.disabledReason || "Token disabled"
+          }`,
+        });
+      }
+    }
+  }
+
+  res.status(401).json({ error: "Unauthorized" });
+};

src/proxy/auth/user-store.ts

@@ -0,0 +1,123 @@
+/**
+ * Basic user management. Handles creation and tracking of proxy users, personal
+ * access tokens, and quota management. No persistence is provided, users must
+ * be re-created on each proxy start via the /admin API.
+ *
+ * Users are identified solely by their personal access token. The token is
+ * used to authenticate the user for all proxied requests.
+ */
+
+import { v4 as uuid } from "uuid";
+
+export interface User {
+  /** The user's personal access token. */
+  token: string;
+  /** The IP addresses the user has connected from. */
+  ip: string[];
+  /** The user's privilege level. */
+  type: UserType;
+  /** The number of prompts the user has made. */
+  promptCount: number;
+  /** The number of tokens the user has consumed. Not yet implemented. */
+  tokenCount: number;
+  /** The time at which the user was created. */
+  createdAt: number;
+  /** The time at which the user last connected. */
+  lastUsedAt?: number;
+  /** The time at which the user was disabled, if applicable. */
+  disabledAt?: number;
+  /** The reason for which the user was disabled, if applicable. */
+  disabledReason?: string;
+}
+
+/**
+ * Possible privilege levels for a user.
+ * - `normal`: Default role. Subject to usual rate limits and quotas.
+ * - `special`: Special role. Higher quotas and exempt from auto-ban/lockout.
+ * TODO: implement auto-ban/lockout for normal users when they do naughty shit
+ */
+export type UserType = "normal" | "special";
+
+type UserUpdate = Partial<User> & Pick<User, "token">;
+
+const users: Map<string, User> = new Map();
+
+/** Creates a new user and returns their token. */
+export function createUser() {
+  const token = uuid();
+  users.set(token, {
+    token,
+    ip: [],
+    type: "normal",
+    promptCount: 0,
+    tokenCount: 0,
+    createdAt: Date.now(),
+  });
+  return token;
+}
+
+/** Returns the user with the given token if they exist. */
+export function getUser(token: string) {
+  return users.get(token);
+}
+
+/** Returns a list of all users. */
+export function getUsers() {
+  return Array.from(users.values()).map((user) => ({ ...user }));
+}
+
+/**
+ * Upserts the given user. Intended for use with the /admin API for updating
+ * user information via JSON. Use other functions for more specific operations.
+ */
+export function upsertUser(user: UserUpdate) {
+  const existing: User = users.get(user.token) ?? {
+    token: user.token,
+    ip: [],
+    type: "normal",
+    promptCount: 0,
+    tokenCount: 0,
+    createdAt: Date.now(),
+  };
+
+  users.set(user.token, {
+    ...existing,
+    ...user,
+  });
+  return users.get(user.token);
+}
+
+/** Increments the prompt count for the given user. */
+export function incrementPromptCount(token: string) {
+  const user = users.get(token);
+  if (!user) return;
+  user.promptCount++;
+}
+
+/** Increments the token count for the given user by the given amount. */
+export function incrementTokenCount(token: string, amount = 1) {
+  const user = users.get(token);
+  if (!user) return;
+  user.tokenCount += amount;
+}
+
+/**
+ * Given a user's token and IP address, authenticates the user and adds the IP
+ * to the user's list of IPs. Returns the user if they exist and are not
+ * disabled, otherwise returns undefined.
+ */
+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);
+  user.lastUsedAt = Date.now();
+  return user;
+}
+
+/** Disables the given user, optionally providing a reason. */
+export function disableUser(token: string, reason?: string) {
+  const user = users.get(token);
+  if (!user) return;
+  user.disabledAt = Date.now();
+  user.disabledReason = reason;
+}

src/proxy/middleware/response/index.ts

@@ -9,6 +9,7 @@ import { keyPool } from "../../../key-management";
 import { buildFakeSseMessage, enqueue, trackWaitTime } from "../../queue";
 import { handleStreamedResponse } from "./handle-streamed-response";
 import { logPrompt } from "./log-prompt";
+import { incrementPromptCount } from "../../auth/user-store";
 
 export const QUOTA_ROUTES = ["/v1/chat/completions"];
 const DECODER_MAP = {
@@ -369,6 +370,9 @@ export const handleInternalError: httpProxy.ErrorCallback = (
 const incrementKeyUsage: ProxyResHandlerWithBody = async (_proxyRes, req) => {
   if (QUOTA_ROUTES.includes(req.path)) {
     keyPool.incrementPrompt(req.key?.hash);
+    if (req.user) {
+      incrementPromptCount(req.user.token);
+    }
   }
 };
 

src/proxy/routes.ts

@@ -5,13 +5,13 @@ subset of the API is supported. Kobold requests must be transformed into
 equivalent OpenAI requests. */
 
 import * as express from "express";
-import { auth } from "./auth";
+import { gatekeeper } from "./auth/gatekeeper";
 import { kobold } from "./kobold";
 import { openai } from "./openai";
 
 const router = express.Router();
 
-router.use(auth);
+router.use(gatekeeper);
 router.use("/kobold", kobold);
 router.use("/openai", openai);
 

src/server.ts

@@ -1,4 +1,4 @@
-import { config } from "./config";
+import { assertConfigIsValid, config } from "./config";
 import "source-map-support/register";
 import express from "express";
 import cors from "cors";
@@ -6,6 +6,7 @@ import pinoHttp from "pino-http";
 import childProcess from "child_process";
 import { logger } from "./logger";
 import { keyPool } from "./key-management";
+import { adminRouter } from "./admin/routes";
 import { proxyRouter, rewriteTavernRequests } from "./proxy/routes";
 import { handleInfoPage } from "./info-page";
 import { logQueue } from "./prompt-logging";
@@ -28,6 +29,7 @@ app.use(
         'res.headers["set-cookie"]',
         "req.headers.authorization",
         'req.headers["x-forwarded-for"]',
+        'req.headers["x-real-ip"]',
       ],
       censor: "********",
     },
@@ -50,6 +52,7 @@ app.set("trust proxy", true);
 
 // routes
 app.get("/", handleInfoPage);
+app.use("/admin", adminRouter);
 app.use("/proxy", proxyRouter);
 
 // 500 and 404
@@ -74,6 +77,8 @@ app.use((_req: unknown, res: express.Response) => {
 
 // start server and load keys
 app.listen(PORT, async () => {
+  assertConfigIsValid();
+
   try {
     // Huggingface seems to have changed something about how they deploy Spaces
     // and git commands fail because of some ownership issue with the .git

src/types/custom.d.ts

@@ -1,11 +1,13 @@
 import { Express } from "express-serve-static-core";
 import { Key } from "../key-management/key-pool";
+import { User } from "../proxy/auth/user-store";
 
 declare global {
   namespace Express {
     interface Request {
       key?: Key;
       api: "kobold" | "openai" | "anthropic";
+      user: User;
       isStreaming?: boolean;
       startTime: number;
       retryCount: number;