diff --git a/.sqlx/query-2e5bfa8b38514ea1eb638c158622d5345cb1bcd943017773493f3ce555c09710.json b/.sqlx/query-0a3cb405100f3695b10d4da60284970b194ec70409def4135fe02cc560001ccc.json similarity index 61% rename from .sqlx/query-2e5bfa8b38514ea1eb638c158622d5345cb1bcd943017773493f3ce555c09710.json rename to .sqlx/query-0a3cb405100f3695b10d4da60284970b194ec70409def4135fe02cc560001ccc.json index 9671a70..9166857 100644 --- a/.sqlx/query-2e5bfa8b38514ea1eb638c158622d5345cb1bcd943017773493f3ce555c09710.json +++ b/.sqlx/query-0a3cb405100f3695b10d4da60284970b194ec70409def4135fe02cc560001ccc.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n INSERT INTO permission_registry (\n permission_id, signing_key_id, subject_identity, subject_public_key_spki_der, destination,\n not_before, not_after, signature\n )\n VALUES ($1, $2, $3, $4, $5, $6, $7, $8)\n ", + "query": "\n INSERT INTO permission_registry (\n permission_id, signing_key_id, subject_identity, subject_public_key_spki_der, destination,\n not_before, not_after, signature, capacity_bytes, refill_bytes_per_sec\n )\n VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)\n ", "describe": { "columns": [], "parameters": { @@ -12,10 +12,12 @@ "Text", "Timestamptz", "Timestamptz", - "Bytea" + "Bytea", + "Int8", + "Int8" ] }, "nullable": [] }, - "hash": "2e5bfa8b38514ea1eb638c158622d5345cb1bcd943017773493f3ce555c09710" + "hash": "0a3cb405100f3695b10d4da60284970b194ec70409def4135fe02cc560001ccc" } diff --git a/.sqlx/query-a4afff0f8d466ec2eb6c8e0dfd0e893ff27ce48c5d1989dc65f0b2256cb32065.json b/.sqlx/query-a4afff0f8d466ec2eb6c8e0dfd0e893ff27ce48c5d1989dc65f0b2256cb32065.json new file mode 100644 index 0000000..db93e52 --- /dev/null +++ b/.sqlx/query-a4afff0f8d466ec2eb6c8e0dfd0e893ff27ce48c5d1989dc65f0b2256cb32065.json @@ -0,0 +1,20 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT permission_id FROM permission_registry WHERE revoked_at IS NOT NULL", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "permission_id", + "type_info": "Text" + } + ], + "parameters": { + "Left": [] + }, + "nullable": [ + false + ] + }, + "hash": "a4afff0f8d466ec2eb6c8e0dfd0e893ff27ce48c5d1989dc65f0b2256cb32065" +} diff --git a/.sqlx/query-e3b8aebb3bd68bb84e6d09cb0b214c5ba1e25a847aea0e877f2a6682a760efa6.json b/.sqlx/query-c0dd29c7fd2b1f455de36a23100caf7582fc316d6f3237f543f11c0e7e5a4f5f.json similarity index 61% rename from .sqlx/query-e3b8aebb3bd68bb84e6d09cb0b214c5ba1e25a847aea0e877f2a6682a760efa6.json rename to .sqlx/query-c0dd29c7fd2b1f455de36a23100caf7582fc316d6f3237f543f11c0e7e5a4f5f.json index c4f6a9e..ffacca5 100644 --- a/.sqlx/query-e3b8aebb3bd68bb84e6d09cb0b214c5ba1e25a847aea0e877f2a6682a760efa6.json +++ b/.sqlx/query-c0dd29c7fd2b1f455de36a23100caf7582fc316d6f3237f543f11c0e7e5a4f5f.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n SELECT\n p.permission_id,\n p.subject_identity,\n p.subject_public_key_spki_der,\n p.destination,\n p.signing_key_id,\n p.not_before AS \"permission_not_before!\",\n p.not_after AS \"permission_not_after!\",\n p.signature,\n s.algorithm AS \"signer_algorithm!\",\n s.public_key_spki_der AS \"signer_public_key_spki_der!\",\n s.not_before AS \"signer_not_before!\",\n s.not_after AS \"signer_not_after!\",\n s.revoked_at AS signer_revoked_at,\n (\n s.revoked_at IS NULL\n AND s.not_before <= now()\n AND s.not_after > now()\n ) AS \"signer_active_now!\"\n FROM permission_registry p\n JOIN principal_signing_keys s ON s.key_id = p.signing_key_id\n WHERE p.subject_identity = $1\n AND p.destination = $2\n AND p.subject_public_key_spki_der = $3\n AND p.revoked_at IS NULL\n AND p.not_before <= now()\n AND p.not_after > now()\n ORDER BY p.not_after DESC\n LIMIT 16\n ", + "query": "\n SELECT\n p.permission_id,\n p.subject_identity,\n p.subject_public_key_spki_der,\n p.destination,\n p.signing_key_id,\n p.not_before AS \"permission_not_before!\",\n p.not_after AS \"permission_not_after!\",\n p.capacity_bytes,\n p.refill_bytes_per_sec,\n p.signature,\n s.algorithm AS \"signer_algorithm!\",\n s.public_key_spki_der AS \"signer_public_key_spki_der!\",\n s.not_before AS \"signer_not_before!\",\n s.not_after AS \"signer_not_after!\",\n s.revoked_at AS signer_revoked_at,\n (\n s.revoked_at IS NULL\n AND s.not_before <= now()\n AND s.not_after > now()\n ) AS \"signer_active_now!\"\n FROM permission_registry p\n JOIN principal_signing_keys s ON s.key_id = p.signing_key_id\n WHERE p.subject_identity = $1\n AND p.destination = $2\n AND p.subject_public_key_spki_der = $3\n AND p.revoked_at IS NULL\n AND p.not_before <= now()\n AND p.not_after > now()\n ORDER BY p.not_after DESC\n LIMIT 16\n ", "describe": { "columns": [ { @@ -40,36 +40,46 @@ }, { "ordinal": 7, + "name": "capacity_bytes", + "type_info": "Int8" + }, + { + "ordinal": 8, + "name": "refill_bytes_per_sec", + "type_info": "Int8" + }, + { + "ordinal": 9, "name": "signature", "type_info": "Bytea" }, { - "ordinal": 8, + "ordinal": 10, "name": "signer_algorithm!", "type_info": "Text" }, { - "ordinal": 9, + "ordinal": 11, "name": "signer_public_key_spki_der!", "type_info": "Bytea" }, { - "ordinal": 10, + "ordinal": 12, "name": "signer_not_before!", "type_info": "Timestamptz" }, { - "ordinal": 11, + "ordinal": 13, "name": "signer_not_after!", "type_info": "Timestamptz" }, { - "ordinal": 12, + "ordinal": 14, "name": "signer_revoked_at", "type_info": "Timestamptz" }, { - "ordinal": 13, + "ordinal": 15, "name": "signer_active_now!", "type_info": "Bool" } @@ -94,9 +104,11 @@ false, false, false, + false, + false, true, null ] }, - "hash": "e3b8aebb3bd68bb84e6d09cb0b214c5ba1e25a847aea0e877f2a6682a760efa6" + "hash": "c0dd29c7fd2b1f455de36a23100caf7582fc316d6f3237f543f11c0e7e5a4f5f" } diff --git a/.sqlx/query-fdeb6c7618cfa0f19e33cc99b0ec4e8f815e850bf446f198223d7318127656c6.json b/.sqlx/query-fdeb6c7618cfa0f19e33cc99b0ec4e8f815e850bf446f198223d7318127656c6.json new file mode 100644 index 0000000..f1df2c9 --- /dev/null +++ b/.sqlx/query-fdeb6c7618cfa0f19e33cc99b0ec4e8f815e850bf446f198223d7318127656c6.json @@ -0,0 +1,14 @@ +{ + "db_name": "PostgreSQL", + "query": "UPDATE permission_registry SET capacity_bytes = capacity_bytes * 10 WHERE permission_id = $1", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [] + }, + "hash": "fdeb6c7618cfa0f19e33cc99b0ec4e8f815e850bf446f198223d7318127656c6" +} diff --git a/README.md b/README.md index 6042084..0ecf340 100644 --- a/README.md +++ b/README.md @@ -89,7 +89,7 @@ The authorization registry has three main tables: The signed bytes are the following UTF-8 text, with fields in this exact order and timestamps formatted as UTC RFC 3339 with six fractional digits: ```text -agent-gateway-permission-v1 +agent-gateway-permission-v2 permission_id=perm-1 signing_key_id=org-alice subject_identity=agent-alpha @@ -97,10 +97,23 @@ subject_public_key_spki_der=3059301306072a8648ce3d020106082a8648ce3d030107034200 destination=api.example.com:443 not_before=2026-05-01T00:00:00.000000Z not_after=2026-06-01T00:00:00.000000Z +capacity_bytes=10485760 +refill_bytes_per_sec=1048576 ``` 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. +`capacity_bytes` and `refill_bytes_per_sec` define a per-permission token bucket. `capacity_bytes` is the maximum burst the agent can emit in a single window; `refill_bytes_per_sec` is the sustained throughput. The bucket is enforced at CONNECT time (refuse with `429` if empty) and continuously during the tunnel (the upstream-side write path is metered, and a write that would drain the bucket terminates the tunnel cleanly). Because the limits are part of the signed canonical bytes, a compromised gateway with read-write access to `permission_registry` still cannot raise them without the principal's signing key. + +### In-flight propagation of expiry and revocation + +The rate-limit bucket also acts as the kill switch for two other reasons an in-flight tunnel might need to end early: + +- **Permission expiry.** The bucket carries the signed `not_after` and treats itself as dead when `now() >= not_after`. The next write through the metered upstream returns `io::Error`, the tunnel closes, and the next CONNECT for the same permission is refused by the policy engine in the usual way. +- **Explicit revocation.** A background task in the gateway polls `permission_registry` every 30 seconds for rows whose `revoked_at` has been set. For each such row, it marks the corresponding in-process bucket as revoked. Active tunnels using that bucket then fail their next write through the same path. Propagation is bounded by the 30-second polling interval; throughout that window the bucket continues to enforce the signed rate limit at full strength, so a just-revoked agent's total transfer between revocation and the next poll is capped at one burst plus up to 30 seconds of sustained throughput at the signed rate. + +These two paths use the same `MeteredStream` error code path as the rate-limit deny; there is no special-cased control flow per termination reason. + ## Client requirements Clients must: @@ -126,6 +139,7 @@ Clients must: | `400` | Malformed request (missing/invalid authority) | | `403` | Policy denied the connection | | `405` | Non-CONNECT method used | +| `429` | Rate limit exceeded (the signed `capacity_bytes` for the matching permission has been depleted); response includes a `Retry-After` header in seconds when the bucket has a non-zero refill rate | | `502` | Could not reach the destination | ### Client certificate extension @@ -164,6 +178,10 @@ 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 ``` +## Demo repository: registry-cli must be updated to sign v2 + +This gateway expects every permission row to be signed in the v2 canonical-bytes format (with `capacity_bytes` and `refill_bytes_per_sec`). The sister demo repository `agent-gateway-demo` ships a `registry-cli/agent-permissions.sh` that currently signs v1. To use this gateway with that demo end-to-end, the script needs a small parallel update: two extra lines in the canonical-bytes here-doc and two corresponding `--set=` flags on the `psql` invocation that inserts the row, plus two new CLI parameters surfacing the values. The diff is small and mechanical; this gateway repository's PR does not modify the demo repo. + ## License This project is licensed under the MIT License. See [LICENSE](LICENSE). diff --git a/migrations/0002_permission_rate_limit.sql b/migrations/0002_permission_rate_limit.sql new file mode 100644 index 0000000..c3bc3ca --- /dev/null +++ b/migrations/0002_permission_rate_limit.sql @@ -0,0 +1,8 @@ +-- Adds the byte-budget fields signed in v2 of the permission canonical bytes. +-- The CHECK constraints guarantee non-negative values, which lets the gateway +-- safely cast Postgres BIGINT (i64) to Rust u64 at the rate-limit boundary. +ALTER TABLE permission_registry + ADD COLUMN capacity_bytes BIGINT NOT NULL CHECK (capacity_bytes >= 0), + ADD COLUMN refill_bytes_per_sec BIGINT NOT NULL CHECK (refill_bytes_per_sec >= 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/main.rs b/src/main.rs index bfb0ddc..d97b2bc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,13 +1,22 @@ use std::path::PathBuf; use std::sync::Arc; +use std::time::Duration; use agent_gateway::proxy::MakeProxyService; +use agent_gateway::rate_limit::BucketStore; use agent_gateway::{config, observability, policy, proxy, tls}; use anyhow::Context; use clap::Parser; use hyper_util::rt::TokioExecutor; +use sqlx::PgPool; use tokio::net::TcpListener; -use tracing::{error, info}; +use tokio::time::MissedTickBehavior; +use tracing::{error, info, warn}; + +/// How often the background task polls `permission_registry` for newly +/// revoked rows and propagates their `revoked` flag to the matching +/// in-process buckets. Bounded latency for explicit revocation. +const REVOCATION_POLL_INTERVAL: Duration = Duration::from_secs(30); #[derive(Parser)] #[command(name = "agent_gateway", about = "mTLS HTTP/2 CONNECT proxy")] @@ -37,7 +46,14 @@ async fn serve(config: config::Config) -> anyhow::Result<()> { let tls_acceptor = tls::TlsAcceptor::from(server_tls); let policy_engine = policy::build_engine(&config.policy).await?; - let make_service = Arc::new(MakeProxyService::new(policy_engine)); + let bucket_store = Arc::new(BucketStore::new()); + let make_service = Arc::new(MakeProxyService::new( + policy_engine, + bucket_store.clone(), + )); + + let revocation_pool = policy::build_pool(&config.policy).await?; + let revocation_task = tokio::spawn(run_revocation_poll(revocation_pool, bucket_store)); let listen_addr: std::net::SocketAddr = config.server.listen_addr.parse()?; let listener = TcpListener::bind(listen_addr).await?; @@ -52,10 +68,42 @@ async fn serve(config: config::Config) -> anyhow::Result<()> { } } + revocation_task.abort(); observability::shutdown(); Ok(()) } +/// Background task that propagates explicit permission revocations +/// (`revoked_at` set on `permission_registry`) to the matching in-process +/// buckets. Bounded latency = `REVOCATION_POLL_INTERVAL`. +async fn run_revocation_poll(pool: PgPool, store: Arc) { + let mut ticker = tokio::time::interval(REVOCATION_POLL_INTERVAL); + ticker.set_missed_tick_behavior(MissedTickBehavior::Skip); + loop { + ticker.tick().await; + match query_revoked_permission_ids(&pool).await { + Ok(ids) => { + for id in ids { + store.mark_revoked(&id); + } + } + Err(e) => warn!(error = ?e, "revocation poll query failed"), + } + } +} + +/// Read-only query that returns every `permission_id` with `revoked_at` set. +/// The gateway calls this on a 30-second timer; integration tests call it +/// directly so they don't have to wait on the timer. +async fn query_revoked_permission_ids(pool: &PgPool) -> anyhow::Result> { + let rows = sqlx::query_scalar!( + "SELECT permission_id FROM permission_registry WHERE revoked_at IS NOT NULL" + ) + .fetch_all(pool) + .await?; + Ok(rows) +} + async fn serve_loop( listener: &TcpListener, tls_acceptor: &tls::TlsAcceptor, diff --git a/src/policy.rs b/src/policy.rs index 08cf163..0460376 100644 --- a/src/policy.rs +++ b/src/policy.rs @@ -22,6 +22,7 @@ pub struct RequestContext { pub enum PolicyDecision { Allow { source_identity: String, + permission: AllowedPermission, }, Deny { source_identity: Option, @@ -29,6 +30,18 @@ pub enum PolicyDecision { }, } +/// Fields from the matched permission row that the proxy needs to set up the +/// per-permission rate-limit bucket. `capacity_bytes` and `refill_bytes_per_sec` +/// come from the signed v2 canonical bytes (so they cannot be raised without +/// the principal's signing key); `not_after` is the bucket's expiry deadline. +#[derive(Clone)] +pub struct AllowedPermission { + pub permission_id: String, + pub capacity_bytes: u64, + pub refill_bytes_per_sec: u64, + pub not_after: chrono::DateTime, +} + #[async_trait] pub trait PolicyEngine: Send + Sync + 'static { async fn evaluate(&self, ctx: &RequestContext) -> PolicyDecision; @@ -67,6 +80,18 @@ pub async fn build_engine(config: &PolicyConfig) -> anyhow::Result anyhow::Result { + build_pg_pool(config).await +} + async fn build_pg_pool(policy: &PolicyConfig) -> anyhow::Result { let database_url = policy.database_url()?; let connect_options = PgConnectOptions::from_str(&database_url) @@ -193,7 +218,23 @@ 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(()) => { + // capacity_bytes and refill_bytes_per_sec are i64 in the + // database row but the CHECK (>= 0) constraint in + // migration 0002 guarantees they are non-negative, so the + // cast to u64 is safe. + #[allow(clippy::cast_sign_loss)] + let permission = AllowedPermission { + permission_id: candidate.permission_id.clone(), + capacity_bytes: candidate.capacity_bytes as u64, + refill_bytes_per_sec: candidate.refill_bytes_per_sec as u64, + not_after: candidate.permission_not_after, + }; + return PolicyDecision::Allow { + source_identity, + permission, + }; + } Err(reason) => last_denial = Some(reason), } } @@ -314,7 +355,7 @@ fn verify_signature(candidate: &CandidatePermission) -> anyhow::Result<()> { fn canonical_permission_bytes(candidate: &CandidatePermission) -> Vec { format!( - "agent-gateway-permission-v1\npermission_id={}\nsigning_key_id={}\nsubject_identity={}\nsubject_public_key_spki_der={}\ndestination={}\nnot_before={}\nnot_after={}\n", + "agent-gateway-permission-v2\npermission_id={}\nsigning_key_id={}\nsubject_identity={}\nsubject_public_key_spki_der={}\ndestination={}\nnot_before={}\nnot_after={}\ncapacity_bytes={}\nrefill_bytes_per_sec={}\n", candidate.permission_id, candidate.signing_key_id, candidate.subject_identity, @@ -326,6 +367,8 @@ fn canonical_permission_bytes(candidate: &CandidatePermission) -> Vec { candidate .permission_not_after .to_rfc3339_opts(SecondsFormat::Micros, true), + candidate.capacity_bytes, + candidate.refill_bytes_per_sec, ) .into_bytes() } @@ -462,6 +505,8 @@ mod tests { signing_key_id: "org-alice".to_owned(), permission_not_before: Utc.with_ymd_and_hms(2026, 5, 1, 0, 0, 0).single().unwrap(), permission_not_after: Utc.with_ymd_and_hms(2026, 6, 1, 0, 0, 0).single().unwrap(), + capacity_bytes: 1024, + refill_bytes_per_sec: 256, signature: vec![], signer_algorithm: "ecdsa_p256_sha256".to_owned(), signer_public_key_spki_der: vec![], @@ -473,7 +518,7 @@ mod tests { assert_eq!( canonical_permission_bytes(&candidate), - b"agent-gateway-permission-v1\npermission_id=perm-1\nsigning_key_id=org-alice\nsubject_identity=agent-alpha\nsubject_public_key_spki_der=305901\ndestination=api.example.com:443\nnot_before=2026-05-01T00:00:00.000000Z\nnot_after=2026-06-01T00:00:00.000000Z\n" + b"agent-gateway-permission-v2\npermission_id=perm-1\nsigning_key_id=org-alice\nsubject_identity=agent-alpha\nsubject_public_key_spki_der=305901\ndestination=api.example.com:443\nnot_before=2026-05-01T00:00:00.000000Z\nnot_after=2026-06-01T00:00:00.000000Z\ncapacity_bytes=1024\nrefill_bytes_per_sec=256\n" ); } } diff --git a/src/proxy.rs b/src/proxy.rs index ef3beba..23338e9 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -21,6 +21,7 @@ use tracing::{Instrument, error, info, warn}; use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::policy::{self, PolicyDecision, PolicyEngine, RequestContext}; +use crate::rate_limit::{BucketStore, MeteredStream, TokenBucket}; type ProxyBody = BoxBody; @@ -43,6 +44,7 @@ fn extract_trace_context(headers: &HeaderMap) -> opentelemetry::Context { #[derive(Clone)] pub struct ProxyService { policy_engine: Arc, + bucket_store: Arc, peer_certs: Vec>, source_peer_addr: SocketAddr, } @@ -50,11 +52,13 @@ pub struct ProxyService { impl ProxyService { fn new( policy_engine: Arc, + bucket_store: Arc, peer_certs: Vec>, source_peer_addr: SocketAddr, ) -> Self { Self { policy_engine, + bucket_store, peer_certs, source_peer_addr, } @@ -84,8 +88,11 @@ impl ProxyService { destination: dest.authority.clone(), }; - let source_identity = match self.policy_engine.evaluate(&ctx).await { - PolicyDecision::Allow { source_identity } => { + let (source_identity, permission) = match self.policy_engine.evaluate(&ctx).await { + PolicyDecision::Allow { + source_identity, + permission, + } => { info!( source_identity = %source_identity, source_peer_addr = %self.source_peer_addr, @@ -93,7 +100,7 @@ impl ProxyService { policy_decision = "allow", "CONNECT allowed" ); - source_identity + (source_identity, permission) } PolicyDecision::Deny { source_identity, @@ -109,6 +116,32 @@ impl ProxyService { } }; + let bucket = self.bucket_store.get_or_create( + &permission.permission_id, + permission.capacity_bytes, + permission.refill_bytes_per_sec, + permission.not_after, + ); + if bucket.available_bytes() == 0 { + let retry_after = bucket.seconds_until_one_byte(); + warn!( + source_identity = %source_identity, + source_peer_addr = %self.source_peer_addr, + dest_authority = %dest.authority, + permission_id = %permission.permission_id, + policy_decision = "rate_limit_deny", + retry_after_seconds = retry_after.unwrap_or(0), + "CONNECT rate-limited" + ); + let mut resp = response(StatusCode::TOO_MANY_REQUESTS, "rate limit exceeded"); + if let Some(secs) = retry_after + && let Ok(value) = http::HeaderValue::from_str(&secs.to_string()) + { + resp.headers_mut().insert("retry-after", value); + } + return resp; + } + // Connect to destination BEFORE returning 200 so the client knows // the tunnel is actually established. let upstream = match TcpStream::connect((&*dest.host, dest.port)).await { @@ -129,6 +162,7 @@ impl ProxyService { spawn_tunnel( on_upgrade, upstream, + bucket, source_identity, self.source_peer_addr, dest.authority, @@ -154,11 +188,15 @@ impl Service> for ProxyService { pub struct MakeProxyService { policy_engine: Arc, + bucket_store: Arc, } impl MakeProxyService { - pub fn new(policy_engine: Arc) -> Self { - Self { policy_engine } + pub fn new(policy_engine: Arc, bucket_store: Arc) -> Self { + Self { + policy_engine, + bucket_store, + } } #[must_use] @@ -167,7 +205,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.bucket_store.clone(), + peer_certs, + source_peer_addr, + ) } } @@ -206,7 +249,8 @@ fn log_denial( fn spawn_tunnel( on_upgrade: hyper::upgrade::OnUpgrade, - mut upstream: TcpStream, + upstream: TcpStream, + bucket: Arc, source_identity: String, source_peer_addr: SocketAddr, dest_authority: String, @@ -229,8 +273,9 @@ fn spawn_tunnel( }; let mut downstream = hyper_util::rt::TokioIo::new(upgraded); + let mut metered_upstream = MeteredStream::new(upstream, bucket); - match copy_bidirectional(&mut downstream, &mut upstream).await { + match copy_bidirectional(&mut downstream, &mut metered_upstream).await { Ok((up, down)) => { info!( source_identity = %source_identity, diff --git a/src/rate_limit.rs b/src/rate_limit.rs new file mode 100644 index 0000000..843f5b9 --- /dev/null +++ b/src/rate_limit.rs @@ -0,0 +1,466 @@ +//! Byte-denominated token bucket for per-permission egress limiting. +//! Internally "tokens" are bytes (1 token = 1 byte); the field name +//! follows the standard token-bucket algorithm vocabulary. + +use std::collections::HashMap; +use std::io; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; +use std::time::Instant; + +use chrono::{DateTime, Utc}; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + +struct Inner { + tokens: f64, + last_refill: Instant, +} + +pub struct TokenBucket { + inner: Mutex, + capacity_bytes: u64, + refill_bytes_per_sec: u64, + permission_not_after: DateTime, + revoked: AtomicBool, +} + +impl TokenBucket { + /// Construct a bucket starting at full capacity. The signed permission + /// values flow directly into this constructor; the bucket's parameters + /// are immutable after this point. + #[must_use] + #[allow(clippy::cast_precision_loss)] + pub fn new( + capacity_bytes: u64, + refill_bytes_per_sec: u64, + permission_not_after: DateTime, + ) -> Self { + Self { + inner: Mutex::new(Inner { + tokens: capacity_bytes as f64, + last_refill: Instant::now(), + }), + capacity_bytes, + refill_bytes_per_sec, + permission_not_after, + revoked: AtomicBool::new(false), + } + } + + /// Mark this bucket as revoked. Subsequent consume and peek operations + /// return 0. Idempotent: calling on an already-revoked bucket is a + /// no-op. The Release ordering on the store pairs with the Acquire + /// load in `is_dead` so the revocation is visible to every subsequent + /// dead-check across threads. + pub fn mark_revoked(&self) { + self.revoked.store(true, Ordering::Release); + } + + /// Atomically grants up to `max` bytes of budget. Returns the number + /// of bytes actually consumed from the bucket. Returns 0 if the bucket + /// is empty or expired. The returned value is the amount the caller is + /// committed to either using or refunding. + #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation, clippy::cast_sign_loss)] + pub fn try_consume_up_to(&self, max: u64) -> u64 { + if self.is_dead() { + return 0; + } + let mut inner = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + self.refill_locked(&mut inner); + // `tokens` is bounded above by `capacity_bytes` (a u64) after every + // refill, and clamped at 0 below. The floor-then-cast to u64 is safe. + let available = inner.tokens.max(0.0).floor() as u64; + let granted = std::cmp::min(available, max); + if granted > 0 { + inner.tokens -= granted as f64; + } + granted + } + + /// Return budget previously consumed (e.g. because the inner write was + /// partial, errored, or returned Pending). Refund never raises tokens + /// above capacity. + #[allow(clippy::cast_precision_loss)] + pub fn refund(&self, bytes: u64) { + let mut inner = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + let cap = self.capacity_bytes as f64; + inner.tokens = (inner.tokens + bytes as f64).min(cap); + } + + /// How many whole bytes the bucket currently holds. Returns 0 if the + /// bucket is expired. Useful for the CONNECT-time peek. + #[must_use] + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + pub fn available_bytes(&self) -> u64 { + if self.is_dead() { + return 0; + } + let mut inner = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + self.refill_locked(&mut inner); + inner.tokens.max(0.0).floor() as u64 + } + + /// Seconds until at least one byte is available, or `None` if the bucket + /// already has budget, refill is zero, or the permission has expired + /// (in any of these cases no amount of waiting will make budget appear). + #[must_use] + #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation, clippy::cast_sign_loss)] + pub fn seconds_until_one_byte(&self) -> Option { + if self.is_dead() { + return None; + } + let mut inner = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + self.refill_locked(&mut inner); + if inner.tokens >= 1.0 || self.refill_bytes_per_sec == 0 { + return None; + } + let deficit = 1.0 - inner.tokens; + let seconds = (deficit / self.refill_bytes_per_sec as f64).ceil() as u64; + Some(seconds.max(1)) + } + + fn is_dead(&self) -> bool { + self.revoked.load(Ordering::Acquire) || Utc::now() >= self.permission_not_after + } + + #[allow(clippy::cast_precision_loss)] + fn refill_locked(&self, inner: &mut Inner) { + let now = Instant::now(); + let elapsed = now + .saturating_duration_since(inner.last_refill) + .as_secs_f64(); + let cap = self.capacity_bytes as f64; + inner.tokens = (inner.tokens + elapsed * (self.refill_bytes_per_sec as f64)).min(cap); + inner.last_refill = now; + } +} + +/// Per-process store of `TokenBucket`s keyed on `permission_id`. The store is +/// shared between the proxy (which calls `get_or_create` per CONNECT) and +/// the revocation poll task (added in a later commit). +pub struct BucketStore { + map: Mutex>>, +} + +impl BucketStore { + #[must_use] + pub fn new() -> Self { + Self { + map: Mutex::new(HashMap::new()), + } + } + + /// Return the bucket for `permission_id`, creating it with the supplied + /// parameters if absent. The supplied parameters are ignored on + /// subsequent calls; a bucket's parameters are fixed at construction. + pub fn get_or_create( + &self, + permission_id: &str, + capacity_bytes: u64, + refill_bytes_per_sec: u64, + permission_not_after: DateTime, + ) -> Arc { + let mut map = self.map.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(existing) = map.get(permission_id) { + return existing.clone(); + } + let bucket = Arc::new(TokenBucket::new( + capacity_bytes, + refill_bytes_per_sec, + permission_not_after, + )); + map.insert(permission_id.to_owned(), bucket.clone()); + bucket + } + + /// Mark the bucket for `permission_id` as revoked, if one exists. + /// No-op for `permission_id`s the gateway has not seen yet. Called by + /// the background revocation poll task in `main.rs`. + pub fn mark_revoked(&self, permission_id: &str) { + let map = self.map.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(bucket) = map.get(permission_id) { + bucket.mark_revoked(); + } + } +} + +impl Default for BucketStore { + fn default() -> Self { + Self::new() + } +} + +/// IO wrapper that meters writes against a `TokenBucket` while delegating +/// reads straight through. Designed for the upstream side of the +/// CONNECT-tunnel `copy_bidirectional` call. +pub struct MeteredStream { + inner: S, + bucket: Arc, +} + +impl MeteredStream { + pub fn new(inner: S, bucket: Arc) -> Self { + Self { inner, bucket } + } +} + +impl AsyncRead for MeteredStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let me = self.get_mut(); + Pin::new(&mut me.inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for MeteredStream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let me = self.get_mut(); + // Empty write is not a rate-limit event; delegate. Required because + // `try_consume_up_to(0)` returns 0, which would otherwise trip the + // empty-bucket branch below. + if buf.is_empty() { + return Pin::new(&mut me.inner).poll_write(cx, buf); + } + // `buf.len()` is usize; on 64-bit platforms (the only ones rustls + // supports here) this fits a u64 without loss. + #[allow(clippy::cast_possible_truncation)] + let granted = me.bucket.try_consume_up_to(buf.len() as u64); + if granted == 0 { + return Poll::Ready(Err(io::Error::other("rate_limit_exceeded"))); + } + #[allow(clippy::cast_possible_truncation)] + let granted_usize = granted as usize; + match Pin::new(&mut me.inner).poll_write(cx, &buf[..granted_usize]) { + Poll::Ready(Ok(n)) => { + if n < granted_usize { + me.bucket.refund((granted_usize - n) as u64); + } + Poll::Ready(Ok(n)) + } + Poll::Ready(Err(e)) => { + me.bucket.refund(granted); + Poll::Ready(Err(e)) + } + Poll::Pending => { + me.bucket.refund(granted); + Poll::Pending + } + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let me = self.get_mut(); + Pin::new(&mut me.inner).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let me = self.get_mut(); + Pin::new(&mut me.inner).poll_shutdown(cx) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Duration as ChronoDuration; + use std::sync::atomic::{AtomicU64, Ordering}; + + fn far_future() -> DateTime { + Utc::now() + ChronoDuration::hours(1) + } + + fn far_past() -> DateTime { + Utc::now() - ChronoDuration::hours(1) + } + + #[test] + fn try_consume_up_to_grants_at_most_max() { + let bucket = TokenBucket::new(1000, 0, far_future()); + let granted = bucket.try_consume_up_to(500); + assert_eq!(granted, 500); + assert_eq!(bucket.available_bytes(), 500); + } + + #[test] + fn try_consume_up_to_grants_at_most_available() { + let bucket = TokenBucket::new(1000, 0, far_future()); + bucket.try_consume_up_to(800); + let granted = bucket.try_consume_up_to(500); + assert_eq!(granted, 200); + assert_eq!(bucket.available_bytes(), 0); + } + + #[test] + fn try_consume_up_to_zero_when_empty() { + let bucket = TokenBucket::new(100, 0, far_future()); + bucket.try_consume_up_to(100); + assert_eq!(bucket.try_consume_up_to(50), 0); + } + + #[test] + fn bucket_refills_over_time() { + let bucket = TokenBucket::new(1000, 1000, far_future()); + bucket.try_consume_up_to(900); + assert_eq!(bucket.available_bytes(), 100); + std::thread::sleep(std::time::Duration::from_millis(200)); + let available = bucket.available_bytes(); + // After ~200ms with refill=1000 bytes/sec from 100 starting tokens + // we expect roughly 300 tokens. Use a broad tolerance to avoid CI + // flakiness on slow runners; the property under test is "some + // refill happened, still capped at capacity". + assert!( + available > 100 && available <= 1000, + "expected some refill after sleep, got {available}" + ); + } + + #[test] + fn concurrent_consume_no_overshoot() { + use std::thread; + + let bucket = Arc::new(TokenBucket::new(1000, 0, far_future())); + let total = Arc::new(AtomicU64::new(0)); + let mut handles = Vec::new(); + for _ in 0..100 { + let bucket = bucket.clone(); + let total = total.clone(); + handles.push(thread::spawn(move || { + let granted = bucket.try_consume_up_to(10); + total.fetch_add(granted, Ordering::Relaxed); + })); + } + for h in handles { + h.join().unwrap(); + } + let final_total = total.load(Ordering::Relaxed); + assert!( + final_total <= 1000, + "total granted {final_total} must not exceed capacity 1000" + ); + assert_eq!(final_total, 1000); + } + + #[test] + fn refund_does_not_exceed_capacity() { + let bucket = TokenBucket::new(1000, 0, far_future()); + bucket.try_consume_up_to(500); + bucket.refund(1000); + assert_eq!(bucket.available_bytes(), 1000); + } + + #[test] + fn refund_after_partial_write_pattern() { + let bucket = TokenBucket::new(1000, 0, far_future()); + let granted = bucket.try_consume_up_to(1000); + assert_eq!(granted, 1000); + bucket.refund(200); + assert_eq!(bucket.available_bytes(), 200); + } + + #[test] + fn seconds_until_one_byte_when_empty_with_refill() { + let bucket = TokenBucket::new(100, 10, far_future()); + bucket.try_consume_up_to(100); + assert_eq!(bucket.seconds_until_one_byte(), Some(1)); + } + + #[test] + fn seconds_until_one_byte_when_empty_no_refill() { + let bucket = TokenBucket::new(100, 0, far_future()); + bucket.try_consume_up_to(100); + assert_eq!(bucket.seconds_until_one_byte(), None); + } + + #[test] + fn seconds_until_one_byte_when_already_full() { + let bucket = TokenBucket::new(100, 10, far_future()); + assert_eq!(bucket.seconds_until_one_byte(), None); + } + + #[test] + fn expired_bucket_returns_zero() { + let bucket = TokenBucket::new(1000, 100, far_past()); + assert_eq!(bucket.try_consume_up_to(100), 0); + assert_eq!(bucket.available_bytes(), 0); + assert_eq!(bucket.seconds_until_one_byte(), None); + } + + #[test] + fn revoked_bucket_returns_zero() { + let bucket = TokenBucket::new(1000, 100, far_future()); + assert!(bucket.available_bytes() > 0); + bucket.mark_revoked(); + assert_eq!(bucket.try_consume_up_to(100), 0); + assert_eq!(bucket.available_bytes(), 0); + assert_eq!(bucket.seconds_until_one_byte(), None); + } + + #[test] + fn mark_revoked_is_idempotent() { + let bucket = TokenBucket::new(1000, 100, far_future()); + bucket.mark_revoked(); + bucket.mark_revoked(); + assert_eq!(bucket.try_consume_up_to(1), 0); + } + + #[test] + fn bucket_store_mark_revoked_revokes_existing_bucket() { + let store = BucketStore::new(); + let bucket = store.get_or_create("perm-1", 1000, 100, far_future()); + assert!(bucket.available_bytes() > 0); + store.mark_revoked("perm-1"); + assert_eq!(bucket.available_bytes(), 0); + } + + #[test] + fn bucket_store_mark_revoked_is_noop_for_unknown_permission_id() { + let store = BucketStore::new(); + store.mark_revoked("never-seen"); + // No panic, no error. Acceptable; the poll task may query + // permission_ids the gateway has not yet served. + } + + #[test] + fn bucket_store_returns_same_instance_for_same_permission_id() { + let store = BucketStore::new(); + let a = store.get_or_create("perm-1", 1000, 100, far_future()); + let b = store.get_or_create("perm-1", 9999, 9999, far_future()); + assert!(Arc::ptr_eq(&a, &b)); + // The second call's parameters are ignored; the original bucket's + // capacity stands. + assert_eq!(a.available_bytes(), 1000); + } + + #[tokio::test] + async fn metered_stream_caps_writes_at_capacity_no_overshoot() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let bucket = Arc::new(TokenBucket::new(1024, 0, far_future())); + let (sender, mut receiver) = tokio::io::duplex(8192); + let mut metered = MeteredStream::new(sender, bucket.clone()); + + let payload = vec![0u8; 4096]; + // write_all will fail when the bucket drains; we don't assert on the + // result type, only on the bytes that actually arrived. + let _ = metered.write_all(&payload).await; + drop(metered); + + let mut received_bytes = Vec::new(); + let _ = receiver.read_to_end(&mut received_bytes).await; + + assert!( + received_bytes.len() <= 1024, + "receiver got {} bytes; bucket capacity was 1024 — overshoot detected!", + received_bytes.len() + ); + } +} diff --git a/src/registry.rs b/src/registry.rs index 41d4441..0259e52 100644 --- a/src/registry.rs +++ b/src/registry.rs @@ -4,7 +4,7 @@ use anyhow::Context; use chrono::{DateTime, Utc}; use sqlx::postgres::PgPool; -const EXPECTED_SCHEMA_VERSION: i32 = 1; +const EXPECTED_SCHEMA_VERSION: i32 = 2; #[derive(Clone)] pub(crate) struct RegistryStore { @@ -21,6 +21,8 @@ pub(crate) struct CandidatePermission { pub(crate) signing_key_id: String, pub(crate) permission_not_before: DateTime, pub(crate) permission_not_after: DateTime, + pub(crate) capacity_bytes: i64, + pub(crate) refill_bytes_per_sec: i64, pub(crate) signature: Vec, pub(crate) signer_algorithm: String, pub(crate) signer_public_key_spki_der: Vec, @@ -70,6 +72,8 @@ impl RegistryStore { p.signing_key_id, p.not_before AS "permission_not_before!", p.not_after AS "permission_not_after!", + p.capacity_bytes, + p.refill_bytes_per_sec, p.signature, s.algorithm AS "signer_algorithm!", s.public_key_spki_der AS "signer_public_key_spki_der!", diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 4b33d8b..41d15d7 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -21,6 +21,10 @@ use agent_gateway::policy::PolicyEngine; use agent_gateway::proxy::MakeProxyService; const CLIENT_EXTENSION_OID: &[u64] = &[1, 3, 6, 1, 4, 1, 57264, 1, 1]; +/// "Effectively unlimited" byte budget for tests that don't care about +/// the rate limit. Uses `i64::MAX as u64` so the value fits Postgres BIGINT +/// without wrapping when we cast back at insert time. ~9 exabytes. +const UNLIMITED_BYTES: u64 = i64::MAX as u64; static TEST_ID: AtomicU64 = AtomicU64::new(1); static TEST_MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("./migrations"); @@ -356,6 +360,7 @@ pub struct TestAuthzRegistry { pub struct SeededPermission { pub permission_id: String, pub normalized_destination: String, + pub not_after: DateTime, } impl TestAuthzRegistry { @@ -458,6 +463,8 @@ impl TestAuthzRegistry { subject_public_key_spki_der, destination, true, + UNLIMITED_BYTES, + UNLIMITED_BYTES, ) .await } @@ -473,6 +480,44 @@ impl TestAuthzRegistry { .await } + pub async fn allow_with_limits( + &self, + subject_identity: &str, + subject_public_key_spki_der: &[u8], + destination: &str, + capacity_bytes: u64, + refill_bytes_per_sec: u64, + ) -> SeededPermission { + self.allow_inner( + subject_identity, + subject_public_key_spki_der, + destination, + true, + capacity_bytes, + refill_bytes_per_sec, + ) + .await + } + + pub async fn allow_with_limits_for_pki( + &self, + pki: &TestPki, + subject_identity: &str, + destination: &str, + capacity_bytes: u64, + refill_bytes_per_sec: u64, + ) -> SeededPermission { + let subject_public_key_spki_der = pki.client_spki_der(); + self.allow_with_limits( + subject_identity, + &subject_public_key_spki_der, + destination, + capacity_bytes, + refill_bytes_per_sec, + ) + .await + } + pub async fn allow_without_signer_scope( &self, subject_identity: &str, @@ -484,6 +529,8 @@ impl TestAuthzRegistry { subject_public_key_spki_der, destination, false, + UNLIMITED_BYTES, + UNLIMITED_BYTES, ) .await } @@ -530,12 +577,15 @@ impl TestAuthzRegistry { .expect("tamper permission destination"); } + #[allow(clippy::cast_possible_wrap)] async fn allow_inner( &self, subject_identity: &str, subject_public_key_spki_der: &[u8], destination: &str, include_scope: bool, + capacity_bytes: u64, + refill_bytes_per_sec: u64, ) -> SeededPermission { let permission_id = unique_id("test-permission"); let (not_before, not_after) = active_window(); @@ -566,17 +616,25 @@ impl TestAuthzRegistry { destination, not_before, not_after, + capacity_bytes, + refill_bytes_per_sec, ); let signature: p256::ecdsa::Signature = self.signing_key.sign(&signed_bytes); let signature_der = signature.to_der(); + // The CHECK constraint in migration 0002 enforces non-negative on + // capacity_bytes and refill_bytes_per_sec, and tests stay below + // `i64::MAX as u64`, so the cast back to i64 cannot wrap. + let capacity_for_db = capacity_bytes.min(i64::MAX as u64) as i64; + let refill_for_db = refill_bytes_per_sec.min(i64::MAX as u64) as i64; + sqlx::query!( r" INSERT INTO permission_registry ( permission_id, signing_key_id, subject_identity, subject_public_key_spki_der, destination, - not_before, not_after, signature + not_before, not_after, signature, capacity_bytes, refill_bytes_per_sec ) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) ", &permission_id, &self.key_id, @@ -586,6 +644,8 @@ impl TestAuthzRegistry { not_before, not_after, signature_der.as_bytes(), + capacity_for_db, + refill_for_db, ) .execute(&self.pool) .await @@ -594,6 +654,7 @@ impl TestAuthzRegistry { SeededPermission { permission_id, normalized_destination: destination.to_owned(), + not_after, } } } @@ -652,6 +713,7 @@ fn active_window() -> (DateTime, DateTime) { (now - Duration::hours(1), now + Duration::hours(1)) } +#[allow(clippy::too_many_arguments)] fn test_canonical_permission_bytes( permission_id: &str, signing_key_id: &str, @@ -660,9 +722,11 @@ fn test_canonical_permission_bytes( destination: &str, not_before: DateTime, not_after: DateTime, + capacity_bytes: u64, + refill_bytes_per_sec: u64, ) -> Vec { format!( - "agent-gateway-permission-v1\npermission_id={permission_id}\nsigning_key_id={signing_key_id}\nsubject_identity={subject_identity}\nsubject_public_key_spki_der={}\ndestination={destination}\nnot_before={}\nnot_after={}\n", + "agent-gateway-permission-v2\npermission_id={permission_id}\nsigning_key_id={signing_key_id}\nsubject_identity={subject_identity}\nsubject_public_key_spki_der={}\ndestination={destination}\nnot_before={}\nnot_after={}\ncapacity_bytes={capacity_bytes}\nrefill_bytes_per_sec={refill_bytes_per_sec}\n", lower_hex(subject_public_key_spki_der), not_before.to_rfc3339_opts(SecondsFormat::Micros, true), not_after.to_rfc3339_opts(SecondsFormat::Micros, true), @@ -701,6 +765,17 @@ pub async fn start_proxy( pki: &TestPki, policy_engine: Arc, ) -> (SocketAddr, ServerGuard) { + let (addr, guard, _store) = start_proxy_with_store(pki, policy_engine).await; + (addr, guard) +} + +/// Like `start_proxy` but also returns a handle to the in-process +/// `BucketStore` so tests can inspect bucket state, drain buckets directly, +/// or call `mark_revoked` without waiting on a background timer. +pub async fn start_proxy_with_store( + pki: &TestPki, + policy_engine: Arc, +) -> (SocketAddr, ServerGuard, Arc) { install_test_crypto_provider(); let mut server_config = rustls::ServerConfig::builder() @@ -711,7 +786,8 @@ pub async fn start_proxy( let tls_acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(server_config)); - let make_service = Arc::new(MakeProxyService::new(policy_engine)); + let bucket_store = Arc::new(agent_gateway::rate_limit::BucketStore::new()); + let make_service = Arc::new(MakeProxyService::new(policy_engine, bucket_store.clone())); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -748,7 +824,7 @@ pub async fn start_proxy( } }); - (addr, ServerGuard { task }) + (addr, ServerGuard { task }, bucket_store) } /// Connect an HTTP/2 mTLS client to the proxy. Returns a `SendRequest` handle. diff --git a/tests/e2e.rs b/tests/e2e.rs index 46dca67..55b6ce3 100644 --- a/tests/e2e.rs +++ b/tests/e2e.rs @@ -523,3 +523,108 @@ async fn mtls_accepts_untrusted_ca_but_policy_denies_unregistered_key() { ); policy.cleanup().await; } + +#[tokio::test] +async fn tunnel_closes_on_bucket_drain_no_overshoot() { + let _guard = serial_test_lock().await; + let log = init_tracing_capture(); + drain_events(&log); + + let (echo_addr, _echo_guard) = start_echo_server().await; + let dest = format!("127.0.0.1:{}", echo_addr.port()); + + let subject = unique_test_identity("agent-alpha"); + let pki = TestPki::new(&subject); + let registry = common::TestAuthzRegistry::new().await; + registry + .allow_with_limits_for_pki(&pki, &subject, &dest, 1024, 0) + .await; + let engine = registry.engine(EXT_OID).await; + let (proxy_addr, _proxy_guard) = start_proxy(&pki, engine).await; + + let mut send_req = common::connect_client(proxy_addr, &pki).await; + let req = Request::connect(&dest) + .body(Empty::::new()) + .unwrap(); + let resp = send_req.send_request(req).await.unwrap(); + assert_eq!(resp.status(), 200); + + let upgraded = hyper::upgrade::on(resp).await.unwrap(); + let mut io = hyper_util::rt::TokioIo::new(upgraded); + + // Try to push 4 KiB through a tunnel whose bucket holds 1 KiB. + // write_all will fail mid-stream once the bucket drains; we do not + // assert on its result, only on the bytes that actually arrived + // at the destination-side echo server (echoed back as the tunnel + // download). The echo server bounces bytes back as we send them, + // so whatever we read here is what got through the metered side. + let payload = vec![0u8; 4096]; + let _ = io.write_all(&payload).await; + + let mut received_back = Vec::new(); + let _ = io.read_to_end(&mut received_back).await; + + assert!( + received_back.len() <= 1024, + "echo received {} bytes; bucket capacity was 1024 \u{2014} overshoot detected!", + received_back.len() + ); + + let events = wait_for_event(&log, "tunnel error", EVENT_TIMEOUT).await; + let err_evt = find_event(&events, "tunnel error").expect("expected a tunnel error event"); + let err_msg = err_evt + .fields + .get("error") + .map_or("", String::as_str); + assert!( + err_msg.contains("rate_limit_exceeded"), + "tunnel error message should mention rate_limit_exceeded; got {err_msg:?}" + ); + + registry.cleanup().await; +} + +#[tokio::test] +async fn subsequent_connect_after_drain_returns_429() { + let _guard = serial_test_lock().await; + let log = init_tracing_capture(); + drain_events(&log); + + let (echo_addr, _echo_guard) = start_echo_server().await; + let dest = format!("127.0.0.1:{}", echo_addr.port()); + + let subject = unique_test_identity("agent-alpha"); + let pki = TestPki::new(&subject); + let registry = common::TestAuthzRegistry::new().await; + registry + .allow_with_limits_for_pki(&pki, &subject, &dest, 512, 0) + .await; + let engine = registry.engine(EXT_OID).await; + let (proxy_addr, _proxy_guard) = start_proxy(&pki, engine).await; + + // First CONNECT: drain the bucket through the tunnel. + let mut send_req = common::connect_client(proxy_addr, &pki).await; + let req = Request::connect(&dest) + .body(Empty::::new()) + .unwrap(); + let resp = send_req.send_request(req).await.unwrap(); + assert_eq!(resp.status(), 200); + + let upgraded = hyper::upgrade::on(resp).await.unwrap(); + let mut io = hyper_util::rt::TokioIo::new(upgraded); + let _ = io.write_all(&[0u8; 2048]).await; + drop(io); + + // Wait for the proxy to register the close. + let _ = wait_for_event(&log, "tunnel error", EVENT_TIMEOUT).await; + + // Second CONNECT: should return 429. + let mut send_req2 = common::connect_client(proxy_addr, &pki).await; + let req2 = Request::connect(&dest) + .body(Empty::::new()) + .unwrap(); + let resp2 = send_req2.send_request(req2).await.unwrap(); + assert_eq!(resp2.status(), 429, "second CONNECT should be rate-limited"); + + registry.cleanup().await; +} diff --git a/tests/integration.rs b/tests/integration.rs index d47c0c1..06adb46 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -541,3 +541,172 @@ async fn proxy_dest_ipv6_matches_policy() { assert_deny(&eval(engine.as_ref(), &pki, "[::1]:443").await); registry.cleanup().await; } + +#[tokio::test] +async fn policy_rejects_tampered_capacity_column() { + let subject = unique_test_identity("agent-alpha"); + let registry = TestAuthzRegistry::new().await; + let pki = TestPki::new(&subject); + let seeded = registry + .allow_with_limits_for_pki(&pki, &subject, "api.example.com:443", 1000, 100) + .await; + + sqlx::query!( + "UPDATE permission_registry SET capacity_bytes = capacity_bytes * 10 WHERE permission_id = $1", + &seeded.permission_id + ) + .execute(®istry.pool) + .await + .expect("tamper capacity_bytes"); + + let engine = registry.engine(EXT_OID).await; + let ctx = agent_gateway::policy::RequestContext { + peer_certificates: pki.client_cert_chain(), + destination: "api.example.com:443".into(), + }; + match engine.evaluate(&ctx).await { + PolicyDecision::Deny { reason, .. } => { + assert!( + reason.contains("invalid permission signature"), + "expected signature verification failure, got: {reason}" + ); + } + PolicyDecision::Allow { .. } => { + panic!("tampered row should not verify, but policy allowed it") + } + } + registry.cleanup().await; +} + +#[tokio::test] +async fn policy_with_rate_limit_exhausts_then_denies() { + let _guard = common::serial_test_lock().await; + let log = common::init_tracing_capture(); + common::drain_events(&log); + + let subject = unique_test_identity("agent-alpha"); + let registry = TestAuthzRegistry::new().await; + let pki = TestPki::new(&subject); + let (echo_addr, _echo_guard) = common::start_echo_server().await; + let dest = format!("127.0.0.1:{}", echo_addr.port()); + + let seeded = registry + .allow_with_limits_for_pki(&pki, &subject, &dest, 100, 0) + .await; + + let engine = registry.engine(EXT_OID).await; + let (proxy_addr, _proxy_guard, store) = + common::start_proxy_with_store(&pki, engine).await; + + // First CONNECT: should succeed because the bucket starts at full capacity. + let mut send_req = common::connect_client(proxy_addr, &pki).await; + let req = hyper::Request::connect(&dest) + .body(http_body_util::Empty::::new()) + .unwrap(); + let resp = send_req.send_request(req).await.unwrap(); + assert_eq!(resp.status(), 200, "first CONNECT should be allowed"); + drop(resp); + + // Drain the bucket directly via the test-held store. Mid-tunnel + // enforcement lands in the next commit; at this stage the tunnel + // writes do not yet flow through MeteredStream, so we drain + // explicitly. get_or_create returns the existing bucket because + // the proxy already created it during the first CONNECT. + let bucket = store.get_or_create(&seeded.permission_id, 100, 0, seeded.not_after); + assert_eq!( + bucket.try_consume_up_to(100), + 100, + "drain should consume exactly the bucket's capacity" + ); + + // Second CONNECT on a fresh client: bucket is now empty, expect 429. + let mut send_req2 = common::connect_client(proxy_addr, &pki).await; + let req2 = hyper::Request::connect(&dest) + .body(http_body_util::Empty::::new()) + .unwrap(); + let resp2 = send_req2.send_request(req2).await.unwrap(); + assert_eq!( + resp2.status(), + 429, + "second CONNECT should be rate-limited" + ); + + registry.cleanup().await; +} + +#[tokio::test] +async fn tunnel_closes_when_permission_revoked() { + let _guard = common::serial_test_lock().await; + let log = common::init_tracing_capture(); + common::drain_events(&log); + + let (echo_addr, _echo_guard) = common::start_echo_server().await; + let dest = format!("127.0.0.1:{}", echo_addr.port()); + + let subject = unique_test_identity("agent-alpha"); + let pki = TestPki::new(&subject); + let registry = TestAuthzRegistry::new().await; + let seeded = registry + .allow_with_limits_for_pki(&pki, &subject, &dest, 1_000_000, 1_000_000) + .await; + let engine = registry.engine(EXT_OID).await; + let (proxy_addr, _proxy_guard, bucket_store) = + common::start_proxy_with_store(&pki, engine).await; + + let mut send_req = common::connect_client(proxy_addr, &pki).await; + let req = hyper::Request::connect(&dest) + .body(http_body_util::Empty::::new()) + .unwrap(); + let resp = send_req.send_request(req).await.unwrap(); + assert_eq!(resp.status(), 200); + + let upgraded = hyper::upgrade::on(resp).await.unwrap(); + let mut io = hyper_util::rt::TokioIo::new(upgraded); + + // Confirm the tunnel works for a small write before revocation. + tokio::io::AsyncWriteExt::write_all(&mut io, b"hello") + .await + .expect("first write should succeed"); + + // Revoke the permission via the same SQL the registry-cli would issue. + sqlx::query!( + "UPDATE permission_registry SET revoked_at = now() WHERE permission_id = $1", + &seeded.permission_id + ) + .execute(®istry.pool) + .await + .expect("revoke permission"); + + // Simulate the revocation-poll task's work directly. Keeps the test + // deterministic without waiting on the 30-second timer. + let revoked_ids: Vec = sqlx::query_scalar!( + "SELECT permission_id FROM permission_registry WHERE revoked_at IS NOT NULL" + ) + .fetch_all(®istry.pool) + .await + .expect("list revoked"); + assert!( + revoked_ids.contains(&seeded.permission_id), + "revoked list should include our permission" + ); + for id in &revoked_ids { + bucket_store.mark_revoked(id); + } + + // The next write should fail because the bucket is now dead. We loop + // a small number of times because the tunnel write path may buffer + // a tiny amount in the h2 layer before the error surfaces. + let mut closed = false; + for _ in 0..16 { + if tokio::io::AsyncWriteExt::write_all(&mut io, b"more data") + .await + .is_err() + { + closed = true; + break; + } + } + assert!(closed, "write after revocation should eventually fail"); + + registry.cleanup().await; +}