Add per-identity rate limit on data proxied (with revocation propagation) - #6
Closed
jude-sph wants to merge 7 commits into
Closed
Add per-identity rate limit on data proxied (with revocation propagation)#6jude-sph wants to merge 7 commits into
jude-sph wants to merge 7 commits into
Conversation
…registry Adds the two byte-budget columns that will be signed in v2 of the canonical permission bytes. Bumps the schema version to 2 and updates the verifier to emit and check the new format. Test helpers are migrated to v2 with effectively-unlimited defaults (i64::MAX) so the existing test suite keeps passing without per-test rewrites. PolicyDecision::Allow now carries an AllowedPermission struct bundling the matched row's permission_id, capacity_bytes, refill_bytes_per_sec, and not_after. Subsequent commits use these to set up the rate-limit bucket. Adds a tampering test that proves the v2 signature catches unsigned changes to capacity_bytes.
Adds a new module exposing TokenBucket, BucketStore, and MeteredStream. The bucket holds capacity_bytes, refill_bytes_per_sec, and the matched permission's not_after as immutable u64/DateTime fields. try_consume_up_to is atomic and short-circuits to zero when the permission has expired, making the bucket the natural mid-tunnel kill switch. MeteredStream wraps an AsyncWrite, reserving tokens before each inner write and refunding any tail the inner stream did not accept. This makes overshoot impossible: tokens consumed equals bytes the inner TCP stack confirmed accepted. Reads are delegated straight through; only the upstream egress direction is metered. Unit tests cover the bucket math, concurrent consume (no overshoot), refund semantics on partial writes, retry-after computation, expiry, the store's get-or-create contract, and the algorithmic no-overshoot property of MeteredStream against an in-memory duplex pipe. Mutex locks use the soft-recovery pattern (unwrap_or_else with PoisonError::into_inner) so a panicking task cannot poison the bucket state and bring down unrelated tunnels.
ProxyService now consults an in-process BucketStore (keyed by permission_id) before allowing each CONNECT. If the bucket for the matched permission is empty, return 429 Too Many Requests with a Retry-After header computed from the bucket's refill rate. The bucket is created lazily on first use with the signed capacity, refill rate, and not_after from the permission row. MakeProxyService accepts the BucketStore at construction; main.rs instantiates it alongside the policy engine. The test harness exposes start_proxy_with_store so integration tests can hold a reference to the same store the proxy uses, which lets them drain or revoke buckets without going through the tunnel. Integration test seeds a v2 row with capacity 100, refill 0, exhausts the bucket via the test-held store, and asserts the next CONNECT returns 429.
spawn_tunnel now wraps the upstream TCP socket in a MeteredStream that reserves bytes from the per-permission bucket before each inner write and refunds anything the inner stream did not accept. When the bucket drains mid-stream, the next write returns io::Error and the tunnel collapses through the existing copy_bidirectional error path. No new control flow added. Two end-to-end tests prove the property: the first writes 4 KiB through a tunnel with a 1 KiB bucket and asserts the destination-side echo server received at most 1 KiB (the no-overshoot proof), then inspects the tunnel-error structured log to confirm the rate_limit error string. The second drains the bucket through the tunnel and asserts a fresh CONNECT returns 429.
TokenBucket gains a revoked AtomicBool, set once via mark_revoked. The bucket's is_dead helper now returns true for any of: tokens exhausted, permission expired, or revoked flag set. All three flow through the same MeteredStream::poll_write error path, so there is no special-cased control flow per reason; the bucket is the unified mid-tunnel kill switch. A small background task in main.rs polls permission_registry every 30 seconds for rows with revoked_at set and calls mark_revoked on the matching bucket. The query is read-only and respects the gateway's read-only database role. The query function is factored out so integration tests can invoke it directly without waiting on the timer, which keeps the test deterministic. Memory ordering on the revoked flag is Release on store / Acquire on load, which makes the revocation visible to every subsequent is_dead check across threads. Adds unit tests for the revoked-bucket behaviour, the BucketStore::mark_revoked lookup, and idempotency. Adds an integration test that revokes a permission mid-tunnel and asserts the next write fails.
…dependency Updates the README to cover the new 429 status code with the Retry-After header, the two signed byte-budget fields in the v2 canonical permission format, and the kill-switch behaviour now shared between cap exhaustion, permission expiry, and explicit revocation. Adds a section explaining the 30-second revocation poll and the bounded latency it implies. Adds a note about the parallel update needed in the sister demo repo's registry-cli.
The previous phrasing said the agent "cannot exfiltrate faster than capacity_bytes + refill_bytes_per_sec × 30s after a revocation" -- strictly the units don't match (faster is a rate, not a quantity). Rephrased to make the bound a transferred-bytes quantity and to be explicit that the 30s is the upper bound on poll-to-mark latency, not a fixed latency.
Author
|
Sorry, PR'd here instead of my fork by accident |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add per-identity rate limit on data proxied (with revocation propagation)
What this adds
A per-identity, byte-denominated rate limit on data proxied through the
gateway. Each permission row carries a signed
capacity_bytesandrefill_bytes_per_sec; the gateway enforces them both at CONNECT time(refuse with 429) and continuously during the tunnel (drain → close).
While building the bucket infrastructure I realised the same mechanism
naturally addresses two existing gaps: permission expiry and explicit
revocation today only affect new CONNECTs, not in-flight tunnels.
The bucket now also closes the tunnel when
not_afterpasses (inline,on every write) or when a background poll learns the permission has
been revoked (within 30 seconds).
Design choices and tradeoffs
gateway cannot raise them. (See
policy.rs::canonical_permission_bytes,bumped to v2.)
framing. Doubles the design surface for marginal additional
protection if we counted downloads.
verification; the demo repo's
registry-cliwill need a smallparallel update (see note at the end). Simpler than a v1/v2
conditional verifier.
signed
not_afterbounds total harm per permission at(not_after - not_before) × refill + capacity.permission_id. Reissuance is the explicitreset mechanism. Bucket parameters are immutable after creation —
in-place permission updates do not propagate to existing buckets
until restart, by design.
MeteredStreamreserves tokens beforethe inner write and refunds whatever the inner stream did not
accept. Tokens consumed = bytes the TCP stack confirmed accepted.
Verified by a unit test (100 threads racing on a 1 KiB bucket sum to
exactly 1000 bytes) and an end-to-end test (echo server asserts ≤
capacity bytes received).
std::sync::Mutexand
AtomicBool. Consistent with the existing pattern (X.509parsing, canonical bytes, hex encoding all hand-rolled).
Known limitations (also in README)
not_after.permission_ids ever seen.limit continues to enforce throughout the window, so total transfer
between revocation and the next poll is capped at one burst plus 30
seconds of sustained throughput at the signed rate.
Required parallel change in agent-gateway-demo
The
registry-cli/agent-permissions.shscript inSL5TaskForce/agent-gateway-democurrently signs v1 canonical bytes.After this PR is merged, two extra lines need to be added to the
<<EOFblock incmd_grant:with corresponding
--set=flags on the psql invocation that insertsthe row, and two additional CLI parameters surfacing the values.
Without this change, the demo's permissions will fail verification
against this gateway.
Verification
cargo build --release— greencargo clippy --all-targets --all-features -- -D warnings— greencargo test— all 80 tests pass, including:src/rate_limit.rsfor bucket math, refund,concurrent consume (no overshoot), retry-after, expiry, and
revocation (14 tests).
policy_rejects_tampered_capacity_column(proves the signed-not-config security property).
policy_with_rate_limit_exhausts_then_denies.tunnel_closes_when_permission_revoked.tunnel_closes_on_bucket_drain_no_overshoot(thealgorithmic no-overshoot proof at the IO-wrapper layer).
subsequent_connect_after_drain_returns_429.