mirror of
https://github.com/fosrl/pangolin.git
synced 2026-01-28 22:00:51 +00:00
Chungus 2.0
This commit is contained in:
Owen
454
server/private/lib/rateLimit.ts
Normal file
454
server/private/lib/rateLimit.ts
Normal file
@@ -0,0 +1,454 @@
|
||||
/*
|
||||
* This file is part of a proprietary work.
|
||||
*
|
||||
* Copyright (c) 2025 Fossorial, Inc.
|
||||
* All rights reserved.
|
||||
*
|
||||
* This file is licensed under the Fossorial Commercial License.
|
||||
* You may not use this file except in compliance with the License.
|
||||
* Unauthorized use, copying, modification, or distribution is strictly prohibited.
|
||||
*
|
||||
* This file is not licensed under the AGPLv3.
|
||||
*/
|
||||
|
||||
import logger from "@server/logger";
|
||||
import redisManager from "@server/private/lib/redis";
|
||||
import { build } from "@server/build";
|
||||
|
||||
// Rate limiting configuration
|
||||
export const RATE_LIMIT_WINDOW = 60; // 1 minute in seconds
|
||||
export const RATE_LIMIT_MAX_REQUESTS = 100;
|
||||
export const RATE_LIMIT_PER_MESSAGE_TYPE = 20; // Per message type limit within the window
|
||||
|
||||
// Configuration for batched Redis sync
|
||||
export const REDIS_SYNC_THRESHOLD = 15; // Sync to Redis every N messages
|
||||
export const REDIS_SYNC_FORCE_INTERVAL = 30000; // Force sync every 30 seconds as backup
|
||||
|
||||
interface RateLimitTracker {
|
||||
count: number;
|
||||
windowStart: number;
|
||||
pendingCount: number;
|
||||
lastSyncedCount: number;
|
||||
}
|
||||
|
||||
interface RateLimitResult {
|
||||
isLimited: boolean;
|
||||
reason?: string;
|
||||
totalHits?: number;
|
||||
resetTime?: Date;
|
||||
}
|
||||
|
||||
export class RateLimitService {
|
||||
private localRateLimitTracker: Map<string, RateLimitTracker> = new Map();
|
||||
private localMessageTypeRateLimitTracker: Map<string, RateLimitTracker> = new Map();
|
||||
private cleanupInterval: NodeJS.Timeout | null = null;
|
||||
private forceSyncInterval: NodeJS.Timeout | null = null;
|
||||
|
||||
constructor() {
|
||||
if (build == "oss") {
|
||||
return;
|
||||
}
|
||||
|
||||
// Start cleanup and sync intervals
|
||||
this.cleanupInterval = setInterval(() => {
|
||||
this.cleanupLocalRateLimit().catch((error) => {
|
||||
logger.error("Error during rate limit cleanup:", error);
|
||||
});
|
||||
}, 60000); // Run cleanup every minute
|
||||
|
||||
this.forceSyncInterval = setInterval(() => {
|
||||
this.forceSyncAllPendingData().catch((error) => {
|
||||
logger.error("Error during force sync:", error);
|
||||
});
|
||||
}, REDIS_SYNC_FORCE_INTERVAL);
|
||||
}
|
||||
|
||||
// Redis keys
|
||||
private getRateLimitKey(clientId: string): string {
|
||||
return `ratelimit:${clientId}`;
|
||||
}
|
||||
|
||||
private getMessageTypeRateLimitKey(clientId: string, messageType: string): string {
|
||||
return `ratelimit:${clientId}:${messageType}`;
|
||||
}
|
||||
|
||||
// Helper function to sync local rate limit data to Redis
|
||||
private async syncRateLimitToRedis(
|
||||
clientId: string,
|
||||
tracker: RateLimitTracker
|
||||
): Promise<void> {
|
||||
if (!redisManager.isRedisEnabled() || tracker.pendingCount === 0) return;
|
||||
|
||||
try {
|
||||
const currentTime = Math.floor(Date.now() / 1000);
|
||||
const globalKey = this.getRateLimitKey(clientId);
|
||||
|
||||
// Get current value and add pending count
|
||||
const currentValue = await redisManager.hget(
|
||||
globalKey,
|
||||
currentTime.toString()
|
||||
);
|
||||
const newValue = (
|
||||
parseInt(currentValue || "0") + tracker.pendingCount
|
||||
).toString();
|
||||
await redisManager.hset(globalKey, currentTime.toString(), newValue);
|
||||
|
||||
// Set TTL using the client directly
|
||||
if (redisManager.getClient()) {
|
||||
await redisManager
|
||||
.getClient()
|
||||
.expire(globalKey, RATE_LIMIT_WINDOW + 10);
|
||||
}
|
||||
|
||||
// Update tracking
|
||||
tracker.lastSyncedCount = tracker.count;
|
||||
tracker.pendingCount = 0;
|
||||
|
||||
logger.debug(`Synced global rate limit to Redis for client ${clientId}`);
|
||||
} catch (error) {
|
||||
logger.error("Failed to sync global rate limit to Redis:", error);
|
||||
}
|
||||
}
|
||||
|
||||
private async syncMessageTypeRateLimitToRedis(
|
||||
clientId: string,
|
||||
messageType: string,
|
||||
tracker: RateLimitTracker
|
||||
): Promise<void> {
|
||||
if (!redisManager.isRedisEnabled() || tracker.pendingCount === 0) return;
|
||||
|
||||
try {
|
||||
const currentTime = Math.floor(Date.now() / 1000);
|
||||
const messageTypeKey = this.getMessageTypeRateLimitKey(clientId, messageType);
|
||||
|
||||
// Get current value and add pending count
|
||||
const currentValue = await redisManager.hget(
|
||||
messageTypeKey,
|
||||
currentTime.toString()
|
||||
);
|
||||
const newValue = (
|
||||
parseInt(currentValue || "0") + tracker.pendingCount
|
||||
).toString();
|
||||
await redisManager.hset(
|
||||
messageTypeKey,
|
||||
currentTime.toString(),
|
||||
newValue
|
||||
);
|
||||
|
||||
// Set TTL using the client directly
|
||||
if (redisManager.getClient()) {
|
||||
await redisManager
|
||||
.getClient()
|
||||
.expire(messageTypeKey, RATE_LIMIT_WINDOW + 10);
|
||||
}
|
||||
|
||||
// Update tracking
|
||||
tracker.lastSyncedCount = tracker.count;
|
||||
tracker.pendingCount = 0;
|
||||
|
||||
logger.debug(
|
||||
`Synced message type rate limit to Redis for client ${clientId}, type ${messageType}`
|
||||
);
|
||||
} catch (error) {
|
||||
logger.error("Failed to sync message type rate limit to Redis:", error);
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize local tracker from Redis data
|
||||
private async initializeLocalTracker(clientId: string): Promise<RateLimitTracker> {
|
||||
const currentTime = Math.floor(Date.now() / 1000);
|
||||
const windowStart = currentTime - RATE_LIMIT_WINDOW;
|
||||
|
||||
if (!redisManager.isRedisEnabled()) {
|
||||
return {
|
||||
count: 0,
|
||||
windowStart: currentTime,
|
||||
pendingCount: 0,
|
||||
lastSyncedCount: 0
|
||||
};
|
||||
}
|
||||
|
||||
try {
|
||||
const globalKey = this.getRateLimitKey(clientId);
|
||||
const globalRateLimitData = await redisManager.hgetall(globalKey);
|
||||
|
||||
let count = 0;
|
||||
for (const [timestamp, countStr] of Object.entries(globalRateLimitData)) {
|
||||
const time = parseInt(timestamp);
|
||||
if (time >= windowStart) {
|
||||
count += parseInt(countStr);
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
count,
|
||||
windowStart: currentTime,
|
||||
pendingCount: 0,
|
||||
lastSyncedCount: count
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error("Failed to initialize global tracker from Redis:", error);
|
||||
return {
|
||||
count: 0,
|
||||
windowStart: currentTime,
|
||||
pendingCount: 0,
|
||||
lastSyncedCount: 0
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
private async initializeMessageTypeTracker(
|
||||
clientId: string,
|
||||
messageType: string
|
||||
): Promise<RateLimitTracker> {
|
||||
const currentTime = Math.floor(Date.now() / 1000);
|
||||
const windowStart = currentTime - RATE_LIMIT_WINDOW;
|
||||
|
||||
if (!redisManager.isRedisEnabled()) {
|
||||
return {
|
||||
count: 0,
|
||||
windowStart: currentTime,
|
||||
pendingCount: 0,
|
||||
lastSyncedCount: 0
|
||||
};
|
||||
}
|
||||
|
||||
try {
|
||||
const messageTypeKey = this.getMessageTypeRateLimitKey(clientId, messageType);
|
||||
const messageTypeRateLimitData = await redisManager.hgetall(messageTypeKey);
|
||||
|
||||
let count = 0;
|
||||
for (const [timestamp, countStr] of Object.entries(messageTypeRateLimitData)) {
|
||||
const time = parseInt(timestamp);
|
||||
if (time >= windowStart) {
|
||||
count += parseInt(countStr);
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
count,
|
||||
windowStart: currentTime,
|
||||
pendingCount: 0,
|
||||
lastSyncedCount: count
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error("Failed to initialize message type tracker from Redis:", error);
|
||||
return {
|
||||
count: 0,
|
||||
windowStart: currentTime,
|
||||
pendingCount: 0,
|
||||
lastSyncedCount: 0
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Main rate limiting function
|
||||
async checkRateLimit(
|
||||
clientId: string,
|
||||
messageType?: string,
|
||||
maxRequests: number = RATE_LIMIT_MAX_REQUESTS,
|
||||
messageTypeLimit: number = RATE_LIMIT_PER_MESSAGE_TYPE,
|
||||
windowMs: number = RATE_LIMIT_WINDOW * 1000
|
||||
): Promise<RateLimitResult> {
|
||||
const currentTime = Math.floor(Date.now() / 1000);
|
||||
const windowStart = currentTime - Math.floor(windowMs / 1000);
|
||||
|
||||
// Check global rate limit
|
||||
let globalTracker = this.localRateLimitTracker.get(clientId);
|
||||
|
||||
if (!globalTracker || globalTracker.windowStart < windowStart) {
|
||||
// New window or first request - initialize from Redis if available
|
||||
globalTracker = await this.initializeLocalTracker(clientId);
|
||||
globalTracker.windowStart = currentTime;
|
||||
this.localRateLimitTracker.set(clientId, globalTracker);
|
||||
}
|
||||
|
||||
// Increment global counters
|
||||
globalTracker.count++;
|
||||
globalTracker.pendingCount++;
|
||||
this.localRateLimitTracker.set(clientId, globalTracker);
|
||||
|
||||
// Check if global limit would be exceeded
|
||||
if (globalTracker.count >= maxRequests) {
|
||||
return {
|
||||
isLimited: true,
|
||||
reason: "global",
|
||||
totalHits: globalTracker.count,
|
||||
resetTime: new Date((globalTracker.windowStart + Math.floor(windowMs / 1000)) * 1000)
|
||||
};
|
||||
}
|
||||
|
||||
// Sync to Redis if threshold reached
|
||||
if (globalTracker.pendingCount >= REDIS_SYNC_THRESHOLD) {
|
||||
this.syncRateLimitToRedis(clientId, globalTracker);
|
||||
}
|
||||
|
||||
// Check message type specific rate limit if messageType is provided
|
||||
if (messageType) {
|
||||
const messageTypeKey = `${clientId}:${messageType}`;
|
||||
let messageTypeTracker = this.localMessageTypeRateLimitTracker.get(messageTypeKey);
|
||||
|
||||
if (!messageTypeTracker || messageTypeTracker.windowStart < windowStart) {
|
||||
// New window or first request for this message type - initialize from Redis if available
|
||||
messageTypeTracker = await this.initializeMessageTypeTracker(clientId, messageType);
|
||||
messageTypeTracker.windowStart = currentTime;
|
||||
this.localMessageTypeRateLimitTracker.set(messageTypeKey, messageTypeTracker);
|
||||
}
|
||||
|
||||
// Increment message type counters
|
||||
messageTypeTracker.count++;
|
||||
messageTypeTracker.pendingCount++;
|
||||
this.localMessageTypeRateLimitTracker.set(messageTypeKey, messageTypeTracker);
|
||||
|
||||
// Check if message type limit would be exceeded
|
||||
if (messageTypeTracker.count >= messageTypeLimit) {
|
||||
return {
|
||||
isLimited: true,
|
||||
reason: `message_type:${messageType}`,
|
||||
totalHits: messageTypeTracker.count,
|
||||
resetTime: new Date((messageTypeTracker.windowStart + Math.floor(windowMs / 1000)) * 1000)
|
||||
};
|
||||
}
|
||||
|
||||
// Sync to Redis if threshold reached
|
||||
if (messageTypeTracker.pendingCount >= REDIS_SYNC_THRESHOLD) {
|
||||
this.syncMessageTypeRateLimitToRedis(clientId, messageType, messageTypeTracker);
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
isLimited: false,
|
||||
totalHits: globalTracker.count,
|
||||
resetTime: new Date((globalTracker.windowStart + Math.floor(windowMs / 1000)) * 1000)
|
||||
};
|
||||
}
|
||||
|
||||
// Decrement function for skipSuccessfulRequests/skipFailedRequests functionality
|
||||
async decrementRateLimit(clientId: string, messageType?: string): Promise<void> {
|
||||
// Decrement global counter
|
||||
const globalTracker = this.localRateLimitTracker.get(clientId);
|
||||
if (globalTracker && globalTracker.count > 0) {
|
||||
globalTracker.count--;
|
||||
// We need to account for this in pending count to sync correctly
|
||||
globalTracker.pendingCount--;
|
||||
}
|
||||
|
||||
// Decrement message type counter if provided
|
||||
if (messageType) {
|
||||
const messageTypeKey = `${clientId}:${messageType}`;
|
||||
const messageTypeTracker = this.localMessageTypeRateLimitTracker.get(messageTypeKey);
|
||||
if (messageTypeTracker && messageTypeTracker.count > 0) {
|
||||
messageTypeTracker.count--;
|
||||
messageTypeTracker.pendingCount--;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Reset key function
|
||||
async resetKey(clientId: string): Promise<void> {
|
||||
// Remove from local tracking
|
||||
this.localRateLimitTracker.delete(clientId);
|
||||
|
||||
// Remove all message type entries for this client
|
||||
for (const [key] of this.localMessageTypeRateLimitTracker) {
|
||||
if (key.startsWith(`${clientId}:`)) {
|
||||
this.localMessageTypeRateLimitTracker.delete(key);
|
||||
}
|
||||
}
|
||||
|
||||
// Remove from Redis if enabled
|
||||
if (redisManager.isRedisEnabled()) {
|
||||
const globalKey = this.getRateLimitKey(clientId);
|
||||
await redisManager.del(globalKey);
|
||||
|
||||
// Get all message type keys for this client and delete them
|
||||
const client = redisManager.getClient();
|
||||
if (client) {
|
||||
const messageTypeKeys = await client.keys(`ratelimit:${clientId}:*`);
|
||||
if (messageTypeKeys.length > 0) {
|
||||
await Promise.all(messageTypeKeys.map(key => redisManager.del(key)));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Cleanup old local rate limit entries and force sync pending data
|
||||
private async cleanupLocalRateLimit(): Promise<void> {
|
||||
const currentTime = Math.floor(Date.now() / 1000);
|
||||
const windowStart = currentTime - RATE_LIMIT_WINDOW;
|
||||
|
||||
// Clean up global rate limit tracking and sync pending data
|
||||
for (const [clientId, tracker] of this.localRateLimitTracker.entries()) {
|
||||
if (tracker.windowStart < windowStart) {
|
||||
// Sync any pending data before cleanup
|
||||
if (tracker.pendingCount > 0) {
|
||||
await this.syncRateLimitToRedis(clientId, tracker);
|
||||
}
|
||||
this.localRateLimitTracker.delete(clientId);
|
||||
}
|
||||
}
|
||||
|
||||
// Clean up message type rate limit tracking and sync pending data
|
||||
for (const [key, tracker] of this.localMessageTypeRateLimitTracker.entries()) {
|
||||
if (tracker.windowStart < windowStart) {
|
||||
// Sync any pending data before cleanup
|
||||
if (tracker.pendingCount > 0) {
|
||||
const [clientId, messageType] = key.split(":", 2);
|
||||
await this.syncMessageTypeRateLimitToRedis(clientId, messageType, tracker);
|
||||
}
|
||||
this.localMessageTypeRateLimitTracker.delete(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Force sync all pending rate limit data to Redis
|
||||
async forceSyncAllPendingData(): Promise<void> {
|
||||
if (!redisManager.isRedisEnabled()) return;
|
||||
|
||||
logger.debug("Force syncing all pending rate limit data to Redis...");
|
||||
|
||||
// Sync all pending global rate limits
|
||||
for (const [clientId, tracker] of this.localRateLimitTracker.entries()) {
|
||||
if (tracker.pendingCount > 0) {
|
||||
await this.syncRateLimitToRedis(clientId, tracker);
|
||||
}
|
||||
}
|
||||
|
||||
// Sync all pending message type rate limits
|
||||
for (const [key, tracker] of this.localMessageTypeRateLimitTracker.entries()) {
|
||||
if (tracker.pendingCount > 0) {
|
||||
const [clientId, messageType] = key.split(":", 2);
|
||||
await this.syncMessageTypeRateLimitToRedis(clientId, messageType, tracker);
|
||||
}
|
||||
}
|
||||
|
||||
logger.debug("Completed force sync of pending rate limit data");
|
||||
}
|
||||
|
||||
// Cleanup function for graceful shutdown
|
||||
async cleanup(): Promise<void> {
|
||||
if (build == "oss") {
|
||||
return;
|
||||
}
|
||||
|
||||
// Clear intervals
|
||||
if (this.cleanupInterval) {
|
||||
clearInterval(this.cleanupInterval);
|
||||
}
|
||||
if (this.forceSyncInterval) {
|
||||
clearInterval(this.forceSyncInterval);
|
||||
}
|
||||
|
||||
// Force sync all pending data
|
||||
await this.forceSyncAllPendingData();
|
||||
|
||||
// Clear local data
|
||||
this.localRateLimitTracker.clear();
|
||||
this.localMessageTypeRateLimitTracker.clear();
|
||||
|
||||
logger.info("Rate limit service cleanup completed");
|
||||
}
|
||||
}
|
||||
|
||||
// Export singleton instance
|
||||
export const rateLimitService = new RateLimitService();
|
||||
Reference in New Issue
Block a user