diff --git a/.env.example b/.env.example index b8196a1..f8c1416 100644 --- a/.env.example +++ b/.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 diff --git a/src/config.ts b/src/config.ts index 46c4ece..7104484 100644 --- a/src/config.ts +++ b/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), diff --git a/src/proxy/auth/user-store.ts b/src/proxy/auth/user-store.ts index bc2dbca..9e27447 100644 --- a/src/proxy/auth/user-store.ts +++ b/src/proxy/auth/user-store.ts @@ -43,6 +43,8 @@ export type UserType = "normal" | "special"; type UserUpdate = Partial & Pick; +const MAX_IPS_PER_USER = config.maxIpsPerUser; + const users: Map = new Map(); const usersToFlush = new Set(); @@ -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;