@beignet/core/src/server/hooks/rate-limit.ts

Core framework primitives for Beignet

/**
 * Rate limit hooks for @beignet/core/server
 */

import type { RateLimitScope } from "../../contracts/index.js";
import { AppError, httpErrors } from "../../errors/index.js";
import {
  type ActivityActor,
  AuthUnauthorizedError,
  type RateLimitPort,
} from "../../ports/index.js";
import {
  createProviderInstrumentation,
  type ProviderInstrumentationTarget,
} from "../../providers/index.js";
import {
  resolveTrustedClientIp,
  type TrustedProxyClientIpSource,
  type TrustedProxyConfig,
  type TrustedRequestInfo,
} from "../trusted-proxy.js";
import type { HttpRequestLike, ServerHook } from "../types.js";

/**
 * Ports required by rate-limit hooks.
 */
export type RateLimitPorts = {
  rateLimit: RateLimitPort;
};

/**
 * Minimal context shape required for user-scoped rate limits.
 */
export type CtxWithRateLimit = {
  ports: RateLimitPorts;
  actor?: ActivityActor;
};

type EarlyRateLimitScope = Exclude<RateLimitScope, "user">;

/**
 * Strategy for resolving the client IP used by `ip`-scoped limits.
 *
 * - `"none"`: do not trust request headers for IP resolution. Every request
 *   to one contract shares that contract's unknown-client bucket; this is the
 *   explicit opt-out for apps that declare `ip` scopes without a trusted
 *   client-IP source.
 * - `"x-forwarded-for-last"`: the last `x-forwarded-for` entry. Use this only
 *   when the app is always behind a trusted reverse proxy that appends the
 *   socket address.
 * - `"x-forwarded-for-first"`: the first `x-forwarded-for` entry. This value
 *   is client-controlled, so only use it when a trusted edge normalizes the
 *   header before it reaches the app.
 * - `"x-real-ip"` and `"cf-connecting-ip"`: dedicated platform headers.
 * - A function receives the raw request and returns the client IP, for
 *   platform-specific resolution.
 */
export type RateLimitIpSource = "none" | TrustedProxyClientIpSource;

/**
 * Options for `createRateLimitHooks(...)`.
 */
export interface RateLimitOptions<Ctx> {
  /**
   * Build a rate-limit key after context exists.
   *
   * This is used for user-scoped limits and any late key strategy. The
   * returned value is the complete key and is not automatically namespaced by
   * contract.
   */
  key?: (args: {
    ctx: Ctx;
    req: HttpRequestLike;
    scope: RateLimitScope;
  }) => string;
  /**
   * Build a rate-limit key before request parsing and context creation.
   *
   * This is used for global and IP-scoped limits. The returned value is the
   * complete key and is not automatically namespaced by contract.
   */
  earlyKey?: (args: {
    req: HttpRequestLike;
    scope: EarlyRateLimitScope;
  }) => string;
  /**
   * Resolve the client IP for `ip`-scoped limits.
   *
   * There is no default: when any contract declares an `ip`-scoped rate limit
   * and neither the server-level `trustedProxy.clientIp`,
   * this hook's `trustedProxy.clientIp`, `ipSource`, nor a custom `earlyKey`
   * is configured, the hook fails at startup. Prefer the server-level policy
   * for production proxy headers; keep these hook-local options for an
   * intentional override or custom keying.
   */
  ipSource?: RateLimitIpSource;
  /**
   * Hook-local trusted-proxy policy used to resolve client IPs when
   * `ipSource` is not set. This overrides the server-level policy.
   *
   * Configure this only when the app is always behind a platform or reverse
   * proxy that strips or normalizes forwarding headers. For `ip`-scoped rate
   * limits, set `trustedProxy.clientIp` to the header source written by that
   * trusted edge.
   */
  trustedProxy?: TrustedProxyConfig;
}

function resolveClientIp(
  req: HttpRequestLike,
  ipSource: RateLimitIpSource,
): string | undefined {
  if (ipSource === "none") {
    return undefined;
  }
  return resolveTrustedClientIp(req, ipSource);
}

function hasTrustedProxyClientIp(
  config: TrustedProxyConfig | undefined,
): boolean {
  if (!config) return false;
  return Boolean(config.clientIp);
}

