Mirrors
/
oai-reverse-proxy
mirror of https://gitgud.io/khanon/oai-reverse-proxy.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

adds configurable PoW timeout and iteration count

This commit is contained in:
nai-degen 2024-05-21 12:38:41 -05:00
parent 63ab1a7685
commit b76db652e0
10 changed files with 228 additions and 84 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
.env.example
View File

@ -80,7 +80,8 @@ NODE_ENV=production
# CAPTCHA_MODE=none
# POW_TOKEN_HOURS=24
# POW_TOKEN_MAX_IPS=2
# POW-DIFFICULTY_LEVEL=low
# POW_DIFFICULTY_LEVEL=low
# POW_CHALLENGE_TIMEOUT=30
# ------------------------------------------------------------------------------
# Optional settings for user management, access control, and quota enforcement:

82
docs/pow-captcha.md
View File

@ -4,11 +4,13 @@ You can require users to complete a proof-of-work before they can access the
proxy. This can increase the cost of denial of service attacks and slow down
automated abuse.
When configured, users access the challenge UI and request a proof of work. The
server will generate a challenge according to the difficulty level you have set.
The user can then start the worker to solve the challenge. Once the challenge is
solved, the user can submit the solution to the server. The server will verify
the solution and issue a temporary token for that user.
When configured, users access the challenge UI and request a token. The server
sends a challenge to the client, which asks the user's browser to find a
solution to the challenge that meets a certain constraint (the difficulty
level). Once the user has found a solution, they can submit it to the server
and get a user token valid for a period you specify.
The proof-of-work challenge uses the argon2id hash function.
## Configuration
@ -21,37 +23,73 @@ CAPTCHA_MODE=proof_of_work
POW_TOKEN_HOURS=24
# Max number of IPs that can use a user_token issued via proof-of-work
POW_TOKEN_MAX_IPS=2
# The difficulty level of the proof-of-work challenge
# The difficulty level of the proof-of-work challenge. You can use one of the
# predefined levels specified below, or you can specify a custom number of
# expected hash iterations.
POW_DIFFICULTY_LEVEL=low
```
## Difficulty Levels
The difficulty level controls how long it takes to solve the proof-of-work,
specifically by adjusting the average number of iterations required to find a
valid solution. Due to randomness, the actual number of iterations required can
vary significantly.
The difficulty level controls how long, on average, it will take for a user to
solve the proof-of-work challenge. Due to randomness, the actual time can very
significantly; lucky users may solve the challenge in a fraction of the average
time, while unlucky users may take much longer.
You can adjust the difficulty while the proxy is running from the admin interface.
The difficulty level doesn't affect the speed of the hash function itself, only
the number of hashes that will need to be computed. Therefore, the time required
to complete the challenge scales linearly with the difficulty level's iteration
count.
### Extreme
You can adjust the difficulty level while the proxy is running from the admin
interface.
- Average of 4000 iterations required
- Not recommended unless you are expecting very high levels of abuse
### High
- Average of 1900 iterations required
### Medium
- Average of 900 iterations required
Be aware that there is a time limit for solving the challenge, by default set to
30 minutes. Above 'high' difficulty, you will probably need to increase the time
limit or it will be very hard for users with slow devices to find a solution
within the time limit.
### Low
- Average of 200 iterations required
- Default setting.
### Medium
- Average of 900 iterations required
### High
- Average of 1900 iterations required
### Extreme
- Average of 4000 iterations required
- Not recommended unless you are expecting very high levels of abuse
- May require increasing `POW_CHALLENGE_TIMEOUT`
### Custom
Setting `POW_DIFFICULTY_LEVEL` to an integer will use that number of iterations
as the difficulty level.
## Other challenge settings
- `POW_CHALLENGE_TIMEOUT`: The time limit for solving the challenge, in minutes.
Default is 30.
- `POW_TOKEN_HOURS`: The period of time for which a user token issued via proof-
of-work can be used. Default is 24 hours. Starts when the challenge is solved.
- `POW_TOKEN_MAX_IPS`: The maximum number of unique IPs that can use a single
user token issued via proof-of-work. Default is 2.
- `POW_TOKEN_PURGE_HOURS`: The period of time after which an expired user token
issued via proof-of-work will be removed from the database. Until it is
purged, users can refresh expired tokens by completing a half-difficulty
challenge. Default is 48 hours.
- `POW_MAX_TOKENS_PER_IP`: The maximum number of active user tokens that can
be associated with a single IP address. After this limit is reached, the
oldest token will be forcibly expired when a new token is issued. Set to 0
to disable this feature. Default is 0.
## Custom argon2id parameters
You can set custom argon2id parameters for the proof-of-work challenge.

3
src/admin/web/manage.ts
View File

@ -17,6 +17,7 @@ import {
} from "../../shared/users/schema";
import { getLastNImages } from "../../shared/file-storage/image-history";
import { blacklists, parseCidrs, whitelists } from "../../shared/cidr";
import { invalidatePowHmacKey } from "../../user/web/pow-captcha";
const router = Router();
@ -313,8 +314,10 @@ router.post("/maintenance", (req, res) => {
const temps = users.filter((u) => u.type === "temporary");
temps.forEach((user) => {
user.expiresAt = Date.now();
user.disabledReason = "Admin forced expiration."
userStore.upsertUser(user);
});
invalidatePowHmacKey()
flash.type = "success";
flash.message = `${temps.length} temporary users marked for expiration.`;
break;

5
src/admin/web/views/admin_anti-abuse.ejs
View File

@ -12,13 +12,12 @@
}
#token-manage {
width: 100%;
display: flex;
width: 100%;
}
#token-manage button {
padding: 0.5em;
margin: 0 0.5em;
flex-grow: 1;
margin: 0 0.5em;
}
</style>

5
src/admin/web/views/admin_create-user.ejs
View File

@ -51,9 +51,8 @@
<legend>Temporary User Options</legend>
<div class="temporary-user-fieldset">
<p class="full-width">
Temporary users will be disabled after the specified duration, and their records will be deleted 72 hours after that.
These options apply only to new
temporary users; existing ones use whatever options were in effect when they were created.
Temporary users will be disabled after the specified duration, and their records will be permanently deleted after some time.
These options apply only to new temporary users; existing ones use whatever options were in effect when they were created.
</p>
<label for="temporaryUserDuration" class="full-width">Access duration (in minutes)</label>
<input type="number" name="temporaryUserDuration" id="temporaryUserDuration" value="60" class="full-width" />

2
src/admin/web/views/admin_index.ejs
View File

@ -32,7 +32,7 @@
<form id="maintenanceForm" action="/admin/manage/maintenance" method="post">
<input id="_csrf" type="hidden" name="_csrf" value="<%= csrfToken %>" />
<input id="hiddenAction" type="hidden" name="action" value="" />
<div display="flex" flex-direction="column">
<div>
<fieldset>
<legend>Key Recheck</legend>
<button id="recheck-keys" type="button" onclick="submitForm('recheck')">Force Key Recheck</button>

76
src/config.ts
View File

@ -117,25 +117,58 @@ type Config = {
*/
captchaMode: "none" | "proof_of_work";
/**
* Duration in hours for which a PoW-issued temporary user token is valid.
* Duration (in hours) for which a PoW-issued temporary user token is valid.
*/
powTokenHours: number;
/** Maximum number of IPs allowed per PoW-issued temporary user token. */
/**
* The maximum number of IPs from which a single temporary user token can be
* used. Upon reaching the limit, the `maxIpsAutoBan` behavior is triggered.
*/
powTokenMaxIps: number;
/**
* Difficulty level for the proof-of-work. Refer to docs/pow-captcha.md for
* details on the available modes.
* Difficulty level for the proof-of-work challenge.
* - `low`: 200 iterations
* - `medium`: 900 iterations
* - `high`: 1900 iterations
* - `extreme`: 4000 iterations
* - `number`: A custom number of iterations to use.
*
* Difficulty level only affects the number of iterations used in the PoW,
* not the complexity of the hash itself. Therefore, the average time-to-solve
* will scale linearly with the number of iterations.
*
* Refer to docs/proof-of-work.md for guidance and hashrate benchmarks.
*/
powDifficultyLevel: "low" | "medium" | "high" | "extreme";
powDifficultyLevel: "low" | "medium" | "high" | "extreme" | number;
/**
* Duration in minutes before a PoW challenge expires. Users' browsers must
* solve the challenge within this time frame or it will be rejected.
* Defaults to 30 minutes. It should be kept somewhat low to prevent abusive
* clients from working on many challenges in parallel, but you may need to
* increase this value for higher difficulty levels or older devices will not
* be able to solve the challenge in time.
* Duration (in minutes) before a PoW challenge expires. Users' browsers must
* solve the challenge within this time frame or it will be rejected. Should
* be kept somewhat low to prevent abusive clients from working on many
* challenges in parallel, but you may need to increase this value for higher
* difficulty levels or older devices will not be able to solve the challenge
* in time.
*
* Defaults to 30 minutes.
*/
powChallengeTimeout: number;
/**
* Duration (in hours) before expired temporary user tokens are purged from
* the user database. Users can refresh expired tokens by solving a faster PoW
* challenge as long as the original token has not been purged. Once purged,
* the user must solve a full PoW challenge to obtain a new token.
*
* Defaults to 48 hours. At 0, tokens are purged immediately upon expiry.
*/
powTokenPurgeHours: number;
/**
* Maximum number of active temporary user tokens that can be associated with
* a single IP address. Note that this may impact users sending requests from
* hosted AI chat clients such as Agnaistic or RisuAI, as they may share IPs.
*
* When the limit is reached, the oldest token with the same IP will be
* expired. At 0, no limit is enforced. Defaults to 0.
*/
// powMaxTokensPerIp: number;
/** Per-user limit for requests per minute to text and chat models. */
textModelRateLimit: number;
/** Per-user limit for requests per minute to image generation models. */
@ -332,6 +365,7 @@ export const config: Config = {
powTokenMaxIps: getEnvWithDefault("POW_TOKEN_MAX_IPS", 2),
powDifficultyLevel: getEnvWithDefault("POW_DIFFICULTY_LEVEL", "low"),
powChallengeTimeout: getEnvWithDefault("POW_CHALLENGE_TIMEOUT", 30),
powTokenPurgeHours: getEnvWithDefault("POW_TOKEN_PURGE_HOURS", 48),
firebaseRtdbUrl: getEnvWithDefault("FIREBASE_RTDB_URL", undefined),
firebaseKey: getEnvWithDefault("FIREBASE_KEY", undefined),
textModelRateLimit: getEnvWithDefault("TEXT_MODEL_RATE_LIMIT", 4),
@ -486,12 +520,29 @@ export async function assertConfigIsValid() {
);
}
if (config.captchaMode === "proof_of_work" && config.gatekeeper !== "user_token") {
if (
config.captchaMode === "proof_of_work" &&
config.gatekeeper !== "user_token"
) {
throw new Error(
"Captcha mode 'proof_of_work' requires gatekeeper mode 'user_token'."
);
}
if (config.captchaMode === "proof_of_work") {
const val = config.powDifficultyLevel;
const isDifficulty =
typeof val === "string" &&
["low", "medium", "high", "extreme"].includes(val);
const isIterations =
typeof val === "number" && Number.isInteger(val) && val > 0;
if (!isDifficulty && !isIterations) {
throw new Error(
"Invalid POW_DIFFICULTY_LEVEL. Must be one of: low, medium, high, extreme, or a positive integer."
);
}
}
if (config.gatekeeper === "proxy_key" && !config.proxyKey) {
throw new Error(
"`proxy_key` gatekeeper mode requires a `PROXY_KEY` to be set."
@ -569,6 +620,7 @@ export const OMITTED_KEYS = [
"proxyEndpointRoute",
"adminWhitelist",
"ipBlacklist",
"powTokenPurgeHours",
] satisfies (keyof Config)[];
type OmitKeys = (typeof OMITTED_KEYS)[number];

7
src/server.ts
View File

@ -32,7 +32,12 @@ app.use(
pinoHttp({
quietReqLogger: true,
logger,
autoLogging: { ignore: ({ url }) => ["/health"].includes(url as string) },
autoLogging: {
ignore: ({ url }) => {
const ignoreList = ["/health", "/res", "/user_content"];
return ignoreList.some((path) => (url as string).startsWith(path));
},
},
redact: {
paths: [
"req.headers.cookie",

3
src/shared/users/user-store.ts
View File

@ -307,7 +307,8 @@ function cleanupExpiredTokens() {
user.meta.refreshable = config.captchaMode !== "none";
disabled++;
}
if (user.disabledAt && user.disabledAt + 72 * 60 * 60 * 1000 < now) {
const purgeTimeout = config.powTokenPurgeHours * 60 * 60 * 1000;
if (user.disabledAt && user.disabledAt + purgeTimeout < now) {
users.delete(user.token);
usersToFlush.add(user.token);
deleted++;

122
src/user/web/pow-captcha.ts
View File

@ -2,14 +2,28 @@ import crypto from "crypto";
import express from "express";
import argon2 from "@node-rs/argon2";
import { z } from "zod";
import { createUser, getUser, upsertUser } from "../../shared/users/user-store";
import {
authenticate,
createUser,
getUser,
upsertUser,
} from "../../shared/users/user-store";
import { config } from "../../config";
/** HMAC key for signing challenges; regenerated on startup */
const HMAC_KEY = crypto.randomBytes(32).toString("hex");
/** Lockout time after verification in milliseconds */
const LOCKOUT_TIME = 1000 * 60; // 60 seconds
/** HMAC key for signing challenges; regenerated on startup */
let hmacSecret = crypto.randomBytes(32).toString("hex");
/**
* Regenerate the HMAC key used for signing challenges. Calling this function
* will invalidate all existing challenges.
*/
export function invalidatePowHmacKey() {
hmacSecret = crypto.randomBytes(32).toString("hex");
}
const argon2Params = {
ARGON2_TIME_COST: parseInt(process.env.ARGON2_TIME_COST || "8"),
ARGON2_MEMORY_KB: parseInt(process.env.ARGON2_MEMORY_KB || String(1024 * 64)),
@ -100,11 +114,16 @@ setInterval(() => {
}, 1000);
function generateChallenge(clientIp?: string, token?: string): Challenge {
let workFactor = workFactors[config.powDifficultyLevel];
let workFactor =
(typeof config.powDifficultyLevel === "number"
? config.powDifficultyLevel
: workFactors[config.powDifficultyLevel]) || 1000;
// If this is a token refresh, halve the work factor
if (token) {
// Challenge difficulty is reduced for token refreshes
workFactor = Math.floor(workFactor / 2);
}
const hashBits = BigInt(argon2Params.ARGON2_HASH_LENGTH) * 8n;
const hashMax = 2n ** hashBits;
const targetValue = hashMax / BigInt(workFactor);
@ -123,7 +142,7 @@ function generateChallenge(clientIp?: string, token?: string): Challenge {
}
function signMessage(msg: any): string {
const hmac = crypto.createHmac("sha256", HMAC_KEY);
const hmac = crypto.createHmac("sha256", hmacSecret);
if (typeof msg === "object") {
hmac.update(JSON.stringify(msg));
} else {
@ -154,26 +173,33 @@ async function verifySolution(
return result;
}
function verifyTokenRefreshable(token?: string, logger?: any): boolean {
if (!token) {
logger?.warn("No token provided for refresh");
return false;
}
function verifyTokenRefreshable(token: string, req: express.Request) {
const ip = req.ip;
const user = getUser(token);
if (!user) {
logger?.warn({ token }, "No user found for token");
req.log.warn({ token }, "Cannot refresh token - not found");
return false;
}
if (user.type !== "temporary") {
logger?.warn({ token }, "User is not temporary");
req.log.warn({ token }, "Cannot refresh token - wrong token type");
return false;
}
logger?.info(
{ token, refreshable: user.meta?.refreshable },
"Token refreshable"
);
return user.meta?.refreshable;
if (!user.meta?.refreshable) {
req.log.warn({ token }, "Cannot refresh token - not refreshable");
return false;
}
if (!user.ip.includes(ip)) {
// If there are available slots, add the IP to the list
const { result } = authenticate(token, ip);
if (result === "limited") {
req.log.warn({ token, ip }, "Cannot refresh token - IP limit reached");
return false;
}
}
req.log.info({ token }, "Allowing token refresh");
return true;
}
const router = express.Router();
@ -192,10 +218,8 @@ router.post("/challenge", (req, res) => {
}
if (action === "refresh") {
if (!verifyTokenRefreshable(refreshToken, req.log)) {
res
.status(400)
.json({
if (!refreshToken || !verifyTokenRefreshable(refreshToken, req)) {
res.status(400).json({
error: "Not allowed to refresh that token; request a new one",
});
return;
@ -262,7 +286,7 @@ router.post("/verify", async (req, res) => {
return;
}
if (challenge.token && !verifyTokenRefreshable(challenge.token, req.log)) {
if (challenge.token && !verifyTokenRefreshable(challenge.token, req)) {
res.status(400).json({ error: "Not allowed to refresh that usertoken" });
return;
}
@ -271,6 +295,7 @@ router.post("/verify", async (req, res) => {
try {
const success = await verifySolution(challenge, solution, req.log);
if (!success) {
recentAttempts.set(ip, Date.now() + 1000 * 60 * 60 * 6);
req.log.warn("Solution failed verification");
res.status(400).json({ error: "Solution failed verification" });
return;
@ -294,21 +319,8 @@ router.post("/verify", async (req, res) => {
return res.json({ success: true, token: challenge.token });
}
} else {
const token = createUser({
type: "temporary",
expiresAt: Date.now() + config.powTokenHours * 60 * 60 * 1000,
});
upsertUser({
token,
ip: [ip],
maxIps: config.powTokenMaxIps,
meta: { refreshable: true },
});
req.log.info(
{ ip, token: `...${token.slice(-5)}` },
"Proof-of-work token issued"
);
return res.json({ success: true, token });
const newToken = issueToken(req);
return res.json({ success: true, token: newToken });
}
});
@ -318,7 +330,41 @@ router.get("/", (_req, res) => {
difficultyLevel: config.powDifficultyLevel,
tokenLifetime: config.powTokenHours,
tokenMaxIps: config.powTokenMaxIps,
challengeTimeout: config.powChallengeTimeout,
});
});
// const ipTokenCache = new Map<string, Set<string>>();
//
// function buildIpTokenCountCache() {
// ipTokenCache.clear();
// const users = getUsers().filter((u) => u.type === "temporary");
// for (const user of users) {
// for (const ip of user.ip) {
// if (!ipTokenCache.has(ip)) {
// ipTokenCache.set(ip, new Set());
// }
// ipTokenCache.get(ip)?.add(user.token);
// }
// }
// }
function issueToken(req: express.Request) {
const token = createUser({
type: "temporary",
expiresAt: Date.now() + config.powTokenHours * 60 * 60 * 1000,
});
upsertUser({
token,
ip: [req.ip],
maxIps: config.powTokenMaxIps,
meta: { refreshable: true },
});
req.log.info(
{ ip: req.ip, token: `...${token.slice(-5)}` },
"Proof-of-work token issued"
);
return token;
}
export { router as powRouter };