adds option to not disable keys when reaching IP limit

This commit is contained in:
nai-degen 2023-11-05 11:49:20 -06:00
parent 5a8fb3aff6
commit a27163a629
5 changed files with 50 additions and 24 deletions

View File

@ -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

View File

@ -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),

View File

@ -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}` });
}
}
}

View File

@ -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}`);

View File

@ -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 })