feat: enhance rate limiting with new public request budgets and client IP validation

src/handlers/sends.ts

@@ -1,5 +1,6 @@
 import { Env, Send, SendAuthType, SendResponse, SendType, DEFAULT_DEV_SECRET } from '../types';
 import { StorageService } from '../services/storage';
+import { RateLimitService, getClientIdentifier } from '../services/ratelimit';
 import { jsonResponse, errorResponse } from '../utils/response';
 import { generateUUID } from '../utils/uuid';
 import { parsePagination, encodeContinuationToken } from '../utils/pagination';
@@ -13,6 +14,7 @@ import {
 
 const SEND_INACCESSIBLE_MSG = 'Send does not exist or is no longer available';
 const SEND_PASSWORD_ITERATIONS = 100_000;
+const SEND_PASSWORD_LIMIT_SCOPE = 'send-password';
 
 function getAliasedProp(source: unknown, aliases: string[]): { present: boolean; value: unknown } {
   if (!source || typeof source !== 'object') return { present: false, value: undefined };
@@ -383,12 +385,44 @@ async function getCreatorIdentifier(storage: StorageService, send: Send): Promis
   return owner?.email ?? null;
 }
 
-async function validatePublicSendAccess(send: Send, body: unknown): Promise<Response | null> {
+type PublicSendAccessValidationResult =
+  | { ok: true }
+  | { ok: false; response: Response; reason: 'email_auth_unsupported' | 'password_missing' | 'invalid_password' };
+
+function sendPasswordLimitKey(clientIdentifier: string): string {
+  return `${clientIdentifier}:${SEND_PASSWORD_LIMIT_SCOPE}`;
+}
+
+function sendPasswordLockMessage(retryAfterSeconds: number): string {
+  return `Too many failed send password attempts. Try again in ${Math.ceil(retryAfterSeconds / 60)} minutes.`;
+}
+
+function sendPasswordLockedErrorResponse(retryAfterSeconds: number): Response {
+  return errorResponse(sendPasswordLockMessage(retryAfterSeconds), 429);
+}
+
+function sendPasswordLockedOAuthResponse(retryAfterSeconds: number): Response {
+  const message = sendPasswordLockMessage(retryAfterSeconds);
+  return jsonResponse(
+    {
+      error: 'invalid_grant',
+      error_description: message,
+      send_access_error_type: 'too_many_password_attempts',
+      ErrorModel: {
+        Message: message,
+        Object: 'error',
+      },
+    },
+    429
+  );
+}
+
+async function validatePublicSendAccess(send: Send, body: unknown): Promise<PublicSendAccessValidationResult> {
   if (hasEmailAuth(send)) {
-    return errorResponse(SEND_INACCESSIBLE_MSG, 404);
+    return { ok: false, response: errorResponse(SEND_INACCESSIBLE_MSG, 404), reason: 'email_auth_unsupported' };
   }
 
-  if (!send.passwordHash) return null;
+  if (!send.passwordHash) return { ok: true };
 
   const passwordRaw = getAliasedProp(body, ['password', 'Password']);
   const passwordHashB64Raw = getAliasedProp(body, [
@@ -401,7 +435,7 @@ async function validatePublicSendAccess(send: Send, body: unknown): Promise<Resp
   let validPassword = false;
   if (send.passwordSalt && send.passwordIterations) {
     if (typeof passwordRaw.value !== 'string') {
-      return errorResponse('Password not provided', 401);
+      return { ok: false, response: errorResponse('Password not provided', 401), reason: 'password_missing' };
     }
     validPassword = await verifySendPassword(send, passwordRaw.value);
   } else {
@@ -411,12 +445,14 @@ async function validatePublicSendAccess(send: Send, body: unknown): Promise<Resp
         : typeof passwordRaw.value === 'string'
           ? passwordRaw.value
           : '';
-    if (!candidate) return errorResponse('Password not provided', 401);
+    if (!candidate) return { ok: false, response: errorResponse('Password not provided', 401), reason: 'password_missing' };
     validPassword = verifySendPasswordHashB64(send, candidate);
   }
-  if (!validPassword) return errorResponse('Invalid password', 400);
+  if (!validPassword) {
+    return { ok: false, response: errorResponse('Invalid password', 400), reason: 'invalid_password' };
+  }
 
-  return null;
+  return { ok: true };
 }
 
 // GET /api/sends
@@ -1016,9 +1052,34 @@ export async function handleAccessSend(request: Request, env: Env, accessId: str
     body = {};
   }
 
-  const validationErr = await validatePublicSendAccess(send, body);
-  if (validationErr) {
-    return validationErr;
+  let sendPasswordLimitIpKey: string | null = null;
+  let sendPasswordRateLimit: RateLimitService | null = null;
+  if (send.passwordHash) {
+    const clientIdentifier = getClientIdentifier(request);
+    if (!clientIdentifier) {
+      return errorResponse('Client IP is required', 403);
+    }
+    sendPasswordLimitIpKey = sendPasswordLimitKey(clientIdentifier);
+    sendPasswordRateLimit = new RateLimitService(env.DB);
+    const sendPasswordCheck = await sendPasswordRateLimit.checkLoginAttempt(sendPasswordLimitIpKey);
+    if (!sendPasswordCheck.allowed) {
+      return sendPasswordLockedErrorResponse(sendPasswordCheck.retryAfterSeconds || 60);
+    }
+  }
+
+  const validation = await validatePublicSendAccess(send, body);
+  if (!validation.ok) {
+    if (validation.reason === 'invalid_password' && sendPasswordRateLimit && sendPasswordLimitIpKey) {
+      const failed = await sendPasswordRateLimit.recordFailedLogin(sendPasswordLimitIpKey);
+      if (failed.locked) {
+        return sendPasswordLockedErrorResponse(failed.retryAfterSeconds || 60);
+      }
+    }
+    return validation.response;
+  }
+
+  if (send.passwordHash && sendPasswordRateLimit && sendPasswordLimitIpKey) {
+    await sendPasswordRateLimit.clearLoginAttempts(sendPasswordLimitIpKey);
   }
 
   if (send.type === SendType.Text) {
@@ -1065,9 +1126,34 @@ export async function handleAccessSendFile(
     body = {};
   }
 
-  const validationErr = await validatePublicSendAccess(send, body);
-  if (validationErr) {
-    return validationErr;
+  let sendPasswordLimitIpKey: string | null = null;
+  let sendPasswordRateLimit: RateLimitService | null = null;
+  if (send.passwordHash) {
+    const clientIdentifier = getClientIdentifier(request);
+    if (!clientIdentifier) {
+      return errorResponse('Client IP is required', 403);
+    }
+    sendPasswordLimitIpKey = sendPasswordLimitKey(clientIdentifier);
+    sendPasswordRateLimit = new RateLimitService(env.DB);
+    const sendPasswordCheck = await sendPasswordRateLimit.checkLoginAttempt(sendPasswordLimitIpKey);
+    if (!sendPasswordCheck.allowed) {
+      return sendPasswordLockedErrorResponse(sendPasswordCheck.retryAfterSeconds || 60);
+    }
+  }
+
+  const validation = await validatePublicSendAccess(send, body);
+  if (!validation.ok) {
+    if (validation.reason === 'invalid_password' && sendPasswordRateLimit && sendPasswordLimitIpKey) {
+      const failed = await sendPasswordRateLimit.recordFailedLogin(sendPasswordLimitIpKey);
+      if (failed.locked) {
+        return sendPasswordLockedErrorResponse(failed.retryAfterSeconds || 60);
+      }
+    }
+    return validation.response;
+  }
+
+  if (send.passwordHash && sendPasswordRateLimit && sendPasswordLimitIpKey) {
+    await sendPasswordRateLimit.clearLoginAttempts(sendPasswordLimitIpKey);
   }
 
   const updated = await storage.incrementSendAccessCount(send.id);
@@ -1221,7 +1307,9 @@ export async function issueSendAccessToken(
   env: Env,
   sendIdOrAccessId: string,
   passwordHashB64?: string | null,
-  password?: string | null
+  password?: string | null,
+  rateLimit?: RateLimitService,
+  sendPasswordLimitIpKey?: string
 ): Promise<{ token: string } | { error: Response }> {
   const jwt = getSafeJwtSecret(env);
   if (!jwt.ok) {
@@ -1267,6 +1355,15 @@ export async function issueSendAccessToken(
   }
 
   if (send.passwordHash) {
+    if (rateLimit && sendPasswordLimitIpKey) {
+      const sendPasswordCheck = await rateLimit.checkLoginAttempt(sendPasswordLimitIpKey);
+      if (!sendPasswordCheck.allowed) {
+        return {
+          error: sendPasswordLockedOAuthResponse(sendPasswordCheck.retryAfterSeconds || 60),
+        };
+      }
+    }
+
     let ok = false;
     if (passwordHashB64) {
       ok = verifySendPasswordHashB64(send, passwordHashB64);
@@ -1275,6 +1372,14 @@ export async function issueSendAccessToken(
     }
 
     if (!ok) {
+      if (rateLimit && sendPasswordLimitIpKey) {
+        const failed = await rateLimit.recordFailedLogin(sendPasswordLimitIpKey);
+        if (failed.locked) {
+          return {
+            error: sendPasswordLockedOAuthResponse(failed.retryAfterSeconds || 60),
+          };
+        }
+      }
       return {
         error: jsonResponse(
           {
@@ -1290,6 +1395,10 @@ export async function issueSendAccessToken(
         ),
       };
     }
+
+    if (rateLimit && sendPasswordLimitIpKey) {
+      await rateLimit.clearLoginAttempts(sendPasswordLimitIpKey);
+    }
   }
 
   const token = await createSendAccessToken(send.id, jwt.secret);