misty.cafe/mistykey
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases Packages Wiki Activity Actions

manage caches in SkRateLimiterService

Browse source
This commit is contained in:
Hazelnoot 2025-10-01 11:30:35 -04:00
parent 338dad9a87
commit b32f3b5019
2 changed files with 32 additions and 32 deletions
Download patch file Download diff file Expand all files Collapse all files

24
packages/backend/src/server/SkRateLimiterService.ts
View file

@ -5,14 +5,14 @@
import { Inject, Injectable } from '@nestjs/common';
import Redis from 'ioredis';
import type { TimeService } from '@/core/TimeService.js';
import type { EnvService } from '@/core/EnvService.js';
import type { MiUser } from '@/models/_.js';
import { TimeService } from '@/core/TimeService.js';
import { EnvService } from '@/core/EnvService.js';
import { BucketRateLimit, LegacyRateLimit, LimitInfo, RateLimit, hasMinLimit, isLegacyRateLimit, Keyed, hasMaxLimit, disabledLimitInfo, MaxLegacyLimit, MinLegacyLimit } from '@/misc/rate-limit-utils.js';
import { RoleService } from '@/core/RoleService.js';
import { CacheManagementService, type ManagedMemoryKVCache } from '@/core/CacheManagementService.js';
import { ConflictError } from '@/misc/errors/ConflictError.js';
import { DI } from '@/di-symbols.js';
import { MemoryKVCache } from '@/misc/cache.js';
import type { MiUser } from '@/models/_.js';
import type { RoleService } from '@/core/RoleService.js';
// Sentinel value used for caching the default role template.
// Required because MemoryKVCache doesn't support null keys.
@ -30,26 +30,24 @@ interface ParsedLimit {
@Injectable()
export class SkRateLimiterService {
// 1-minute cache interval
private readonly factorCache = new MemoryKVCache<number>(1000 * 60);
// 10-second cache interval
private readonly lockoutCache = new MemoryKVCache<number>(1000 * 10);
private readonly factorCache: ManagedMemoryKVCache<number>;
private readonly lockoutCache: ManagedMemoryKVCache<number>;
private readonly requestCounts = new Map<string, number>();
private readonly disabled: boolean;
constructor(
@Inject('TimeService')
private readonly timeService: TimeService,
@Inject(DI.redisForRateLimit)
private readonly redisClient: Redis.Redis,
@Inject('RoleService')
private readonly roleService: RoleService,
private readonly timeService: TimeService,
@Inject('EnvService')
envService: EnvService,
cacheManagementService: CacheManagementService,
) {
this.factorCache = cacheManagementService.createMemoryKVCache<number>(1000 * 60); // 1m
this.lockoutCache = cacheManagementService.createMemoryKVCache<number>(1000 * 10); // 10s
this.disabled = envService.env.NODE_ENV === 'test';
}

40
packages/backend/test/unit/server/api/SkRateLimiterServiceTests.ts
View file

@ -3,30 +3,27 @@
* SPDX-License-Identifier: AGPL-3.0-only
*/
import { GodOfTimeService } from '../../../misc/GodOfTimeService.js';
import type Redis from 'ioredis';
import type { MiUser } from '@/models/User.js';
import type { RolePolicies, RoleService } from '@/core/RoleService.js';
import type { Config } from '@/config.js';
import { SkRateLimiterService } from '@/server/SkRateLimiterService.js';
import { BucketRateLimit, Keyed, LegacyRateLimit } from '@/misc/rate-limit-utils.js';
/* eslint-disable @typescript-eslint/no-non-null-assertion */
import { CacheManagementService } from '@/core/CacheManagementService.js';
import { InternalEventService } from '@/core/InternalEventService.js';
describe(SkRateLimiterService, () => {
let mockTimeService: { now: number, date: Date } = null!;
let mockRedis: Array<(command: [string, ...unknown[]]) => [Error | null, unknown] | null> = null!;
let mockRedisExec: (batch: [string, ...unknown[]][]) => Promise<[Error | null, unknown][] | null> = null!;
let mockEnvironment: Record<string, string | undefined> = null!;
let serviceUnderTest: () => SkRateLimiterService = null!;
let mockDefaultUserPolicies: Partial<RolePolicies> = null!;
let mockUserPolicies: Record<string, Partial<RolePolicies>> = null!;
let mockTimeService: GodOfTimeService;
let mockRedis: Array<(command: [string, ...unknown[]]) => [Error | null, unknown] | null>;
let mockRedisExec: (batch: [string, ...unknown[]][]) => Promise<[Error | null, unknown][] | null>;
let mockEnvironment: Record<string, string | undefined>;
let serviceUnderTest: () => SkRateLimiterService;
let mockDefaultUserPolicies: Partial<RolePolicies>;
let mockUserPolicies: Record<string, Partial<RolePolicies>>;
beforeEach(() => {
mockTimeService = {
now: 0,
get date() {
return new Date(mockTimeService.now);
},
};
mockTimeService = new GodOfTimeService();
function callMockRedis(command: [string, ...unknown[]]) {
const handlerResults = mockRedis.map(handler => handler(command));
@ -62,9 +59,10 @@ describe(SkRateLimiterService, () => {
},
};
},
reset() {
return Promise.resolve();
},
reset: () => Promise.resolve(),
publish: () => Promise.resolve(),
on() {},
off() {},
} as unknown as Redis.Redis;
mockEnvironment = Object.create(process.env);
@ -82,9 +80,13 @@ describe(SkRateLimiterService, () => {
},
} as unknown as RoleService;
const fakeConfig = { host: 'example.com' } as unknown as Config;
const internalEventService = new InternalEventService(mockRedisClient, mockRedisClient, fakeConfig);
const cacheManagementService = new CacheManagementService(mockRedisClient, mockTimeService, internalEventService);
let service: SkRateLimiterService | undefined = undefined;
serviceUnderTest = () => {
return service ??= new SkRateLimiterService(mockTimeService, mockRedisClient, mockRoleService, mockEnvService);
return service ??= new SkRateLimiterService(mockRedisClient, mockRoleService, mockTimeService, mockEnvService, cacheManagementService);
};
});