pkgs/bot/src/utils/rateLimit.ts (view raw)
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import { db, uwuifyRateLimit } from "@/database";
import { eq } from "drizzle-orm";
const LIMIT_PER_HOUR = 25;
const HOUR_MS = 60 * 60 * 1000;
export async function checkUwuifyRateLimit(userId: string): Promise<boolean> {
const now = new Date();
const existing = db
.select()
.from(uwuifyRateLimit)
.where(eq(uwuifyRateLimit.userId, userId))
.get();
if (!existing) {
await db.insert(uwuifyRateLimit).values({
userId,
count: 1,
resetAt: new Date(now.getTime() + HOUR_MS),
});
return true;
}
const now_ms = now.getTime();
const resetAt_ms = existing.resetAt.getTime();
if (now_ms >= resetAt_ms) {
await db
.update(uwuifyRateLimit)
.set({
count: 1,
resetAt: new Date(now_ms + HOUR_MS),
})
.where(eq(uwuifyRateLimit.userId, userId));
return true;
}
if (existing.count >= LIMIT_PER_HOUR) {
return false;
}
await db
.update(uwuifyRateLimit)
.set({ count: existing.count + 1 })
.where(eq(uwuifyRateLimit.userId, userId));
return true;
}
export function getRateLimitInfo(
userId: string,
): Promise<{ remaining: number; resetAt: Date } | null> {
return db
.select()
.from(uwuifyRateLimit)
.where(eq(uwuifyRateLimit.userId, userId))
.then((res) => {
if (!res || res.length === 0) return null;
const record = res[0];
const remaining = Math.max(0, LIMIT_PER_HOUR - record.count);
return { remaining, resetAt: record.resetAt };
});
}
|