R524809
132 lines
4.3 KiB
TypeScript
132 lines
4.3 KiB
TypeScript
import {
|
|
Inject,
|
|
Injectable,
|
|
OnModuleDestroy,
|
|
OnModuleInit,
|
|
} from '@nestjs/common';
|
|
import { createHmac } from 'node:crypto';
|
|
import { Redis } from 'ioredis';
|
|
import { ApiError } from '../common/api-error.js';
|
|
import { APP_CONFIG, type AppConfig } from '../config/app-config.module.js';
|
|
|
|
type AuthAction = 'login' | 'register' | 'refresh' | 'change_password';
|
|
type Limit = { ip: number; email?: number; user?: number };
|
|
|
|
@Injectable()
|
|
export class AuthRateLimiter implements OnModuleInit, OnModuleDestroy {
|
|
private readonly redis: Redis;
|
|
|
|
public constructor(@Inject(APP_CONFIG) private readonly config: AppConfig) {
|
|
// Redis 让限流计数在多个 API 实例间共享，不会随单个实例重启清零。
|
|
this.redis = new Redis(config.REDIS_URL, {
|
|
enableOfflineQueue: false,
|
|
lazyConnect: true,
|
|
maxRetriesPerRequest: 1,
|
|
});
|
|
}
|
|
|
|
public async onModuleInit(): Promise<void> {
|
|
// 在开始接收认证请求前完成连接，避免首个请求命中 Redis 连接竞态而被误判为 503。
|
|
await this.redis.connect();
|
|
}
|
|
|
|
public async consume(
|
|
action: AuthAction,
|
|
identifiers: { ip: string | null; email?: string; userId?: string },
|
|
): Promise<void> {
|
|
const limit = this.limitFor(action);
|
|
await this.consumeBucket(
|
|
action,
|
|
'ip',
|
|
identifiers.ip ?? 'unknown',
|
|
limit.ip,
|
|
);
|
|
if (identifiers.email && limit.email) {
|
|
await this.consumeBucket(
|
|
action,
|
|
'email',
|
|
identifiers.email,
|
|
limit.email,
|
|
);
|
|
}
|
|
if (identifiers.userId && limit.user) {
|
|
await this.consumeBucket(
|
|
action,
|
|
'user',
|
|
identifiers.userId,
|
|
limit.user,
|
|
);
|
|
}
|
|
}
|
|
|
|
public async onModuleDestroy(): Promise<void> {
|
|
if (this.redis.status !== 'ready') {
|
|
this.redis.disconnect();
|
|
return;
|
|
}
|
|
await this.redis.quit();
|
|
}
|
|
|
|
private async consumeBucket(
|
|
action: AuthAction,
|
|
dimension: 'ip' | 'email' | 'user',
|
|
identifier: string,
|
|
limit: number,
|
|
): Promise<void> {
|
|
const digest = createHmac(
|
|
'sha256',
|
|
this.config.AUTH_RATE_LIMIT_KEY_SECRET,
|
|
)
|
|
.update(identifier)
|
|
.digest('base64url');
|
|
// Redis key 只保存 IP/邮箱的 HMAC，不暴露可被日志或运维工具读取的原文。
|
|
const key = `atlas:auth-rate-limit:${action}:${dimension}:${digest}`;
|
|
try {
|
|
// INCR 与首次 EXPIRE 必须原子执行，避免并发下产生永不过期的限流 key。
|
|
const count = await this.redis.eval(
|
|
"local value = redis.call('INCR', KEYS[1]); if value == 1 then redis.call('EXPIRE', KEYS[1], ARGV[1]); end; return value;",
|
|
1,
|
|
key,
|
|
this.config.AUTH_RATE_LIMIT_WINDOW_SECONDS,
|
|
);
|
|
if (Number(count) > limit) {
|
|
throw new ApiError(
|
|
429,
|
|
'AUTH_RATE_LIMITED',
|
|
'请求过于频繁，请稍后重试',
|
|
);
|
|
}
|
|
} catch (error) {
|
|
if (error instanceof ApiError) throw error;
|
|
// 认证限流故障时必须拒绝请求；静默跳过会扩大暴力破解风险。
|
|
throw new ApiError(
|
|
503,
|
|
'AUTH_RATE_LIMIT_UNAVAILABLE',
|
|
'认证服务暂时不可用',
|
|
);
|
|
}
|
|
}
|
|
|
|
private limitFor(action: AuthAction): Limit {
|
|
if (action === 'login') {
|
|
return {
|
|
ip: this.config.AUTH_LOGIN_IP_LIMIT,
|
|
email: this.config.AUTH_LOGIN_EMAIL_LIMIT,
|
|
};
|
|
}
|
|
if (action === 'register') {
|
|
return {
|
|
ip: this.config.AUTH_REGISTER_IP_LIMIT,
|
|
email: this.config.AUTH_REGISTER_EMAIL_LIMIT,
|
|
};
|
|
}
|
|
if (action === 'change_password') {
|
|
return {
|
|
ip: this.config.AUTH_CHANGE_PASSWORD_IP_LIMIT,
|
|
user: this.config.AUTH_CHANGE_PASSWORD_USER_LIMIT,
|
|
};
|
|
}
|
|
return { ip: this.config.AUTH_REFRESH_IP_LIMIT };
|
|
}
|
|
}
|