mirror of
https://github.com/fosrl/pangolin.git
synced 2025-12-16 13:06:27 +00:00
454 lines
16 KiB
TypeScript
454 lines
16 KiB
TypeScript
/*
|
|
* 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 "#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(); |