TallyNote/tests/rate-limit.test.ts

import { afterEach, describe, expect, it } from "vitest";
import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import { buildApp } from "../server/app.js";
import { loadConfig, prepareDataDirectories } from "../server/config.js";
import { openDatabase } from "../server/db/index.js";
import { createRateLimiter } from "../server/rate-limit.js";

const configKeys = [
  "TALLYNOTE_DATA_DIR",
  "TALLYNOTE_PUBLIC_ORIGIN",
  "TALLYNOTE_COOKIE_SECURE",
  "TALLYNOTE_ALLOW_INSECURE_HTTP",
  "TALLYNOTE_RATE_LIMIT_PER_MINUTE",
  "TALLYNOTE_TRUST_PROXY",
  "NODE_ENV",
  "TALLYNOTE_ENV",
];

afterEach(() => { for (const key of configKeys) delete process.env[key]; });

describe("内存滑动窗口限流器", () => {
  it("窗口内未超限时放行", () => {
    let clock = 1_000;
    const limiter = createRateLimiter({ limit: 3, windowMs: 60_000, now: () => clock });
    expect(limiter.check("a")).toEqual({ allowed: true, retryAfterSeconds: 0 });
    clock += 1_000;
    expect(limiter.check("a")).toEqual({ allowed: true, retryAfterSeconds: 0 });
    clock += 1_000;
    expect(limiter.check("a")).toEqual({ allowed: true, retryAfterSeconds: 0 });
  });

  it("达到上限后拒绝并给出 Retry-After 秒数", () => {
    let clock = 10_000;
    const limiter = createRateLimiter({ limit: 2, windowMs: 30_000, now: () => clock });
    expect(limiter.check("a").allowed).toBe(true);
    expect(limiter.check("a").allowed).toBe(true);
    clock += 5_000;
    const denied = limiter.check("a");
    expect(denied.allowed).toBe(false);
    expect(denied.retryAfterSeconds).toBeGreaterThan(0);
    // The window started at t=10000 and lasts 30s, so at t=15000 the caller
    // must wait the remaining 25 seconds.
    expect(denied.retryAfterSeconds).toBe(25);
  });

  it("窗口过期后计数重置并重新放行", () => {
    let clock = 0;
    const limiter = createRateLimiter({ limit: 1, windowMs: 1_000, now: () => clock });
    expect(limiter.check("a").allowed).toBe(true);
    expect(limiter.check("a").allowed).toBe(false);
    // One millisecond before the window closes the key is still limited.
    clock = 999;
    expect(limiter.check("a").allowed).toBe(false);
    clock = 1_000;
    expect(limiter.check("a")).toEqual({ allowed: true, retryAfterSeconds: 0 });
    // The reset key starts a brand-new window from the reset moment, so the
    // same key is limited again until that new window also elapses.
    clock = 1_500;
    expect(limiter.check("a").allowed).toBe(false);
    clock = 2_000;
    expect(limiter.check("a")).toEqual({ allowed: true, retryAfterSeconds: 0 });
  });

  it("不同键互不影响", () => {
    const limiter = createRateLimiter({ limit: 1, windowMs: 60_000, now: () => 0 });
    expect(limiter.check("1.2.3.4").allowed).toBe(true);
    expect(limiter.check("1.2.3.4").allowed).toBe(false);
    expect(limiter.check("5.6.7.8").allowed).toBe(true);
    expect(limiter.check("5.6.7.8").allowed).toBe(false);
    expect(limiter.size()).toBe(2);
  });

  it("惰性清理过期桶，避免长期运行内存增长", () => {
    let clock = 0;
    const limiter = createRateLimiter({ limit: 10, windowMs: 1_000, now: () => clock });
    for (let index = 0; index < 999; index += 1) limiter.check(`stale-${index}`);
    expect(limiter.size()).toBe(999);
    clock = 5_000;
    // The sweep is amortized: only a periodic full pass removes dead keys, so
    // the count must drop back to just the key currently receiving traffic.
    for (let index = 0; index < 1_000; index += 1) limiter.check("noisy");
    expect(limiter.size()).toBe(1);
  });

  it("拒绝无效的限流参数", () => {
    expect(() => createRateLimiter({ limit: 0, windowMs: 1_000 })).toThrow(/limit/);
    expect(() => createRateLimiter({ limit: 1.5, windowMs: 1_000 })).toThrow(/limit/);
    expect(() => createRateLimiter({ limit: 1, windowMs: 0 })).toThrow(/窗口/);
  });
});

describe("全局限流配置", () => {
  function validConfigEnv() {
    process.env.TALLYNOTE_PUBLIC_ORIGIN = "http://127.0.0.1:3996";
    process.env.TALLYNOTE_COOKIE_SECURE = "false";
  }

  it("默认每分钟 600 次，并支持显式覆盖", () => {
    validConfigEnv();
    expect(loadConfig().apiRateLimitPerMinute).toBe(600);
    process.env.TALLYNOTE_RATE_LIMIT_PER_MINUTE = "120";
    expect(loadConfig().apiRateLimitPerMinute).toBe(120);
    process.env.TALLYNOTE_RATE_LIMIT_PER_MINUTE = "1";
    expect(loadConfig().apiRateLimitPerMinute).toBe(1);
  });

  it("拒绝非整数或小于 1 的限流值", () => {
    validConfigEnv();
    process.env.TALLYNOTE_RATE_LIMIT_PER_MINUTE = "0";
    expect(() => loadConfig()).toThrow(/TALLYNOTE_RATE_LIMIT_PER_MINUTE/);
    process.env.TALLYNOTE_RATE_LIMIT_PER_MINUTE = "-10";
    expect(() => loadConfig()).toThrow(/TALLYNOTE_RATE_LIMIT_PER_MINUTE/);
    process.env.TALLYNOTE_RATE_LIMIT_PER_MINUTE = "abc";
    expect(() => loadConfig()).toThrow(/TALLYNOTE_RATE_LIMIT_PER_MINUTE/);
    process.env.TALLYNOTE_RATE_LIMIT_PER_MINUTE = "12.5";
    expect(() => loadConfig()).toThrow(/TALLYNOTE_RATE_LIMIT_PER_MINUTE/);
  });
});

describe("全局限流接入 HTTP 层", () => {
  const dataDirs: string[] = [];

  afterEach(() => {
    while (dataDirs.length > 0) rmSync(dataDirs.pop()!, { recursive: true, force: true });
  });

  async function buildLimitedApp(limit: string, withWeb = false) {
    const dataDir = mkdtempSync(path.join(tmpdir(), "tallynote-rate-limit-"));
    dataDirs.push(dataDir);
    process.env.TALLYNOTE_DATA_DIR = dataDir;
    process.env.TALLYNOTE_PUBLIC_ORIGIN = "http://127.0.0.1:3996";
    process.env.TALLYNOTE_COOKIE_SECURE = "false";
    process.env.TALLYNOTE_RATE_LIMIT_PER_MINUTE = limit;
    const config = loadConfig();
    // By default keep the test independent from the locally generated dist/web
    // tree. When a real asset graph is requested the exemption must still hold,
    // which proves it is path-based rather than an artefact of a missing webDir.
    config.webDir = withWeb ? path.join(dataDir, "web") : path.join(dataDir, "missing-web");
    prepareDataDirectories(config);
    if (withWeb) {
      mkdirSync(path.join(config.webDir, "assets"), { recursive: true });
      writeFileSync(path.join(config.webDir, "index.html"), "<!doctype html><title>tallynote-test-index</title>");
      writeFileSync(path.join(config.webDir, "assets", "probe.js"), "console.log('tallynote-test-asset');");
    }
    const database = openDatabase(config);
    const app = await buildApp(database, config);
    return { app, database };
  }

  it("超过配置的 /api/* 配额后返回 429 与 Retry-After，非 API 路径不受影响", async () => {
    const { app, database } = await buildLimitedApp("2");
    try {
      // Static/health traffic is exempt: the limiter only owns /api/*.
      for (let index = 0; index < 5; index += 1) {
        const health = await app.inject({ method: "GET", url: "/health" });
        expect(health.statusCode).toBe(200);
      }
      const first = await app.inject({ method: "GET", url: "/api/auth/status" });
      expect(first.statusCode).toBe(200);
      const second = await app.inject({ method: "GET", url: "/api/auth/status" });
      expect(second.statusCode).toBe(200);
      const limited = await app.inject({ method: "GET", url: "/api/auth/status" });
      expect(limited.statusCode).toBe(429);
      expect(limited.json().error.code).toBe("RATE_LIMITED");
      expect(limited.json().error.message).toBe("请求过于频繁，请稍后再试");
      expect(limited.json().error.requestId).toBeTruthy();
      expect(Number(limited.headers["retry-after"])).toBeGreaterThan(0);
      // The limiter must not touch the database: no new table, no writes to
      // the login lockout table used by the stricter login protection.
      const attempts = database.sqlite.prepare("SELECT COUNT(*) AS count FROM login_attempts").get() as { count: number };
      expect(attempts.count).toBe(0);
    } finally {
      await app.close();
      database.sqlite.close();
    }
  });

  it("配额耗尽后静态资源与 SPA 回退仍可访问", async () => {
    const { app, database } = await buildLimitedApp("1", true);
    try {
      // Spend the whole /api/* quota for this client.
      const first = await app.inject({ method: "GET", url: "/api/auth/status" });
      expect(first.statusCode).toBe(200);
      const limited = await app.inject({ method: "GET", url: "/api/auth/status" });
      expect(limited.statusCode).toBe(429);

      // A limited client must still be able to load the page and its assets,
      // otherwise recovery from the limit is impossible without a cache purge.
      const index = await app.inject({ method: "GET", url: "/" });
      expect(index.statusCode).toBe(200);
      expect(index.body).toContain("tallynote-test-index");

      const asset = await app.inject({ method: "GET", url: "/assets/probe.js" });
      expect(asset.statusCode).toBe(200);
      expect(asset.body).toContain("tallynote-test-asset");

      // An unknown non-API path falls back to the SPA entry and stays exempt.
      const fallback = await app.inject({ method: "GET", url: "/expenses" });
      expect(fallback.statusCode).toBe(200);
      expect(fallback.body).toContain("tallynote-test-index");
    } finally {
      await app.close();
      database.sqlite.close();
    }
  });
});