function formatContractNames(names: readonly string[]): string {
  return names.map((name) => `"${name}"`).join(", ");
}

function ipSourceConfigurationError(contractNames: string): Error {
  return new Error(
    `createRateLimitHooks(...) has no client IP source configured, but contract(s) ${contractNames} declare an "ip"-scoped rate limit. ` +
      `Set trustedProxy.clientIp or ipSource to a header source written by a trusted edge, ` +
      `for example "x-forwarded-for-last", "x-forwarded-for-first", "x-real-ip", "cf-connecting-ip", or a custom function, ` +
      `or "none" to explicitly accept one shared unknown-client bucket per contract.`,
  );
}

function namespaceContractKey(contractName: string, key: string): string {
  const encodedContractName = contractName
    .replaceAll("%", "%25")
    .replaceAll(":", "%3A");
  return `contract:${encodedContractName}:${key}`;
}

function emitUserKey(contractName: string, userId: string): string {
  return namespaceContractKey(contractName, `user:${userId}`);
}

function emitIpKey(contractName: string, ip: string): string {
  return namespaceContractKey(contractName, `ip:${ip}`);
}

function defaultRateLimitKey<Ctx extends CtxWithRateLimit>(
  args: {
    ctx: Ctx;
    req: HttpRequestLike;
    scope: RateLimitScope;
    contractName: string;
  },
  getClientIp: (req: HttpRequestLike) => string | undefined,
): string {
  const { contractName, ctx, req, scope } = args;

  if (scope === "user") {
    if (ctx.actor?.type !== "user" || !ctx.actor.id) {
      throw new AuthUnauthorizedError();
    }
    return emitUserKey(contractName, ctx.actor.id);
  }

  if (scope === "ip") {
    const ip = getClientIp(req) || "unknown";
    return emitIpKey(contractName, ip);
  }

  return namespaceContractKey(contractName, "global");
}

function defaultEarlyRateLimitKey(
  args: {
    req: HttpRequestLike;
    scope: EarlyRateLimitScope;
    contractName: string;
  },
  getClientIp: (req: HttpRequestLike) => string | undefined,
): string {
  if (args.scope === "ip") {
    const ip = getClientIp(args.req) || "unknown";
    return emitIpKey(args.contractName, ip);
  }

  return namespaceContractKey(args.contractName, "global");
}

async function enforceRateLimit(
  ports: RateLimitPorts,
  args: {
    key: string;
    limit: number;
    windowSec: number;
    scope: RateLimitScope;
  },
): Promise<void> {
  const result = await ports.rateLimit.hit({
    key: args.key,
    limit: args.limit,
    windowSec: args.windowSec,
  });

  if (result.allowed) {
    return;
  }

  // App ports may carry an `instrumentation` or `devtools` sink alongside the
  // typed rate-limit port; the helper resolves them at runtime.
  const instrumentation = createProviderInstrumentation(
    ports as ProviderInstrumentationTarget,
    {
      providerName: "rate-limit",
      watcher: "rateLimit",
    },
  );
  instrumentation.custom({
    name: "rateLimit.denied",
    label: "Rate limit denied",
    summary: `Rate limit denied for ${args.key}`,
    details: {
      key: args.key,
      scope: args.scope,
      limit: args.limit,
      windowSec: args.windowSec,
    },
  });

  throw new AppError(
    httpErrors.TooManyRequests,
    {
      scope: args.scope,
      retryAfterSeconds: result.retryAfterSeconds,
      resetAt: result.resetAt?.toISOString() ?? null,
    },
    "Rate limit exceeded",
    result.retryAfterSeconds !== null
      ? { headers: { "Retry-After": String(result.retryAfterSeconds) } }
      : undefined,
  );
}

