diff --git a/.agent/PLANS.md b/.agent/PLANS.md new file mode 100644 index 0000000..15d9583 --- /dev/null +++ b/.agent/PLANS.md @@ -0,0 +1,150 @@ +# Codex Execution Plans (ExecPlans): + +This document describes the requirements for an execution plan ("ExecPlan"), a design document that a coding agent can follow to deliver a working feature or system change. Treat the reader as a complete beginner to this repository: they have only the current working tree and the single ExecPlan file you provide. There is no memory of prior plans and no external context. + +## How to use ExecPlans and PLANS.md + +When authoring an executable specification (ExecPlan), follow PLANS.md _to the letter_. If it is not in your context, refresh your memory by reading the entire PLANS.md file. Be thorough in reading (and re-reading) source material to produce an accurate specification. When creating a spec, start from the skeleton and flesh it out as you do your research. + +When implementing an executable specification (ExecPlan), do not prompt the user for "next steps"; simply proceed to the next milestone. Keep all sections up to date, add or split entries in the list at every stopping point to affirmatively state the progress made and next steps. Resolve ambiguities autonomously, and commit frequently. + +When discussing an executable specification (ExecPlan), record decisions in a log in the spec for posterity; it should be unambiguously clear why any change to the specification was made. ExecPlans are living documents, and it should always be possible to restart from _only_ the ExecPlan and no other work. + +When researching a design with challenging requirements or significant unknowns, use milestones to implement proof of concepts, "toy implementations", etc., that allow validating whether the user's proposal is feasible. Read the source code of libraries by finding or acquiring them, research deeply, and include prototypes to guide a fuller implementation. + +## Requirements + +NON-NEGOTIABLE REQUIREMENTS: + +* Every ExecPlan must be fully self-contained. Self-contained means that in its current form it contains all knowledge and instructions needed for a novice to succeed. +* Every ExecPlan is a living document. Contributors are required to revise it as progress is made, as discoveries occur, and as design decisions are finalized. Each revision must remain fully self-contained. +* Every ExecPlan must enable a complete novice to implement the feature end-to-end without prior knowledge of this repo. +* Every ExecPlan must produce a demonstrably working behavior, not merely code changes to "meet a definition". +* Every ExecPlan must define every term of art in plain language or do not use it. + +Purpose and intent come first. Begin by explaining, in a few sentences, why the work matters from a user's perspective: what someone can do after this change that they could not do before, and how to see it working. Then guide the reader through the exact steps to achieve that outcome, including what to edit, what to run, and what they should observe. + +The agent executing your plan can list files, read files, search, run the project, and run tests. It does not know any prior context and cannot infer what you meant from earlier milestones. Repeat any assumption you rely on. Do not point to external blogs or docs; if knowledge is required, embed it in the plan itself in your own words. If an ExecPlan builds upon a prior ExecPlan and that file is checked in, incorporate it by reference. If it is not, you must include all relevant context from that plan. + +## Formatting + +Format and envelope are simple and strict. Each ExecPlan must be one single fenced code block labeled as `md` that begins and ends with triple backticks. Do not nest additional triple-backtick code fences inside; when you need to show commands, transcripts, diffs, or code, present them as indented blocks within that single fence. Use indentation for clarity rather than code fences inside an ExecPlan to avoid prematurely closing the ExecPlan's code fence. Use two newlines after every heading, use # and ## and so on, and correct syntax for ordered and unordered lists. + +When writing an ExecPlan to a Markdown (.md) file where the content of the file *is only* the single ExecPlan, you should omit the triple backticks. + +Write in plain prose. Prefer sentences over lists. Avoid checklists, tables, and long enumerations unless brevity would obscure meaning. Checklists are permitted only in the `Progress` section, where they are mandatory. Narrative sections must remain prose-first. + +## Guidelines + +Self-containment and plain language are paramount. If you introduce a phrase that is not ordinary English ("daemon", "middleware", "RPC gateway", "filter graph"), define it immediately and remind the reader how it manifests in this repository (for example, by naming the files or commands where it appears). Do not say "as defined previously" or "according to the architecture doc." Include the needed explanation here, even if you repeat yourself. + +Avoid common failure modes. Do not rely on undefined jargon. Do not describe "the letter of a feature" so narrowly that the resulting code compiles but does nothing meaningful. Do not outsource key decisions to the reader. When ambiguity exists, resolve it in the plan itself and explain why you chose that path. Err on the side of over-explaining user-visible effects and under-specifying incidental implementation details. + +Anchor the plan with observable outcomes. State what the user can do after implementation, the commands to run, and the outputs they should see. Acceptance should be phrased as behavior a human can verify ("after starting the server, navigating to [http://localhost:8080/health](http://localhost:8080/health) returns HTTP 200 with body OK") rather than internal attributes ("added a HealthCheck struct"). If a change is internal, explain how its impact can still be demonstrated (for example, by running tests that fail before and pass after, and by showing a scenario that uses the new behavior). + +Specify repository context explicitly. Name files with full repository-relative paths, name functions and modules precisely, and describe where new files should be created. If touching multiple areas, include a short orientation paragraph that explains how those parts fit together so a novice can navigate confidently. When running commands, show the working directory and exact command line. When outcomes depend on environment, state the assumptions and provide alternatives when reasonable. + +Be idempotent and safe. Write the steps so they can be run multiple times without causing damage or drift. If a step can fail halfway, include how to retry or adapt. If a migration or destructive operation is necessary, spell out backups or safe fallbacks. Prefer additive, testable changes that can be validated as you go. + +Validation is not optional. Include instructions to run tests, to start the system if applicable, and to observe it doing something useful. Describe comprehensive testing for any new features or capabilities. Include expected outputs and error messages so a novice can tell success from failure. Where possible, show how to prove that the change is effective beyond compilation (for example, through a small end-to-end scenario, a CLI invocation, or an HTTP request/response transcript). State the exact test commands appropriate to the project’s toolchain and how to interpret their results. + +Capture evidence. When your steps produce terminal output, short diffs, or logs, include them inside the single fenced block as indented examples. Keep them concise and focused on what proves success. If you need to include a patch, prefer file-scoped diffs or small excerpts that a reader can recreate by following your instructions rather than pasting large blobs. + +## Milestones + +Milestones are narrative, not bureaucracy. If you break the work into milestones, introduce each with a brief paragraph that describes the scope, what will exist at the end of the milestone that did not exist before, the commands to run, and the acceptance you expect to observe. Keep it readable as a story: goal, work, result, proof. Progress and milestones are distinct: milestones tell the story, progress tracks granular work. Both must exist. Never abbreviate a milestone merely for the sake of brevity, do not leave out details that could be crucial to a future implementation. + +Each milestone must be independently verifiable and incrementally implement the overall goal of the execution plan. + +## Living plans and design decisions + +* ExecPlans are living documents. As you make key design decisions, update the plan to record both the decision and the thinking behind it. Record all decisions in the `Decision Log` section. +* ExecPlans must contain and maintain a `Progress` section, a `Surprises & Discoveries` section, a `Decision Log`, and an `Outcomes & Retrospective` section. These are not optional. +* When you discover optimizer behavior, performance tradeoffs, unexpected bugs, or inverse/unapply semantics that shaped your approach, capture those observations in the `Surprises & Discoveries` section with short evidence snippets (test output is ideal). +* If you change course mid-implementation, document why in the `Decision Log` and reflect the implications in `Progress`. Plans are guides for the next contributor as much as checklists for you. +* At completion of a major task or the full plan, write an `Outcomes & Retrospective` entry summarizing what was achieved, what remains, and lessons learned. + +# Prototyping milestones and parallel implementations + +It is acceptable—-and often encouraged—-to include explicit prototyping milestones when they de-risk a larger change. Examples: adding a low-level operator to a dependency to validate feasibility, or exploring two composition orders while measuring optimizer effects. Keep prototypes additive and testable. Clearly label the scope as “prototyping”; describe how to run and observe results; and state the criteria for promoting or discarding the prototype. + +Prefer additive code changes followed by subtractions that keep tests passing. Parallel implementations (e.g., keeping an adapter alongside an older path during migration) are fine when they reduce risk or enable tests to continue passing during a large migration. Describe how to validate both paths and how to retire one safely with tests. When working with multiple new libraries or feature areas, consider creating spikes that evaluate the feasibility of these features _independently_ of one another, proving that the external library performs as expected and implements the features we need in isolation. + +## Skeleton of a Good ExecPlan + + # + + This ExecPlan is a living document. The sections `Progress`, `Surprises & Discoveries`, `Decision Log`, and `Outcomes & Retrospective` must be kept up to date as work proceeds. + + If PLANS.md file is checked into the repo, reference the path to that file here from the repository root and note that this document must be maintained in accordance with PLANS.md. + + ## Purpose / Big Picture + + Explain in a few sentences what someone gains after this change and how they can see it working. State the user-visible behavior you will enable. + + ## Progress + + Use a list with checkboxes to summarize granular steps. Every stopping point must be documented here, even if it requires splitting a partially completed task into two (“done” vs. “remaining”). This section must always reflect the actual current state of the work. + + - [x] (2025-10-01 13:00Z) Example completed step. + - [ ] Example incomplete step. + - [ ] Example partially completed step (completed: X; remaining: Y). + + Use timestamps to measure rates of progress. + + ## Surprises & Discoveries + + Document unexpected behaviors, bugs, optimizations, or insights discovered during implementation. Provide concise evidence. + + - Observation: … + Evidence: … + + ## Decision Log + + Record every decision made while working on the plan in the format: + + - Decision: … + Rationale: … + Date/Author: … + + ## Outcomes & Retrospective + + Summarize outcomes, gaps, and lessons learned at major milestones or at completion. Compare the result against the original purpose. + + ## Context and Orientation + + Describe the current state relevant to this task as if the reader knows nothing. Name the key files and modules by full path. Define any non-obvious term you will use. Do not refer to prior plans. + + ## Plan of Work + + Describe, in prose, the sequence of edits and additions. For each edit, name the file and location (function, module) and what to insert or change. Keep it concrete and minimal. + + ## Concrete Steps + + State the exact commands to run and where to run them (working directory). When a command generates output, show a short expected transcript so the reader can compare. This section must be updated as work proceeds. + + ## Validation and Acceptance + + Describe how to start or exercise the system and what to observe. Phrase acceptance as behavior, with specific inputs and outputs. If tests are involved, say "run and expect passed; the new test fails before the change and passes after>". + + ## Idempotence and Recovery + + If steps can be repeated safely, say so. If a step is risky, provide a safe retry or rollback path. Keep the environment clean after completion. + + ## Artifacts and Notes + + Include the most important transcripts, diffs, or snippets as indented examples. Keep them concise and focused on what proves success. + + ## Interfaces and Dependencies + + Be prescriptive. Name the libraries, modules, and services to use and why. Specify the types, traits/interfaces, and function signatures that must exist at the end of the milestone. Prefer stable names and paths such as `crate::module::function` or `package.submodule.Interface`. E.g.: + + In crates/foo/planner.rs, define: + + pub trait Planner { + fn plan(&self, observed: &Observed) -> Vec; + } + +If you follow the guidance above, a single, stateless agent -- or a human novice -- can read your ExecPlan from top to bottom and produce a working, observable result. That is the bar: SELF-CONTAINED, SELF-SUFFICIENT, NOVICE-GUIDING, OUTCOME-FOCUSED. + +When you revise a plan, you must ensure your changes are comprehensively reflected across all sections, including the living document sections, and you must write a note at the bottom of the plan describing the change and the reason why. ExecPlans must describe not just the what but the why for almost everything. \ No newline at end of file diff --git a/.agent/execplans/per-identity-rate-limit-poc.md b/.agent/execplans/per-identity-rate-limit-poc.md new file mode 100644 index 0000000..39901db --- /dev/null +++ b/.agent/execplans/per-identity-rate-limit-poc.md @@ -0,0 +1,370 @@ +# Add Per-Identity Proxied-Data Rate Limits + +This ExecPlan is a living document. The sections `Progress`, `Surprises & Discoveries`, `Decision Log`, and `Outcomes & Retrospective` must be kept up to date as work proceeds. + +This plan follows `.agent/PLANS.md` from the repository root. The implementation must not begin until the user has reviewed and approved this plan. + + +## Purpose / Big Picture + +The gateway currently authorizes whether an identity may open a CONNECT tunnel to a destination, but once a tunnel is established there is no per-identity data-volume throttling. After this proof of concept, an operator can store an optional bytes-per-second limit for a source identity in the authorization registry. When that identity opens one or more tunnels through a single gateway process, bytes flowing in both directions across all of that identity's active tunnels share one in-memory token bucket and are slowed to that configured rate. + +The visible behavior is that an identity with no rate-limit row behaves exactly as it does today, while an identity with a rate limit can still connect to approved destinations but large transfers take measurably longer. This is intentionally process-local: it assumes only one running gateway process and does not coordinate buckets across multiple gateway instances. + + +## Progress + +- [x] (2026-05-17T04:42Z) Read `.agent/PLANS.md`, `src/proxy.rs`, `src/policy.rs`, `src/registry.rs`, `src/main.rs`, `src/config.rs`, `migrations/0001_signed_authorization_registry.sql`, and the integration/e2e test helpers to ground the design in the current repository. +- [x] (2026-05-17T04:42Z) Created this ExecPlan for user review. +- [x] (2026-05-17T04:47Z) Updated the plan to reflect the user decision to introduce a `SubjectIdentity` struct carrying identity metadata instead of threading a raw identity string and separate settings. +- [x] (2026-05-17T04:51Z) Updated the plan so `RegistryStore` returns a complete `SubjectIdentity`, with absent identity metadata represented as an unlimited identity rather than as `None`. +- [x] (2026-05-17T04:54Z) Updated the limiter design to use per-bucket `std::sync::Mutex` locking, while keeping a short-lived map lock only for bucket lookup and creation. +- [x] (2026-05-17T05:54Z) Updated the limiter design to cache a `RateLimitHandle` per tunnel so forwarded chunks do not perform a map lookup. +- [x] (2026-05-17T05:26Z) Implement the registry schema and SQLx metadata changes. +- [x] (2026-05-17T05:31Z) Add `SubjectIdentity` and thread it through policy evaluation. +- [x] (2026-05-17T05:31Z) Add a process-local shared token bucket limiter. +- [x] (2026-05-17T05:31Z) Replace the tunnel copy path with a rate-aware bidirectional copy loop. +- [x] (2026-05-17T05:31Z) Add focused unit, integration, and e2e tests. +- [x] (2026-05-17T05:39Z) Run formatting, SQLx metadata checks, clippy, and the full test suite. +- [x] (2026-05-17T05:39Z) Update README/configuration documentation and record final outcomes. + + +## Surprises & Discoveries + +- Observation: `src/proxy.rs` currently uses `tokio::io::copy_bidirectional`, which reports final byte counts but does not expose a hook before each chunk is forwarded. + Evidence: `spawn_tunnel` calls `copy_bidirectional(&mut downstream, &mut upstream).await` and logs the returned `(up, down)` byte counts. + +- Observation: the policy engine returns only `PolicyDecision::Allow { source_identity }`, so policy evaluation needs a new structured allow result that carries both the identity string and the identity metadata loaded from the registry. + Evidence: `src/policy.rs` defines `PolicyDecision::Allow { source_identity: String }` and the proxy stores that identity before opening the upstream TCP connection. + +- Observation: schema version checking is strict, so adding a migration requires bumping the expected registry schema version. + Evidence: `src/registry.rs` has `const EXPECTED_SCHEMA_VERSION: i32 = 1` and `verify_schema_version` rejects any other latest version. + +- Observation: This checkout has committed `.sqlx` metadata, but `TEST_DATABASE_URL` is not set and `cargo sqlx` is not installed in the current environment. + Evidence: `cargo test --locked --no-run` failed on the new SQLx macro with `SQLX_OFFLINE=true` and no cached data; `cargo sqlx --version` returned `error: no such command: sqlx`; `printenv TEST_DATABASE_URL` returned no value. + + +## Decision Log + +- Decision: Store limits in a new `identity_metadata` registry table keyed by `subject_identity`, with an optional positive `rate_limit_bytes_per_second` column. + Rationale: The task asks for a per-identity limit, not a per-permission or per-destination limit. A separate table keeps metadata independent from permission grants and allows an identity to have no row, or a row with no limit, without affecting authorization semantics. + Date/Author: 2026-05-17 / Codex + +- Decision: Create a `SubjectIdentity` struct to represent an identity and its associated settings in the `identity_metadata` table, rather than just representing identities as a `String`. + Rationale: Designing for future extensibility, and indicating a clear separation of concerns when we embed identity metadata in other structs/enums. + Date/Author: 2026-05-17 / ajd + +- Decision: Interpret the rate limit as one aggregate limit per identity across all active tunnels in one gateway process, with both client-to-destination and destination-to-client bytes charged to the same bucket. + Rationale: This matches the user's high-level design and is the simplest useful behavior. It avoids destination-level buckets and avoids distributed coordination. + Date/Author: 2026-05-17 / Codex + +- Decision: Fetch the identity's current rate limit during policy evaluation and attach it to the `SubjectIdentity` returned inside `PolicyDecision::Allow`. + Rationale: The policy engine already has the registry pool and already establishes the authenticated source identity. Returning structured identity metadata with the allow decision avoids adding a second database dependency to the proxy layer. + Date/Author: 2026-05-17 / Codex + +- Decision: `RegistryStore` should fetch identity metadata and return a complete `SubjectIdentity`, not just a standalone rate-limit value. + Rationale: `SubjectIdentity` is the representation of an identity plus its registry-backed settings. A missing `identity_metadata` row means an authenticated identity with no configured rate limit, so the registry method should return `Ok(SubjectIdentity { rate_limit_bytes_per_second: None })` rather than `Ok(None)`. + Date/Author: 2026-05-17 / ajd and Codex + +- Decision: Enforce limits in the proxy layer with a shared in-memory token bucket map owned by `MakeProxyService`. + Rationale: The proxy layer sees every proxied byte and already creates per-connection services from a shared factory. Keeping the limiter there avoids mixing byte-forwarding concerns into authorization code. + Date/Author: 2026-05-17 / Codex + +- Decision: Use `std::sync::Mutex` for both the bucket map and each individual bucket, but put token accounting behind per-bucket locks rather than one global lock. + Rationale: The proof of concept does not need Tokio's `sync` feature as long as no lock is held across `.await`. A short map lock for lookup or insertion is simple, while per-bucket locks reduce unnecessary contention between different identities during high-throughput copying. + Date/Author: 2026-05-17 / ajd and Codex + +- Decision: Cache a per-tunnel `RateLimitHandle` after policy authorization and identity configuration, rather than looking up the bucket in the map for every copied chunk. + Rationale: The bucket map is only needed when configuring the identity at CONNECT time. Once a tunnel has an `Arc` to the relevant bucket, each copy loop can use that handle directly and avoid global map lock traffic in the hot path. Unlimited identities use an empty handle that returns immediately. + Date/Author: 2026-05-17 / ajd and Codex + +- Decision: The gateway refreshes the process-local bucket configuration for an identity when a new CONNECT request for that identity is authorized; it does not poll the database or read metadata per chunk. + Rationale: A shared bucket must have one current configuration per identity. Refreshing on CONNECT keeps implementation simple and avoids database work in the forwarding loop. Existing tunnels for the same identity may observe the refreshed process-local limit after a later CONNECT for that identity, which is acceptable for this proof of concept and must be documented. + Date/Author: 2026-05-17 / Codex + + +## Outcomes & Retrospective + +This section is intentionally empty until implementation begins. At completion, record what was implemented, what validation passed, and which proof-of-concept limitations remain. + +2026-05-17 / Codex: Implemented the proof of concept. The registry now has `identity_metadata` schema version 2, policy evaluation returns `SubjectIdentity`, the proxy configures a shared process-local `RateLimiter` and passes a cached `RateLimitHandle` into each tunnel, and tunnel forwarding uses rate-aware copy loops for both directions. README documents the table and limitations. Validation passed with `cargo fmt -- --check`, `SQLX_OFFLINE=false DATABASE_URL=postgres://agent_gateway_admin:agent_gateway_dev@localhost:5432/agent_gateway cargo sqlx prepare --check -- --all-targets --locked`, `rustup run stable cargo-clippy --locked --all-targets`, and `SQLX_OFFLINE=true TEST_DATABASE_URL=postgres://agent_gateway_admin:agent_gateway_dev@localhost:5432/agent_gateway cargo test --locked`. + + +## Context and Orientation + +This repository is a Rust 2024 project named `agent_gateway`. The entry point in `src/main.rs` starts an mTLS HTTP/2 CONNECT proxy. A CONNECT proxy accepts an HTTP CONNECT request, opens a TCP connection to the requested destination, and then forwards raw bytes between the client and that destination. + +The current startup path is: + + src/main.rs loads config, builds TLS config, builds a policy engine with policy::build_engine, constructs MakeProxyService::new(policy_engine), accepts TLS connections, extracts peer certificates, and serves HTTP/2 requests using ProxyService. + +The current authorization path is: + + src/proxy.rs parses the CONNECT authority into a canonical destination string. + src/proxy.rs builds a RequestContext containing the mTLS peer certificates and destination. + src/policy.rs extracts the source identity from a configured X.509 certificate extension. + src/policy.rs normalizes the destination. + src/registry.rs queries permission_registry and principal_signing_keys for active candidate permissions. + src/policy.rs verifies the permission signature and checks signer destination scope. + src/policy.rs returns PolicyDecision::Allow with a SubjectIdentity, or PolicyDecision::Deny with a reason. + +The current tunnel path is: + + src/proxy.rs connects to the upstream TCP destination before returning HTTP 200. + src/proxy.rs waits for the HTTP upgrade. + src/proxy.rs wraps the upgraded stream in hyper_util::rt::TokioIo. + src/proxy.rs calls tokio::io::copy_bidirectional to forward bytes both ways until EOF or error. + src/proxy.rs logs bytes_client_to_dest and bytes_dest_to_client when the tunnel closes. + +The current registry schema is in `migrations/0001_signed_authorization_registry.sql`. It has `agent_gateway_schema_version`, `principal_signing_keys`, `principal_key_permissions`, and `permission_registry`. The schema version check in `src/registry.rs` expects version 1. Tests use `tests/common/mod.rs`, which runs migrations against `TEST_DATABASE_URL` and provides helpers for seeding signed permissions. + +The new term "token bucket" means a small in-memory state object that holds a number of available byte tokens for one identity. Tokens refill over time at the configured bytes-per-second rate. Before the proxy forwards a chunk of N bytes for that identity, it waits until the bucket has N tokens, then subtracts those tokens. This makes short bursts possible up to the bucket capacity while keeping the long-term average near the configured rate. + +The new term `SubjectIdentity` means the Rust struct that represents an authenticated source identity after policy evaluation has extracted it from the client certificate and loaded any associated identity metadata from the registry. Its required fields for this proof of concept are the identity value itself and the optional `rate_limit_bytes_per_second` value. Code should use this struct when passing an allowed identity between policy and proxy layers instead of passing a bare `String` plus separate metadata fields. + + +## Plan of Work + +First, add registry storage for identity metadata. Create a new migration `migrations/0002_identity_metadata.sql`. It should create a table named `identity_metadata` with `subject_identity TEXT PRIMARY KEY`, `rate_limit_bytes_per_second BIGINT`, `created_at TIMESTAMPTZ NOT NULL DEFAULT now()`, `updated_at TIMESTAMPTZ NOT NULL DEFAULT now()`, `CHECK (subject_identity <> '')`, and `CHECK (rate_limit_bytes_per_second IS NULL OR rate_limit_bytes_per_second > 0)`. Insert schema version 2 into `agent_gateway_schema_version`. Then change `EXPECTED_SCHEMA_VERSION` in `src/registry.rs` to 2. For this proof of concept, a missing row and a row with `rate_limit_bytes_per_second = NULL` both mean unlimited. + +Next, define the structured identity type in `src/policy.rs`. Add `pub struct SubjectIdentity` with private fields for `value: String` and `rate_limit_bytes_per_second: Option`. Implement a constructor usable inside the policy module, plus public accessor methods `value(&self) -> &str` and `rate_limit_bytes_per_second(&self) -> Option`. Keeping fields private lets future identity metadata be added without turning the struct into an unstructured bag of public fields. + +Then add a registry method for loading identity metadata into the structured identity type. In `src/registry.rs`, add a method on `RegistryStore` named `subject_identity(&self, subject_identity: String) -> anyhow::Result`. The SQL should select `rate_limit_bytes_per_second` from `identity_metadata` for the identity value. If no row exists or the value is NULL, return a `SubjectIdentity` with `rate_limit_bytes_per_second: None`; absence of metadata does not mean absence of the authenticated identity. If a positive value exists, convert it to `u64`, then to `NonZeroU64`, and return a `SubjectIdentity` with that limit. Because the database check constraint prevents zero and negative values, conversion failure should be treated as a registry data error with context. Wrap the query in the same `self.query_timeout` pattern used by `candidate_permissions` and `signer_has_scope`. + +Then thread the structured identity through policy evaluation. In `src/policy.rs`, change `PolicyDecision::Allow` to `Allow { subject_identity: SubjectIdentity }`. In `PostgresPolicyEngine::evaluate`, after a candidate permission is accepted, call `self.registry.subject_identity(source_identity).await` and return the resulting `SubjectIdentity` in the allow decision. If the metadata lookup fails, deny the request with the original source identity in `source_identity: Some(...)` and a reason such as `identity metadata lookup failed: ...`; this mirrors existing registry lookup failure behavior. Existing tests that match `Allow { .. }` can continue using wildcard fields, but any exact construction or match must be updated. + +After that, add the limiter implementation. Prefer a new module `src/rate_limit.rs`, exported from `src/lib.rs` as `pub mod rate_limit;`, because the logic is independent enough to unit test without TLS or Postgres. Define a cloneable `RateLimiter` type that owns an `Arc>>>>`. Use `std::sync::Mutex`, not `tokio::sync::Mutex`, so no new Tokio feature is required; do not hold any mutex guard across `.await`. The outer map mutex should be held only long enough to insert, remove, or clone an `Arc` to a per-identity bucket during identity configuration. Token refill and subtraction should happen while holding only that identity bucket's mutex, so unrelated identities do not block each other while copying data. + +Define a cloneable `RateLimitHandle` that contains `Option>>`. Define a synchronous method `configure_identity(&self, identity: &SubjectIdentity) -> RateLimitHandle` that creates or updates the per-identity bucket when `identity.rate_limit_bytes_per_second()` returns a positive limit, removes the bucket from the map when it returns `None`, and returns the handle that the tunnel should cache. Define an async method `RateLimitHandle::acquire(&self, bytes: usize)`. If the handle has no bucket or `bytes == 0`, return immediately. If a bucket exists, refill and inspect the bucket under the bucket lock. If enough tokens are present, subtract them and return. If not, compute a sleep duration, release the bucket lock, sleep with `tokio::time::sleep`, and retry. Bucket capacity should be one second of traffic: `limit.get()` tokens. If a read chunk is larger than capacity, acquire it in repeated chunks no larger than capacity so low configured rates can still make progress. + +Use integer math for refill to keep clippy clean. One workable representation is `Bucket { limit_bytes_per_second: NonZeroU64, capacity: u64, tokens: u64, refill_remainder: u128, last_refill: Instant }`, with nanoseconds as the time unit. On refill, compute `total = elapsed_nanos * u128::from(limit) + refill_remainder`, add `total / 1_000_000_000` tokens up to capacity, and retain `total % 1_000_000_000` as the remainder. When the configured limit for an identity changes, update that identity's bucket under its bucket lock, clamp tokens to the new capacity, and keep the bucket otherwise intact. If the identity becomes unlimited, remove its bucket from the map under the map lock. For a couple-hour proof of concept, do not implement garbage collection of idle buckets beyond removing buckets for identities explicitly configured as unlimited. + +Then wire the limiter into the proxy. In `src/proxy.rs`, add `rate_limiter: RateLimiter` to `MakeProxyService` and `ProxyService`. `MakeProxyService::new(policy_engine)` should create `RateLimiter::default()` internally so most callers do not need a new argument. `make_service` should clone the limiter into each `ProxyService`. In `handle`, bind `subject_identity` from the allow decision. After the policy decision allows the CONNECT and before spawning the tunnel, call `let rate_limit = rate_limiter.configure_identity(&subject_identity)`. Clone or extract the identity value string from `subject_identity` for logging, and pass the `RateLimitHandle` into `spawn_tunnel`. + +Replace `copy_bidirectional` with an explicit rate-aware tunnel copy. One straightforward design is to split both streams with `tokio::io::split`, then run two `rate_limited_copy` futures concurrently with `tokio::try_join!`. The client-to-destination copy reads from the upgraded downstream stream, calls `rate_limit.acquire(n).await`, writes the bytes to the upstream writer, and accumulates the byte count. The destination-to-client copy does the same in reverse and uses a clone of the same `RateLimitHandle`, so both directions charge the same bucket. On EOF, call `shutdown()` on the writer for that direction. Preserve the existing final log fields `bytes_client_to_dest` and `bytes_dest_to_client` so existing observability expectations continue to hold. + +Finally, update tests and docs. Add unit tests for the token bucket behavior in `src/rate_limit.rs`. Add registry or policy integration tests showing that an identity without metadata is allowed with no limit, an identity with metadata is allowed with the expected limit, and an invalid metadata lookup failure denies authorization if such a failure can be induced cleanly. Add an e2e test that configures a very low limit, transfers enough bytes through the echo tunnel, and asserts elapsed time is greater than a conservative lower bound. Keep timing thresholds loose to avoid flakes. Update `README.md` to describe the `identity_metadata` table and the proof-of-concept limitations. + + +## Concrete Steps + +Work from the repository root: `/home/ajd/projects/sl5_work_test/agent-gateway`. + +1. Create `migrations/0002_identity_metadata.sql` with the new table and schema version insert. + +2. Edit `src/registry.rs`: + + - Change `EXPECTED_SCHEMA_VERSION` from 1 to 2. + - Import `std::num::NonZeroU64`. + - Add `RegistryStore::subject_identity` returning `anyhow::Result`. + +3. Edit `src/policy.rs`: + + - Import `std::num::NonZeroU64`. + - Add `pub struct SubjectIdentity` with private identity value and rate-limit fields. + - Add accessor methods for the identity value and optional rate limit. + - Change `PolicyDecision::Allow` to contain `subject_identity: SubjectIdentity`. + - Query identity metadata through `RegistryStore::subject_identity` after a permission candidate is accepted. + - Deny on metadata lookup errors with an explicit reason. + +4. Add `src/rate_limit.rs`: + + - Define `RateLimiter`. + - Define `RateLimitHandle`. + - Define internal `Bucket`. + - Store buckets as `HashMap>>` behind a short-lived map mutex. + - Implement `Default`, `Clone`, `configure_identity`, and `acquire`. + - Have `configure_identity` accept `&SubjectIdentity` and return `RateLimitHandle`. + - Have `RateLimitHandle::acquire` accept only the byte count. + - Hold the map mutex only during identity configuration; hold the bucket mutex only for token accounting; hold no lock across `.await`. + - Add unit tests using short real-time sleeps with generous assertions. Avoid adding Tokio's `test-util` feature unless the timing tests become too slow or flaky. + +5. Edit `src/lib.rs` to export `pub mod rate_limit;`. + +6. Edit `src/proxy.rs`: + + - Remove the `copy_bidirectional` import. + - Add `use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};` or equivalent imports. + - Add the limiter to `MakeProxyService` and `ProxyService`. + - Configure the limiter from the allow decision's `SubjectIdentity` before spawning the tunnel and pass the returned `RateLimitHandle` into the tunnel. + - Use `SubjectIdentity::value()` for log fields. + - Add `rate_limited_copy`. + - Use `tokio::try_join!` to run both directions concurrently. + +7. Edit `tests/common/mod.rs`: + + - Add a helper on `TestAuthzRegistry` to insert or clear an identity rate limit, for example `set_identity_rate_limit(&self, subject_identity: &str, bytes_per_second: Option)`. + - Extend cleanup to delete rows from `identity_metadata` for test identities touched by the registry, or have the helper return enough information for test-local cleanup. + +8. Edit `tests/integration.rs`: + + - Add tests for policy allow decisions with absent metadata and configured metadata. + - Adjust existing matches for the new `PolicyDecision::Allow` shape if needed. + +9. Edit `tests/e2e.rs`: + + - Add a conservative timing test using an echo server and a low rate limit. + - Transfer enough data to observe throttling without making the suite slow. For example, 4096 bytes through an echo server at 2048 bytes per second accounts for 8192 total bytes across both directions; with a one-second initial bucket capacity, that should take roughly three seconds, but the assertion should use a conservative lower bound to allow scheduling variability. + +10. Edit `README.md`: + + - Document the new identity metadata table. + - State that rate limits are process-local, are refreshed when a new CONNECT is authorized for an identity, aggregate across active tunnels by identity, and count both directions. + +11. Regenerate SQLx metadata if the repository contains `.sqlx` metadata or if CI expects it. Use the exact commands from README/CI: + + cargo install sqlx-cli --version 0.8.6 --locked --no-default-features --features postgres + SQLX_OFFLINE=false DATABASE_URL="$TEST_DATABASE_URL" cargo sqlx database setup + SQLX_OFFLINE=false DATABASE_URL="$TEST_DATABASE_URL" cargo sqlx prepare -- --all-targets --locked + +12. Run validation: + + cargo fmt + SQLX_OFFLINE=false DATABASE_URL="$TEST_DATABASE_URL" cargo sqlx prepare --check -- --all-targets --locked + cargo clippy --locked --all-targets + cargo test --locked + + +## Validation and Acceptance + +The implementation is acceptable when all of the following are true. + +An identity without a row in `identity_metadata` can still connect and tunnel data exactly as before. The existing `tunnel_echoes_data` e2e test should continue to pass, and the final `tunnel closed` event should still include positive `bytes_client_to_dest` and `bytes_dest_to_client` values. + +An identity with `identity_metadata.rate_limit_bytes_per_second = N` receives `PolicyDecision::Allow { subject_identity }` when it otherwise has a valid signed permission, where `subject_identity.value()` is the certificate identity and `subject_identity.rate_limit_bytes_per_second()` is `Some(N)`. This should be proven with a database-backed integration test. + +All active tunnels for the same identity share one limiter in the same process. This can be proven with either a focused unit test that uses cloned `RateLimitHandle` values for the same identity or an e2e test with two tunnels if there is enough time. For the proof of concept, the unit test is sufficient if the e2e timing test proves real tunnel throttling. + +Both traffic directions count against the same configured limit. The e2e timing test should use the echo server so bytes written by the client are also returned by the destination; the observed elapsed time should reflect charging both the outbound and inbound chunks. + +All commands below should complete successfully: + + cargo fmt + SQLX_OFFLINE=false DATABASE_URL="$TEST_DATABASE_URL" cargo sqlx prepare --check -- --all-targets --locked + cargo clippy --locked --all-targets + cargo test --locked + + +## Idempotence and Recovery + +The migration is additive and safe to run once through SQLx migrations. If a local test database has already applied only version 1, rerunning `cargo sqlx database setup` with the new migration should apply version 2. If local test data causes conflicts, use a disposable test database because the test harness expects to write and clean up registry rows. + +The limiter is process-local state only. Restarting the gateway clears all buckets. That is acceptable for this proof of concept and should be documented as a limitation. + +If the e2e timing test is flaky, first loosen its lower bound or reduce scheduler sensitivity. Do not replace it with only a pure unit test unless the user agrees, because the feature must be demonstrated on actual proxied data. + + +## Artifacts and Notes + +Current tunnel forwarding code in `src/proxy.rs`: + + match copy_bidirectional(&mut downstream, &mut upstream).await { + Ok((up, down)) => { + info!( + bytes_client_to_dest = up, + bytes_dest_to_client = down, + "tunnel closed" + ); + } + Err(e) => { ... } + } + +Target policy decision shape: + + pub enum PolicyDecision { + Allow { + subject_identity: SubjectIdentity, + }, + Deny { + source_identity: Option, + reason: String, + }, + } + +Target migration shape: + + CREATE TABLE identity_metadata ( + subject_identity TEXT PRIMARY KEY, + rate_limit_bytes_per_second BIGINT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + CHECK (subject_identity <> ''), + CHECK ( + rate_limit_bytes_per_second IS NULL + OR rate_limit_bytes_per_second > 0 + ) + ); + + INSERT INTO agent_gateway_schema_version (version) VALUES (2); + + +## Interfaces and Dependencies + +Use only existing project dependencies for the proof of concept. Do not add a rate-limiting crate unless implementation uncovers a concrete blocker. + +In `src/policy.rs`, define: + + #[derive(Debug, Clone)] + pub struct SubjectIdentity { ... } + + impl SubjectIdentity { + pub(crate) fn new( + value: String, + rate_limit_bytes_per_second: Option, + ) -> Self; + + pub fn value(&self) -> &str; + + pub fn rate_limit_bytes_per_second( + &self, + ) -> Option; + } + +In `src/rate_limit.rs`, define: + + #[derive(Clone, Default)] + pub struct RateLimiter { + buckets: Arc>>>>, + } + + #[derive(Clone, Default)] + pub struct RateLimitHandle { ... } + + impl RateLimiter { + pub fn configure_identity( + &self, + identity: &crate::policy::SubjectIdentity, + ) -> RateLimitHandle; + } + + impl RateLimitHandle { + pub async fn acquire(&self, bytes: usize); + } + +In `src/registry.rs`, define: + + pub(crate) async fn subject_identity( + &self, + subject_identity: String, + ) -> anyhow::Result; + +In `src/proxy.rs`, keep the external constructor shape: + + impl MakeProxyService { + pub fn new(policy_engine: Arc) -> Self; + } + +This preserves existing call sites in `src/main.rs` and `tests/common/mod.rs`. + + +## Plan Change Notes + +2026-05-17 / Codex: Initial plan created for user review. The plan deliberately scopes rate limiting to a single gateway process and refreshes the process-local bucket configuration when an identity opens a new CONNECT because the task is a simple proof of concept with a couple-hour implementation budget. + +2026-05-17 / Codex: Updated the plan after ajd added the `SubjectIdentity` decision. The plan now makes `SubjectIdentity` the policy-to-proxy carrier for the authenticated identity value and its optional rate-limit metadata. + +2026-05-17 / Codex: Updated the registry lookup design so `RegistryStore::subject_identity` returns a complete `SubjectIdentity`. Missing `identity_metadata` rows now map to an unlimited `SubjectIdentity` instead of an optional result. + +2026-05-17 / Codex: Updated the limiter design to keep `std::sync::Mutex` but use per-bucket locking. The global map lock is now only for managing bucket handles, and each identity's token accounting happens under that identity's own bucket lock. + +2026-05-17 / Codex: Updated the limiter design to cache a per-tunnel `RateLimitHandle` returned by `RateLimiter::configure_identity`. Forwarding no longer performs a bucket-map lookup per copied chunk. diff --git a/.sqlx/query-48633e01bfb078fd16a8ea365bed0307a01844fa132ae470ac7a4d7df05e670a.json b/.sqlx/query-48633e01bfb078fd16a8ea365bed0307a01844fa132ae470ac7a4d7df05e670a.json new file mode 100644 index 0000000..c6a5148 --- /dev/null +++ b/.sqlx/query-48633e01bfb078fd16a8ea365bed0307a01844fa132ae470ac7a4d7df05e670a.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "\n INSERT INTO identity_metadata (\n subject_identity, rate_limit_bytes_per_second\n )\n VALUES ($1, $2)\n ON CONFLICT (subject_identity)\n DO UPDATE SET\n rate_limit_bytes_per_second = EXCLUDED.rate_limit_bytes_per_second,\n updated_at = now()\n ", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Text", + "Int8" + ] + }, + "nullable": [] + }, + "hash": "48633e01bfb078fd16a8ea365bed0307a01844fa132ae470ac7a4d7df05e670a" +} diff --git a/.sqlx/query-59c2c346e5f1572cb528413cabf0896cf235a47242f2c82ccc12dc0fc9e2708b.json b/.sqlx/query-59c2c346e5f1572cb528413cabf0896cf235a47242f2c82ccc12dc0fc9e2708b.json new file mode 100644 index 0000000..609e3fa --- /dev/null +++ b/.sqlx/query-59c2c346e5f1572cb528413cabf0896cf235a47242f2c82ccc12dc0fc9e2708b.json @@ -0,0 +1,14 @@ +{ + "db_name": "PostgreSQL", + "query": "DELETE FROM identity_metadata WHERE subject_identity = $1", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [] + }, + "hash": "59c2c346e5f1572cb528413cabf0896cf235a47242f2c82ccc12dc0fc9e2708b" +} diff --git a/.sqlx/query-924126efaaf9b1d50c24ae577537a7fae2b6c1d6992fb59c11a5a20ff0d875f3.json b/.sqlx/query-924126efaaf9b1d50c24ae577537a7fae2b6c1d6992fb59c11a5a20ff0d875f3.json new file mode 100644 index 0000000..d98e418 --- /dev/null +++ b/.sqlx/query-924126efaaf9b1d50c24ae577537a7fae2b6c1d6992fb59c11a5a20ff0d875f3.json @@ -0,0 +1,22 @@ +{ + "db_name": "PostgreSQL", + "query": "\n SELECT rate_limit_bytes_per_second\n FROM identity_metadata\n WHERE subject_identity = $1\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "rate_limit_bytes_per_second", + "type_info": "Int8" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + true + ] + }, + "hash": "924126efaaf9b1d50c24ae577537a7fae2b6c1d6992fb59c11a5a20ff0d875f3" +} diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..6b32215 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,75 @@ +# Agent Instructions + +## Project Overview + +This repository contains `agent_gateway`, a Rust 2024 mTLS HTTP/2 CONNECT proxy. It authorizes tunnel requests from custom X.509 client certificate extensions against a PostgreSQL-backed signed permission registry, then proxies raw TCP bytes to approved destinations. + +Key paths: + +- `src/main.rs` starts the gateway binary and wires configuration, policy, TLS, proxying, and observability together. +- `src/lib.rs` exposes the library modules used by tests and the binary. +- `src/config.rs`, `src/tls.rs`, `src/policy.rs`, `src/registry.rs`, `src/proxy.rs`, and `src/observability.rs` hold the main implementation. +- `migrations/0001_signed_authorization_registry.sql` defines the authorization registry schema. +- `tests/integration.rs`, `tests/e2e.rs`, and `tests/common/mod.rs` cover policy behavior, TLS material, database registry setup, and proxy flows. + +## Planning Larger Changes + +When writing complex features or significant refactors, use an ExecPlan (as described in `.agent/PLANS.md`) from design to implementation. + +Keep ExecPlans current while working. Record decisions, discoveries, progress, validation evidence, and any change in direction in the plan before relying on that context later. + +## Build and Test Commands + +Use the repository root as the working directory for all commands below. + +- Format Rust code with `cargo fmt`. +- Run static checks with `cargo clippy --locked --all-targets`. The project denies `clippy::pedantic`, so address warnings instead of suppressing them unless there is a narrow, documented reason. +- Run the full test suite with `cargo test --locked`. +- For a release build, use `cargo build --release`. + +The CI test job also checks SQLx query metadata: + +```bash +SQLX_OFFLINE=false DATABASE_URL="$TEST_DATABASE_URL" cargo sqlx database setup +SQLX_OFFLINE=false DATABASE_URL="$TEST_DATABASE_URL" cargo sqlx prepare --check -- --all-targets --locked +``` + +Run the SQLx check when changing migrations, SQL query text, or registry/policy code that relies on checked SQLx macros. + +## Test Suite Notes + +Most tests are ordinary Rust tests run by `cargo test --locked`, but database-backed policy and end-to-end tests require `TEST_DATABASE_URL` to point at a writable PostgreSQL database. The test process applies migrations and writes test data, so use a disposable test database. + +The GitHub Actions workflow uses Postgres 16 with: + +```text +TEST_DATABASE_URL=postgres://agent_gateway_admin:agent_gateway_dev@localhost:5432/agent_gateway_test +``` + +The tests cover destination normalization, config validation, TLS PKI generation, request parsing, signed permission verification, signer delegation scope enforcement, and database-backed authorization behavior. When changing authorization semantics, add or update tests close to the affected behavior rather than relying only on broad e2e coverage. + +## SQLx and Migrations + +Checked SQLx query macros compile against committed metadata in normal builds. If you change migration SQL or SQL query text, regenerate or check the metadata using `sqlx-cli` version `0.8.6`, matching the README and CI workflow. + +Do not treat SQLx metadata as a replacement for migrations. Schema changes belong in `migrations/`, with tests updated to prove the migrated database supports the intended behavior. + +## Commit Discipline + +Make regular git commits while working, especially before and after risky changes. Each commit should represent one logical change: for example, keep a behavior change separate from formatting, dependency updates, SQLx metadata refreshes, or test-only follow-ups. + +For all Codex-initiated commits, include a co-author trailer: + +```text +Co-authored-by: Codex +``` + +Before committing, review `git diff --stat` and `git diff` to ensure the commit only contains the intended files. Do not revert unrelated user changes in a dirty working tree. + +## Implementation Guidelines + +Prefer the existing module boundaries and patterns. Keep authorization and registry logic explicit and well tested, because small changes can affect security-sensitive decisions about which client identities may connect to which destinations. + +Preserve locked dependency behavior by using `--locked` for checks where practical. Avoid introducing new dependencies unless they clearly reduce complexity or match an established project need. + +Update `README.md`, `config.example.toml`, migrations, or tests whenever a code change alters user-visible behavior, configuration, database requirements, or operational expectations. diff --git a/README.md b/README.md index 6042084..d054c20 100644 --- a/README.md +++ b/README.md @@ -76,13 +76,14 @@ The gateway authorizes a CONNECT only when all of these checks pass: ### Database Structure -The authorization registry has three main tables: +The authorization registry has four main tables: | Table | Key Columns | Purpose | |---|---|---| | `principal_signing_keys` | `key_id`, `algorithm`, `public_key_spki_der`, `not_before`, `not_after`, `revoked_at` | Stores trusted P-256 public keys that may sign permissions. | | `principal_key_permissions` | `signing_key_id`, `destination`, `not_before`, `not_after`, `revoked_at` | Defines which destinations each signing key is allowed to delegate. | | `permission_registry` | `permission_id`, `signing_key_id`, `subject_identity`, `subject_public_key_spki_der`, `destination`, `not_before`, `not_after`, `revoked_at`, `signature` | Stores signed permissions that authorize a subject identity and exact subject key to reach a normalized destination. | +| `identity_metadata` | `subject_identity`, `rate_limit_bytes_per_second` | Stores optional per-identity settings that are independent of destination grants. | `principal_key_permissions.signing_key_id` and `permission_registry.signing_key_id` both reference `principal_signing_keys.key_id`. A permission is usable only when the permission row is active, the signing key is active, the signature verifies over the canonical row fields, and the signing key has a matching destination delegation row. @@ -101,6 +102,12 @@ not_after=2026-06-01T00:00:00.000000Z Destination strings are normalized with the same rules used for CONNECT requests: hostnames are lowercased, omitted ports default to `443`, and IPv6 destinations use bracketed `host:port` form. +### Per-identity rate limits + +`identity_metadata.rate_limit_bytes_per_second` optionally configures a process-local byte rate limit for an identity. Missing rows and `NULL` values mean unlimited. + +When a CONNECT request is authorized, the gateway refreshes the in-memory bucket for that identity from `identity_metadata`. All active tunnels for the same identity in that gateway process share the same bucket, and both client-to-destination and destination-to-client bytes count against it. The proof-of-concept limiter is not distributed across multiple gateway processes, and restarting the gateway resets in-memory buckets. + ## Client requirements Clients must: diff --git a/migrations/0002_identity_metadata.sql b/migrations/0002_identity_metadata.sql new file mode 100644 index 0000000..262bc83 --- /dev/null +++ b/migrations/0002_identity_metadata.sql @@ -0,0 +1,13 @@ +CREATE TABLE identity_metadata ( + subject_identity TEXT PRIMARY KEY, + rate_limit_bytes_per_second BIGINT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + CHECK (subject_identity <> ''), + CHECK ( + rate_limit_bytes_per_second IS NULL + OR rate_limit_bytes_per_second > 0 + ) +); + +INSERT INTO agent_gateway_schema_version (version) VALUES (2); diff --git a/src/lib.rs b/src/lib.rs index e8ac142..96504de 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,5 +2,6 @@ pub mod config; pub mod observability; pub mod policy; pub mod proxy; +pub mod rate_limit; mod registry; pub mod tls; diff --git a/src/policy.rs b/src/policy.rs index 08cf163..9db14d5 100644 --- a/src/policy.rs +++ b/src/policy.rs @@ -1,3 +1,4 @@ +use std::num::NonZeroU64; use std::str::FromStr; use std::sync::Arc; @@ -21,7 +22,7 @@ pub struct RequestContext { pub enum PolicyDecision { Allow { - source_identity: String, + subject_identity: SubjectIdentity, }, Deny { source_identity: Option, @@ -29,6 +30,31 @@ pub enum PolicyDecision { }, } +#[derive(Debug, Clone)] +pub struct SubjectIdentity { + value: String, + rate_limit_bytes_per_second: Option, +} + +impl SubjectIdentity { + pub(crate) fn new(value: String, rate_limit_bytes_per_second: Option) -> Self { + Self { + value, + rate_limit_bytes_per_second, + } + } + + #[must_use] + pub fn value(&self) -> &str { + &self.value + } + + #[must_use] + pub fn rate_limit_bytes_per_second(&self) -> Option { + self.rate_limit_bytes_per_second + } +} + #[async_trait] pub trait PolicyEngine: Send + Sync + 'static { async fn evaluate(&self, ctx: &RequestContext) -> PolicyDecision; @@ -193,7 +219,19 @@ impl PolicyEngine for PostgresPolicyEngine { let mut last_denial = None; for candidate in candidates { match self.evaluate_candidate(&candidate, &normalized_dest).await { - Ok(()) => return PolicyDecision::Allow { source_identity }, + Ok(()) => { + return match self + .registry + .subject_identity(source_identity.clone()) + .await + { + Ok(subject_identity) => PolicyDecision::Allow { subject_identity }, + Err(e) => PolicyDecision::Deny { + source_identity: Some(source_identity), + reason: format!("identity metadata lookup failed: {e:#}"), + }, + }; + } Err(reason) => last_denial = Some(reason), } } diff --git a/src/proxy.rs b/src/proxy.rs index ef3beba..21e3c9a 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -15,12 +15,13 @@ use opentelemetry::global; use opentelemetry::propagation::Extractor; use rustls::ServerConnection; use rustls_pki_types::CertificateDer; -use tokio::io::copy_bidirectional; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::net::TcpStream; use tracing::{Instrument, error, info, warn}; use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::policy::{self, PolicyDecision, PolicyEngine, RequestContext}; +use crate::rate_limit::{RateLimitHandle, RateLimiter}; type ProxyBody = BoxBody; @@ -43,6 +44,7 @@ fn extract_trace_context(headers: &HeaderMap) -> opentelemetry::Context { #[derive(Clone)] pub struct ProxyService { policy_engine: Arc, + rate_limiter: RateLimiter, peer_certs: Vec>, source_peer_addr: SocketAddr, } @@ -50,11 +52,13 @@ pub struct ProxyService { impl ProxyService { fn new( policy_engine: Arc, + rate_limiter: RateLimiter, peer_certs: Vec>, source_peer_addr: SocketAddr, ) -> Self { Self { policy_engine, + rate_limiter, peer_certs, source_peer_addr, } @@ -84,16 +88,16 @@ impl ProxyService { destination: dest.authority.clone(), }; - let source_identity = match self.policy_engine.evaluate(&ctx).await { - PolicyDecision::Allow { source_identity } => { + let subject_identity = match self.policy_engine.evaluate(&ctx).await { + PolicyDecision::Allow { subject_identity } => { info!( - source_identity = %source_identity, + source_identity = %subject_identity.value(), source_peer_addr = %self.source_peer_addr, dest_authority = %dest.authority, policy_decision = "allow", "CONNECT allowed" ); - source_identity + subject_identity } PolicyDecision::Deny { source_identity, @@ -108,6 +112,7 @@ impl ProxyService { return response(StatusCode::FORBIDDEN, "forbidden"); } }; + let source_identity = subject_identity.value().to_owned(); // Connect to destination BEFORE returning 200 so the client knows // the tunnel is actually established. @@ -126,9 +131,11 @@ impl ProxyService { }; let on_upgrade = hyper::upgrade::on(req); + let rate_limit = self.rate_limiter.configure_identity(&subject_identity); spawn_tunnel( on_upgrade, upstream, + rate_limit, source_identity, self.source_peer_addr, dest.authority, @@ -154,11 +161,15 @@ impl Service> for ProxyService { pub struct MakeProxyService { policy_engine: Arc, + rate_limiter: RateLimiter, } impl MakeProxyService { pub fn new(policy_engine: Arc) -> Self { - Self { policy_engine } + Self { + policy_engine, + rate_limiter: RateLimiter::default(), + } } #[must_use] @@ -167,7 +178,12 @@ impl MakeProxyService { peer_certs: Vec>, source_peer_addr: SocketAddr, ) -> ProxyService { - ProxyService::new(self.policy_engine.clone(), peer_certs, source_peer_addr) + ProxyService::new( + self.policy_engine.clone(), + self.rate_limiter.clone(), + peer_certs, + source_peer_addr, + ) } } @@ -206,7 +222,8 @@ fn log_denial( fn spawn_tunnel( on_upgrade: hyper::upgrade::OnUpgrade, - mut upstream: TcpStream, + upstream: TcpStream, + rate_limit: RateLimitHandle, source_identity: String, source_peer_addr: SocketAddr, dest_authority: String, @@ -229,8 +246,14 @@ fn spawn_tunnel( }; let mut downstream = hyper_util::rt::TokioIo::new(upgraded); + let (downstream_read, downstream_write) = tokio::io::split(&mut downstream); + let (upstream_read, upstream_write) = tokio::io::split(upstream); - match copy_bidirectional(&mut downstream, &mut upstream).await { + let client_to_dest = + rate_limited_copy(downstream_read, upstream_write, rate_limit.clone()); + let dest_to_client = rate_limited_copy(upstream_read, downstream_write, rate_limit); + + match tokio::try_join!(client_to_dest, dest_to_client) { Ok((up, down)) => { info!( source_identity = %source_identity, @@ -256,6 +279,31 @@ fn spawn_tunnel( ); } +async fn rate_limited_copy( + mut reader: R, + mut writer: W, + rate_limit: RateLimitHandle, +) -> std::io::Result +where + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, +{ + let mut buf = [0_u8; 16 * 1024]; + let mut total = 0_u64; + + loop { + let bytes_read = reader.read(&mut buf).await?; + if bytes_read == 0 { + writer.shutdown().await?; + return Ok(total); + } + + rate_limit.acquire(bytes_read).await; + writer.write_all(&buf[..bytes_read]).await?; + total = total.saturating_add(u64::try_from(bytes_read).unwrap_or(u64::MAX)); + } +} + fn response(status: StatusCode, message: &str) -> Response { let body: ProxyBody = if message.is_empty() { Empty::::new().boxed() diff --git a/src/rate_limit.rs b/src/rate_limit.rs new file mode 100644 index 0000000..8976d4d --- /dev/null +++ b/src/rate_limit.rs @@ -0,0 +1,230 @@ +use std::collections::HashMap; +use std::num::NonZeroU64; +use std::sync::{Arc, Mutex, PoisonError}; +use std::time::{Duration, Instant}; + +use crate::policy::SubjectIdentity; + +const NANOS_PER_SECOND: u128 = 1_000_000_000; + +#[derive(Clone, Default)] +pub struct RateLimiter { + buckets: Arc>>>>, +} + +#[derive(Clone, Default)] +pub struct RateLimitHandle { + bucket: Option>>, +} + +struct Bucket { + limit_bytes_per_second: NonZeroU64, + capacity: u64, + tokens: u64, + refill_remainder: u128, + last_refill: Instant, +} + +impl Bucket { + fn new(limit_bytes_per_second: NonZeroU64) -> Self { + let capacity = limit_bytes_per_second.get(); + Self { + limit_bytes_per_second, + capacity, + tokens: capacity, + refill_remainder: 0, + last_refill: Instant::now(), + } + } + + fn update_limit(&mut self, limit_bytes_per_second: NonZeroU64) { + self.refill(Instant::now()); + self.limit_bytes_per_second = limit_bytes_per_second; + self.capacity = limit_bytes_per_second.get(); + self.tokens = self.tokens.min(self.capacity); + } + + fn refill(&mut self, now: Instant) { + let elapsed = now.saturating_duration_since(self.last_refill); + if elapsed.is_zero() { + return; + } + + let elapsed_nanos = elapsed.as_nanos(); + let total = + elapsed_nanos * u128::from(self.limit_bytes_per_second.get()) + self.refill_remainder; + let added = (total / NANOS_PER_SECOND).min(u128::from(u64::MAX)); + self.refill_remainder = total % NANOS_PER_SECOND; + self.last_refill = now; + + if added == 0 { + return; + } + + let added = u64::try_from(added).unwrap_or(u64::MAX); + self.tokens = self.tokens.saturating_add(added).min(self.capacity); + } + + fn try_acquire(&mut self, requested: u64, now: Instant) -> Option { + self.refill(now); + if self.tokens >= requested { + self.tokens -= requested; + None + } else { + let missing = requested - self.tokens; + let nanos = (u128::from(missing) * NANOS_PER_SECOND) + .div_ceil(u128::from(self.limit_bytes_per_second.get())); + Some(duration_from_nanos(nanos)) + } + } +} + +impl RateLimiter { + pub fn configure_identity(&self, identity: &SubjectIdentity) -> RateLimitHandle { + if let Some(limit) = identity.rate_limit_bytes_per_second() { + let bucket = { + let mut buckets = self.buckets.lock().unwrap_or_else(PoisonError::into_inner); + buckets + .entry(identity.value().to_owned()) + .or_insert_with(|| Arc::new(Mutex::new(Bucket::new(limit)))) + .clone() + }; + bucket + .lock() + .unwrap_or_else(PoisonError::into_inner) + .update_limit(limit); + RateLimitHandle { + bucket: Some(bucket), + } + } else { + self.buckets + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remove(identity.value()); + RateLimitHandle::default() + } + } +} + +impl RateLimitHandle { + pub async fn acquire(&self, bytes: usize) { + let Some(bucket) = self.bucket.as_ref() else { + return; + }; + let mut remaining = u64::try_from(bytes).unwrap_or(u64::MAX); + if remaining == 0 { + return; + } + + while remaining > 0 { + let requested = { + let bucket = bucket.lock().unwrap_or_else(PoisonError::into_inner); + remaining.min(bucket.capacity) + }; + + loop { + let wait = { + let mut bucket = bucket.lock().unwrap_or_else(PoisonError::into_inner); + bucket.try_acquire(requested, Instant::now()) + }; + if let Some(wait) = wait { + tokio::time::sleep(wait).await; + } else { + remaining -= requested; + break; + } + } + } + } +} + +fn duration_from_nanos(nanos: u128) -> Duration { + let nanos = u64::try_from(nanos).unwrap_or(u64::MAX); + Duration::from_nanos(nanos) +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroU64; + use std::time::{Duration, Instant}; + + use super::*; + + fn subject(value: &str, limit: Option) -> SubjectIdentity { + SubjectIdentity::new(value.to_owned(), limit.and_then(NonZeroU64::new)) + } + + #[tokio::test] + async fn unlimited_identity_returns_immediately() { + let limiter = RateLimiter::default(); + let handle = limiter.configure_identity(&subject("agent-alpha", None)); + handle + .acquire(usize::try_from(u64::MAX).unwrap_or(usize::MAX)) + .await; + } + + #[tokio::test] + async fn configured_identity_waits_for_refill() { + let limiter = RateLimiter::default(); + let handle = limiter.configure_identity(&subject("agent-alpha", Some(10_000))); + handle.acquire(10_000).await; + + let start = Instant::now(); + handle.acquire(1_000).await; + + assert!( + start.elapsed() >= Duration::from_millis(50), + "second acquire should wait for bucket refill" + ); + } + + #[tokio::test] + async fn different_identities_have_independent_buckets() { + let limiter = RateLimiter::default(); + let alpha = limiter.configure_identity(&subject("agent-alpha", Some(10_000))); + let beta = limiter.configure_identity(&subject("agent-beta", Some(10_000))); + + alpha.acquire(10_000).await; + + let start = Instant::now(); + beta.acquire(10_000).await; + + assert!( + start.elapsed() < Duration::from_millis(50), + "agent-beta should have its own full bucket" + ); + } + + #[tokio::test] + async fn cloned_handles_share_the_same_bucket() { + let limiter = RateLimiter::default(); + let first = limiter.configure_identity(&subject("agent-alpha", Some(10_000))); + let second = first.clone(); + + first.acquire(10_000).await; + + let start = Instant::now(); + second.acquire(1_000).await; + + assert!( + start.elapsed() >= Duration::from_millis(50), + "second handle should wait for the shared bucket to refill" + ); + } + + #[tokio::test] + async fn removing_limit_makes_identity_unlimited() { + let limiter = RateLimiter::default(); + let rate_limit = limiter.configure_identity(&subject("agent-alpha", Some(10_000))); + rate_limit.acquire(10_000).await; + let unlimited = limiter.configure_identity(&subject("agent-alpha", None)); + + let start = Instant::now(); + unlimited.acquire(10_000).await; + + assert!( + start.elapsed() < Duration::from_millis(50), + "unlimited identity should not wait" + ); + } +} diff --git a/src/registry.rs b/src/registry.rs index 41d4441..a01bc69 100644 --- a/src/registry.rs +++ b/src/registry.rs @@ -1,10 +1,13 @@ +use std::num::NonZeroU64; use std::time::Duration; use anyhow::Context; use chrono::{DateTime, Utc}; use sqlx::postgres::PgPool; -const EXPECTED_SCHEMA_VERSION: i32 = 1; +use crate::policy::SubjectIdentity; + +const EXPECTED_SCHEMA_VERSION: i32 = 2; #[derive(Clone)] pub(crate) struct RegistryStore { @@ -135,4 +138,33 @@ impl RegistryStore { .context("authorization registry signer scope lookup timed out")? .context("querying authorization registry signer scope") } + + pub(crate) async fn subject_identity( + &self, + subject_identity: String, + ) -> anyhow::Result { + let query = sqlx::query_scalar!( + r#" + SELECT rate_limit_bytes_per_second + FROM identity_metadata + WHERE subject_identity = $1 + "#, + &subject_identity, + ); + + let rate_limit = tokio::time::timeout(self.query_timeout, query.fetch_optional(&self.pool)) + .await + .context("authorization registry identity metadata lookup timed out")? + .context("querying authorization registry identity metadata")? + .flatten() + .map(Self::rate_limit_from_db) + .transpose()?; + + Ok(SubjectIdentity::new(subject_identity, rate_limit)) + } + + fn rate_limit_from_db(value: i64) -> anyhow::Result { + let value = u64::try_from(value).context("identity rate limit must be positive")?; + NonZeroU64::new(value).context("identity rate limit must be non-zero") + } } diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 4b33d8b..2bf356b 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -350,6 +350,7 @@ pub struct TestAuthzRegistry { pub pool: sqlx::PgPool, signing_key: SigningKey, key_id: String, + identity_metadata_subjects: Arc>>, } #[derive(Debug, Clone)] @@ -403,6 +404,7 @@ impl TestAuthzRegistry { pool, signing_key, key_id, + identity_metadata_subjects: Arc::new(Mutex::new(Vec::new())), } } @@ -424,6 +426,17 @@ impl TestAuthzRegistry { } pub async fn cleanup(&self) { + let identity_metadata_subjects = + { self.identity_metadata_subjects.lock().unwrap().clone() }; + for subject_identity in identity_metadata_subjects { + sqlx::query!( + "DELETE FROM identity_metadata WHERE subject_identity = $1", + &subject_identity + ) + .execute(&self.pool) + .await + .expect("delete test identity metadata"); + } sqlx::query!( "DELETE FROM permission_registry WHERE signing_key_id = $1", &self.key_id @@ -447,6 +460,34 @@ impl TestAuthzRegistry { .expect("delete test signing key"); } + pub async fn set_identity_rate_limit( + &self, + subject_identity: &str, + bytes_per_second: Option, + ) { + self.identity_metadata_subjects + .lock() + .unwrap() + .push(subject_identity.to_owned()); + sqlx::query!( + r" + INSERT INTO identity_metadata ( + subject_identity, rate_limit_bytes_per_second + ) + VALUES ($1, $2) + ON CONFLICT (subject_identity) + DO UPDATE SET + rate_limit_bytes_per_second = EXCLUDED.rate_limit_bytes_per_second, + updated_at = now() + ", + subject_identity, + bytes_per_second, + ) + .execute(&self.pool) + .await + .expect("upsert test identity metadata"); + } + pub async fn allow( &self, subject_identity: &str, @@ -608,6 +649,16 @@ impl TestPolicyEngine { self.engine.clone() } + pub async fn set_identity_rate_limit( + &self, + subject_identity: &str, + bytes_per_second: Option, + ) { + self.registry + .set_identity_rate_limit(subject_identity, bytes_per_second) + .await; + } + pub async fn cleanup(&self) { self.registry.cleanup().await; } diff --git a/tests/e2e.rs b/tests/e2e.rs index 46dca67..c0fd0e4 100644 --- a/tests/e2e.rs +++ b/tests/e2e.rs @@ -138,6 +138,61 @@ async fn tunnel_echoes_data() { policy.cleanup().await; } +#[tokio::test] +async fn tunnel_rate_limits_echo_data_per_identity() { + let _guard = serial_test_lock().await; + let log = init_tracing_capture(); + drain_events(&log); + + let (echo_addr, _echo_guard) = start_echo_server().await; + + let subject = unique_test_identity("agent-alpha"); + let pki = TestPki::new(&subject); + let policy = policy_allowing( + &pki, + &subject, + vec![format!("127.0.0.1:{}", echo_addr.port())], + ) + .await; + policy.set_identity_rate_limit(&subject, Some(4_096)).await; + let (proxy_addr, _proxy_guard) = start_proxy(&pki, policy.engine()).await; + + let mut send_req = common::connect_client(proxy_addr, &pki).await; + + let dest = format!("127.0.0.1:{}", echo_addr.port()); + let req = Request::connect(&dest) + .body(Empty::::new()) + .unwrap(); + + let resp = send_req.send_request(req).await.unwrap(); + assert_eq!(resp.status(), 200, "expected 200 OK for allowed CONNECT"); + + let upgraded = hyper::upgrade::on(resp).await.unwrap(); + let mut io = hyper_util::rt::TokioIo::new(upgraded); + let payload = vec![0x5a; 4_096]; + + let start = std::time::Instant::now(); + io.write_all(&payload).await.unwrap(); + io.shutdown().await.unwrap(); + + let mut buf = Vec::new(); + io.read_to_end(&mut buf).await.unwrap(); + let elapsed = start.elapsed(); + assert_eq!(buf, payload, "echo server should return the same data"); + assert!( + elapsed >= std::time::Duration::from_millis(500), + "round trip should be throttled by the shared identity bucket, elapsed {elapsed:?}" + ); + + let events = wait_for_event(&log, "tunnel closed", EVENT_TIMEOUT).await; + let closed_evt = find_event(&events, "tunnel closed").expect("expected tunnel closed event"); + assert_field_eq(closed_evt, "source_identity", &subject); + assert_field_eq(closed_evt, "dest_authority", &dest); + assert_field_eq(closed_evt, "bytes_client_to_dest", "4096"); + assert_field_eq(closed_evt, "bytes_dest_to_client", "4096"); + policy.cleanup().await; +} + #[tokio::test] async fn tunnel_policy_deny() { let _guard = serial_test_lock().await; diff --git a/tests/integration.rs b/tests/integration.rs index d47c0c1..aa0335b 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -46,6 +46,56 @@ async fn policy_allows_matching_cert_and_destination() { registry.cleanup().await; } +#[tokio::test] +async fn policy_allow_without_identity_metadata_has_no_rate_limit() { + let subject = unique_test_identity("agent-alpha"); + let registry = TestAuthzRegistry::new().await; + let pki = TestPki::new(&subject); + registry + .allow_for_pki(&pki, &subject, "api.example.com:443") + .await; + let engine = registry.engine(EXT_OID).await; + + match eval(engine.as_ref(), &pki, "api.example.com:443").await { + PolicyDecision::Allow { subject_identity } => { + assert_eq!(subject_identity.value(), subject); + assert!(subject_identity.rate_limit_bytes_per_second().is_none()); + } + PolicyDecision::Deny { reason, .. } => panic!("expected Allow, got Deny: {reason}"), + } + + registry.cleanup().await; +} + +#[tokio::test] +async fn policy_allow_with_identity_metadata_has_rate_limit() { + let subject = unique_test_identity("agent-alpha"); + let registry = TestAuthzRegistry::new().await; + let pki = TestPki::new(&subject); + registry + .allow_for_pki(&pki, &subject, "api.example.com:443") + .await; + registry + .set_identity_rate_limit(&subject, Some(2_048)) + .await; + let engine = registry.engine(EXT_OID).await; + + match eval(engine.as_ref(), &pki, "api.example.com:443").await { + PolicyDecision::Allow { subject_identity } => { + assert_eq!(subject_identity.value(), subject); + assert_eq!( + subject_identity + .rate_limit_bytes_per_second() + .map(std::num::NonZeroU64::get), + Some(2_048) + ); + } + PolicyDecision::Deny { reason, .. } => panic!("expected Allow, got Deny: {reason}"), + } + + registry.cleanup().await; +} + #[tokio::test] async fn policy_allows_explicit_non_default_port() { let subject = unique_test_identity("agent-alpha");