Implement user persistence via Firebase (khanon/oai-reverse-proxy!8)

.env.example

@@ -14,7 +14,10 @@
 # Note: CHECK_KEYS is disabled by default in local development mode, but enabled
 # by default in production mode.
 
-# Optional settings for prompt logging to Google Sheets
+# Optional settings for user management. See docs/user-management.md.
+# GATEKEEPER=none
+
+# Optional settings for prompt logging. See docs/logging-sheets.md.
 # PROMPT_LOGGING=false
 
 # ------------------------------------------------------------------------------
@@ -26,9 +29,17 @@
 # You can add multiple keys by separating them with a comma.
 OPENAI_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
 
-# You can require a bearer token to access the proxy by setting the key below.
+# You can require a Bearer token for requests when using proxy_token gatekeeper.
 # PROXY_KEY=your-secret-key
 
+# You can set an admin key for user management when using user_token gatekeeper.
+# ADMIN_KEY=your-very-secret-key
+
+# These are used for various persistence features. Refer to the docs for more
+# info.
+# FIREBASE_KEY=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+# FIREBASE_RTDB_URL=https://xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx.firebaseio.com
+
 # This is only relevant if you want to use the prompt logging feature.
 # GOOGLE_SHEETS_SPREADSHEET_ID=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
 # GOOGLE_SHEETS_KEY=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx

docs/user-management.md

