Skip to content

Add opt-in per-request target-host retargeting to the reflector - #118

Open
aszarama wants to merge 2 commits into
fix/reflector-race-and-key-stabilityfrom
feat/reflector-dynamic-target-host
Open

Add opt-in per-request target-host retargeting to the reflector#118
aszarama wants to merge 2 commits into
fix/reflector-race-and-key-stabilityfrom
feat/reflector-dynamic-target-host

Conversation

@aszarama

@aszarama aszarama commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Lets one proxy entry forward to many upstream hosts, chosen per request by an X-Cortex-Relay-Target-Host header and validated against a per-entry allowlist.

Stacked on #117 — review that one first. This PR targets fix/reflector-race-and-key-stability, so the diff here is feature-only; retarget to main once #117 merges.

Accept files switch it on with a dynamicTargetHosts key; the reflector mechanism and that wiring are both in this PR, so it is reviewable end to end.

Why

Routing GCP traffic through the relay means one agent must reach ~25 *.googleapis.com hosts. Enumerating them as static origins does not work once Vertex is in scope: every Vertex client sets a zone-derived endpoint (us-central1-aiplatform.googleapis.com), so the host set is unbounded in the zone dimension rather than a fixed list. A wildcard allowlist is required, not merely more general.

Today each rule bakes in a single fixed origin at render time. NewSingleHostReverseProxy's Director is free to rewrite the destination per request, and the shared http.Transport pools connections per host, so nothing downstream cares that consecutive requests through one entry land on different hosts.

Design

Opt-in per entry. An entry with no patterns ignores the header entirely and passes it through untouched. An ambient implementation would retarget every rule's traffic — including the injected /__axon/* route and the default registration entry — so the opt-in is a safety property, not ergonomics.

Allowlist check in ServeHTTP, not the Director. A Director cannot cleanly abort a request. The check runs after entry resolution and before any forwarding; a rejected host gets a 403 and a log naming it. It never fails open.

The validated host travels in the request context. ServeHTTP stashes it, the Director reads it. A Director rewrite is therefore unreachable unless the allowlist check ran first — the ordering is enforced by the data flow rather than by convention. The routing header is stripped before the request goes upstream.

WebSocket upgrades are refused with a 400 when the header is present. The upgrade branch uses entry.TargetURI directly and bypasses the Director, so a retarget would silently not apply. GCP needs no WS; better to reject loudly than to forward to the wrong host.