/**
 * Create metadata-driven rate-limit hooks.
 *
 * The hook reads `contract.metadata.rateLimit`. Global and IP-scoped limits run
 * in `onRequest` before context creation; user-scoped limits run in
 * `beforeHandle` after route hooks have resolved identity and `ctx.actor` is
 * available. Default keys include the contract name so unrelated contracts do
 * not share counters. A user-scoped limit without a resolved user actor fails
 * with `AuthUnauthorizedError` instead of falling back to a global bucket.
 * Exceeded limits throw the framework `TooManyRequests` app error
 * with `scope`, `retryAfterSeconds`, and `resetAt` details, and the 429
 * response carries a `Retry-After` header when the limiter reports a reset
 * time. The bucket key is
 * never sent to clients; denials emit a `rateLimit.denied` instrumentation
 * event that carries the key for operators.
 *
 * `ip`-scoped limits require an explicit `trustedProxy.clientIp`, `ipSource`,
 * or custom `earlyKey`: the hook's `validate` phase fails `createServer(...)`
 * startup when a registered contract declares an `ip` scope without one,
 * instead of silently collapsing all clients into one shared bucket.
 * Contracts added later through `server.route(...)` are not visible to
 * `validate`, so enforcing an `ip`-scoped limit without a client-IP source
 * throws the same configuration error at request time as a backstop. Pass
 * `ipSource: "none"` to explicitly opt in to one unknown-client bucket per
 * contract.
 *
 * @param options - Optional key builders and client-IP source.
 * @returns A server hook backed by `ctx.ports.rateLimit`.
 */
export function createRateLimitHooks<Ctx extends CtxWithRateLimit>(
  options: RateLimitOptions<Ctx> = {},
): ServerHook<Ctx, RateLimitPorts> {
  const { ipSource, trustedProxy } = options;
  const getClientIp = (
    req: HttpRequestLike,
    requestInfo: TrustedRequestInfo | undefined,
    contractName: string,
  ): string | undefined => {
    if (ipSource !== undefined) {
      return resolveClientIp(req, ipSource);
    }
    if (trustedProxy !== undefined) {
      if (hasTrustedProxyClientIp(trustedProxy) && trustedProxy) {
        return resolveTrustedClientIp(req, trustedProxy.clientIp);
      }
      throw ipSourceConfigurationError(formatContractNames([contractName]));
    }
    if (requestInfo?.clientIpTrusted) {
      return requestInfo.clientIp;
    }
    throw ipSourceConfigurationError(formatContractNames([contractName]));
  };

  return {
    name: "rate-limit",
    validate: ({ contracts, trustedProxy: serverTrustedProxy }) => {
      if (
        ipSource !== undefined ||
        hasTrustedProxyClientIp(
          trustedProxy !== undefined ? trustedProxy : serverTrustedProxy,
        ) ||
        options.earlyKey
      ) {
        return;
      }
      const ipScoped = contracts
        .filter((contract) => contract.metadata?.rateLimit?.scope === "ip")
        .map((contract) => contract.name);
      if (ipScoped.length === 0) {
        return;
      }
      throw ipSourceConfigurationError(formatContractNames(ipScoped));
    },
    onRequest: async ({ contract, ports, req, requestInfo }) => {
      const rlMeta = contract.metadata?.rateLimit;
      if (!rlMeta) {
        return undefined;
      }

      const scope: RateLimitScope = rlMeta.scope ?? "global";
      if (scope === "user") {
        return undefined;
      }

      const key =
        options.earlyKey?.({ req, scope }) ??
        defaultEarlyRateLimitKey(
          { req, scope, contractName: contract.name },
          (r) => getClientIp(r, requestInfo, contract.name),
        );

      await enforceRateLimit(ports, {
        key,
        limit: rlMeta.max,
        windowSec: rlMeta.windowSec,
        scope,
      });

      return undefined;
    },
    beforeHandle: async ({ ctx, contract, req, requestInfo }) => {
      const rlMeta = contract.metadata?.rateLimit;
      if (!rlMeta) {
        return undefined;
      }

      const scope: RateLimitScope = rlMeta.scope ?? "global";
      if (scope !== "user") {
        return undefined;
      }

      const key =
        options.key?.({ ctx, req, scope }) ??
        defaultRateLimitKey(
          { ctx, req, scope, contractName: contract.name },
          (r) => getClientIp(r, requestInfo, contract.name),
        );

      await enforceRateLimit(ctx.ports, {
        key,
        limit: rlMeta.max,
        windowSec: rlMeta.windowSec,
        scope,
      });

      return undefined;
    },
  };
}