duocthu/apps/web/middleware.ts

import { NextResponse } from "next/server";
import type { NextRequest } from "next/server";
import { jwtVerify } from "jose";
import { SESSION_COOKIE } from "./app/api/auth/session";

/**
 * Rate limiting for the public API surface.
 *
 * The architecture assigns rate limiting to `api-gateway`, which is not built
 * yet, so until it exists this is where limits for the remaining bounded API
 * routes live. `/api/chat` is intentionally unlimited.
 *
 * It runs here rather than inside the route handlers because the edge
 * middleware rejects an abusive request before any handler work — and because
 * it covers every current and future `/api/*` route by default rather than
 * one endpoint at a time.
 *
 * Known limitations, stated rather than hidden:
 * - Counters are per process and in memory. Production runs a single `web`
 *   container, so this is a real limit today; the moment that scales to more
 *   than one replica each replica gets its own allowance, and this needs to
 *   move to Redis (already reserved for exactly this in `docs/architecture.md`)
 *   or to the gateway.
 * - It is keyed by client IP, so it throttles a shared NAT as one caller. That
 *   is the correct trade for a cost guard with no authentication; per-user
 *   limits need auth, which does not exist yet.
 * - It is a cost and abuse guard, not a security control. It does not
 *   authenticate anyone and must not be described as if it does.
 */

interface Bucket {
  hits: number[];
}

interface Rule {
  windowMs: number;
  max: number;
}

// `/api/chat` is intentionally absent, so chat requests pass through without
// rate limiting. Autocomplete is a local catalog lookup with no model call.
const RULES: Array<{ prefix: string; rules: Rule[] }> = [
  {
    prefix: "/api/suggest",
    rules: [{ windowMs: 60_000, max: 120 }],
  },
  // `/api/pdf` streams the whole 37MB source PDF from disk on every request,
  // with no range support and no caching headers. It is unauthenticated, so
  // repeated fetches are a bandwidth and memory cost on a single small EC2
  // host. Sized against real UI behaviour rather than guessed: `tra-cuu`
  // uses it as an iframe `src` whose `#page=` fragment changes per citation
  // click, and `CitationCard` links to it, so a clinician working through a
  // long evidence list can legitimately fetch it repeatedly. 30/min is far
  // above that and still bounds an automated puller.
  {
    prefix: "/api/pdf",
    rules: [{ windowMs: 60_000, max: 30 }],
  },
  // `/api/feedback` writes one row per answer (upsert keyed on trace_id), so
  // the ceiling only needs to exceed how fast a human can rate answers. It
  // reaches PostgreSQL on every call, which is why it is bounded at all.
  {
    prefix: "/api/feedback",
    rules: [{ windowMs: 60_000, max: 60 }],
  },
];

const buckets = new Map<string, Bucket>();
const LONGEST_WINDOW_MS = 3_600_000;
let lastSweep = 0;

/** Drop entries no rule can still be counting, so the map cannot grow without bound. */
function sweep(now: number) {
  if (now - lastSweep < 60_000) return;
  lastSweep = now;
  for (const [key, bucket] of buckets) {
    const live = bucket.hits.filter((t) => now - t < LONGEST_WINDOW_MS);
    if (live.length === 0) buckets.delete(key);
    else bucket.hits = live;
  }
}

/**
 * Caddy sits in front and sets `X-Forwarded-For`; the left-most entry is the
 * original client. Falling back to a shared key rather than to "unlimited"
 * matters: an unknown IP must not become a way to opt out of the limit.
 */
function clientKey(request: NextRequest): string {
  const forwarded = request.headers.get("x-forwarded-for");
  if (forwarded) {
    const first = forwarded.split(",")[0]?.trim();
    if (first) return first;
  }
  return request.headers.get("x-real-ip")?.trim() || "unknown";
}

function matchRules(pathname: string) {
  return RULES.find((entry) => pathname.startsWith(entry.prefix))?.rules;
}

