Skip to content

feat(proxy): inject session ID and sequence number headers on matching requests#198

Open
SasSwart wants to merge 7 commits intomainfrom
sasswart/feat-proxy-inject-session-headers
Open

feat(proxy): inject session ID and sequence number headers on matching requests#198
SasSwart wants to merge 7 commits intomainfrom
sasswart/feat-proxy-inject-session-headers

Conversation

@SasSwart
Copy link
Copy Markdown
Contributor

@SasSwart SasSwart commented Apr 30, 2026

PR Map

  1. feat(audit): session ID generation, sequence counter, and BoundaryLog wiring #196
  2. feat(config): session correlation header injection configuration #197
  3. 👉🏼 feat(proxy): inject session ID and sequence number headers on matching requests #198

RFC: Bridge ↔ Boundaries Correlation

When session correlation is enabled and the outgoing request matches a configured inject target, set X-Coder-Agent-Firewall-Session-Id and X-Coder-Agent-Firewall-Sequence-Number headers on the forwarded request. Any values the jailed client may have set are overwritten so the upstream always sees boundary's authoritative session ID and sequence number.

Depends on #197.

Changes

  • audit/request.go: Add SequenceNumber int32 field to Request. The proxy assigns this value before auditing so the audit log and the injected HTTP header carry the same sequence number.
  • audit/socket_auditor.go: Remove the seq *SequenceCounter field. The auditor now reads req.SequenceNumber directly instead of generating its own value. NewSocketAuditor no longer accepts a *SequenceCounter parameter.
  • audit/multi_auditor.go: Remove SequenceCounter creation in SetupAuditor; call NewSocketAuditor with three arguments (logger, socket path, session ID).
  • audit/socket_auditor_test.go: Remove seq field from test auditor structs. Tests that exercise sequence numbers now set SequenceNumber on the Request directly.
  • cli/cli.go: Minor punctuation fix in the --session-id-inject-target flag description.
  • proxy/proxy.go:
    • Add sessionCorrelation, sessionID, seqCounter, forwardTransport fields to Server.
    • Add SessionCorrelation, SessionID, ForwardTransport to Config (the sequence counter is an internal value-type field on Server, not user-configurable).
    • New shouldInjectHeaders(host, path) method: matches request host (case-insensitive, port-stripped) and optional path glob against configured InjectTargets.
    • processHTTPRequest allocates a sequence number from p.seqCounter and passes it to both AuditRequest and forwardRequest.
    • forwardRequest stamps both headers on matching requests after copying client headers (overwriting any client-supplied values). Uses ForwardTransport when set.
  • proxy/proxy_audit_test.go: New test TestSequenceNumberIncrementsAcrossRequestTypes — verifies sequence numbers increment correctly across plain HTTP, implicit-CONNECT HTTPS, and explicit-CONNECT tunnel requests.
  • proxy/proxy_framework_test.go: Add WithSessionCorrelation, WithSessionID, WithForwardTransport test options; wire them into the proxy Config.
  • proxy/proxy_session_correlation_test.go: New test file with seven tests:
    • TestSessionCorrelation_MatchedDomain: headers injected on matching domain.
    • TestSessionCorrelation_MatchedDomainWithPort: headers injected when inject target domain includes a port.
    • TestSessionCorrelation_UnmatchedDomain: headers absent on non-matching domain.
    • TestSessionCorrelation_Disabled: headers absent when correlation is disabled.
    • TestSessionCorrelation_OverwritesClientValue: proxy overwrites spoofed client headers.
    • TestSessionCorrelation_PathMatching: headers injected only when path glob matches.
    • TestSessionCorrelation_SequenceNumberIncrements: sequence numbers increment across requests.
Implementation notes
  • The sequence number is allocated once in processHTTPRequest via p.seqCounter.Next(), then threaded through to both the auditor and the header injection site. This guarantees the audit event and the HTTP header always agree on the sequence number for a given request.
  • shouldInjectHeaders strips the port from both the request host and the configured target domain using net.SplitHostPort before comparison, and uses path.Match for glob matching, consistent with the RFC's path=/api/v2/aibridge/* syntax.
  • The SequenceCounter on Server is a value type (audit.SequenceCounter) backed by an atomic.Int32. It is always present and starts at 0 — sequence numbers are allocated for every request regardless of whether session correlation is enabled.

Note

This PR was authored by Coder Agents.

