import { RateLimitError } from '../utils/errors.js';
import { RateLimitConfig } from '../utils/types.js';
interface RateLimitBucket {
tokens: number;
lastRefill: number;
}
export class RateLimiter {
private globalBucket: RateLimitBucket;
private vaultBuckets: Map<string, RateLimitBucket> = new Map();
private operationBuckets: Map<string, RateLimitBucket> = new Map();
constructor(private config: RateLimitConfig) {
this.globalBucket = {
tokens: config.global.burstSize,
lastRefill: Date.now(),
};
}
checkLimit(operation: string, vault: string): void {
const now = Date.now();
// Check global limit
this.refillBucket(this.globalBucket, this.config.global.requestsPerMinute, now);
if (this.globalBucket.tokens < 1) {
throw new RateLimitError('Global rate limit exceeded');
}
this.globalBucket.tokens--;
// Check per-vault limit
const vaultKey = vault;
if (!this.vaultBuckets.has(vaultKey)) {
this.vaultBuckets.set(vaultKey, {
tokens: this.config.perVault.requestsPerMinute,
lastRefill: now,
});
}
const vaultBucket = this.vaultBuckets.get(vaultKey)!;
this.refillBucket(vaultBucket, this.config.perVault.requestsPerMinute, now);
if (vaultBucket.tokens < 1) {
throw new RateLimitError(`Rate limit exceeded for vault: ${vault}`);
}
vaultBucket.tokens--;
// Check per-operation limit
const opType = this.getOperationType(operation);
const opConfig = this.config.perOperation[opType];
if (opConfig) {
const opKey = `${vault}:${opType}`;
if (!this.operationBuckets.has(opKey)) {
this.operationBuckets.set(opKey, {
tokens: opConfig.requestsPerMinute,
lastRefill: now,
});
}
const opBucket = this.operationBuckets.get(opKey)!;
this.refillBucket(opBucket, opConfig.requestsPerMinute, now);
if (opBucket.tokens < 1) {
throw new RateLimitError(`Rate limit exceeded for operation: ${operation}`);
}
opBucket.tokens--;
}
}
private refillBucket(bucket: RateLimitBucket, tokensPerMinute: number, now: number): void {
const timePassed = now - bucket.lastRefill;
const tokensToAdd = (timePassed / 60000) * tokensPerMinute;
bucket.tokens = Math.min(tokensPerMinute, bucket.tokens + tokensToAdd);
bucket.lastRefill = now;
}
private getOperationType(operation: string): 'read' | 'write' | 'command' {
if (operation.startsWith('get_') || operation.startsWith('list_') || operation.startsWith('search_')) {
return 'read';
} else if (operation.startsWith('write_') || operation.startsWith('append_')) {
return 'write';
} else if (operation.startsWith('execute_') || operation.startsWith('open_')) {
return 'command';
}
return 'read'; // default
}
}