@@ -0,0 +1,65 @@
+# User Management
+
+The proxy supports several different user management strategies.  You can choose the one that best fits your needs by setting the `GATEKEEPER` environment variable.
+
+Several of these features require you to set secrets in your environment.  If using Huggingface Spaces to deploy, do not set these in your `.env` file because that file is public and anyone can see it.
+
+## Table of Contents
+- [No user management](#no-user-management-gatekeepernone)
+- [Single-password authentication](#single-password-authentication-gatekeeperproxy_key)
+- [Per-user authentication](#per-user-authentication-gatekeeperuser_token)
+  - [Memory](#memory)
+  - [Firebase Realtime Database](#firebase-realtime-database)
+    - [Firebase setup instructions](#firebase-setup-instructions)
+
+## No user management (`GATEKEEPER=none`)
+
+This is the default mode. The proxy will not require any authentication to access the server and offers basic IP-based rate limiting and anti-abuse features.
+
+## Single-password authentication (`GATEKEEPER=proxy_key`)
+
+This mode allows you to set a password that must be passed in the `Authentication` header of every request to the server as a bearer token.  This is useful if you want to restrict access to the server, but don't want to create a separate account for every user.
+
+To set the password, create a `PROXY_KEY` secret in your environment.
+
+## Per-user authentication (`GATEKEEPER=user_token`)
+
+This mode allows you to provision separate Bearer tokens for each user. You can manage users via the /admin/users REST API, which itself requires an admin Bearer token.
+
+To begin, set `ADMIN_KEY` to a secret value.  This will be used to authenticate requests to the /admin/users REST API.
+
+[You can find an OpenAPI specification for the /admin/users REST API here.](openapi-admin-users.yaml)
+ 
+By default, the proxy will store user data in memory. Naturally, this means that user data will be lost when the proxy is restarted, though you can use the bulk user import/export feature to save and restore user data manually or via a script. However, the proxy also supports persisting user data to an external data store with some additional configuration.
+
+Below are the supported data stores and their configuration options.
+
+### Memory
+
+This is the default data store (`GATEKEEPER_STORE=memory`)  User data will be stored in memory and will be lost when the proxy is restarted. You are responsible for downloading and re-uploading user data via the REST API if you want to persist it.
+
+### Firebase Realtime Database
+
+To use Firebase Realtime Database to persist user data, set the following environment variables:
+- `GATEKEEPER_STORE`: Set this to `firebase_rtdb`
+- **Secret** `FIREBASE_RTDB_URL`: The URL of your Firebase Realtime Database, e.g. `https://my-project-default-rtdb.firebaseio.com`
+- **Secret** `FIREBASE_KEY`: A base-64 encoded service account key for your Firebase project. Refer to the instructions below for how to create this key.
+
+**Firebase setup instructions**
+
+1. Go to the [Firebase console](https://console.firebase.google.com/) and click "Add project", then follow the prompts to create a new project.
+2. From the **Project Overview** page, click **All products** in the left sidebar, then click **Realtime Database**.
+3. Click **Create database** and choose **Start in test mode**.  Click **Enable**.
+    - Test mode is fine for this use case as it still requires authentication to access the database. You may wish to set up more restrictive rules if you plan to use the database for other purposes.
+    - The reference URL for the database will be displayed on the page. You will need this later.
+4. Click the gear icon next to **Project Overview** in the left sidebar, then click **Project settings**.
+5. Click the **Service accounts** tab, then click **Generate new private key**.
+6. The downloaded file contains your key. Encode it as base64 and set it as the `FIREBASE_KEY` secret in your environment.
+7. Set `FIREBASE_RTDB_URL` to the reference URL of your Firebase Realtime Database, e.g. `https://my-project-default-rtdb.firebaseio.com`.
+8. Set `GATEKEEPER_STORE` to `firebase_rtdb` in your environment if you haven't already.
+
+The proxy will attempt to connect to your Firebase Realtime Database at startup and will throw an error if it cannot connect.  If you see this error, check that your `FIREBASE_RTDB_URL` and `FIREBASE_KEY` secrets are set correctly.
+
+---
+
+Users are loaded from the database and changes are flushed periodically.  You can use the PUT /admin/users API to bulk import users and force a flush to the database.

package.json

@@ -22,6 +22,7 @@
     "cors": "^2.8.5",
     "dotenv": "^16.0.3",
     "express": "^4.18.2",
+    "firebase-admin": "^11.8.0",
     "googleapis": "^117.0.0",
     "http-proxy-middleware": "^3.0.0-beta.1",
     "pino": "^8.11.0",

src/config.ts

@@ -1,4 +1,5 @@
 import dotenv from "dotenv";
+import type firebase from "firebase-admin";
 dotenv.config();
 
 const isDev = process.env.NODE_ENV !== "production";
@@ -17,7 +18,7 @@ type Config = {
    **/
   proxyKey?: string;
   /**
-   * The admin key to used for accessing the /admin API. Required if the user
+   * The admin key used to access the /admin API. Required if the user
    * management mode is set to 'user_token'.
    **/
   adminKey?: string;
@@ -35,6 +36,19 @@ type Config = {
    *  Configure this function and add users via the /admin API.
    */
   gatekeeper: "none" | "proxy_key" | "user_token";
+  /**
+   * Persistence layer to use for user management.
+   *
+   * `memory`: Users are stored in memory and are lost on restart (default)
+   *
+   * `firebase_rtdb`: Users are stored in a Firebase Realtime Database; requires
+   *  `firebaseKey` and `firebaseRtdbUrl` to be set.
+   **/
+  gatekeeperStore: "memory" | "firebase_rtdb";
+  /** URL of the Firebase Realtime Database if using the Firebase RTDB store. */
+  firebaseRtdbUrl?: string;
+  /** Base64-encoded Firebase service account key if using the Firebase RTDB store. */
+  firebaseKey?: string;
   /** 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. */
@@ -58,21 +72,21 @@ type Config = {
   /**
    * How to display quota information on the info page.
    *
-   * `none` - Hide quota information
+   * `none`: Hide quota information
    *
-   * `partial` - Display quota information only as a percentage
+   * `partial`: Display quota information only as a percentage
    *
-   * `full` - Display quota information as usage against total capacity
+   * `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)
+   * `fair`: Requests are serviced in the order they were received (default)
    *
-   * `random` - Requests are serviced randomly
+   * `random`: Requests are serviced randomly
    *
-   * `none` - Requests are not queued and users have to retry manually
+   * `none`: Requests are not queued and users have to retry manually
    */
   queueMode: DequeueMode;
 };
@@ -85,6 +99,9 @@ export const config: Config = {
   proxyKey: getEnvWithDefault("PROXY_KEY", ""),
   adminKey: getEnvWithDefault("ADMIN_KEY", ""),
   gatekeeper: getEnvWithDefault("GATEKEEPER", "none"),
+  gatekeeperStore: getEnvWithDefault("GATEKEEPER_STORE", "memory"),
+  firebaseRtdbUrl: getEnvWithDefault("FIREBASE_RTDB_URL", undefined),
+  firebaseKey: getEnvWithDefault("FIREBASE_KEY", undefined),
   modelRateLimit: getEnvWithDefault("MODEL_RATE_LIMIT", 4),
   maxOutputTokens: getEnvWithDefault("MAX_OUTPUT_TOKENS", 300),
   rejectDisallowed: getEnvWithDefault("REJECT_DISALLOWED", false),
@@ -106,7 +123,7 @@ export const config: Config = {
 } as const;
 
 /** Prevents the server from starting if config state is invalid. */
-export function assertConfigIsValid(): void {
+export async function assertConfigIsValid() {
   // Ensure gatekeeper mode is valid.
   if (!["none", "proxy_key", "user_token"].includes(config.gatekeeper)) {
     throw new Error(
@@ -134,12 +151,29 @@ export function assertConfigIsValid(): void {
       "`PROXY_KEY` is set, but gatekeeper mode is not `proxy_key`. Make sure to set `GATEKEEPER=proxy_key`."
     );
   }
+
+  // Require appropriate firebase config if using firebase store.
+  if (
+    config.gatekeeperStore === "firebase_rtdb" &&
+    (!config.firebaseKey || !config.firebaseRtdbUrl)
+  ) {
+    throw new Error(
+      "Firebase RTDB store requires `FIREBASE_KEY` and `FIREBASE_RTDB_URL` to be set."
+    );
+  }
+
+  await maybeInitializeFirebase();
 }
 
-/** Masked, but not omitted as users may wish to see if they're set. */
+/**
+ * Masked, but not omitted as users may wish to see if they're set due to their
+ * implications on privacy.
+ */
 export const SENSITIVE_KEYS: (keyof Config)[] = [
   "googleSheetsKey",
   "googleSheetsSpreadsheetId",
+  "firebaseRtdbUrl",
+  "firebaseKey",
 ];
 
 /** Omitted as they're not useful to display, masked or not. */
@@ -184,3 +218,29 @@ function getEnvWithDefault<T>(name: string, defaultValue: T): T {
     return value as unknown as T;
   }
 }
+
+let firebaseApp: firebase.app.App | undefined;
+
+async function maybeInitializeFirebase() {
+  if (!config.gatekeeperStore.startsWith("firebase")) {
+    return;
+  }
+
+  const firebase = await import("firebase-admin");
+  const firebaseKey = Buffer.from(config.firebaseKey!, "base64").toString();
+  const app = firebase.initializeApp({
+    credential: firebase.credential.cert(JSON.parse(firebaseKey)),
+    databaseURL: config.firebaseRtdbUrl,
+  });
+
+  await app.database().ref("connection-test").set(Date.now());
+
+  firebaseApp = app;
+}
+
+export function getFirebaseApp(): firebase.app.App {
+  if (!firebaseApp) {
+    throw new Error("Firebase app not initialized.");
+  }
+  return firebaseApp;
+}

src/key-management/key-checker.ts

@@ -41,7 +41,7 @@ export class KeyChecker {
   }
 
   public start() {
-    this.log.info("Starting key checker");
+    this.log.info("Starting key checker...");
     this.scheduleNextCheck();
   }
 

src/proxy/auth/user-store.ts

@@ -1,13 +1,16 @@
 /**
  * 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.
+ * access tokens, and quota management. Supports in-memory and Firebase Realtime
+ * Database persistence stores.
  *
  * Users are identified solely by their personal access token. The token is
  * used to authenticate the user for all proxied requests.
  */
 
+import admin from "firebase-admin";
 import { v4 as uuid } from "uuid";
+import { config, getFirebaseApp } from "../../config";
+import { logger } from "../../logger";
 
 export interface User {
   /** The user's personal access token. */
@@ -41,6 +44,15 @@ export type UserType = "normal" | "special";
 type UserUpdate = Partial<User> & Pick<User, "token">;
 
 const users: Map<string, User> = new Map();
+const usersToFlush = new Set<string>();
+
+export async function init() {
+  logger.info({ store: config.gatekeeperStore }, "Initializing user store...");
+  if (config.gatekeeperStore === "firebase_rtdb") {
+    await initFirebase();
+  }
+  logger.info("User store initialized.");
+}
 
 /** Creates a new user and returns their token. */
 export function createUser() {
@@ -84,6 +96,13 @@ export function upsertUser(user: UserUpdate) {
     ...existing,
     ...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);
 }
 
@@ -92,6 +111,7 @@ export function incrementPromptCount(token: string) {
   const user = users.get(token);
   if (!user) return;
   user.promptCount++;
+  usersToFlush.add(token);
 }
 
 /** Increments the token count for the given user by the given amount. */
@@ -99,6 +119,7 @@ export function incrementTokenCount(token: string, amount = 1) {
   const user = users.get(token);
   if (!user) return;
   user.tokenCount += amount;
+  usersToFlush.add(token);
 }
 
 /**
@@ -111,6 +132,7 @@ export function authenticate(token: string, ip: string) {
   if (!user || user.disabledAt) return;
   if (!user.ip.includes(ip)) user.ip.push(ip);
   user.lastUsedAt = Date.now();
+  usersToFlush.add(token);
   return user;
 }
 
@@ -120,4 +142,58 @@ export function disableUser(token: string, reason?: string) {
   if (!user) return;
   user.disabledAt = Date.now();
   user.disabledReason = reason;
+  usersToFlush.add(token);
+}
+
+// TODO: Firebase persistence is pretend right now and just polls the in-memory
+// store to sync it with Firebase when it changes. Will refactor to abstract
+// persistence layer later so we can support multiple stores.
+let firebaseTimeout: NodeJS.Timeout | undefined;
+
+async function initFirebase() {
+  logger.info("Connecting to Firebase...");
+  const app = getFirebaseApp();
+  const db = admin.database(app);
+  const usersRef = db.ref("users");
+  const snapshot = await usersRef.once("value");
+  const users: Record<string, User> | null = snapshot.val();
+  firebaseTimeout = setInterval(flushUsers, 20 * 1000);
+  if (!users) {
+    logger.info("No users found in Firebase.");
+    return;
+  }
+  for (const token in users) {
+    upsertUser(users[token]);
+  }
+  usersToFlush.clear();
+  const numUsers = Object.keys(users).length;
+  logger.info({ users: numUsers }, "Loaded users from Firebase");
+}
+
+async function flushUsers() {
+  const app = getFirebaseApp();
+  const db = admin.database(app);
+  const usersRef = db.ref("users");
+  const updates: Record<string, User> = {};
+
+  for (const token of usersToFlush) {
+    const user = users.get(token);
+    if (!user) {
+      continue;
+    }
+    updates[token] = user;
+  }
+
+  usersToFlush.clear();
+
+  const numUpdates = Object.keys(updates).length;
+  if (numUpdates === 0) {
+    return;
+  }
+  
+  await usersRef.update(updates);
+  logger.info(
+    { users: Object.keys(updates).length },
+    "Flushed users to Firebase"
+  );
 }

src/proxy/queue.ts

@@ -191,7 +191,7 @@ function cleanQueue() {
     (waitTime) => now - waitTime.end > 90 * 1000
   );
   const removed = waitTimes.splice(0, index + 1);
-  log.info(
+  log.debug(
     { stalledRequests: oldRequests.length, prunedWaitTimes: removed.length },
     `Cleaning up request queue.`
   );

src/server.ts

@@ -11,22 +11,10 @@ import { proxyRouter, rewriteTavernRequests } from "./proxy/routes";
 import { handleInfoPage } from "./info-page";
 import { logQueue } from "./prompt-logging";
 import { start as startRequestQueue } from "./proxy/queue";
+import { init as initUserStore } from "./proxy/auth/user-store";
 
 const PORT = config.port;
 
-process.on("uncaughtException", (err: any) => {
-  logger.error(
-    { err, stack: err?.stack },
-    "UNCAUGHT EXCEPTION. Please report this error trace."
-  );
-});
-process.on("unhandledRejection", (err: any) => {
-  logger.error(
-    { err, stack: err?.stack },
-    "UNCAUGHT PROMISE REJECTION. Please report this error trace."
-  );
-});
-
 const app = express();
 // middleware
 app.use("/", rewriteTavernRequests);
@@ -88,10 +76,56 @@ app.use((_req: unknown, res: express.Response) => {
   res.status(404).json({ error: "Not found" });
 });
 
-// start server and load keys
-app.listen(PORT, async () => {
-  assertConfigIsValid();
+async function start() {
+  logger.info("Server starting up...");
+  setGitSha();
 
+  logger.info("Checking configs and external dependencies...");
+  await assertConfigIsValid();
+
+  keyPool.init();
+
+  if (config.gatekeeper === "user_token") {
+    await initUserStore();
+  }
+
+  if (config.promptLogging) {
+    logger.info("Starting prompt logging...");
+    logQueue.start();
+  }
+
+  if (config.queueMode !== "none") {
+    logger.info("Starting request queue...");
+    startRequestQueue();
+  }
+
+  app.listen(PORT, async () => {
+    logger.info({ port: PORT }, "Now listening for connections.");
+    registerUncaughtExceptionHandler();
+  });
+
+  logger.info(
+    { sha: process.env.COMMIT_SHA, nodeEnv: process.env.NODE_ENV },
+    "Startup complete."
+  );
+}
+
+function registerUncaughtExceptionHandler() {
+  process.on("uncaughtException", (err: any) => {
+    logger.error(
+      { err, stack: err?.stack },
+      "UNCAUGHT EXCEPTION. Please report this error trace."
+    );
+  });
+  process.on("unhandledRejection", (err: any) => {
+    logger.error(
+      { err, stack: err?.stack },
+      "UNCAUGHT PROMISE REJECTION. Please report this error trace."
+    );
+  });
+}
+
+function setGitSha() {
   try {
     // Huggingface seems to have changed something about how they deploy Spaces
     // and git commands fail because of some ownership issue with the .git
@@ -132,19 +166,6 @@ app.listen(PORT, async () => {
     );
     process.env.COMMIT_SHA = "unknown";
   }
+}
 
-  logger.info(
-    { sha: process.env.COMMIT_SHA },
-    `Server listening on port ${PORT}`
-  );
-  keyPool.init();
-
-  if (config.promptLogging) {
-    logger.info("Starting prompt logging...");
-    logQueue.start();
-  }
-  if (config.queueMode !== "none") {
-    logger.info("Starting request queue...");
-    startRequestQueue();
-  }
-});
+start();