writeonce/docs/examples/porch/middleware/limiter.wo

-- porch/middleware/limiter.wo — rate limiter middleware. All counting is
-- delegated to the key pool (keypool.wo): this file never reads or writes
-- RateLimitCounter and holds no window arithmetic of its own. That is what
-- makes the pool's per-key serialization guarantee actually apply — a
-- store call here would be a second, uncoordinated writer.
-- Iteration 1 of the porch track, porch-store task 3.

use time
use http
use net

-- Limiter counts requests per key per window by calling into the shared
-- pool and acting on the Verdict it returns.
-- On limit exceeded: 429 with Retry-After and X-RateLimit-* headers.
-- On a saturated pool (the key's actor mailbox is full under load): 503
-- with Retry-After. The request is refused, never let through — a limiter
-- that stops limiting under load is worse than no limiter, since
-- saturating the pool would otherwise be the bypass.
-- Key selection: req.principal wins when non-empty. Otherwise, trust_proxy
-- false (default) keys on net.peer(req.conn), which cannot be forged;
-- trust_proxy true keys on client_ip(req) (the left-most X-Forwarded-For
-- entry) — the app author's assertion that a proxy they control overwrites
-- that header.
-- The refused/saturated paths stamp their own headers directly on the Resp
-- they return. The allowed path has no Resp yet to stamp — before() stashes
-- the numbers on req.ctx, and `after` (same shape as Cors: register the one
-- value as both Mw and Aw) copies them onto whatever response the chain
-- eventually produces.
pub class Limiter {
  pool:        Pool
  limit:       Int
  window:      Int   -- window size in µs (e.g., 60_000_000 = 60s)
  trust_proxy: Bool = false

  fn before(mut req: Req) -> ?Resp {
    let key = limiter_key(self, req);
    let v = try pool_count(self.pool, key, self.limit, self.window) catch (e) {
      print_err("limiter: pool_count trapped: ${e.msg}");
      nil
    };

    if v == nil {
      let r = Resp { status: 503, headers: {}, body: "{\"error\":\"rate limiter saturated\"}" };
      set_header(r, "content-type", "application/json");
      set_header(r, "retry-after", "1");
      return r;
    }

    let limit_hdr = "${v.limit}";
    let remaining = v.limit - v.count;
    if remaining < 0 { remaining = 0; }
    let reset_hdr = "${v.reset_at / 1000}"; -- wall-clock ms -> Unix seconds

    if v.allowed == false {
      let retry_sec = ((v.reset_at - time.now()) / 1000) + 1;
      let r = Resp { status: 429, headers: {}, body: "{\"error\":\"rate limit exceeded\"}" };
      set_header(r, "content-type", "application/json");
      set_header(r, "x-ratelimit-limit", limit_hdr);
      set_header(r, "x-ratelimit-remaining", "0");
      set_header(r, "x-ratelimit-reset", reset_hdr);
      set_header(r, "retry-after", "${retry_sec}");
      return r;
    }

    -- Allowed: attach headers to request for after-chain to stamp on response
    req.ctx["ratelimit_limit"] = limit_hdr;
    req.ctx["ratelimit_remaining"] = "${remaining}";
    req.ctx["ratelimit_reset"] = reset_hdr;

    return nil;
  }

  -- Stamps the allowed-path numbers before() stashed. A refused/saturated
  -- request never reaches here with anything to stamp (before() only
  -- writes ctx on the allowed path), so this is a no-op for those.
  fn after(req: Req, mut r: Resp) {
    let limit_hdr = req.ctx["ratelimit_limit"];
    let remaining_hdr = req.ctx["ratelimit_remaining"];
    let reset_hdr = req.ctx["ratelimit_reset"];
    if limit_hdr == nil { return; }
    if remaining_hdr == nil { return; }
    if reset_hdr == nil { return; }
    r.headers["x-ratelimit-limit"] = limit_hdr;
    r.headers["x-ratelimit-remaining"] = remaining_hdr;
    r.headers["x-ratelimit-reset"] = reset_hdr;
  }
}

-- Key selection: identity first, then the trust_proxy-gated peer address.
-- An absent X-Forwarded-For under trust_proxy means there is nothing to
-- trust, not an empty identity: client_ip(req) == "" falls through to
-- net.peer(req.conn) rather than keying every such client on the literal
-- "ip:" bucket (that collapse was a real bug -- one client omitting the
-- header could exhaust the shared bucket and deny/hide the rest).
pub fn limiter_key(self: Limiter, req: Req) -> Text {
  if req.principal != "" { return "principal:${req.principal}"; }
  if self.trust_proxy {
    let ip = client_ip(req);
    if ip != "" { return "ip:${ip}"; }
  }
  let peer = net.peer(req.conn);
  if peer != "" { return "ip:${peer}"; }
  return "unknown";
}

-- Helper to build a Limiter with defaults
pub fn make_limiter(pool: Pool, limit: Int, window_sec: Int) -> Limiter {
  return Limiter { pool: pool, limit: limit, window: window_sec * 1_000_000 };
}