@SasSwart SasSwart marked this pull request as draft April 30, 2026 13:00
@SasSwart SasSwart force-pushed the sasswart/feat-config-session-correlation-header-injection branch from c24a0fd to c08f291 Compare May 1, 2026 10:51
@SasSwart SasSwart force-pushed the sasswart/feat-config-session-correlation-header-injection branch from c08f291 to 3ea7428 Compare May 5, 2026 09:20
@SasSwart SasSwart force-pushed the sasswart/feat-proxy-inject-session-headers branch from 6e61348 to d8635fb Compare May 5, 2026 10:53
@SasSwart SasSwart marked this pull request as ready for review May 5, 2026 10:56
Base automatically changed from sasswart/feat-config-session-correlation-header-injection to main May 7, 2026 14:04
SasSwart added 2 commits May 7, 2026 14:24
…g requests

When session correlation is enabled and the outgoing request matches a
configured inject target, set X-Coder-Agent-Firewall-Session-Id and
X-Coder-Agent-Firewall-Sequence-Number headers on the forwarded request.
Any values the jailed client may have set are overwritten so the upstream
always sees boundary's authoritative session ID and sequence number.

The sequence number is pre-allocated before the audit event so both the
audit log and the injected header carry the same value. audit.Request
gains a SequenceNumber pointer field; when non-nil the socket auditor
uses it instead of calling its own counter.

New proxy.Config fields: SessionCorrelation, SessionID, SequenceCounter.
New Server method: shouldInjectHeaders (domain + optional path glob
matching).

Tests cover matched domain, unmatched domain, disabled injection,
client-supplied header overwrite, path glob matching, and sequence
number incrementing.
Updated the SequenceNumber field in the audit.Request struct and related handling in the SocketAuditor and proxy components to use int32 instead of uint64. This change ensures consistency in data types across the application. Additionally, minor adjustments were made to the session correlation tests to reflect the updated header names and improve clarity in assertions.
@SasSwart SasSwart force-pushed the sasswart/feat-proxy-inject-session-headers branch from d8635fb to 79d6485 Compare May 7, 2026 14:33
SasSwart added 5 commits May 7, 2026 14:57
…ber handling

Eliminated the SequenceCounter from the SocketAuditor and related components, simplifying the sequence number management. The SequenceNumber field in the audit.Request struct is now a non-pointer int32, ensuring consistent handling across the application. Adjusted tests to reflect these changes and maintain functionality.
…comments

Updated the host port stripping logic in the shouldInjectHeaders function to utilize net.SplitHostPort for improved clarity and reliability. Removed outdated comments regarding session correlation header injection to enhance code readability.
…nd transport configuration

Added a new test to verify that the sequence number increments correctly across different request types (HTTP and HTTPS) in the proxy. Introduced a `WithForwardTransport` option to allow tests to specify a custom HTTP transport for handling self-signed certificates. Updated the proxy server configuration to support this transport, ensuring proper handling of requests with session correlation and sequence number injection.
@SasSwart SasSwart requested a review from johnstcn May 7, 2026 15:35
@SasSwart
Copy link
Copy Markdown
Contributor Author

SasSwart commented May 7, 2026

I've self-reviewed this agent generated PR. Apart from double-checking TestSequenceNumberIncrementsAcrossRequestTypes, I can vouch for everything in this PR. Its ready for external review.

SasSwart added a commit that referenced this pull request May 7, 2026
…er agreement

Add integration tests that verify the core invariants of session
correlation across the proxy, auditor, and forwarded request headers
working together. These tests fill the gap identified during review
of the session correlation PR stack (#196, #197, #198) where unit
tests verified each component in isolation but did not verify them
in concert.

New test file: proxy/proxy_session_correlation_integration_test.go

Tests added:
- LLMRequestAuditAndHeadersAgree: audit sequence number matches
  the forwarded header value on inject-target requests.
- NonLLMRequestAuditedWithoutHeaders: allowed non-inject-target
  requests are audited but carry no correlation headers.
- DeniedRequestAuditedNeverForwarded: denied requests consume a
  sequence number but are never forwarded.
- MixedRequestsSequenceOrdering: interleaved LLM, non-LLM, and
  denied requests all advance the counter monotonically.
- SequenceGapRevealsAgenticLoop: gap between two LLM sequence
  numbers precisely equals intermediate tool-use requests.
- SpoofedHeadersOverwrittenWithCorrectSequence: client-supplied
  headers are replaced and the audit event still agrees.
- DisabledCorrelationNoHeadersNoPreallocatedSequence: disabled
  correlation means no headers and no pre-allocated sequence.
- ConcurrentRequestsUniqueSequenceNumbers: concurrent requests
  each get a unique, dense sequence number.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant