diff --git a/.cursor/rules/mermaid-github.mdc b/.cursor/rules/mermaid-github.mdc new file mode 100644 index 0000000..2fb00fe --- /dev/null +++ b/.cursor/rules/mermaid-github.mdc @@ -0,0 +1,26 @@ +--- +description: GitHub-compatible Mermaid diagrams in markdown docs +globs: docs/**/*.md +alwaysApply: false +--- + +# Mermaid on GitHub + +GitHub uses a strict Mermaid parser. Diagrams that render in VS Code or mermaid.live may still fail on GitHub. + +## Sequence diagrams (`sequenceDiagram`) + +- **Never use `;` in arrow labels** — GitHub treats `;` as end-of-statement. Use `,`, ` and `, or ` — ` instead. +- **Multi-line labels:** use `
` inside the label, not a real newline. +- **Special characters:** if a label must contain `;`, `:`, or parentheses, wrap the whole label in double quotes: `A->>B: "Parse SIWE; recover (ecrecover)"`. +- **Participant aliases:** keep short; avoid unquoted `:` inside message text when possible. + +## Flowcharts (`flowchart`) + +- Prefer quoted edge labels for paths, colons, or braces: `A -->|"POST /session {msg}"| B`. +- Avoid raw newlines inside node text; use `
` or split into separate nodes. + +## Before committing doc diagrams + +1. Grep mermaid blocks for `;` in `->>` / `-->>` lines. +2. Preview on GitHub (PR) or match [GitHub Mermaid docs](https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/creating-diagrams#creating-mermaid-diagrams). diff --git a/AGENTS.md b/AGENTS.md index f43d683..61cb885 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -127,3 +127,6 @@ Figment order (later overrides earlier): All fields must have sensible defaults so the server starts with zero config. Secrets (signing keys, service-account creds) come from env / a secret manager, never from committed files. + +`APP_SERVER__EXPOSE` selects the HTTP surface: `all` (dev), `public` (wallet + +JWKS), or `internal` (RS `POST /v1/auth/quota/consume`). See `docs/DEPLOYMENT.md`. diff --git a/config/default.toml b/config/default.toml index 14eb672..0ca6c0e 100644 --- a/config/default.toml +++ b/config/default.toml @@ -5,6 +5,8 @@ [server] host = "127.0.0.1" port = 3001 +# Route surface: "all" (dev), "public" (wallet + JWKS), "internal" (RS quota APIs). +expose = "all" [telemetry] # "json" for production, "pretty" for local development. @@ -26,6 +28,19 @@ default_scopes = ["ai:invoke", "mint:request", "scan:submit", "profile:read"] refresh_token_ttl_secs = 604800 issue_refresh_tokens = true +[quota] +enabled = true +window_secs = 86400 +paid_scopes = ["ai:invoke", "mint:request"] + +[quota.limits.new_wallet] +"ai:invoke" = 10 +"mint:request" = 2 + +[quota.limits.established] +"ai:invoke" = 1000 +"mint:request" = 50 + [rate_limit] nonce_per_ip_per_minute = 30 diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index e0bdbb1..8f21255 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -8,6 +8,7 @@ _Version 0.1 — phase-0 wallet-signature session authority_ 2. [Design Principles](#2-design-principles) 3. [The Pattern: OAuth2 AS/RS with a SIWE Grant](#3-the-pattern-oauth2-asrs-with-a-siwe-grant) 4. [Deployment Topology](#4-deployment-topology) + - [4.1 Network exposure (public vs internal)](#41-network-exposure-public-vs-internal) 5. [Build vs. Buy](#5-build-vs-buy) 6. [Component Architecture](#6-component-architecture) 7. [Authentication Flow](#7-authentication-flow) @@ -180,6 +181,50 @@ hot part (signature verification) wants to be everywhere. B splits them correctly; C smears sensitive state across every service. `beacon-relay` becomes the first *consumer*, not the host. +### 4.1 Network exposure (public vs internal) + +One logical service, **two ingress surfaces** on Google Cloud (see +[DEPLOYMENT.md](DEPLOYMENT.md) and [adr/0003-public-internal-deployment.md](adr/0003-public-internal-deployment.md)): + +```mermaid +flowchart TB + subgraph Public["trust-relay-public — APP_SERVER__EXPOSE=public"] + JWKS["/.well-known/jwks.json"] + SIWE_EP["/v1/auth/nonce · session · refresh · logout"] + QGET["GET /v1/auth/quota"] + end + + subgraph Internal["trust-relay-internal — APP_SERVER__EXPOSE=internal"] + QCON["POST /v1/auth/quota/consume"] + end + + subgraph Store["Memorystore Redis (private IP)"] + N[nonces] + Q[quota counters] + V[revocation · refresh] + end + + Wallet((Nodle app)) --> SIWE_EP + Wallet --> QGET + RS[Resource servers] -->|user Bearer forwarded| QCON + RS -.->|cache JWKS| JWKS + Public --> Store + Internal --> Store +``` + +| Surface | Cloud Run ingress | Routes | Authenticated identity (today) | +| --- | --- | --- | --- | +| **Public** | `all` / external HTTPS LB | SIWE ceremony, JWKS, logout, optional quota read | Wallet (SIWE) or user JWT | +| **Internal** | `internal` (VPC only) | `POST /v1/auth/quota/consume`, `/healthz` | **User JWT** in `Authorization` (RS not identified) | + +Resource servers **must not** call quota consume on the public URL. They call the +internal service inside the VPC and **must forward the user's bearer token** — +the same `Authorization: Bearer ` header the client sent to the RS. +trust-relay does not authenticate the RS on this path in phase 0; it verifies the +forwarded user JWT and decrements quota for that wallet (`sub`). Long term, Rust +RSs should use a **`trust-auth`** layer with optional Redis replica reads instead +of per-request HTTP (M3b). + --- ## 5. Build vs. Buy @@ -242,7 +287,8 @@ Signing keys: from a secret manager / KMS, never committed. sequenceDiagram autonumber participant C as Client (wallet) - participant P as trust-relay (AS) + participant P as trust-relay public + participant I as trust-relay internal participant R as Redis participant RS as Resource Server @@ -253,7 +299,7 @@ sequenceDiagram C->>C: Build EIP-4361 message (domain, uri, chainId, nonce, iat, exp) C->>C: Sign message with wallet key (EOA) C->>P: POST /v1/auth/session { message, signature } - P->>P: Parse SIWE; recover address (ecrecover) + P->>P: Parse SIWE, recover address via ecrecover P->>R: validate nonce (exists, unused, unexpired, ip-bound?) P->>R: mark nonce used P->>P: optional wallet-heuristic gate for paid scopes @@ -266,7 +312,8 @@ sequenceDiagram RS->>RS: verify signature with cached JWKS public key RS->>RS: check exp/nbf/iss/aud, required scope, tier alt sensitive route (/ai/*, /mint/*) - RS->>R: jti not revoked? wallet not blocked? quota remaining? + RS->>I: POST /quota/consume (user Bearer) + I->>R: decrement quota / check revocation end RS-->>C: 200 / 401 / 403 / 429 @@ -274,18 +321,31 @@ sequenceDiagram C->>P: POST /v1/auth/session/refresh { refreshToken } P->>R: validate + rotate refresh token P-->>C: { accessToken, expiresIn, refreshToken } + + Note over C,RS: Paid route (sensitive) + C->>RS: Bearer JWT on /ai/* or /mint/* + RS->>RS: verify JWT (cached JWKS from public) + RS->>I: POST /v1/auth/quota/consume (internal ingress) + I->>R: decrement wallet quota + I-->>RS: remaining or 429 + RS-->>C: 200 / 403 / 429 ``` +`P` = trust-relay **public** deploy; `I` = trust-relay **internal** deploy (same +image, `APP_SERVER__EXPOSE=internal`). Both share Redis. + ### Endpoints -| Endpoint | Method | Purpose | -| --- | --- | --- | -| `GET /healthz` | GET | Liveness. | -| `GET /.well-known/jwks.json` | GET | Public verification keys (JWKS). | -| `GET /v1/auth/nonce` | GET | Issue a single-use nonce. Per-IP rate limited. | -| `POST /v1/auth/session` | POST | Verify SIWE message + signature; issue tokens. | -| `POST /v1/auth/session/refresh` | POST | Rotate the bearer without a new signature. | -| `POST /v1/auth/logout` | POST | Revoke the current bearer (`jti`). | +| Endpoint | Method | Exposure | Purpose | +| --- | --- | --- | --- | +| `GET /healthz` | GET | public + internal | Liveness. | +| `GET /.well-known/jwks.json` | GET | public | Public verification keys (JWKS). | +| `GET /v1/auth/nonce` | GET | public | Issue a single-use nonce. Per-IP rate limited. | +| `POST /v1/auth/session` | POST | public | Verify SIWE message + signature; issue tokens. | +| `POST /v1/auth/session/refresh` | POST | public | Rotate the bearer without a new signature. | +| `POST /v1/auth/logout` | POST | public | Revoke the current bearer (`jti`). | +| `GET /v1/auth/quota` | GET | public | Remaining quota buckets for the bearer wallet. | +| `POST /v1/auth/quota/consume` | POST | **internal** | Decrement a paid scope (RS forwards user JWT). | --- @@ -423,11 +483,11 @@ sequenceDiagram B-->>D: attestation_result (signed by beacon-relay) D->>T: SIWE re-auth or session upgrade with resources: attestation:result - T->>T: Verify result via beacon-relay JWKS; wallet match + T->>T: Verify result via beacon-relay JWKS, wallet match T-->>D: attested-tier JWT (tier=attested, att.* claims) D->>B: POST /v2/scan/ble Bearer attested JWT - B->>B: Verify via trust-relay JWKS; attested middleware path + B->>B: Verify via trust-relay JWKS, attested middleware path ``` - **Phase 2 — hardware-backed wallet key.** Wallet key moves to secure enclave; diff --git a/docs/DEPLOYMENT.md b/docs/DEPLOYMENT.md new file mode 100644 index 0000000..b34e3ed --- /dev/null +++ b/docs/DEPLOYMENT.md @@ -0,0 +1,141 @@ +# Trust Relay — Deployment & network exposure + +How to split **public** (wallet + JWKS) and **internal** (resource-server quota) surfaces on Google Cloud, and how the binary selects routes. + +## Route exposure matrix + +| Endpoint | Method | `public` | `internal` | Caller | +| --- | --- | :---: | :---: | --- | +| `/healthz` | GET | yes | yes | Load balancer / Cloud Run probe | +| `/.well-known/jwks.json` | GET | yes | no | All resource servers (cache) | +| `/v1/auth/nonce` | GET | yes | no | Nodle app (wallet) | +| `/v1/auth/session` | POST | yes | no | Nodle app (wallet) | +| `/v1/auth/session/refresh` | POST | yes | no | Nodle app (wallet) | +| `/v1/auth/logout` | POST | yes | no | Nodle app (wallet) | +| `/v1/auth/quota` | GET | yes | no | App (UX) or RS pre-check | +| `/v1/auth/quota/consume` | POST | **no** | yes | **Resource servers only** | + +Set `server.expose` via config or `APP_SERVER__EXPOSE`: + +| Value | Use case | +| --- | --- | +| `all` | Local dev, integration tests (default in `config/default.toml`) | +| `public` | Cloud Run service with `ingress: all` (internet-facing) | +| `internal` | Cloud Run service with `ingress: internal` (VPC-only) | + +## Google Cloud topology (recommended) + +Two **Cloud Run** services from the **same container image**, different env and ingress: + +```mermaid +flowchart TB + subgraph Internet + APP[Nodle app] + end + + subgraph GCP_VPC["VPC"] + subgraph Public_CR["trust-relay-public"] + P_ING["ingress: all"] + P_ROUTES["expose = public"] + end + + subgraph Internal_CR["trust-relay-internal"] + I_ING["ingress: internal"] + I_ROUTES["expose = internal"] + end + + REDIS[(Memorystore for Redis)] + AI[AI agent RS] + MINT[NFT mint RS] + end + + APP -->|SIWE ceremony| Public_CR + AI -->|VPC + user Bearer| Internal_CR + MINT -->|VPC + user Bearer| Internal_CR + Public_CR --> REDIS + Internal_CR --> REDIS + AI -.->|JWKS over HTTPS| Public_CR + MINT -.->|JWKS over HTTPS| Public_CR +``` + +### Checklist + +1. **Memorystore for Redis** — private IP in the VPC; nonces, refresh tokens, revocation, quotas. +2. **Serverless VPC Access** (or Direct VPC egress) on **both** Cloud Run services. +3. **trust-relay-public** — `APP_SERVER__EXPOSE=public`; external HTTPS LB or Cloud Run `ingress: all`. +4. **trust-relay-internal** — `APP_SERVER__EXPOSE=internal`; `ingress: internal` only ([Cloud Run internal ingress](https://cloud.google.com/run/docs/securing/ingress)). +5. **RS** — call internal URL; **forward** the user's `Authorization: Bearer` JWT (see below). +6. **Signing key** — Secret Manager → env `APP_SIGNING__KEY_SEED_B64`; never in the image. +7. **Optional** — Cloud Armor on public service (rate-limit `GET /nonce`); `roles/run.invoker` on internal for RS service accounts. + +## Quota consume: who is authenticated? + +`POST /v1/auth/quota/consume` does **not** authenticate the resource server today. trust-relay: + +1. Verifies the **end-user access JWT** (same as logout). +2. Checks the token's `scope` includes the requested scope. +3. Decrements the wallet counter in Redis (`sub` from the JWT). + +The RS is trusted to **reach the internal URL** (network) and to **forward the +user's bearer token**. Private ingress is the phase-0 control; RS service +identity (mTLS / service JWT) is a follow-up. + +### RS MUST forward the user `Authorization` header + +When the Nodle app calls a resource server (e.g. `POST /ai/...`), it sends +`Authorization: Bearer `. Before serving the route, the RS **MUST** +copy that header onto the internal trust-relay request: + +```http +POST /v1/auth/quota/consume HTTP/1.1 +Host: trust-relay-internal +Authorization: Bearer +Content-Type: application/json + +{"scope":"ai:invoke","amount":1} +``` + +| Do | Don't | +| --- | --- | +| Forward the client's `Authorization: Bearer` unchanged | Call internal trust-relay with no `Authorization` | +| Use the internal Cloud Run / VPC URL | Expose or call consume on the public URL | +| Verify the JWT locally first (JWKS) | Send an RS service-account or workload-identity token instead of the user JWT | + +```mermaid +sequenceDiagram + autonumber + participant C as Nodle app + participant P as trust-relay-public + participant RS as Resource server + participant I as trust-relay-internal + participant R as Redis + + C->>P: POST /session (SIWE) + P->>R: init quota buckets + P-->>C: access JWT + + C->>RS: Bearer JWT (e.g. /ai/*) + RS->>RS: Verify JWT locally (JWKS from public) + RS->>I: POST /quota/consume
Authorization: user JWT + I->>I: Verify JWT + scope + I->>R: decrement quota + I-->>RS: remaining / 429 + RS-->>C: 200 or 429 +``` + +## Local development + +```bash +# Everything on one port (default) +APP_SERVER__EXPOSE=all cargo run + +# Simulate split deploys (two terminals) +APP_SERVER__PORT=3001 APP_SERVER__EXPOSE=public cargo run +APP_SERVER__PORT=3002 APP_SERVER__EXPOSE=internal cargo run +``` + +## Future: `trust-auth` crate + +Rust resource servers should move quota/revocation checks into a **Tower layer** with optional **Redis replica** reads, avoiding a per-request HTTP hop to trust-relay-internal. The internal service remains the **authoritative writer** until then. + +See [adr/0003-public-internal-deployment.md](adr/0003-public-internal-deployment.md) and [ARCHITECTURE.md](ARCHITECTURE.md) §4.1. diff --git a/docs/README.md b/docs/README.md index e9f4ee8..9d8e9c0 100644 --- a/docs/README.md +++ b/docs/README.md @@ -7,14 +7,20 @@ Reading order: attestation upgrade flow with beacon-relay, revocation authority, Redis topology, and the implementation roadmap. -2. **[adr/0001-siwe-wallet-session-auth.md](adr/0001-siwe-wallet-session-auth.md)** +2. **[DEPLOYMENT.md](DEPLOYMENT.md)** — public vs internal route surfaces, + `APP_SERVER__EXPOSE`, Google Cloud (two Cloud Run services + Memorystore). + +3. **[adr/0001-siwe-wallet-session-auth.md](adr/0001-siwe-wallet-session-auth.md)** — standalone SIWE session authority; build vs buy; JWT/JWKS pattern. -3. **[adr/0002-attestation-upgrade-and-sole-issuer.md](adr/0002-attestation-upgrade-and-sole-issuer.md)** +4. **[adr/0002-attestation-upgrade-and-sole-issuer.md](adr/0002-attestation-upgrade-and-sole-issuer.md)** — sole session issuer; beacon-relay attestation results; identity model; revocation authority; Redis split with beacon-relay. -4. **[TOKEN-SPEC.md](TOKEN-SPEC.md)** — bearer JWT contract, scopes, attested +5. **[adr/0003-public-internal-deployment.md](adr/0003-public-internal-deployment.md)** + — split public/internal ingress on GCP; quota consume on VPC-only surface. + +6. **[TOKEN-SPEC.md](TOKEN-SPEC.md)** — bearer JWT contract, scopes, attested `att.*` claims, attestation-result credential, RS verification checklist. ## Two-service relationship diff --git a/docs/TOKEN-SPEC.md b/docs/TOKEN-SPEC.md index c7a8e07..463a612 100644 --- a/docs/TOKEN-SPEC.md +++ b/docs/TOKEN-SPEC.md @@ -162,6 +162,40 @@ request whose required scope is absent (`403`). - Suggested lifetime: ~7 days. Because EOA re-signing is silent, refresh tokens are a latency/UX optimization, not a requirement. +### Quota APIs (resource servers) + +Per-wallet counters for paid scopes (`ai:invoke`, `mint:request` by default) are +initialized on `POST /v1/auth/session` and stored in trust-relay Redis. + +| Endpoint | Exposure | Caller | Auth | +| --- | --- | --- | --- | +| `GET /v1/auth/quota` | **Public** deploy | App (UX) or RS | User access JWT | +| `POST /v1/auth/quota/consume` | **Internal** deploy only | Resource servers | User access JWT (forwarded by RS) | + +**RS forwarding rule (required):** On sensitive routes the resource server calls +`POST /v1/auth/quota/consume` on the **internal** trust-relay URL and **MUST** +pass through the end-user's access JWT in `Authorization: Bearer ` — +the **same** bearer the client sent to the RS. trust-relay authenticates that +user token (not the RS). Do **not** substitute an RS service account, mTLS +identity, or omit `Authorization`; phase 0 has no RS credential on this path. + +`POST /v1/auth/quota/consume` request: + +```json +{ "scope": "ai:invoke", "amount": 1 } +``` + +Success: + +```json +{ "scope": "ai:invoke", "remaining": 9, "limit": 10 } +``` + +Exhausted: `429` with `error: "quota_exceeded"`. + +Deploy split and GCP layout: [DEPLOYMENT.md](DEPLOYMENT.md). Configure the binary +with `APP_SERVER__EXPOSE=public|internal|all`. + --- ## 6. Revocation @@ -202,7 +236,11 @@ is the hybrid revocation check, required only on sensitive routes. 8. Verify the required `scope` is present; if the route requires it, verify `tier == "attested"`. 9. **Sensitive routes only:** check the `jti` is not in the revocation set and - `sub` is not in the wallet blocklist; check and decrement per-wallet quota. + `sub` is not in the wallet blocklist; check and decrement per-wallet quota + (via `POST /v1/auth/quota/consume` on the **internal** trust-relay URL, or a + future `trust-auth` / Redis-replica check). The HTTP call **MUST** forward the + client's `Authorization: Bearer` header unchanged — trust-relay decrements + quota for `sub` in that JWT. On failure, respond: diff --git a/docs/adr/0003-public-internal-deployment.md b/docs/adr/0003-public-internal-deployment.md new file mode 100644 index 0000000..236d4db --- /dev/null +++ b/docs/adr/0003-public-internal-deployment.md @@ -0,0 +1,53 @@ +# ADR 0003: Public vs internal HTTP surfaces on Google Cloud + +**Status:** Accepted (phase 0, M5) +**Date:** 2026-06-05 + +## Context + +trust-relay exposes both **wallet-facing** APIs (SIWE ceremony, JWKS) and **backend-only** APIs (`POST /v1/auth/quota/consume`). A single internet-reachable process would let any holder of a user JWT call quota consume from the public internet, even though only resource servers (RS) should drive that path. + +ARCHITECTURE already separates control plane (rare, centralized) from data plane (JWT verify on RS). We need the same split at **network ingress** on GCP. + +## Decision + +1. **One binary, two route surfaces** selected by `server.expose`: + - `public` — healthz, JWKS, nonce, session, refresh, logout, `GET /quota` + - `internal` — healthz, `POST /quota/consume` + - `all` — union (default for dev/tests) + +2. **Production on Google Cloud:** deploy **two Cloud Run services** from the same image: + - `trust-relay-public` — `APP_SERVER__EXPOSE=public`, ingress **all** (or external HTTPS LB) + - `trust-relay-internal` — `APP_SERVER__EXPOSE=internal`, ingress **internal** (VPC-only) + +3. **Memorystore for Redis** stays **private**; both services connect via VPC connector / Direct VPC egress. + +4. **Quota consume auth (phase 0):** authenticate the **user bearer JWT**, not the RS. The RS **MUST** forward the client's `Authorization: Bearer` header on `POST /v1/auth/quota/consume` (unchanged). Network isolation on the internal service is the first RS guardrail; service-to-service credentials are deferred. + +## Consequences + +**Positive** + +- Clear operational model aligned with TOKEN-SPEC §7 (RS enforces quota on sensitive routes). +- No second codebase; config-only split. +- Internal URL can be locked to VPC without blocking JWKS on the public service. + +**Negative / accepted** + +- Two Cloud Run services to operate (scale, deploy, monitor) — acceptable vs coupling everything to one URL. +- RS HTTP hop to consume remains until `trust-auth` + Redis replica (M3b). +- `GET /v1/auth/quota` stays on **public** so apps can show remaining quota; RS should prefer consume's response or internal metrics. + +## Alternatives considered + +| Alternative | Why not (for now) | +| --- | --- | +| Single public service + path blocking at LB | Default `run.app` URL may still expose blocked paths; weaker than `internal` ingress | +| Separate binaries | Unnecessary duplication | +| mTLS on every consume call | Correct long-term; deferred to avoid blocking M5 | + +## References + +- [DEPLOYMENT.md](../DEPLOYMENT.md) — route matrix and GCP checklist +- [ARCHITECTURE.md](../ARCHITECTURE.md) §4.1 — deployment diagram +- [TOKEN-SPEC.md](../TOKEN-SPEC.md) §7 — RS verification checklist (quota step) diff --git a/src/bootstrap.rs b/src/bootstrap.rs index aa4e11d..80e48e3 100644 --- a/src/bootstrap.rs +++ b/src/bootstrap.rs @@ -7,10 +7,12 @@ use crate::{ error::AppError, middleware::rate_limit::{new_nonce_rate_limiter, RateLimiter}, services::{ - nonce::NonceService, refresh::RefreshService, revocation::RevocationService, - token::TokenService, + heuristics::HeuristicsService, nonce::NonceService, quota::QuotaService, + refresh::RefreshService, revocation::RevocationService, token::TokenService, }, + store::heuristics::{HeuristicsStore, InMemoryHeuristicsStore, RedisHeuristicsStore}, store::nonce::{InMemoryNonceStore, NonceStore, RedisNonceStore}, + store::quota::{InMemoryQuotaStore, QuotaStore, RedisQuotaStore}, store::refresh::{InMemoryRefreshStore, RedisRefreshStore, RefreshStore}, store::revocation::{InMemoryRevocationStore, RedisRevocationStore, RevocationStore}, }; @@ -69,6 +71,32 @@ pub fn build_refresh_service(config: &Config, store: Arc) -> A Arc::new(RefreshService::new(store, config.auth.clone())) } +pub async fn build_heuristics_store(config: &Config) -> Result, AppError> { + if config.uses_redis() { + Ok(Arc::new( + RedisHeuristicsStore::connect(&config.redis.url).await?, + )) + } else { + Ok(Arc::new(InMemoryHeuristicsStore::default())) + } +} + +pub fn build_heuristics_service(store: Arc) -> Arc { + Arc::new(HeuristicsService::new(store)) +} + +pub async fn build_quota_store(config: &Config) -> Result, AppError> { + if config.uses_redis() { + Ok(Arc::new(RedisQuotaStore::connect(&config.redis.url).await?)) + } else { + Ok(Arc::new(InMemoryQuotaStore::default())) + } +} + +pub fn build_quota_service(config: &Config, store: Arc) -> Arc { + Arc::new(QuotaService::new(store, config.quota.clone())) +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/config.rs b/src/config.rs index 128b3c6..8462e76 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,5 +1,7 @@ //! Figment-based configuration loading. +use std::collections::HashMap; + use figment::{ providers::{Env, Format, Toml}, Figment, @@ -14,12 +16,28 @@ pub struct Config { pub auth: AuthConfig, pub rate_limit: RateLimitConfig, pub signing: SigningConfig, + pub quota: QuotaConfig, +} + +/// Which HTTP routes this process exposes (see `docs/DEPLOYMENT.md`). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Default, serde::Serialize)] +#[serde(rename_all = "lowercase")] +pub enum ExposeMode { + /// All routes (local dev and single-process deploys). + #[default] + All, + /// Internet-facing: SIWE ceremony, JWKS, logout; no RS quota APIs. + Public, + /// VPC-only: quota consume and related backend helpers; plus `/healthz`. + Internal, } #[derive(Debug, Clone, Deserialize)] pub struct ServerConfig { pub host: String, pub port: u16, + #[serde(default)] + pub expose: ExposeMode, } #[derive(Debug, Clone, Deserialize)] @@ -73,6 +91,55 @@ pub struct RateLimitConfig { pub nonce_per_ip_per_minute: u32, } +#[derive(Debug, Clone, Deserialize)] +pub struct QuotaConfig { + #[serde(default = "default_quota_enabled")] + pub enabled: bool, + #[serde(default = "default_quota_window_secs")] + pub window_secs: u32, + #[serde(default = "default_paid_scopes")] + pub paid_scopes: Vec, + #[serde(default)] + pub limits: QuotaLimitTiers, +} + +fn default_quota_enabled() -> bool { + true +} + +fn default_quota_window_secs() -> u32 { + 86_400 +} + +fn default_paid_scopes() -> Vec { + vec!["ai:invoke".into(), "mint:request".into()] +} + +#[derive(Debug, Clone, Deserialize)] +pub struct QuotaLimitTiers { + #[serde(default = "default_new_wallet_limits")] + pub new_wallet: HashMap, + #[serde(default = "default_established_limits")] + pub established: HashMap, +} + +fn default_new_wallet_limits() -> HashMap { + HashMap::from([("ai:invoke".into(), 10), ("mint:request".into(), 2)]) +} + +fn default_established_limits() -> HashMap { + HashMap::from([("ai:invoke".into(), 1_000), ("mint:request".into(), 50)]) +} + +impl Default for QuotaLimitTiers { + fn default() -> Self { + Self { + new_wallet: default_new_wallet_limits(), + established: default_established_limits(), + } + } +} + #[derive(Debug, Clone, Deserialize)] pub struct SigningConfig { pub issuer: String, @@ -106,12 +173,25 @@ impl Config { mod tests { use super::*; + #[test] + fn expose_mode_deserializes_lowercase() { + assert_eq!( + serde_json::from_str::("\"public\"").unwrap(), + ExposeMode::Public + ); + assert_eq!( + serde_json::from_str::("\"internal\"").unwrap(), + ExposeMode::Internal + ); + } + #[test] fn uses_redis_when_url_is_non_empty() { let cfg = Config { server: ServerConfig { host: "127.0.0.1".into(), port: 3001, + expose: ExposeMode::All, }, telemetry: TelemetryConfig { format: "pretty".into(), @@ -140,6 +220,12 @@ mod tests { access_token_ttl_secs: 3600, key_seed_b64: String::new(), }, + quota: QuotaConfig { + enabled: true, + window_secs: default_quota_window_secs(), + paid_scopes: default_paid_scopes(), + limits: QuotaLimitTiers::default(), + }, }; assert!(cfg.uses_redis()); let mut empty = cfg; diff --git a/src/routes/auth.rs b/src/routes/auth.rs index 82f09e2..94531e1 100644 --- a/src/routes/auth.rs +++ b/src/routes/auth.rs @@ -10,7 +10,10 @@ use axum::{ use serde::{Deserialize, Serialize}; use time::format_description::well_known::Rfc3339; -use crate::{error::AppError, services::siwe::verify_siwe_session, state::AppState}; +use crate::{ + error::AppError, models::claims::AccessTokenClaims, services::siwe::verify_siwe_session, + state::AppState, +}; #[derive(Serialize)] pub struct NonceResponse { @@ -56,6 +59,29 @@ pub struct RefreshResponse { pub refresh_token: String, } +#[derive(Deserialize)] +pub struct QuotaConsumeRequest { + pub scope: String, + #[serde(default = "default_quota_amount")] + pub amount: u64, +} + +fn default_quota_amount() -> u64 { + 1 +} + +#[derive(Serialize)] +pub struct QuotaConsumeResponse { + pub scope: String, + pub remaining: u64, + pub limit: u64, +} + +#[derive(Serialize)] +pub struct QuotaStatusResponse { + pub scopes: Vec, +} + pub async fn get_nonce( State(state): State, headers: HeaderMap, @@ -88,6 +114,15 @@ pub async fn post_session( .await?; state.nonce.validate_and_consume(&verified.nonce).await?; + let profile = state + .heuristics + .profile_for(&verified.wallet_address) + .await?; + state + .quota + .init_for_wallet(&verified.wallet_address, profile) + .await?; + let scopes = state.config.auth.default_scopes.clone(); let scope_refs: Vec<&str> = scopes.iter().map(String::as_str).collect(); let access_token = state.token.mint_wallet_token( @@ -112,6 +147,11 @@ pub async fn post_session( None }; + state + .heuristics + .mark_established(&verified.wallet_address) + .await?; + Ok(Json(SessionResponse { access_token, token_type: "Bearer", @@ -137,6 +177,13 @@ pub async fn post_session_refresh( .revocation .ensure_wallet_active(&rotated.record.wallet) .await?; + state + .quota + .init_for_wallet( + &rotated.record.wallet, + crate::services::heuristics::WalletProfile::Established, + ) + .await?; let scope_refs: Vec<&str> = rotated.record.scopes.iter().map(String::as_str).collect(); let access_token = state.token.mint_wallet_token( @@ -152,6 +199,45 @@ pub async fn post_session_refresh( })) } +/// `POST /v1/auth/quota/consume` — decrement a paid-scope counter (RS hot-path helper). +pub async fn post_quota_consume( + State(state): State, + headers: HeaderMap, + Json(body): Json, +) -> Result, AppError> { + let claims = verified_claims(&state, &headers)?; + ensure_scope(&claims, &body.scope)?; + let status = state + .quota + .consume(&claims.sub, &body.scope, body.amount) + .await?; + Ok(Json(QuotaConsumeResponse { + scope: status.scope, + remaining: status.remaining, + limit: status.limit, + })) +} + +/// `GET /v1/auth/quota` — remaining paid-scope quota for the bearer wallet. +pub async fn get_quota( + State(state): State, + headers: HeaderMap, +) -> Result, AppError> { + let claims = verified_claims(&state, &headers)?; + let scopes: Vec<&str> = claims.scope.split_whitespace().collect(); + let statuses = state.quota.status_for_scopes(&claims.sub, &scopes).await?; + Ok(Json(QuotaStatusResponse { + scopes: statuses + .into_iter() + .map(|s| QuotaConsumeResponse { + scope: s.scope, + remaining: s.remaining, + limit: s.limit, + }) + .collect(), + })) +} + /// `POST /v1/auth/logout` — revoke the bearer access token's `jti`. pub async fn post_logout( State(state): State, @@ -163,6 +249,21 @@ pub async fn post_logout( Ok(StatusCode::NO_CONTENT) } +fn verified_claims(state: &AppState, headers: &HeaderMap) -> Result { + let token = bearer_token(headers)?; + state.token.verify(&token) +} + +fn ensure_scope(claims: &AccessTokenClaims, scope: &str) -> Result<(), AppError> { + if claims.scope.split_whitespace().any(|s| s == scope) { + Ok(()) + } else { + Err(AppError::InsufficientScope(format!( + "token does not grant scope {scope}" + ))) + } +} + fn bearer_token(headers: &HeaderMap) -> Result { let value = headers .get(axum::http::header::AUTHORIZATION) diff --git a/src/routes/mod.rs b/src/routes/mod.rs index ca10657..8d8531f 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -7,47 +7,76 @@ mod jwks; use axum::{routing::get, routing::post, Router}; use tower_http::{cors::CorsLayer, limit::RequestBodyLimitLayer, trace::TraceLayer}; -use crate::state::AppState; +use crate::{config::ExposeMode, state::AppState}; /// Max request body size for auth/session routes (POST JSON). Not applied to `/healthz`. pub const MAX_BODY_BYTES: usize = 64 * 1024; -/// Liveness, JWKS, and other public endpoints — no body limit or rate limiting. -fn public_routes(state: AppState) -> Router { +/// Liveness and JWKS — exposed on `public` and `all` deploys. +fn health_and_jwks_routes(state: AppState) -> Router { Router::new() .route("/healthz", get(health::healthz)) .route("/.well-known/jwks.json", get(jwks::jwks)) .with_state(state) } -/// Auth and session routes — body size cap and nonce handlers only here. -fn api_routes(state: AppState) -> Router { +/// Wallet SIWE ceremony and session lifecycle — internet-facing. +fn wallet_auth_routes(state: AppState) -> Router { Router::new() .route("/v1/auth/nonce", get(auth::get_nonce)) .route("/v1/auth/session", post(auth::post_session)) .route("/v1/auth/session/refresh", post(auth::post_session_refresh)) .route("/v1/auth/logout", post(auth::post_logout)) + .route("/v1/auth/quota", get(auth::get_quota)) .with_state(state) .layer(RequestBodyLimitLayer::new(MAX_BODY_BYTES)) } -/// Build the HTTP router with shared state and global middleware. -pub fn create_router(state: AppState) -> Router { - let api = api_routes(state.clone()); +/// Resource-server helpers — VPC / internal ingress only in production. +fn internal_routes(state: AppState) -> Router { Router::new() - .merge(public_routes(state)) - .merge(api) + .route("/v1/auth/quota/consume", post(auth::post_quota_consume)) + .with_state(state) + .layer(RequestBodyLimitLayer::new(MAX_BODY_BYTES)) +} + +fn apply_global_layers(router: Router) -> Router { + router .layer(TraceLayer::new_for_http()) .layer(CorsLayer::permissive()) } +/// Build the HTTP router for `cfg.server.expose`. +pub fn create_router(state: AppState) -> Router { + let expose = state.config.server.expose; + create_router_for_expose(state, expose) +} + +/// Build the HTTP router for an explicit exposure mode (tests and docs). +pub fn create_router_for_expose(state: AppState, expose: ExposeMode) -> Router { + let health_state = state.clone(); + let router = match expose { + ExposeMode::All => Router::new() + .merge(health_and_jwks_routes(health_state.clone())) + .merge(wallet_auth_routes(state.clone())) + .merge(internal_routes(state)), + ExposeMode::Public => Router::new() + .merge(health_and_jwks_routes(health_state)) + .merge(wallet_auth_routes(state)), + ExposeMode::Internal => Router::new() + .merge(health_and_jwks_routes(health_state)) + .merge(internal_routes(state)), + }; + apply_global_layers(router) +} + /// Integration-test router with in-memory stores (or Redis when `APP_REDIS__URL` is set). pub async fn test_router() -> Router { let config = crate::config::Config::load().expect("test config should load"); test_router_with_config(config).await } -/// Test helper: router with explicit config. +/// Test helper: router with explicit config (`server.expose` respected). pub async fn test_router_with_config(config: crate::config::Config) -> Router { let state = AppState::build(config) .await diff --git a/src/server.rs b/src/server.rs index 1bd7edb..0a5b396 100644 --- a/src/server.rs +++ b/src/server.rs @@ -26,7 +26,8 @@ pub async fn serve( state: AppState, ) -> Result<(), Box> { let addr = listener.local_addr()?; - tracing::info!(%addr, "trust-relay listening"); + let expose = state.config.server.expose; + tracing::info!(%addr, ?expose, "trust-relay listening"); axum::serve(listener, routes::create_router(state)).await?; Ok(()) } diff --git a/src/services/heuristics.rs b/src/services/heuristics.rs new file mode 100644 index 0000000..97d4223 --- /dev/null +++ b/src/services/heuristics.rs @@ -0,0 +1,53 @@ +//! Phase-0 wallet heuristics (first-seen vs established). + +use std::sync::Arc; + +use crate::{error::AppError, store::heuristics::HeuristicsStore}; + +/// Phase-0 profile derived before minting paid-scope quotas. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WalletProfile { + /// First successful session for this wallet. + New, + Established, +} + +pub struct HeuristicsService { + store: Arc, +} + +impl HeuristicsService { + pub fn new(store: Arc) -> Self { + Self { store } + } + + pub async fn profile_for(&self, wallet: &str) -> Result { + let wallet = normalize_wallet(wallet)?; + if self.store.is_established(&wallet).await? { + Ok(WalletProfile::Established) + } else { + Ok(WalletProfile::New) + } + } + + pub async fn mark_established(&self, wallet: &str) -> Result<(), AppError> { + let wallet = normalize_wallet(wallet)?; + self.store.mark_established(&wallet).await?; + Ok(()) + } +} + +fn normalize_wallet(wallet: &str) -> Result { + let w = wallet.trim().to_ascii_lowercase(); + if !w.starts_with("0x") || w.len() != 42 { + return Err(AppError::InvalidRequest( + "wallet address must be 0x-prefixed 20-byte hex".into(), + )); + } + if !w[2..].chars().all(|c| c.is_ascii_hexdigit()) { + return Err(AppError::InvalidRequest( + "wallet address must be hex".into(), + )); + } + Ok(w) +} diff --git a/src/services/mod.rs b/src/services/mod.rs index a8bf03f..d04560d 100644 --- a/src/services/mod.rs +++ b/src/services/mod.rs @@ -1,6 +1,8 @@ //! Business logic (SIWE verification, token minting, quotas). +pub mod heuristics; pub mod nonce; +pub mod quota; pub mod refresh; pub mod revocation; pub mod siwe; diff --git a/src/services/quota.rs b/src/services/quota.rs new file mode 100644 index 0000000..f325066 --- /dev/null +++ b/src/services/quota.rs @@ -0,0 +1,191 @@ +//! Per-wallet quota counters for paid scopes. + +use std::collections::HashMap; +use std::sync::Arc; + +use time::OffsetDateTime; + +use crate::{ + config::QuotaConfig, + error::AppError, + services::heuristics::WalletProfile, + store::quota::{current_window_id, QuotaStore}, +}; + +pub struct QuotaService { + store: Arc, + config: QuotaConfig, +} + +#[derive(Debug)] +pub struct QuotaStatus { + pub scope: String, + pub remaining: u64, + pub limit: u64, +} + +impl QuotaService { + pub fn new(store: Arc, config: QuotaConfig) -> Self { + Self { store, config } + } + + pub fn enabled(&self) -> bool { + self.config.enabled + } + + /// Initialize quota buckets for paid scopes using heuristic tier limits. + pub async fn init_for_wallet( + &self, + wallet: &str, + profile: WalletProfile, + ) -> Result<(), AppError> { + if !self.enabled() { + return Ok(()); + } + let wallet = normalize_wallet(wallet)?; + let limits = limits_for_profile(&self.config, profile); + let window_id = self.window_id(); + let ttl = self.config.window_secs as u64; + for scope in &self.config.paid_scopes { + let limit = limits.get(scope).copied().unwrap_or(0); + if limit == 0 { + continue; + } + self.store + .init_scope(&wallet, scope, window_id, limit, ttl) + .await?; + } + Ok(()) + } + + pub async fn consume( + &self, + wallet: &str, + scope: &str, + amount: u64, + ) -> Result { + if !self.enabled() { + return Err(AppError::InvalidRequest( + "quota enforcement is disabled".into(), + )); + } + if !self.config.paid_scopes.iter().any(|s| s == scope) { + return Err(AppError::InvalidRequest(format!( + "scope {scope} is not quota-gated" + ))); + } + let wallet = normalize_wallet(wallet)?; + let window_id = self.window_id(); + let bucket = self + .store + .try_consume(&wallet, scope, window_id, amount.max(1)) + .await?; + Ok(QuotaStatus { + scope: scope.to_string(), + remaining: bucket.limit.saturating_sub(bucket.used), + limit: bucket.limit, + }) + } + + pub async fn status_for_scopes( + &self, + wallet: &str, + scopes: &[&str], + ) -> Result, AppError> { + if !self.enabled() { + return Ok(Vec::new()); + } + let wallet = normalize_wallet(wallet)?; + let window_id = self.window_id(); + let mut out = Vec::new(); + for scope in scopes { + if !self.config.paid_scopes.iter().any(|s| s == scope) { + continue; + } + if let Some(bucket) = self.store.get_bucket(&wallet, scope, window_id).await? { + out.push(QuotaStatus { + scope: scope.to_string(), + remaining: bucket.limit.saturating_sub(bucket.used), + limit: bucket.limit, + }); + } + } + Ok(out) + } + + fn window_id(&self) -> u64 { + let now = OffsetDateTime::now_utc().unix_timestamp(); + current_window_id(now, self.config.window_secs as u64) + } +} + +fn limits_for_profile(config: &QuotaConfig, profile: WalletProfile) -> &HashMap { + match profile { + WalletProfile::New => &config.limits.new_wallet, + WalletProfile::Established => &config.limits.established, + } +} + +fn normalize_wallet(wallet: &str) -> Result { + let w = wallet.trim().to_ascii_lowercase(); + if !w.starts_with("0x") || w.len() != 42 { + return Err(AppError::InvalidRequest( + "wallet address must be 0x-prefixed 20-byte hex".into(), + )); + } + if !w[2..].chars().all(|c| c.is_ascii_hexdigit()) { + return Err(AppError::InvalidRequest( + "wallet address must be hex".into(), + )); + } + Ok(w) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::QuotaLimitTiers; + use crate::store::quota::InMemoryQuotaStore; + + fn test_config() -> QuotaConfig { + let mut new_wallet = HashMap::new(); + new_wallet.insert("ai:invoke".into(), 2); + let mut established = HashMap::new(); + established.insert("ai:invoke".into(), 100); + QuotaConfig { + enabled: true, + window_secs: 3600, + paid_scopes: vec!["ai:invoke".into()], + limits: QuotaLimitTiers { + new_wallet, + established, + }, + } + } + + #[tokio::test] + async fn new_wallet_gets_lower_limit() { + let store = Arc::new(InMemoryQuotaStore::default()); + let svc = QuotaService::new(store, test_config()); + svc.init_for_wallet( + "0xf39fd6e51aad88f6f4ce6ab8827279cfffb92266", + WalletProfile::New, + ) + .await + .unwrap(); + let s1 = svc + .consume("0xf39fd6e51aad88f6f4ce6ab8827279cfffb92266", "ai:invoke", 1) + .await + .unwrap(); + assert_eq!(s1.remaining, 1); + let _ = svc + .consume("0xf39fd6e51aad88f6f4ce6ab8827279cfffb92266", "ai:invoke", 1) + .await + .unwrap(); + let err = svc + .consume("0xf39fd6e51aad88f6f4ce6ab8827279cfffb92266", "ai:invoke", 1) + .await + .unwrap_err(); + assert!(matches!(err, AppError::QuotaExceeded(_))); + } +} diff --git a/src/state.rs b/src/state.rs index e016f97..37beb4d 100644 --- a/src/state.rs +++ b/src/state.rs @@ -8,8 +8,8 @@ use crate::{ error::AppError, middleware::rate_limit::RateLimiter, services::{ - nonce::NonceService, refresh::RefreshService, revocation::RevocationService, - token::TokenService, + heuristics::HeuristicsService, nonce::NonceService, quota::QuotaService, + refresh::RefreshService, revocation::RevocationService, token::TokenService, }, }; @@ -22,6 +22,8 @@ pub struct AppState { pub token: Arc, pub revocation: Arc, pub refresh: Arc, + pub heuristics: Arc, + pub quota: Arc, } impl AppState { @@ -34,6 +36,10 @@ impl AppState { let revocation = bootstrap::build_revocation_service(revocation_store); let refresh_store = bootstrap::build_refresh_store(&config).await?; let refresh = bootstrap::build_refresh_service(&config, refresh_store); + let heuristics_store = bootstrap::build_heuristics_store(&config).await?; + let heuristics = bootstrap::build_heuristics_service(heuristics_store); + let quota_store = bootstrap::build_quota_store(&config).await?; + let quota = bootstrap::build_quota_service(&config, quota_store); Ok(Self { config: Arc::new(config), nonce, @@ -41,6 +47,8 @@ impl AppState { token, revocation, refresh, + heuristics, + quota, }) } } diff --git a/src/store/heuristics.rs b/src/store/heuristics.rs new file mode 100644 index 0000000..0f64f34 --- /dev/null +++ b/src/store/heuristics.rs @@ -0,0 +1,35 @@ +//! Wallet heuristic persistence (phase-0: first-seen / established). + +mod memory; +mod redis; + +pub use memory::InMemoryHeuristicsStore; +pub use redis::RedisHeuristicsStore; + +use async_trait::async_trait; +use thiserror::Error; + +use crate::error::AppError; + +#[derive(Debug, Error)] +pub enum HeuristicsStoreError { + #[error("store unavailable")] + Unavailable(#[from] ::redis::RedisError), + #[error("{0}")] + Other(String), +} + +impl From for AppError { + fn from(err: HeuristicsStoreError) -> Self { + Self::Internal(err.to_string()) + } +} + +#[async_trait] +pub trait HeuristicsStore: Send + Sync { + /// Whether the wallet has completed at least one prior session. + async fn is_established(&self, wallet: &str) -> Result; + + /// Record that the wallet completed a session (called after successful mint). + async fn mark_established(&self, wallet: &str) -> Result<(), HeuristicsStoreError>; +} diff --git a/src/store/heuristics/memory.rs b/src/store/heuristics/memory.rs new file mode 100644 index 0000000..5279e14 --- /dev/null +++ b/src/store/heuristics/memory.rs @@ -0,0 +1,32 @@ +//! In-memory wallet heuristic store. + +use std::collections::HashSet; +use std::sync::Mutex; + +use async_trait::async_trait; + +use super::{HeuristicsStore, HeuristicsStoreError}; + +#[derive(Default)] +pub struct InMemoryHeuristicsStore { + established: Mutex>, +} + +#[async_trait] +impl HeuristicsStore for InMemoryHeuristicsStore { + async fn is_established(&self, wallet: &str) -> Result { + let guard = self + .established + .lock() + .map_err(|_| HeuristicsStoreError::Other("lock poisoned".into()))?; + Ok(guard.contains(wallet)) + } + + async fn mark_established(&self, wallet: &str) -> Result<(), HeuristicsStoreError> { + self.established + .lock() + .map_err(|_| HeuristicsStoreError::Other("lock poisoned".into()))? + .insert(wallet.to_string()); + Ok(()) + } +} diff --git a/src/store/heuristics/redis.rs b/src/store/heuristics/redis.rs new file mode 100644 index 0000000..09d48b9 --- /dev/null +++ b/src/store/heuristics/redis.rs @@ -0,0 +1,42 @@ +//! Redis-backed wallet heuristic store. + +use async_trait::async_trait; +use redis::AsyncCommands; + +use super::{HeuristicsStore, HeuristicsStoreError}; + +const SEEN_PREFIX: &str = "heuristic:seen:"; + +#[derive(Clone)] +pub struct RedisHeuristicsStore { + client: redis::aio::ConnectionManager, +} + +impl RedisHeuristicsStore { + pub async fn connect(url: &str) -> Result { + let client = redis::Client::open(url)?; + let conn = client.get_connection_manager().await?; + Ok(Self { client: conn }) + } + + fn key(wallet: &str) -> String { + format!("{SEEN_PREFIX}{wallet}") + } +} + +#[async_trait] +impl HeuristicsStore for RedisHeuristicsStore { + async fn is_established(&self, wallet: &str) -> Result { + let key = Self::key(wallet); + let mut conn = self.client.clone(); + let exists: bool = conn.exists(key).await?; + Ok(exists) + } + + async fn mark_established(&self, wallet: &str) -> Result<(), HeuristicsStoreError> { + let key = Self::key(wallet); + let mut conn = self.client.clone(); + conn.set::<_, _, ()>(key, "1").await?; + Ok(()) + } +} diff --git a/src/store/mod.rs b/src/store/mod.rs index fad0b00..1150e9f 100644 --- a/src/store/mod.rs +++ b/src/store/mod.rs @@ -1,5 +1,7 @@ //! Redis and in-memory store implementations. +pub mod heuristics; pub mod nonce; +pub mod quota; pub mod refresh; pub mod revocation; diff --git a/src/store/quota.rs b/src/store/quota.rs new file mode 100644 index 0000000..d44bd18 --- /dev/null +++ b/src/store/quota.rs @@ -0,0 +1,111 @@ +//! Per-wallet quota counters for paid scopes. + +mod memory; +mod redis; + +pub use memory::InMemoryQuotaStore; +pub use redis::RedisQuotaStore; + +use async_trait::async_trait; +use thiserror::Error; + +use crate::error::AppError; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QuotaBucket { + pub used: u64, + pub limit: u64, +} + +#[derive(Debug, Error)] +pub enum QuotaStoreError { + #[error("store unavailable")] + Unavailable(#[from] ::redis::RedisError), + #[error("quota bucket missing")] + NotFound, + #[error("quota exceeded")] + Exceeded, + #[error("{0}")] + Other(String), +} + +impl From for AppError { + fn from(err: QuotaStoreError) -> Self { + match err { + QuotaStoreError::NotFound => { + Self::InvalidRequest("quota bucket not initialized for scope".into()) + } + QuotaStoreError::Exceeded => Self::QuotaExceeded("quota exhausted for scope".into()), + other => Self::Internal(other.to_string()), + } + } +} + +#[async_trait] +pub trait QuotaStore: Send + Sync { + /// Create buckets for the current window when absent (`limit` per scope). + async fn init_scope( + &self, + wallet: &str, + scope: &str, + window_id: u64, + limit: u64, + ttl_secs: u64, + ) -> Result<(), QuotaStoreError>; + + async fn get_bucket( + &self, + wallet: &str, + scope: &str, + window_id: u64, + ) -> Result, QuotaStoreError>; + + /// Atomically consume `amount` units; returns updated bucket or `NotFound`. + async fn try_consume( + &self, + wallet: &str, + scope: &str, + window_id: u64, + amount: u64, + ) -> Result; +} + +pub fn encode_bucket(bucket: &QuotaBucket) -> String { + format!("{},{}", bucket.used, bucket.limit) +} + +pub fn decode_bucket(raw: &str) -> Result { + let (used, limit) = raw + .split_once(',') + .ok_or_else(|| QuotaStoreError::Other("invalid quota bucket encoding".into()))?; + let used = used + .parse() + .map_err(|_| QuotaStoreError::Other("invalid used count".into()))?; + let limit = limit + .parse() + .map_err(|_| QuotaStoreError::Other("invalid limit".into()))?; + Ok(QuotaBucket { used, limit }) +} + +pub fn current_window_id(now_unix: i64, window_secs: u64) -> u64 { + (now_unix.max(0) as u64) / window_secs.max(1) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn bucket_roundtrip() { + let raw = encode_bucket(&QuotaBucket { used: 3, limit: 10 }); + let bucket = decode_bucket(&raw).unwrap(); + assert_eq!(bucket.used, 3); + assert_eq!(bucket.limit, 10); + } + + #[test] + fn window_id_buckets_time() { + assert_eq!(current_window_id(0, 3600), 0); + assert_eq!(current_window_id(3600, 3600), 1); + } +} diff --git a/src/store/quota/memory.rs b/src/store/quota/memory.rs new file mode 100644 index 0000000..e07d39f --- /dev/null +++ b/src/store/quota/memory.rs @@ -0,0 +1,73 @@ +//! In-memory quota store. + +use std::collections::HashMap; +use std::sync::Mutex; + +use async_trait::async_trait; + +use super::{QuotaBucket, QuotaStore, QuotaStoreError}; + +fn bucket_key(wallet: &str, scope: &str, window_id: u64) -> String { + format!("{wallet}:{scope}:{window_id}") +} + +#[derive(Default)] +pub struct InMemoryQuotaStore { + buckets: Mutex>, +} + +#[async_trait] +impl QuotaStore for InMemoryQuotaStore { + async fn init_scope( + &self, + wallet: &str, + scope: &str, + window_id: u64, + limit: u64, + _ttl_secs: u64, + ) -> Result<(), QuotaStoreError> { + let mut guard = self + .buckets + .lock() + .map_err(|_| QuotaStoreError::Other("lock poisoned".into()))?; + let key = bucket_key(wallet, scope, window_id); + guard + .entry(key) + .and_modify(|bucket| bucket.limit = limit) + .or_insert(QuotaBucket { used: 0, limit }); + Ok(()) + } + + async fn get_bucket( + &self, + wallet: &str, + scope: &str, + window_id: u64, + ) -> Result, QuotaStoreError> { + let guard = self + .buckets + .lock() + .map_err(|_| QuotaStoreError::Other("lock poisoned".into()))?; + Ok(guard.get(&bucket_key(wallet, scope, window_id)).cloned()) + } + + async fn try_consume( + &self, + wallet: &str, + scope: &str, + window_id: u64, + amount: u64, + ) -> Result { + let mut guard = self + .buckets + .lock() + .map_err(|_| QuotaStoreError::Other("lock poisoned".into()))?; + let key = bucket_key(wallet, scope, window_id); + let bucket = guard.get_mut(&key).ok_or(QuotaStoreError::NotFound)?; + if bucket.used.saturating_add(amount) > bucket.limit { + return Err(QuotaStoreError::Exceeded); + } + bucket.used += amount; + Ok(bucket.clone()) + } +} diff --git a/src/store/quota/redis.rs b/src/store/quota/redis.rs new file mode 100644 index 0000000..9c6ea01 --- /dev/null +++ b/src/store/quota/redis.rs @@ -0,0 +1,112 @@ +//! Redis-backed per-wallet quota counters. + +use async_trait::async_trait; +use redis::AsyncCommands; + +use super::{decode_bucket, encode_bucket, QuotaBucket, QuotaStore, QuotaStoreError}; + +const KEY_PREFIX: &str = "quota:"; + +#[derive(Clone)] +pub struct RedisQuotaStore { + client: redis::aio::ConnectionManager, +} + +impl RedisQuotaStore { + pub async fn connect(url: &str) -> Result { + let client = redis::Client::open(url)?; + let conn = client.get_connection_manager().await?; + Ok(Self { client: conn }) + } + + fn key(wallet: &str, scope: &str, window_id: u64) -> String { + format!("{KEY_PREFIX}{wallet}:{scope}:{window_id}") + } +} + +#[async_trait] +impl QuotaStore for RedisQuotaStore { + async fn init_scope( + &self, + wallet: &str, + scope: &str, + window_id: u64, + limit: u64, + ttl_secs: u64, + ) -> Result<(), QuotaStoreError> { + let key = Self::key(wallet, scope, window_id); + let mut conn = self.client.clone(); + if let Some(raw) = conn.get::<_, Option>(&key).await? { + let mut bucket = decode_bucket(&raw)?; + bucket.limit = limit; + let payload = encode_bucket(&bucket); + let ttl: i64 = conn.ttl(&key).await.unwrap_or(ttl_secs as i64); + if ttl > 0 { + conn.set_ex::<_, _, ()>(&key, payload, ttl as u64).await?; + } else { + conn.set::<_, _, ()>(&key, payload).await?; + } + } else { + let payload = encode_bucket(&QuotaBucket { used: 0, limit }); + conn.set_ex::<_, _, ()>(&key, payload, ttl_secs.max(1)) + .await?; + } + Ok(()) + } + + async fn get_bucket( + &self, + wallet: &str, + scope: &str, + window_id: u64, + ) -> Result, QuotaStoreError> { + let key = Self::key(wallet, scope, window_id); + let mut conn = self.client.clone(); + let raw: Option = conn.get(key).await?; + raw.map(|s| decode_bucket(&s)).transpose() + } + + async fn try_consume( + &self, + wallet: &str, + scope: &str, + window_id: u64, + amount: u64, + ) -> Result { + let key = Self::key(wallet, scope, window_id); + let mut conn = self.client.clone(); + let script = r#" + local raw = redis.call('GET', KEYS[1]) + if not raw then return -1 end + local sep = string.find(raw, ',') + if not sep then return -2 end + local used = tonumber(string.sub(raw, 1, sep - 1)) + local limit = tonumber(string.sub(raw, sep + 1)) + local amount = tonumber(ARGV[1]) + if used + amount > limit then return -3 end + used = used + amount + local updated = tostring(used) .. ',' .. tostring(limit) + local ttl = redis.call('TTL', KEYS[1]) + if ttl > 0 then + redis.call('SET', KEYS[1], updated, 'EX', ttl) + else + redis.call('SET', KEYS[1], updated) + end + return limit - used + "#; + let remaining: i64 = redis::Script::new(script) + .key(key) + .arg(amount) + .invoke_async(&mut conn) + .await?; + match remaining { + -1 => Err(QuotaStoreError::NotFound), + -2 => Err(QuotaStoreError::Other("invalid bucket".into())), + -3 => Err(QuotaStoreError::Exceeded), + _ => self + .get_bucket(wallet, scope, window_id) + .await? + .ok_or(QuotaStoreError::NotFound), + } + } +} diff --git a/tests/common/mod.rs b/tests/common/mod.rs index cfb104b..9f7a165 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -1,8 +1,33 @@ //! Shared helpers for integration tests. +#![allow(dead_code)] // each integration test binary links this module separately + use k256::ecdsa::SigningKey; +use rand::rngs::OsRng; use sha3::{Digest, Keccak256}; +/// Ephemeral wallet + private key for tests that need an isolated Redis identity. +pub struct EphemeralTestAccount { + pub wallet: String, + pub private_key: String, +} + +pub fn ephemeral_test_account() -> EphemeralTestAccount { + let signing_key = SigningKey::random(&mut OsRng); + let wallet = wallet_from_signing_key(&signing_key); + let private_key = format!("0x{}", hex::encode(signing_key.to_bytes())); + EphemeralTestAccount { + wallet, + private_key, + } +} + +fn wallet_from_signing_key(signing_key: &SigningKey) -> String { + let pk = signing_key.verifying_key().to_encoded_point(false); + let hash = Keccak256::new_with_prefix(&pk.as_bytes()[1..]).finalize(); + format!("0x{}", hex::encode(&hash[12..])) +} + /// Well-known Anvil/Hardhat account #0 private key (test only). pub const TEST_PRIVATE_KEY: &str = "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80"; diff --git a/tests/expose.rs b/tests/expose.rs new file mode 100644 index 0000000..49aa893 --- /dev/null +++ b/tests/expose.rs @@ -0,0 +1,64 @@ +//! `server.expose` route surface tests. + +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use tower::ServiceExt; +use trust_relay::{ + config::{Config, ExposeMode}, + routes::create_router_for_expose, + state::AppState, +}; + +async fn router_with_expose(expose: ExposeMode) -> axum::Router { + let mut config = Config::load().expect("config"); + config.server.expose = expose; + config.redis.url.clear(); + let state = AppState::build(config).await.unwrap(); + create_router_for_expose(state, expose) +} + +#[tokio::test] +async fn public_expose_serves_session_but_not_quota_consume() { + let app = router_with_expose(ExposeMode::Public).await; + let nonce = app + .clone() + .oneshot(Request::get("/v1/auth/nonce").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(nonce.status(), StatusCode::OK); + + let missing = app + .oneshot( + Request::post("/v1/auth/quota/consume") + .header("content-type", "application/json") + .body(Body::from(r#"{"scope":"ai:invoke"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(missing.status(), StatusCode::NOT_FOUND); +} + +#[tokio::test] +async fn internal_expose_serves_healthz_but_not_session() { + let app = router_with_expose(ExposeMode::Internal).await; + let health = app + .clone() + .oneshot(Request::get("/healthz").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(health.status(), StatusCode::OK); + + let session = app + .oneshot( + Request::post("/v1/auth/session") + .header("content-type", "application/json") + .body(Body::from("{}")) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(session.status(), StatusCode::NOT_FOUND); +} diff --git a/tests/quota.rs b/tests/quota.rs new file mode 100644 index 0000000..29b8f75 --- /dev/null +++ b/tests/quota.rs @@ -0,0 +1,209 @@ +//! Quota and wallet-heuristic integration tests. + +mod common; + +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use base64::Engine; +use common::{ + build_siwe_message, ephemeral_test_account, sign_siwe_message_with_key, EphemeralTestAccount, +}; +use ed25519_dalek::SigningKey; +use http_body_util::BodyExt; +use rand::rngs::OsRng; +use serde_json::json; +use tower::ServiceExt; +use trust_relay::{config::Config, routes::create_router, state::AppState}; + +fn stable_test_config() -> Config { + let mut config = Config::load().expect("config"); + let key = SigningKey::generate(&mut OsRng); + config.signing.key_seed_b64 = base64::engine::general_purpose::STANDARD.encode(key.to_bytes()); + config.quota.limits.new_wallet.insert("ai:invoke".into(), 2); + config + .quota + .limits + .new_wallet + .insert("mint:request".into(), 0); + config.quota.paid_scopes = vec!["ai:invoke".into()]; + config +} + +async fn session_with_token(app: &axum::Router, account: &EphemeralTestAccount) -> String { + let nonce_resp = app + .clone() + .oneshot( + Request::builder() + .uri("/v1/auth/nonce") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = nonce_resp.into_body().collect().await.unwrap().to_bytes(); + let nonce = serde_json::from_slice::(&body).unwrap()["nonce"] + .as_str() + .unwrap() + .to_string(); + let message = build_siwe_message( + "localhost", + "http://localhost:3001", + 324, + &account.wallet, + &nonce, + ); + let signature = sign_siwe_message_with_key(&message, &account.private_key); + let session_resp = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/auth/session") + .header("content-type", "application/json") + .body(Body::from( + json!({ "message": message, "signature": signature }).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(session_resp.status(), StatusCode::OK); + let body = session_resp.into_body().collect().await.unwrap().to_bytes(); + let access = serde_json::from_slice::(&body).unwrap()["accessToken"] + .as_str() + .unwrap() + .to_string(); + access +} + +async fn router_for(config: Config) -> axum::Router { + create_router(AppState::build(config).await.unwrap()) +} + +#[tokio::test] +async fn new_wallet_gets_lower_quota_than_second_session() { + let account = ephemeral_test_account(); + let app = router_for(stable_test_config()).await; + let token = session_with_token(&app, &account).await; + + let status = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/auth/quota") + .header("authorization", format!("Bearer {token}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(status.status(), StatusCode::OK); + let body = status.into_body().collect().await.unwrap().to_bytes(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let ai = json["scopes"] + .as_array() + .unwrap() + .iter() + .find(|v| v["scope"] == "ai:invoke") + .unwrap(); + assert_eq!(ai["limit"], 2); + assert_eq!(ai["remaining"], 2); + + let consume = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/auth/quota/consume") + .header("authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .body(Body::from(json!({ "scope": "ai:invoke" }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(consume.status(), StatusCode::OK); + let body = consume.into_body().collect().await.unwrap().to_bytes(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(json["remaining"], 1); +} + +#[tokio::test] +async fn quota_consume_returns_429_when_exhausted() { + let account = ephemeral_test_account(); + let app = router_for(stable_test_config()).await; + let token = session_with_token(&app, &account).await; + + for _ in 0..2 { + let resp = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/auth/quota/consume") + .header("authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .body(Body::from(json!({ "scope": "ai:invoke" }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + } + + let exhausted = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/auth/quota/consume") + .header("authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .body(Body::from(json!({ "scope": "ai:invoke" }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(exhausted.status(), StatusCode::TOO_MANY_REQUESTS); + let bytes = exhausted.into_body().collect().await.unwrap().to_bytes(); + let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(json["error"], "quota_exceeded"); +} + +#[tokio::test] +async fn established_wallet_gets_higher_limit_on_second_session() { + let mut config = stable_test_config(); + config + .quota + .limits + .established + .insert("ai:invoke".into(), 50); + let account = ephemeral_test_account(); + let app = router_for(config).await; + let _ = session_with_token(&app, &account).await; + let token2 = session_with_token(&app, &account).await; + + let status = app + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/auth/quota") + .header("authorization", format!("Bearer {token2}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(status.status(), StatusCode::OK); + let body = status.into_body().collect().await.unwrap().to_bytes(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let ai = json["scopes"] + .as_array() + .unwrap() + .iter() + .find(|v| v["scope"] == "ai:invoke") + .unwrap(); + assert_eq!(ai["limit"], 50); +} diff --git a/tests/quota_redis.rs b/tests/quota_redis.rs new file mode 100644 index 0000000..d04b1b8 --- /dev/null +++ b/tests/quota_redis.rs @@ -0,0 +1,39 @@ +//! Redis quota store integration tests (`APP_REDIS__URL` set in CI). + +use std::time::Duration; + +use trust_relay::store::quota::{current_window_id, QuotaStore, RedisQuotaStore}; + +fn redis_url() -> String { + std::env::var("APP_REDIS__URL").unwrap_or_else(|_| "redis://127.0.0.1:6379".into()) +} + +async fn try_connect_redis() -> Option { + if std::env::var("APP_REDIS__URL").is_err() { + return None; + } + let url = redis_url(); + match tokio::time::timeout(Duration::from_millis(800), RedisQuotaStore::connect(&url)).await { + Ok(Ok(store)) => Some(store), + _ => None, + } +} + +#[tokio::test] +async fn redis_quota_init_and_consume() { + let Some(store) = try_connect_redis().await else { + return; + }; + let wallet = format!("0x{:040x}", uuid::Uuid::new_v4().as_u128()); + let window = current_window_id(time::OffsetDateTime::now_utc().unix_timestamp(), 3600); + store + .init_scope(&wallet, "ai:invoke", window, 3, 60) + .await + .unwrap(); + let bucket = store + .try_consume(&wallet, "ai:invoke", window, 1) + .await + .unwrap(); + assert_eq!(bucket.limit.saturating_sub(bucket.used), 2); + assert_eq!(bucket.used, 1); +}