Add per-identity rate limit on data proxied (with revocation propagation)

README.md

@@ -89,18 +89,31 @@ 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
 subject_public_key_spki_der=3059301306072a8648ce3d020106082a8648ce3d03010703420004...
 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).

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);

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;

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<BucketStore>) {
+    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<Vec<String>> {
+    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,

src/policy.rs

@@ -22,13 +22,26 @@ pub struct RequestContext {
 pub enum PolicyDecision {
     Allow {
         source_identity: String,
+        permission: AllowedPermission,
     },
     Deny {
         source_identity: Option<String>,
         reason: String,
     },
 }
 
+/// 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<chrono::Utc>,
+}
+
 #[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<Arc<dyn Polic
     )?))
 }
 
+/// Build a Postgres pool with the same configuration the policy engine uses.
+/// Used by the gateway's revocation poll task to issue read-only queries
+/// against `permission_registry`.
+///
+/// # Errors
+///
+/// Returns an error if the database URL is invalid or the database cannot
+/// be reached.
+pub async fn build_pool(config: &PolicyConfig) -> anyhow::Result<PgPool> {
+    build_pg_pool(config).await
+}
+
 async fn build_pg_pool(policy: &PolicyConfig) -> anyhow::Result<PgPool> {
     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<u8> {
     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<u8> {
         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"
         );
     }
 }