60 lines
1.8 KiB
TypeScript
60 lines
1.8 KiB
TypeScript
import { createHmac } from "node:crypto";
|
|
import { sql } from "drizzle-orm";
|
|
import { getDb } from "@album/database";
|
|
|
|
function rateLimitSecret() {
|
|
const secret =
|
|
process.env.RATE_LIMIT_SECRET ?? process.env.BETTER_AUTH_SECRET;
|
|
if (!secret && process.env.NODE_ENV === "production") {
|
|
throw new Error("RATE_LIMIT_SECRET is required in production");
|
|
}
|
|
return secret ?? "album-development-rate-limit";
|
|
}
|
|
|
|
export function requestClientIdentifier(request: Request) {
|
|
return (
|
|
request.headers.get("x-real-ip") ??
|
|
request.headers.get("x-forwarded-for")?.split(",")[0]?.trim() ??
|
|
"unknown"
|
|
);
|
|
}
|
|
|
|
export async function consumeRateLimit(input: {
|
|
namespace: string;
|
|
identifier: string;
|
|
limit: number;
|
|
windowMs: number;
|
|
}) {
|
|
const key = createHmac("sha256", rateLimitSecret())
|
|
.update(`${input.namespace}:${input.identifier}`)
|
|
.digest("hex");
|
|
const resetAt = new Date(Date.now() + input.windowMs);
|
|
const rows = await getDb().execute<{
|
|
count: number;
|
|
reset_at: Date;
|
|
}>(sql`
|
|
INSERT INTO rate_limit_buckets ("key", "count", "reset_at")
|
|
VALUES (${key}, 1, ${resetAt.toISOString()}::timestamptz)
|
|
ON CONFLICT ("key") DO UPDATE SET
|
|
"count" = CASE
|
|
WHEN rate_limit_buckets."reset_at" <= now() THEN 1
|
|
ELSE rate_limit_buckets."count" + 1
|
|
END,
|
|
"reset_at" = CASE
|
|
WHEN rate_limit_buckets."reset_at" <= now() THEN EXCLUDED."reset_at"
|
|
ELSE rate_limit_buckets."reset_at"
|
|
END,
|
|
"updated_at" = now()
|
|
RETURNING "count", "reset_at"
|
|
`);
|
|
const row = rows[0];
|
|
if (!row) throw new Error("Could not evaluate rate limit");
|
|
return {
|
|
allowed: Number(row.count) <= input.limit,
|
|
retryAfterSeconds: Math.max(
|
|
Math.ceil((new Date(row.reset_at).getTime() - Date.now()) / 1_000),
|
|
1,
|
|
),
|
|
};
|
|
}
|