/** `/admin/**` (except the login page itself) requires a valid session
 * cookie carrying `role: "admin"`. Verified with the same `JWT_SECRET`
 * auth-service signs with — Edge middleware can't call out to auth-service
 * per request without adding real latency to every admin page load, and
 * `jose` (unlike `jsonwebtoken`) works in the Edge runtime this file runs
 * under, so local verification is both correct and the only option here. */
async function guardAdmin(request: NextRequest): Promise<NextResponse | null> {
  const { pathname } = request.nextUrl;
  if (!pathname.startsWith("/admin") || pathname === "/admin/login") return null;

  const token = request.cookies.get(SESSION_COOKIE)?.value;
  const secret = process.env.JWT_SECRET;
  if (!token || !secret) {
    return NextResponse.redirect(new URL("/admin/login", request.url));
  }
  try {
    const { payload } = await jwtVerify(token, new TextEncoder().encode(secret));
    if (payload.role !== "admin") {
      return NextResponse.redirect(new URL("/admin/login", request.url));
    }
    return null;
  } catch {
    return NextResponse.redirect(new URL("/admin/login", request.url));
  }
}

/** `/profile` requires any valid session (no specific role) — it renders
 * unconditionally otherwise (`profile/page.tsx` has no server component
 * data fetch to fail on), so an anonymous visitor would see the full form
 * shell with every save silently 401ing instead of being sent to log in. */
async function guardProfile(request: NextRequest): Promise<NextResponse | null> {
  const { pathname } = request.nextUrl;
  if (pathname !== "/profile") return null;

  const token = request.cookies.get(SESSION_COOKIE)?.value;
  const secret = process.env.JWT_SECRET;
  if (!token || !secret) {
    return NextResponse.redirect(new URL("/login", request.url));
  }
  try {
    await jwtVerify(token, new TextEncoder().encode(secret));
    return null;
  } catch {
    return NextResponse.redirect(new URL("/login", request.url));
  }
}

export async function middleware(request: NextRequest) {
  const adminRedirect = await guardAdmin(request);
  if (adminRedirect) return adminRedirect;

  const profileRedirect = await guardProfile(request);
  if (profileRedirect) return profileRedirect;

  const rules = matchRules(request.nextUrl.pathname);
  if (!rules) return NextResponse.next();

  const now = Date.now();
  sweep(now);

  const key = `${clientKey(request)}:${request.nextUrl.pathname}`;
  const bucket = buckets.get(key) ?? { hits: [] };
  bucket.hits = bucket.hits.filter((t) => now - t < LONGEST_WINDOW_MS);

  for (const rule of rules) {
    const inWindow = bucket.hits.filter((t) => now - t < rule.windowMs);
    if (inWindow.length >= rule.max) {
      const oldest = Math.min(...inWindow);
      const retryAfterSec = Math.max(1, Math.ceil((rule.windowMs - (now - oldest)) / 1000));
      // Record nothing for a rejected request: a client hammering the endpoint
      // should not keep pushing its own window forward and lock itself out for
      // longer than the rule says.
      buckets.set(key, bucket);
      return NextResponse.json(
        {
          error: "rate_limited",
          message:
            "Bạn đang gửi quá nhiều yêu cầu trong thời gian ngắn. Vui lòng đợi một lát rồi thử lại.",
        },
        {
          status: 429,
          headers: {
            "Retry-After": String(retryAfterSec),
            "X-RateLimit-Limit": String(rule.max),
            "X-RateLimit-Remaining": "0",
          },
        }
      );
    }
  }

  bucket.hits.push(now);
  buckets.set(key, bucket);

  const tightest = rules[0];
  const used = bucket.hits.filter((t) => now - t < tightest.windowMs).length;
  const response = NextResponse.next();
  response.headers.set("X-RateLimit-Limit", String(tightest.max));
  response.headers.set("X-RateLimit-Remaining", String(Math.max(0, tightest.max - used)));
  return response;
}

export const config = {
  matcher: ["/api/:path*", "/admin/:path*", "/profile"],
};