Wildcard matching is label-aligned. *.googleapis.com requires at least one leading label on a dot boundary, so it admits compute.googleapis.com and us-central1-aiplatform.googleapis.com but rejects googleapis.com, evilgoogleapis.com and compute.googleapis.com.evil.com. Hosts are lowercased, an optional numeric port is allowed, and anything URL-ish (/, @, ?, #, %, backslash, whitespace) is rejected outright.

key() gains a dynamicTargets suffix so an opt-in entry hashes distinctly from a plain entry on the same origin.

Note key() hashes each header's resolver key — the raw accept-file string such as Bearer ${plugin:gcp-token} — not its resolved value, so a rotating token does not churn keys or spawn duplicate entries. Load-bearing, and easy to break by accident.

Tests

Six new tests, all under -race:

  • retarget across two hosts on one entry, asserting rule headers are still injected and the routing header never reaches the upstream
  • fail-closed on a disallowed host, asserting the backend is never hit
  • header ignored without opt-in, asserting it passes through untouched
  • WebSocket upgrade refused
  • an 18-case resolveDynamicHost table covering the dot-boundary and shape violations above
  • 400 concurrent retargets alternating across two backends

make test-race gates 41 reflector tests, 0 races. Full suite green without -race; go vet and gofmt clean.


Wiring it to the accept file

The mechanism above is reachable from configuration. A private rule declares its own allowlist:

{
  "method": "any",
  "path": "/*",
  "origin": "https://cloudresourcemanager.googleapis.com",
  "headers": { "Authorization": "Bearer ${plugin:gcp-token}" },
  "dynamicTargetHosts": ["*.googleapis.com"]
}

A rule that omits the key cannot be retargeted, so the opt-in stays per rule rather than becoming ambient across the file.

Why a rule key rather than branching on integration type in Render. Retargeting is a general relay capability and GCP is only its first consumer. A rule key means the next integration writes an accept file; integration-type branching would mean editing Render every time.

The broker ignores the key, and that is load-bearing. Rules round-trip through map[string]any precisely so fields Axon does not model survive, and the broker applies no schema validation — its Rule interface is compile-time only and does not declare headers either, which has been shipping through this same path in production. A test pins the round-trip so it cannot regress silently.

Misconfiguration fails loudly. A malformed dynamicTargetHosts panics at render instead of being ignored — silently dropping a typo'd allowlist surfaces as an unexplained 403 on every request, which is much harder to diagnose than a startup failure. Declaring it while the reflector is disabled panics too, matching the existing behaviour for custom headers.

I extracted the render closure into relayInstanceManager.reflectorRenderStep so this can be tested against a real accept file and a real reflector without standing up the broker supervisor. That path was previously only covered by the full manager integration tests, which are among the slowest and flakiest in the package.

Also fixed a latent bug found on the way: AddRule returned a rule wrapper without its acceptFile back-reference, so any accessor that logs would nil-deref. Nothing used the return value in production, so it was unreachable — but the new tests do.

Verification

make test-race gates 45 reflector tests, 0 races, stable over 6 consecutive runs. Full suite green; go vet and gofmt clean.

Lets one proxy entry forward to many upstream hosts, selected per request
by the X-Cortex-Relay-Target-Host header and validated against a
per-entry allowlist. This is the agent-side half of routing GCP traffic
through the relay, where one agent must reach ~25 *.googleapis.com hosts
and the Vertex endpoints are zone-derived, so the host set is unbounded
in the zone dimension and cannot be enumerated as static origins.

Retargeting is opt-in per entry via WithDynamicTargetHosts. An entry with
no patterns ignores the header entirely and passes it through untouched,
so the injected /__axon/* route and the default registration entry are
unaffected.

The allowlist check runs in ServeHTTP, after entry resolution and before
any forwarding, because a Director cannot cleanly abort a request. A
rejected host gets a 403 and a log naming it - it never fails open. The
validated host then travels to the Director in the request context, so a
Director rewrite is unreachable unless the ServeHTTP check ran first.
The routing header is stripped before the request goes upstream.

WebSocket upgrades are refused with a 400 when the header is present,
because the upgrade branch uses entry.TargetURI directly and bypasses the
Director, so a retarget would silently not apply there.

Wildcard patterns match on a dot boundary and require at least one
leading label, so "*.googleapis.com" admits compute.googleapis.com but
rejects googleapis.com, evilgoogleapis.com and
compute.googleapis.com.evil.com. Hosts are lowercased, an optional
numeric port is allowed, and anything URL-ish is rejected outright.

key() gains a dynamicTargets suffix so an opt-in entry hashes distinctly
from a plain entry on the same origin.

Note: WithDynamicTargetHosts has no production callers yet. How the
accept file declares the opt-in and how relay_instance_manager.Render
threads it through is the follow-up that makes this reachable.
Makes per-request retargeting reachable from configuration. Previously
WithDynamicTargetHosts had no production callers, so the mechanism existed
but nothing could switch it on.

A private rule declares its allowlist with a dynamicTargetHosts key:

  {
    "method": "any",
    "path": "/*",
    "origin": "https://cloudresourcemanager.googleapis.com",
    "headers": { "Authorization": "Bearer ${plugin:gcp-token}" },
    "dynamicTargetHosts": ["*.googleapis.com"]
  }

A rule that omits the key cannot be retargeted at all, so this stays
opt-in per rule rather than becoming ambient across the accept file.

Chose a rule key over branching on integration type at render time,
because retargeting is a general relay capability and GCP is only its
first consumer - a rule key means the next integration writes an accept
file instead of editing Render.

The broker has no knowledge of the key and ignores it. Rules round-trip
through map[string]any specifically to preserve fields Axon does not
model, and the broker applies no schema validation to accept files - its
Rule interface is compile-time only and does not declare `headers`
either, which has been shipping through the same path in production.
A test pins the round-trip so this cannot regress silently.

A malformed value panics at render rather than being ignored. Silently
dropping a typo'd allowlist would surface as an unexplained 403 on every
request, which is far harder to diagnose than a startup failure. For the
same reason, declaring dynamicTargetHosts while the reflector is disabled
panics, matching the existing behaviour for custom headers.

Extracted the render closure into relayInstanceManager.reflectorRenderStep
so the wiring can be tested against a real accept file and a real
reflector without standing up the broker supervisor. The existing
coverage for this path only ran inside the full manager integration
tests, which are among the slowest and flakiest in the package.
@aszarama
aszarama force-pushed the feat/reflector-dynamic-target-host branch from 6ba4cc5 to d3bf1b9 Compare August 8, 2026 06:39
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