Merge branch 'main' of https://git.jiong.us/jkjoy/Toot-Worker

src/http.ts

@@ -1,5 +1,5 @@
 export class HttpError extends Error {
-  constructor(readonly status: number, message: string) {
+  constructor(readonly status: number, message: string, readonly headers: HeadersInit = {}) {
     super(message);
   }
 }

src/index.ts

@@ -102,7 +102,7 @@ export default {
       }
       return response;
     } catch (error) {
-      if (error instanceof HttpError) return json({ error: error.message }, error.status);
+      if (error instanceof HttpError) return json({ error: error.message }, error.status, error.headers);
       console.error("unhandled", error);
       return json({ error: "internal_server_error" }, 500);
     }

src/mastodon.ts

@@ -107,6 +107,10 @@ const TOKEN_TTL_SECONDS = 60 * 60 * 24 * 90;
 const MAX_STATUS_CHARS = 5000;
 const MAX_MEDIA_ATTACHMENTS = 20;
 const MAX_MEDIA_BYTES = 10 * 1024 * 1024;
+const AUTH_RATE_LIMIT_WINDOW_SECONDS = 15 * 60;
+const AUTH_RATE_LIMIT_LOCK_SECONDS = 15 * 60;
+const AUTH_RATE_LIMIT_MAX_IP_FAILURES = 20;
+const AUTH_RATE_LIMIT_MAX_USERNAME_FAILURES = 8;
 
 const SUPPORTED_MIME = ["image/jpeg", "image/png", "image/gif", "image/webp"];
 const VALID_STATUS_VISIBILITIES = new Set(["public", "unlisted", "private", "direct"]);
@@ -282,10 +286,14 @@ export async function authorize(request: Request, env: Env): Promise<Response> {
   const redirectUri = selectRedirectUri(app, bodyString(body, "redirect_uri"));
   if (!redirectUri) return json({ error: "invalid_request" }, 400);
 
-  const user = await getUserByUsername(env, bodyString(body, "username"));
+  const username = bodyString(body, "username");
+  await assertAuthNotRateLimited(request, env, username);
+  const user = await getUserByUsername(env, username);
   if (!user || !(await verifyPassword(bodyString(body, "password"), user.password_hash))) {
+    await recordAuthFailure(request, env, username);
     return html("Invalid username or password", 401);
   }
+  await clearAuthFailures(request, env, username);
 
   const code = tokenString(32);
   const scope = requestedScopesWithinApp(bodyString(body, "scope", app.scopes), app.scopes);
@@ -311,8 +319,14 @@ export async function token(request: Request, env: Env): Promise<Response> {
   let scopes = app.scopes;
 
   if (grantType === "password") {
-    const user = await getUserByUsername(env, bodyString(body, "username"));
-    if (!user || !(await verifyPassword(bodyString(body, "password"), user.password_hash))) return json({ error: "invalid_grant" }, 400);
+    const username = bodyString(body, "username");
+    await assertAuthNotRateLimited(request, env, username);
+    const user = await getUserByUsername(env, username);
+    if (!user || !(await verifyPassword(bodyString(body, "password"), user.password_hash))) {
+      await recordAuthFailure(request, env, username);
+      return json({ error: "invalid_grant" }, 400);
+    }
+    await clearAuthFailures(request, env, username);
     userId = user.id;
     scopes = requestedScopesWithinApp(bodyString(body, "scope", app.scopes), app.scopes);
   } else if (grantType === "client_credentials") {
@@ -344,6 +358,51 @@ export async function revoke(request: Request, env: Env): Promise<Response> {
   return json({});
 }
 
+type AuthFailureRecord = {
+  count: number;
+  firstFailureAt: number;
+  lockedUntil?: number;
+};
+
+async function assertAuthNotRateLimited(request: Request, env: Env, username: string): Promise<void> {
+  const keys = authRateLimitKeys(request, username);
+  const now = Math.floor(Date.now() / 1000);
+  const records = await Promise.all(keys.map(({ key }) => env.KV.get<AuthFailureRecord>(key, "json")));
+  const lockedUntil = records.reduce((latest, record) => Math.max(latest, record?.lockedUntil ?? 0), 0);
+  if (lockedUntil > now) {
+    const retryAfter = String(Math.max(1, lockedUntil - now));
+    throw new HttpError(429, "rate_limited", { "retry-after": retryAfter });
+  }
+}
+
+async function recordAuthFailure(request: Request, env: Env, username: string): Promise<void> {
+  const now = Math.floor(Date.now() / 1000);
+  await Promise.all(authRateLimitKeys(request, username).map(async ({ key, limit }) => {
+    const existing = await env.KV.get<AuthFailureRecord>(key, "json");
+    const inWindow = existing && now - existing.firstFailureAt < AUTH_RATE_LIMIT_WINDOW_SECONDS;
+    const next: AuthFailureRecord = {
+      count: inWindow ? existing.count + 1 : 1,
+      firstFailureAt: inWindow ? existing.firstFailureAt : now
+    };
+    if (next.count >= limit) next.lockedUntil = now + AUTH_RATE_LIMIT_LOCK_SECONDS;
+    const ttl = Math.max(AUTH_RATE_LIMIT_WINDOW_SECONDS, AUTH_RATE_LIMIT_LOCK_SECONDS);
+    await env.KV.put(key, JSON.stringify(next), { expirationTtl: ttl });
+  }));
+}
+
+async function clearAuthFailures(request: Request, env: Env, username: string): Promise<void> {
+  await Promise.all(authRateLimitKeys(request, username).map(({ key }) => env.KV.delete(key)));
+}
+
+function authRateLimitKeys(request: Request, username: string): { key: string; limit: number }[] {
+  const ip = request.headers.get("cf-connecting-ip") ?? request.headers.get("x-forwarded-for")?.split(",")[0]?.trim() ?? "unknown";
+  const normalizedUsername = username.trim().toLowerCase() || "empty";
+  return [
+    { key: `auth_fail:ip:${ip}`, limit: AUTH_RATE_LIMIT_MAX_IP_FAILURES },
+    { key: `auth_fail:user:${normalizedUsername}`, limit: AUTH_RATE_LIMIT_MAX_USERNAME_FAILURES }
+  ];
+}
+
 export async function verifyCredentials(request: Request, env: Env): Promise<Response> {
   const user = await requireUser(request, env);
   const account = await accountJson(env, user) as Record<string, unknown>;