112 lines
2.7 KiB
TypeScript
112 lines
2.7 KiB
TypeScript
type WindowCounter = {
|
|
count: number
|
|
resetAt: number
|
|
}
|
|
|
|
const ipRequestStore = new Map<string, WindowCounter>()
|
|
const sessionRequestStore = new Map<string, WindowCounter>()
|
|
const sessionOutputStore = new Map<string, WindowCounter>()
|
|
|
|
function now() {
|
|
return Date.now()
|
|
}
|
|
|
|
function cleanup(store: Map<string, WindowCounter>) {
|
|
const currentTime = now()
|
|
|
|
for (const [key, value] of store.entries()) {
|
|
if (value.resetAt <= currentTime) {
|
|
store.delete(key)
|
|
}
|
|
}
|
|
}
|
|
|
|
function readOrCreate(store: Map<string, WindowCounter>, key: string, windowMs: number) {
|
|
cleanup(store)
|
|
|
|
const currentTime = now()
|
|
const existing = store.get(key)
|
|
|
|
if (!existing || existing.resetAt <= currentTime) {
|
|
const fresh = {
|
|
count: 0,
|
|
resetAt: currentTime + windowMs,
|
|
}
|
|
|
|
store.set(key, fresh)
|
|
return fresh
|
|
}
|
|
|
|
return existing
|
|
}
|
|
|
|
function consume(store: Map<string, WindowCounter>, key: string, windowMs: number, amount: number) {
|
|
const bucket = readOrCreate(store, key, windowMs)
|
|
bucket.count += amount
|
|
store.set(key, bucket)
|
|
return bucket
|
|
}
|
|
|
|
function peek(store: Map<string, WindowCounter>, key: string, windowMs: number) {
|
|
return readOrCreate(store, key, windowMs)
|
|
}
|
|
|
|
export function getChatRateLimitStatus({
|
|
ip,
|
|
maxIpRequests,
|
|
maxSessionRequests,
|
|
maxSessionOutputChars,
|
|
requestWindowMs,
|
|
sessionId,
|
|
outputWindowMs,
|
|
}: {
|
|
ip: string
|
|
maxIpRequests: number
|
|
maxSessionRequests: number
|
|
maxSessionOutputChars: number
|
|
outputWindowMs: number
|
|
requestWindowMs: number
|
|
sessionId: string
|
|
}) {
|
|
const ipBucket = peek(ipRequestStore, ip, requestWindowMs)
|
|
const sessionBucket = peek(sessionRequestStore, sessionId, requestWindowMs)
|
|
const outputBucket = peek(sessionOutputStore, sessionId, outputWindowMs)
|
|
|
|
const ipRemaining = Math.max(0, maxIpRequests - ipBucket.count)
|
|
const sessionRemaining = Math.max(0, maxSessionRequests - sessionBucket.count)
|
|
const outputCharsRemaining = Math.max(0, maxSessionOutputChars - outputBucket.count)
|
|
|
|
return {
|
|
ipRemaining,
|
|
sessionRemaining,
|
|
outputCharsRemaining,
|
|
requestResetAt: new Date(Math.max(ipBucket.resetAt, sessionBucket.resetAt)).toISOString(),
|
|
outputResetAt: new Date(outputBucket.resetAt).toISOString(),
|
|
blocked: ipRemaining <= 0 || sessionRemaining <= 0 || outputCharsRemaining <= 0,
|
|
}
|
|
}
|
|
|
|
export function consumeChatRequest({
|
|
ip,
|
|
requestWindowMs,
|
|
sessionId,
|
|
}: {
|
|
ip: string
|
|
requestWindowMs: number
|
|
sessionId: string
|
|
}) {
|
|
consume(ipRequestStore, ip, requestWindowMs, 1)
|
|
consume(sessionRequestStore, sessionId, requestWindowMs, 1)
|
|
}
|
|
|
|
export function consumeChatOutput({
|
|
chars,
|
|
outputWindowMs,
|
|
sessionId,
|
|
}: {
|
|
chars: number
|
|
outputWindowMs: number
|
|
sessionId: string
|
|
}) {
|
|
consume(sessionOutputStore, sessionId, outputWindowMs, chars)
|
|
}
|