/* * SPDX-FileCopyrightText: hazelnoot and other Sharkey contributors * SPDX-License-Identifier: AGPL-3.0-only */ import { Inject, Injectable } from '@nestjs/common'; import Redis from 'ioredis'; import { LoggerService } from '@/core/LoggerService.js'; import { TimeService } from '@/core/TimeService.js'; import { EnvService } from '@/core/EnvService.js'; import { DI } from '@/di-symbols.js'; import type Logger from '@/logger.js'; import { BucketRateLimit, LegacyRateLimit, LimitInfo, RateLimit, hasMinLimit, isLegacyRateLimit } from '@/misc/rate-limit-utils.js'; @Injectable() export class SkRateLimiterService { private readonly logger: Logger; private readonly disabled: boolean; constructor( @Inject(TimeService) private readonly timeService: TimeService, @Inject(DI.redis) private readonly redisClient: Redis.Redis, @Inject(LoggerService) loggerService: LoggerService, @Inject(EnvService) envService: EnvService, ) { this.logger = loggerService.getLogger('limiter'); this.disabled = envService.env.NODE_ENV !== 'production'; // TODO disable in TEST *only* } public async limit(limit: RateLimit, actor: string, factor = 1): Promise { if (this.disabled || factor === 0) { return { blocked: false, remaining: Number.MAX_SAFE_INTEGER, resetSec: 0, resetMs: 0, fullResetSec: 0, fullResetMs: 0, }; } if (factor < 0) { throw new Error(`Rate limit factor is zero or negative: ${factor}`); } if (isLegacyRateLimit(limit)) { return await this.limitLegacy(limit, actor, factor); } else { return await this.limitBucket(limit, actor, factor); } } private async limitLegacy(limit: LegacyRateLimit, actor: string, factor: number): Promise { const promises: Promise[] = []; // The "min" limit - if present - is handled directly. if (hasMinLimit(limit)) { promises.push( this.limitMin(limit, actor, factor), ); } // Convert the "max" limit into a leaky bucket with 1 drip / second rate. if (limit.max != null && limit.duration != null) { promises.push( this.limitBucket({ type: 'bucket', key: limit.key, size: limit.max, dripRate: Math.max(Math.round(limit.duration / limit.max), 1), }, actor, factor), ); } const [lim1, lim2] = await Promise.all(promises); return { blocked: (lim1?.blocked || lim2?.blocked) ?? false, remaining: Math.min(lim1?.remaining ?? Number.MAX_SAFE_INTEGER, lim2?.remaining ?? Number.MAX_SAFE_INTEGER), resetSec: Math.max(lim1?.resetSec ?? 0, lim2?.resetSec ?? 0), resetMs: Math.max(lim1?.resetMs ?? 0, lim2?.resetMs ?? 0), fullResetSec: Math.max(lim1?.fullResetSec ?? 0, lim2?.fullResetSec ?? 0), fullResetMs: Math.max(lim1?.fullResetMs ?? 0, lim2?.fullResetMs ?? 0), }; } private async limitMin(limit: LegacyRateLimit & { minInterval: number }, actor: string, factor: number): Promise { if (limit.minInterval === 0) return null; if (limit.minInterval < 0) throw new Error(`Invalid rate limit ${limit.key}: minInterval is negative (${limit.minInterval})`); const counter = await this.getLimitCounter(limit, actor, 'min'); const minInterval = Math.max(Math.ceil(limit.minInterval / factor), 0); // Update expiration if (counter.c > 0) { const isCleared = this.timeService.now - counter.t >= minInterval; if (isCleared) { counter.c = 0; } } const blocked = counter.c > 0; if (!blocked) { counter.c++; counter.t = this.timeService.now; } // Calculate limit status const resetMs = Math.max(Math.ceil(minInterval - (this.timeService.now - counter.t)), 0); const resetSec = Math.ceil(resetMs / 1000); const limitInfo: LimitInfo = { blocked, remaining: 0, resetSec, resetMs, fullResetSec: resetSec, fullResetMs: resetMs }; // Update the limit counter, but not if blocked if (!blocked) { // Don't await, or we will slow down the API. this.setLimitCounter(limit, actor, counter, resetSec, 'min') .catch(err => this.logger.error(`Failed to update limit ${limit.key}:min for ${actor}:`, err)); } return limitInfo; } private async limitBucket(limit: BucketRateLimit, actor: string, factor: number): Promise { if (limit.size < 1) throw new Error(`Invalid rate limit ${limit.key}: size is less than 1 (${limit.size})`); if (limit.dripRate != null && limit.dripRate < 1) throw new Error(`Invalid rate limit ${limit.key}: dripRate is less than 1 (${limit.dripRate})`); if (limit.dripSize != null && limit.dripSize < 1) throw new Error(`Invalid rate limit ${limit.key}: dripSize is less than 1 (${limit.dripSize})`); const counter = await this.getLimitCounter(limit, actor, 'bucket'); const bucketSize = Math.max(Math.ceil(limit.size * factor), 1); const dripRate = Math.ceil(limit.dripRate ?? 1000); const dripSize = Math.ceil(limit.dripSize ?? 1); // Update drips if (counter.c > 0) { const dripsSinceLastTick = Math.floor((this.timeService.now - counter.t) / dripRate) * dripSize; counter.c = Math.max(counter.c - dripsSinceLastTick, 0); } const blocked = counter.c >= bucketSize; if (!blocked) { counter.c++; counter.t = this.timeService.now; } // Calculate limit status const remaining = Math.max(bucketSize - counter.c, 0); const resetMs = remaining > 0 ? 0 : Math.max(dripRate - (this.timeService.now - counter.t), 0); const resetSec = Math.ceil(resetMs / 1000); const fullResetMs = Math.ceil(counter.c / dripSize) * dripRate; const fullResetSec = Math.ceil(fullResetMs / 1000); const limitInfo: LimitInfo = { blocked, remaining, resetSec, resetMs, fullResetSec, fullResetMs }; // Update the limit counter, but not if blocked if (!blocked) { // Don't await, or we will slow down the API. this.setLimitCounter(limit, actor, counter, fullResetSec, 'bucket') .catch(err => this.logger.error(`Failed to update limit ${limit.key} for ${actor}:`, err)); } return limitInfo; } private async getLimitCounter(limit: RateLimit, actor: string, subject: string): Promise { const key = createLimitKey(limit, actor, subject); const value = await this.redisClient.get(key); if (value == null) { return { t: 0, c: 0 }; } return JSON.parse(value); } private async setLimitCounter(limit: RateLimit, actor: string, counter: LimitCounter, expiration: number, subject: string): Promise { const key = createLimitKey(limit, actor, subject); const value = JSON.stringify(counter); const expirationSec = Math.max(expiration, 1); await this.redisClient.set(key, value, 'EX', expirationSec); } } function createLimitKey(limit: RateLimit, actor: string, subject: string): string { return `rl_${actor}_${limit.key}_${subject}`; } export interface LimitCounter { /** Timestamp */ t: number; /** Counter */ c: number; }