Skip to content

Commit 418e509

Browse files
pedrojreispedrojreis
committed
refactor(limits): improve Redis command handling and formatting
1 parent 05a2c7f commit 418e509

1 file changed

Lines changed: 167 additions & 137 deletions

File tree

‎service/src/middleware/limits.ts‎

Lines changed: 167 additions & 137 deletions
Original file line numberDiff line numberDiff line change
@@ -11,194 +11,224 @@ import { getExecutionIdentity } from '../execution-identity';
1111
import logger from '../logger';
1212

1313
type RedisCommandTarget = {
14-
call?: (command: string, ...args: (string | number | Buffer)[]) => Promise<unknown>;
14+
call?: (
15+
command: string,
16+
...args: (string | number | Buffer)[]
17+
) => Promise<unknown>;
1518
};
1619

1720
let redisCommands: RedisCommandTarget | undefined;
1821

1922
async function getRedisCommands(): Promise<RedisCommandTarget> {
20-
if (redisCommands) return redisCommands;
21-
const { connection } = await import('../queue');
22-
return connection;
23+
if (redisCommands) return redisCommands;
24+
const { connection } = await import('../queue');
25+
return connection;
2326
}
2427

25-
const sendCommand: SendCommandFn = async (command: string, ...args: (string | number | Buffer)[]): Promise<RedisReply> => {
26-
const target = await getRedisCommands();
27-
if (typeof target.call === 'function') {
28-
const result = await target.call(command, ...args);
28+
const sendCommand: SendCommandFn = async (
29+
command: string,
30+
...args: (string | number | Buffer)[]
31+
): Promise<RedisReply> => {
32+
const target = await getRedisCommands();
33+
const commandMethod = (target as Record<string, unknown>)[
34+
command.toLowerCase()
35+
];
36+
if (typeof commandMethod !== 'function') {
37+
if (typeof target.call === 'function') {
38+
const result = await target.call(command, ...args);
39+
return result as RedisReply;
40+
}
41+
throw new Error(`Redis command method unavailable: ${command}`);
42+
}
43+
const result = await commandMethod.apply(target, args);
2944
return result as RedisReply;
30-
}
31-
const commandMethod = (target as Record<string, unknown>)[command.toLowerCase()];
32-
if (typeof commandMethod !== 'function') {
33-
throw new Error(`Redis command method unavailable: ${command}`);
34-
}
35-
const result = await commandMethod.apply(target, args);
36-
return result as RedisReply;
3745
};
3846

3947
type RequestWithRateLimit = Request & {
40-
rateLimit?: {
41-
limit: number;
42-
used: number;
43-
remaining: number;
44-
resetTime?: Date;
45-
};
48+
rateLimit?: {
49+
limit: number;
50+
used: number;
51+
remaining: number;
52+
resetTime?: Date;
53+
};
4654
};
4755

4856
type RateLimitResponseOptions = {
49-
message: string;
50-
structuredBody?: boolean;
51-
logRejections?: boolean;
57+
message: string;
58+
structuredBody?: boolean;
59+
logRejections?: boolean;
5260
};
5361

5462
const unknownPrincipal = 'unknown';
5563

56-
const keySegment = (value: string | undefined, fallback = unknownPrincipal): string => {
57-
const trimmed = value?.trim();
58-
return trimmed ? trimmed.replace(/:/g, '_') : fallback;
64+
const keySegment = (
65+
value: string | undefined,
66+
fallback = unknownPrincipal,
67+
): string => {
68+
const trimmed = value?.trim();
69+
return trimmed ? trimmed.replace(/:/g, '_') : fallback;
5970
};
6071

6172
const hashLabel = (value: string | undefined): string | undefined => {
62-
const trimmed = value?.trim();
63-
if (!trimmed) return undefined;
64-
return createHash('sha256').update(trimmed).digest('hex').slice(0, 12);
73+
const trimmed = value?.trim();
74+
if (!trimmed) return undefined;
75+
return createHash('sha256').update(trimmed).digest('hex').slice(0, 12);
6576
};
6677

6778
export function setRateLimitRedisForTests(client?: RedisCommandTarget): void {
68-
redisCommands = client;
79+
redisCommands = client;
6980
}
7081

7182
export function retryAfterSeconds(req: Request, windowMs: number): number {
72-
const resetTime = (req as RequestWithRateLimit).rateLimit?.resetTime;
73-
if (resetTime instanceof Date) {
74-
return Math.max(1, Math.ceil((resetTime.getTime() - Date.now()) / 1000));
75-
}
76-
return Math.max(1, Math.ceil(windowMs / 1000));
83+
const resetTime = (req as RequestWithRateLimit).rateLimit?.resetTime;
84+
if (resetTime instanceof Date) {
85+
return Math.max(
86+
1,
87+
Math.ceil((resetTime.getTime() - Date.now()) / 1000),
88+
);
89+
}
90+
return Math.max(1, Math.ceil(windowMs / 1000));
7791
}
7892

79-
export function rateLimitResponseBody(message: string, retryAfter: number): {
80-
error: 'rate_limited';
81-
message: string;
82-
retry_after_seconds: number;
93+
export function rateLimitResponseBody(
94+
message: string,
95+
retryAfter: number,
96+
): {
97+
error: 'rate_limited';
98+
message: string;
99+
retry_after_seconds: number;
83100
} {
84-
const retryAfterLabel = retryAfter === 1 ? '1 second' : `${retryAfter} seconds`;
85-
return {
86-
error: 'rate_limited',
87-
message: `${message} Please retry in ${retryAfterLabel}.`,
88-
retry_after_seconds: retryAfter,
89-
};
101+
const retryAfterLabel =
102+
retryAfter === 1 ? '1 second' : `${retryAfter} seconds`;
103+
return {
104+
error: 'rate_limited',
105+
message: `${message} Please retry in ${retryAfterLabel}.`,
106+
retry_after_seconds: retryAfter,
107+
};
90108
}
91109

92110
export const keyGenerator = (req: Request): string => {
93-
const authReq = req as AuthenticatedRequest;
94-
const identity = getExecutionIdentity(authReq);
95-
if (identity.canonicalUserId) {
96-
return `${keySegment(identity.storageNamespace, 'legacy')}:user:${keySegment(identity.canonicalUserId)}`;
97-
}
98-
return `ip:${keySegment(req.ip)}`;
111+
const authReq = req as AuthenticatedRequest;
112+
const identity = getExecutionIdentity(authReq);
113+
if (identity.canonicalUserId) {
114+
return `${keySegment(identity.storageNamespace, 'legacy')}:user:${keySegment(identity.canonicalUserId)}`;
115+
}
116+
return `ip:${keySegment(req.ip)}`;
99117
};
100118

101119
const buildRateLimiter = (
102-
prefix: string,
103-
windowMs: number,
104-
max: number,
105-
options: RateLimitResponseOptions
120+
prefix: string,
121+
windowMs: number,
122+
max: number,
123+
options: RateLimitResponseOptions,
106124
): RateLimitRequestHandler => {
107-
return rateLimitFactory({
108-
windowMs,
109-
max,
110-
standardHeaders: 'draft-6',
111-
legacyHeaders: false,
112-
store: new RateLimitRedisStore({
113-
sendCommand,
114-
prefix: `${prefix}:`
115-
}),
116-
keyGenerator,
117-
handler: (req: Request, res: Response) => {
118-
const retryAfter = retryAfterSeconds(req, windowMs);
119-
const rateLimit = (req as RequestWithRateLimit).rateLimit;
120-
res.setHeader('Retry-After', retryAfter.toString());
121-
res.setHeader('RateLimit-Limit', (rateLimit?.limit ?? max).toString());
122-
res.setHeader('RateLimit-Remaining', '0');
123-
res.setHeader('RateLimit-Reset', retryAfter.toString());
124-
125-
if (options.logRejections) {
126-
const authReq = req as AuthenticatedRequest;
127-
const principal = authReq.codeApiPrincipal;
128-
const identity = getExecutionIdentity(authReq);
129-
const hasIdentity = Boolean(identity.canonicalUserId);
130-
logger.warn('CodeAPI rate limit rejected', {
131-
limiter: prefix,
132-
path: req.originalUrl || req.path,
133-
retryAfterSeconds: retryAfter,
134-
limit: rateLimit?.limit ?? max,
135-
windowMs,
136-
principalSource: hasIdentity ? identity.principalSource : undefined,
137-
tenantHash: hasIdentity ? hashLabel(identity.storageNamespace) : undefined,
138-
userHash: hasIdentity ? hashLabel(identity.canonicalUserId) : undefined,
139-
credentialHash: hashLabel(principal?.credentialId),
140-
});
141-
}
142-
143-
if (options.structuredBody) {
144-
return res.status(429).json(rateLimitResponseBody(options.message, retryAfter));
145-
}
146-
const retryAfterLabel = retryAfter === 1 ? '1 second' : `${retryAfter} seconds`;
147-
res.status(429).json({
148-
error: `${options.message} Please retry in ${retryAfterLabel}.`,
149-
});
150-
},
151-
});
125+
return rateLimitFactory({
126+
windowMs,
127+
max,
128+
standardHeaders: 'draft-6',
129+
legacyHeaders: false,
130+
store: new RateLimitRedisStore({
131+
sendCommand,
132+
prefix: `${prefix}:`,
133+
}),
134+
keyGenerator,
135+
handler: (req: Request, res: Response) => {
136+
const retryAfter = retryAfterSeconds(req, windowMs);
137+
const rateLimit = (req as RequestWithRateLimit).rateLimit;
138+
res.setHeader('Retry-After', retryAfter.toString());
139+
res.setHeader(
140+
'RateLimit-Limit',
141+
(rateLimit?.limit ?? max).toString(),
142+
);
143+
res.setHeader('RateLimit-Remaining', '0');
144+
res.setHeader('RateLimit-Reset', retryAfter.toString());
145+
146+
if (options.logRejections) {
147+
const authReq = req as AuthenticatedRequest;
148+
const principal = authReq.codeApiPrincipal;
149+
const identity = getExecutionIdentity(authReq);
150+
const hasIdentity = Boolean(identity.canonicalUserId);
151+
logger.warn('CodeAPI rate limit rejected', {
152+
limiter: prefix,
153+
path: req.originalUrl || req.path,
154+
retryAfterSeconds: retryAfter,
155+
limit: rateLimit?.limit ?? max,
156+
windowMs,
157+
principalSource: hasIdentity
158+
? identity.principalSource
159+
: undefined,
160+
tenantHash: hasIdentity
161+
? hashLabel(identity.storageNamespace)
162+
: undefined,
163+
userHash: hasIdentity
164+
? hashLabel(identity.canonicalUserId)
165+
: undefined,
166+
credentialHash: hashLabel(principal?.credentialId),
167+
});
168+
}
169+
170+
if (options.structuredBody) {
171+
return res
172+
.status(429)
173+
.json(rateLimitResponseBody(options.message, retryAfter));
174+
}
175+
const retryAfterLabel =
176+
retryAfter === 1 ? '1 second' : `${retryAfter} seconds`;
177+
res.status(429).json({
178+
error: `${options.message} Please retry in ${retryAfterLabel}.`,
179+
});
180+
},
181+
});
152182
};
153183

154184
export const createRateLimiter = (
155-
prefix: string,
156-
windowMs: number,
157-
max: number,
158-
options: RateLimitResponseOptions
185+
prefix: string,
186+
windowMs: number,
187+
max: number,
188+
options: RateLimitResponseOptions,
159189
): RateLimitRequestHandler => {
160-
let limiter: RateLimitRequestHandler | undefined;
161-
const getLimiter = (): RateLimitRequestHandler => {
162-
limiter ??= buildRateLimiter(prefix, windowMs, max, options);
163-
return limiter;
164-
};
165-
166-
const lazyLimiter = ((req: Request, res: Response, next: NextFunction) => {
167-
return getLimiter()(req, res, next);
168-
}) as RateLimitRequestHandler;
169-
lazyLimiter.resetKey = (key: string) => getLimiter().resetKey(key);
170-
lazyLimiter.getKey = (key: string) => getLimiter().getKey(key);
171-
return lazyLimiter;
190+
let limiter: RateLimitRequestHandler | undefined;
191+
const getLimiter = (): RateLimitRequestHandler => {
192+
limiter ??= buildRateLimiter(prefix, windowMs, max, options);
193+
return limiter;
194+
};
195+
196+
const lazyLimiter = ((req: Request, res: Response, next: NextFunction) => {
197+
return getLimiter()(req, res, next);
198+
}) as RateLimitRequestHandler;
199+
lazyLimiter.resetKey = (key: string) => getLimiter().resetKey(key);
200+
lazyLimiter.getKey = (key: string) => getLimiter().getKey(key);
201+
return lazyLimiter;
172202
};
173203

174204
export const executionLimiter = createRateLimiter(
175-
'exec',
176-
env.EXEC_LIMIT_WINDOW,
177-
env.EXEC_MAX_REQUESTS,
178-
{
179-
message: 'Too many CodeAPI execution requests.',
180-
structuredBody: true,
181-
logRejections: true,
182-
}
205+
'exec',
206+
env.EXEC_LIMIT_WINDOW,
207+
env.EXEC_MAX_REQUESTS,
208+
{
209+
message: 'Too many CodeAPI execution requests.',
210+
structuredBody: true,
211+
logRejections: true,
212+
},
183213
);
184214

185215
export const uploadLimiter = createRateLimiter(
186-
'upload',
187-
env.UPLOAD_LIMIT_WINDOW,
188-
env.UPLOAD_MAX_REQUESTS,
189-
{ message: 'Too many file uploads.' }
216+
'upload',
217+
env.UPLOAD_LIMIT_WINDOW,
218+
env.UPLOAD_MAX_REQUESTS,
219+
{ message: 'Too many file uploads.' },
190220
);
191221

192222
export const downloadLimiter = createRateLimiter(
193-
'download',
194-
env.DOWNLOAD_LIMIT_WINDOW,
195-
env.DOWNLOAD_MAX_REQUESTS,
196-
{ message: 'Too many file downloads.' }
223+
'download',
224+
env.DOWNLOAD_LIMIT_WINDOW,
225+
env.DOWNLOAD_MAX_REQUESTS,
226+
{ message: 'Too many file downloads.' },
197227
);
198228

199229
export const fetchLimiter = createRateLimiter(
200-
'fetch',
201-
env.FETCH_LIMIT_WINDOW,
202-
env.FETCH_MAX_REQUESTS,
203-
{ message: 'Too many file list requests.' }
230+
'fetch',
231+
env.FETCH_LIMIT_WINDOW,
232+
env.FETCH_MAX_REQUESTS,
233+
{ message: 'Too many file list requests.' },
204234
);

0 commit comments

Comments
 (0)