From a27163a6294d749b624df7248d0d0864f95acdb0 Mon Sep 17 00:00:00 2001 From: nai-degen Date: Sun, 5 Nov 2023 11:49:20 -0600 Subject: [PATCH] adds option to not disable keys when reaching IP limit --- .env.example | 2 ++ src/config.ts | 8 +++++++- src/proxy/gatekeeper.ts | 25 +++++++++++++----------- src/shared/users/user-store.ts | 35 ++++++++++++++++++++++------------ src/user/web/self-service.ts | 4 ++++ 5 files changed, 50 insertions(+), 24 deletions(-) diff --git a/.env.example b/.env.example index df5769f..794caa1 100644 --- a/.env.example +++ b/.env.example @@ -66,6 +66,8 @@ # Maximum number of unique IPs a user can connect from. (0 for unlimited) # MAX_IPS_PER_USER=0 +# Whether user_tokens should be automatically disabled when reaching the IP limit. +# MAX_IPS_AUTO_BAN=true # With user_token gatekeeper, whether to allow users to change their nickname. # ALLOW_NICKNAME_CHANGES=true diff --git a/src/config.ts b/src/config.ts index bd04c4e..a2dbbf2 100644 --- a/src/config.ts +++ b/src/config.ts @@ -65,11 +65,16 @@ type Config = { */ firebaseKey?: string; /** - * Maximum number of IPs per user, after which their token is disabled. + * Maximum number of IPs allowed per user token. * Users with the manually-assigned `special` role are exempt from this limit. * - Defaults to 0, which means that users are not IP-limited. */ maxIpsPerUser: number; + /** + * Whether a user token should be automatically disabled if it exceeds the + * `maxIpsPerUser` limit, or if only connections from new IPs are be rejected. + */ + maxIpsAutoBan: boolean; /** Per-IP limit for requests per minute to OpenAI's completions endpoint. */ modelRateLimit: number; /** @@ -172,6 +177,7 @@ export const config: Config = { gatekeeper: getEnvWithDefault("GATEKEEPER", "none"), gatekeeperStore: getEnvWithDefault("GATEKEEPER_STORE", "memory"), maxIpsPerUser: getEnvWithDefault("MAX_IPS_PER_USER", 0), + maxIpsAutoBan: getEnvWithDefault("MAX_IPS_AUTO_BAN", true), firebaseRtdbUrl: getEnvWithDefault("FIREBASE_RTDB_URL", undefined), firebaseKey: getEnvWithDefault("FIREBASE_KEY", undefined), modelRateLimit: getEnvWithDefault("MODEL_RATE_LIMIT", 4), diff --git a/src/proxy/gatekeeper.ts b/src/proxy/gatekeeper.ts index 89b9c95..02bf904 100644 --- a/src/proxy/gatekeeper.ts +++ b/src/proxy/gatekeeper.ts @@ -46,19 +46,22 @@ export const gatekeeper: RequestHandler = (req, res, 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) { + const { user, result } = authenticate(token, req.ip); + + switch (result) { + case "success": + req.user = user; + return next(); + case "limited": return res.status(403).json({ - error: `Forbidden: ${ - maybeBannedUser.disabledReason || "Token disabled" - }`, + error: `Forbidden: no more IPs can authenticate with this token`, }); - } + case "disabled": + const bannedUser = getUser(token); + if (bannedUser?.disabledAt) { + const reason = bannedUser.disabledReason || "Token disabled"; + return res.status(403).json({ error: `Forbidden: ${reason}` }); + } } } diff --git a/src/shared/users/user-store.ts b/src/shared/users/user-store.ts index 8924952..cc3ebbc 100644 --- a/src/shared/users/user-store.ts +++ b/src/shared/users/user-store.ts @@ -180,22 +180,33 @@ export function incrementTokenCount( * 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) { +export function authenticate(token: string, ip: string): + { user?: User; result: "success" | "disabled" | "not_found" | "limited" } + { const user = users.get(token); - if (!user || user.disabledAt) return; - if (!user.ip.includes(ip)) user.ip.push(ip); - - const configIpLimit = user.maxIps ?? config.maxIpsPerUser; - const ipLimit = - user.type === "special" || !configIpLimit ? Infinity : configIpLimit; - if (user.ip.length > ipLimit) { - disableUser(token, "IP address limit exceeded."); - return; + if (!user) return { result: "not_found" }; + if (user.disabledAt) return { result: "disabled" }; + + const newIp = !user.ip.includes(ip); + + const userLimit = user.maxIps ?? config.maxIpsPerUser; + const enforcedLimit = + user.type === "special" || !userLimit ? Infinity : userLimit; + + if (newIp && user.ip.length >= enforcedLimit) { + if (config.maxIpsAutoBan) { + user.ip.push(ip); + disableUser(token, "IP address limit exceeded."); + return { result: "disabled" }; + } + return { result: "limited" }; + } else if (newIp) { + user.ip.push(ip); } user.lastUsedAt = Date.now(); usersToFlush.add(token); - return user; + return { user, result: "success" }; } export function hasAvailableQuota( @@ -366,7 +377,7 @@ function getModelFamilyForQuotaUsage(model: string): ModelFamily { if (model.startsWith("claude")) { return "claude"; } - if(model.startsWith("anthropic.claude")) { + if (model.startsWith("anthropic.claude")) { return "aws-claude"; } throw new Error(`Unknown quota model family for model ${model}`); diff --git a/src/user/web/self-service.ts b/src/user/web/self-service.ts index ca564c6..bd6811d 100644 --- a/src/user/web/self-service.ts +++ b/src/user/web/self-service.ts @@ -46,6 +46,10 @@ router.post("/edit-nickname", (req, res) => { throw new ForbiddenError("Not logged in."); } else if (!config.allowNicknameChanges || existing.disabledAt) { throw new ForbiddenError("Nickname changes are not allowed."); + } else if (!config.maxIpsAutoBan && !existing.ip.includes(req.ip)) { + throw new ForbiddenError( + "Nickname changes are only allowed from registered IPs." + ); } const schema = UserPartialSchema.pick({ nickname: true })