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

lookup and cache rate limit factors directly within SkRateLimiterService

Browse source
This commit is contained in:
Hazelnoot 2025-02-02 22:02:08 -05:00
parent 402933004a
commit 09669d72e7
7 changed files with 124 additions and 71 deletions
Download patch file Download diff file Expand all files Collapse all files

37
packages/backend/src/server/api/ApiCallService.ts
View file

@ -313,35 +313,30 @@ export class ApiCallService implements OnApplicationShutdown {
// eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
if (endpointLimit) {
// koa will automatically load the `X-Forwarded-For` header if `proxy: true` is configured in the app.
let limitActor: string;
let limitActor: string | MiLocalUser;
if (user) {
limitActor = user.id;
limitActor = user;
} else {
limitActor = getIpHash(request.ip);
}
// TODO: 毎リクエスト計算するのもあれだしキャッシュしたい
const factor = user ? (await this.roleService.getUserPolicies(user.id)).rateLimitFactor : 1;
const limit = {
key: ep.name,
...endpointLimit,
};
if (factor > 0) {
const limit = {
key: ep.name,
...endpointLimit,
};
// Rate limit
const info = await this.rateLimiterService.limit(limit, limitActor);
// Rate limit
const info = await this.rateLimiterService.limit(limit, limitActor, factor);
sendRateLimitHeaders(reply, info);
sendRateLimitHeaders(reply, info);
if (info.blocked) {
throw new ApiError({
message: 'Rate limit exceeded. Please try again later.',
code: 'RATE_LIMIT_EXCEEDED',
id: 'd5826d14-3982-4d2e-8011-b9e9f02499ef',
httpStatusCode: 429,
}, info);
}
if (info.blocked) {
throw new ApiError({
message: 'Rate limit exceeded. Please try again later.',
code: 'RATE_LIMIT_EXCEEDED',
id: 'd5826d14-3982-4d2e-8011-b9e9f02499ef',
httpStatusCode: 429,
}, info);
}
}

45
packages/backend/src/server/api/SkRateLimiterService.ts
View file

@ -5,36 +5,59 @@
import { Inject, Injectable } from '@nestjs/common';
import Redis from 'ioredis';
import { TimeService } from '@/core/TimeService.js';
import { EnvService } from '@/core/EnvService.js';
import type { TimeService } from '@/core/TimeService.js';
import type { EnvService } from '@/core/EnvService.js';
import { BucketRateLimit, LegacyRateLimit, LimitInfo, RateLimit, hasMinLimit, isLegacyRateLimit, Keyed, hasMaxLimit, disabledLimitInfo, MaxLegacyLimit, MinLegacyLimit } from '@/misc/rate-limit-utils.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.
const defaultUserKey = '';
@Injectable()
export class SkRateLimiterService {
// 1-minute cache interval
private readonly factorCache = new MemoryKVCache<number>(1000 * 60);
private readonly disabled: boolean;
constructor(
@Inject(TimeService)
@Inject('TimeService')
private readonly timeService: TimeService,
@Inject(DI.redis)
private readonly redisClient: Redis.Redis,
@Inject(EnvService)
@Inject('RoleService')
private readonly roleService: RoleService,
@Inject('EnvService')
envService: EnvService,
) {
this.disabled = envService.env.NODE_ENV === 'test';
}
/**
* Check & increment a rate limit
* Check & increment a rate limit for a client
* @param limit The limit definition
* @param actor Client who is calling this limit
* @param factor Scaling factor - smaller = larger limit (less restrictive)
* @param actorOrUser authenticated client user or IP hash
*/
public async limit(limit: Keyed<RateLimit>, actor: string, factor = 1): Promise<LimitInfo> {
if (this.disabled || factor === 0) {
public async limit(limit: Keyed<RateLimit>, actorOrUser: string | MiUser): Promise<LimitInfo> {
if (this.disabled) {
return disabledLimitInfo;
}
const actor = typeof(actorOrUser) === 'object' ? actorOrUser.id : actorOrUser;
const userCacheKey = typeof(actorOrUser) === 'object' ? actorOrUser.id : defaultUserKey;
const userRoleKey = typeof(actorOrUser) === 'object' ? actorOrUser.id : null;
const factor = this.factorCache.get(userCacheKey) ?? await this.factorCache.fetch(userCacheKey, async () => {
const role = await this.roleService.getUserPolicies(userRoleKey);
return role.rateLimitFactor;
});
if (factor === 0) {
return disabledLimitInfo;
}
@ -42,10 +65,6 @@ export class SkRateLimiterService {
throw new Error(`Rate limit factor is zero or negative: ${factor}`);
}
return await this.tryLimit(limit, actor, factor);
}
private async tryLimit(limit: Keyed<RateLimit>, actor: string, factor: number): Promise<LimitInfo> {
if (isLegacyRateLimit(limit)) {
return await this.limitLegacy(limit, actor, factor);
} else {

16
packages/backend/src/server/api/StreamingApiServerService.ts
View file

@ -18,7 +18,6 @@ import { CacheService } from '@/core/CacheService.js';
import { MiLocalUser } from '@/models/User.js';
import { UserService } from '@/core/UserService.js';
import { ChannelFollowingService } from '@/core/ChannelFollowingService.js';
import { RoleService } from '@/core/RoleService.js';
import { getIpHash } from '@/misc/get-ip-hash.js';
import { LoggerService } from '@/core/LoggerService.js';
import { SkRateLimiterService } from '@/server/api/SkRateLimiterService.js';
@ -49,7 +48,6 @@ export class StreamingApiServerService {
private usersService: UserService,
private channelFollowingService: ChannelFollowingService,
private rateLimiterService: SkRateLimiterService,
private roleService: RoleService,
private loggerService: LoggerService,
) {
}
@ -57,22 +55,18 @@ export class StreamingApiServerService {
@bindThis
private async rateLimitThis(
user: MiLocalUser | null | undefined,
requestIp: string | undefined,
requestIp: string,
limit: IEndpointMeta['limit'] & { key: NonNullable<string> },
) : Promise<boolean> {
let limitActor: string;
let limitActor: string | MiLocalUser;
if (user) {
limitActor = user.id;
limitActor = user;
} else {
limitActor = getIpHash(requestIp || 'wtf');
limitActor = getIpHash(requestIp);
}
const factor = user ? (await this.roleService.getUserPolicies(user.id)).rateLimitFactor : 1;
if (factor <= 0) return false;
// Rate limit
const rateLimit = await this.rateLimiterService.limit(limit, limitActor, factor);
const rateLimit = await this.rateLimiterService.limit(limit, limitActor);
return rateLimit.blocked;
}