feat(mcp-sidecar): implement RSA-OAEP token exchange for dynamic token refresh#1275
feat(mcp-sidecar): implement RSA-OAEP token exchange for dynamic token refresh#1275markturansky merged 1 commit intoalphafrom
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (11)
💤 Files with no reviewable changes (1)
✅ Files skipped from review due to trivial changes (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughAdds a token-exchange subsystem and integrates dynamic token refreshing into the MCP client and mention resolver. MCP sidecar construction became session-aware (injects SESSION_ID). Config loading no longer hardcodes MCP API server URL; manifests removed explicit MCP_API_SERVER_URL. Several tests added/updated for new behavior and concurrency. Changes
Sequence Diagram(s)sequenceDiagram
participant App as Ambient MCP (main)
participant Exchanger as Token Exchanger
participant CP as Control Plane Token Server
participant Client as MCP Client
participant API as Ambient API Server
App->>Exchanger: New(tokenURL, pubKey, sessionID)
Exchanger->>CP: FetchToken() with RSA-OAEP(encrypted sessionID)
CP-->>Exchanger: { "token": access_token }
Exchanger->>Client: OnRefresh(callback) → Client.SetToken(access_token)
Exchanger->>Exchanger: StartBackgroundRefresh()
rect rgba(0, 150, 136, 0.5)
Note over Exchanger,CP: Background refresh loop
loop every refreshPeriod
Exchanger->>CP: FetchToken() with encrypted sessionID
CP-->>Exchanger: { "token": new_token }
Exchanger->>Client: callback(new_token) → Client.SetToken(new_token)
end
end
rect rgba(66, 133, 244, 0.5)
Note over Client,API: Authenticated requests use latest token
Client->>API: GET /... Authorization: Bearer <token>
API-->>Client: response
end
🚥 Pre-merge checks | ✅ 4 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
✨ Simplify code
Comment |
405beea to
3b8600a
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@components/ambient-mcp/tokenexchange/tokenexchange.go`:
- Around line 62-64: OnRefresh currently assigns e.onRefresh without
synchronization while FetchToken reads it unsafely; protect the onRefresh field
with the same mutex used elsewhere (e.g., e.mu) by either setting e.onRefresh
inside a lock in OnRefresh or copying the callback under lock in FetchToken
before invoking it. Locate the Exchanger.OnRefresh method and the
Exchanger.FetchToken method and ensure both reads and writes of the onRefresh
field are performed under the same mutex to eliminate the data race.
- Around line 104-123: StartBackgroundRefresh can spawn multiple goroutines and
Stop can panic on double-close of e.stopCh; modify Exchanger to add two
sync.Once fields (e.startOnce, e.stopOnce) and use startOnce.Do in
StartBackgroundRefresh to ensure only one goroutine is created, and use
stopOnce.Do in Stop to close e.stopCh safely; keep existing logic calling
e.FetchToken() and ticker cleanup but ensure stopCh is initialized when
Exchanger is constructed so the Once guards operate correctly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: fe92cb9b-c89b-4671-b73e-7991900dd72d
📒 Files selected for processing (11)
components/ambient-control-plane/internal/config/config.gocomponents/ambient-control-plane/internal/reconciler/kube_reconciler.gocomponents/ambient-mcp/client/client.gocomponents/ambient-mcp/client/client_test.gocomponents/ambient-mcp/main.gocomponents/ambient-mcp/mention/resolve.gocomponents/ambient-mcp/mention/resolve_test.gocomponents/ambient-mcp/tokenexchange/tokenexchange.gocomponents/ambient-mcp/tokenexchange/tokenexchange_test.gocomponents/ambient-mcp/tools/sessions.gocomponents/manifests/overlays/mpp-openshift/ambient-control-plane.yaml
💤 Files with no reviewable changes (1)
- components/manifests/overlays/mpp-openshift/ambient-control-plane.yaml
…n refresh Replace the static AMBIENT_TOKEN injection with the same RSA-OAEP token exchange protocol the runner uses, so the MCP sidecar can fetch and periodically refresh its own API token from the CP token server. - Add tokenexchange package: RSA-OAEP encrypt session ID, fetch token from CP /token endpoint with retry/backoff, background refresh every 5 minutes - Make client.Client thread-safe (sync.RWMutex) with SetToken() for hot-swapping tokens on refresh - Refactor mention.Resolver to use TokenFunc instead of static string so it always reads the current token - Inject SESSION_ID into MCP sidecar env in buildMCPSidecar() - Default MCP_API_SERVER_URL to AMBIENT_API_SERVER_URL so the sidecar inherits the correct API endpoint automatically - Remove hardcoded MCP_API_SERVER_URL from mpp-openshift overlay - Add unit tests for tokenexchange, client, and mention packages 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
3b8600a to
90b4ecb
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
components/ambient-mcp/tokenexchange/tokenexchange.go (2)
62-64:⚠️ Potential issue | 🔴 CriticalProtect
onRefreshwith the same mutex (current code races).Line 63 writes
e.onRefreshunsafely, while Lines 88-89 read/invoke it unsafely during concurrent refreshes. This is a real data race.Proposed fix
func (e *Exchanger) OnRefresh(fn func(string)) { + e.mu.Lock() + defer e.mu.Unlock() e.onRefresh = fn } @@ - e.mu.Lock() - e.currentToken = token - e.mu.Unlock() - - if e.onRefresh != nil { - e.onRefresh(token) - } + e.mu.Lock() + e.currentToken = token + callback := e.onRefresh + e.mu.Unlock() + + if callback != nil { + callback(token) + }As per coding guidelines, "Flag only errors, security risks, or functionality-breaking problems."
Also applies to: 84-90
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@components/ambient-mcp/tokenexchange/tokenexchange.go` around lines 62 - 64, The OnRefresh setter races with concurrent readers; protect writes to e.onRefresh by acquiring the same Exchanger mutex used when invoking it during refresh (i.e., use the Exchanger's mutex around OnRefresh assignment) and ensure the refresh path (where e.onRefresh is read/invoked in the refresh routine at lines ~84-90) also holds that same mutex while accessing e.onRefresh; update the OnRefresh method to lock/unlock the mutex and keep the existing read/invoke sites locked with the identical mutex to eliminate the data race.
104-123:⚠️ Potential issue | 🟠 MajorGuard lifecycle methods: multiple starts and double-stop are unsafe.
Line 105 starts a new goroutine every call, and Line 122 panics if
Stop()is called twice (closeon closed channel). Addsync.Onceguards.Proposed fix
type Exchanger struct { tokenURL string publicKey *rsa.PublicKey sessionID string httpClient *http.Client mu sync.RWMutex currentToken string onRefresh func(string) stopCh chan struct{} + startOnce sync.Once + stopOnce sync.Once } @@ func (e *Exchanger) StartBackgroundRefresh() { - go func() { + e.startOnce.Do(func() { + go func() { ticker := time.NewTicker(refreshPeriod) defer ticker.Stop() for { select { case <-ticker.C: if _, err := e.FetchToken(); err != nil { fmt.Fprintf(os.Stderr, "background token refresh failed: %v\n", err) } case <-e.stopCh: return } } - }() + }() + }) } func (e *Exchanger) Stop() { - close(e.stopCh) + e.stopOnce.Do(func() { + close(e.stopCh) + }) }As per coding guidelines, "Flag only errors, security risks, or functionality-breaking problems."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@components/ambient-mcp/tokenexchange/tokenexchange.go` around lines 104 - 123, The StartBackgroundRefresh and Stop methods are not guarded and can spawn multiple goroutines or panic on a double close of stopCh; add sync.Once guards: introduce startOnce and stopOnce fields on Exchanger, wrap the goroutine startup in startOnce.Do(...) inside StartBackgroundRefresh to ensure only one background goroutine is started, and wrap close(e.stopCh) in stopOnce.Do(func(){ close(e.stopCh) }) inside Stop to make closing idempotent; ensure stopCh is initialized before startOnce.Do is called (e.g., in constructor) so the goroutine/select uses a valid channel.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@components/ambient-mcp/client/client_test.go`:
- Around line 14-16: TestNew expects trailing slash to be stripped but
client.New currently stores baseURL as-is; update the constructor New(...) in
the client package to normalize the provided baseURL by trimming any trailing
slash before assigning to the baseURL field so BaseURL() returns the normalized
value, and keep BaseURL() behavior unchanged (or alternatively change the test
to expect the raw value) — locate the constructor New, the baseURL field, and
the BaseURL() method to implement the trim logic.
In `@components/ambient-mcp/tokenexchange/tokenexchange.go`:
- Around line 190-193: The token URL validation in tokenexchange.go currently
allows plaintext HTTP by checking parsed.Scheme and permitting "http" or
"https"; change the check to accept only "https" (i.e., require parsed.Scheme ==
"https"), update the error message to reflect that only HTTPS is allowed, and
ensure this logic is applied in the function where parsed.Scheme / scheme is
used for token endpoint validation so non‑TLS endpoints are rejected.
---
Duplicate comments:
In `@components/ambient-mcp/tokenexchange/tokenexchange.go`:
- Around line 62-64: The OnRefresh setter races with concurrent readers; protect
writes to e.onRefresh by acquiring the same Exchanger mutex used when invoking
it during refresh (i.e., use the Exchanger's mutex around OnRefresh assignment)
and ensure the refresh path (where e.onRefresh is read/invoked in the refresh
routine at lines ~84-90) also holds that same mutex while accessing e.onRefresh;
update the OnRefresh method to lock/unlock the mutex and keep the existing
read/invoke sites locked with the identical mutex to eliminate the data race.
- Around line 104-123: The StartBackgroundRefresh and Stop methods are not
guarded and can spawn multiple goroutines or panic on a double close of stopCh;
add sync.Once guards: introduce startOnce and stopOnce fields on Exchanger, wrap
the goroutine startup in startOnce.Do(...) inside StartBackgroundRefresh to
ensure only one background goroutine is started, and wrap close(e.stopCh) in
stopOnce.Do(func(){ close(e.stopCh) }) inside Stop to make closing idempotent;
ensure stopCh is initialized before startOnce.Do is called (e.g., in
constructor) so the goroutine/select uses a valid channel.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 3bd4301f-916f-4e8e-8894-aeaed2be2063
📒 Files selected for processing (11)
components/ambient-control-plane/internal/config/config.gocomponents/ambient-control-plane/internal/reconciler/kube_reconciler.gocomponents/ambient-mcp/client/client.gocomponents/ambient-mcp/client/client_test.gocomponents/ambient-mcp/main.gocomponents/ambient-mcp/mention/resolve.gocomponents/ambient-mcp/mention/resolve_test.gocomponents/ambient-mcp/tokenexchange/tokenexchange.gocomponents/ambient-mcp/tokenexchange/tokenexchange_test.gocomponents/ambient-mcp/tools/sessions.gocomponents/manifests/overlays/mpp-openshift/ambient-control-plane.yaml
💤 Files with no reviewable changes (1)
- components/manifests/overlays/mpp-openshift/ambient-control-plane.yaml
✅ Files skipped from review due to trivial changes (1)
- components/ambient-mcp/tokenexchange/tokenexchange_test.go
🚧 Files skipped from review as they are similar to previous changes (5)
- components/ambient-mcp/tools/sessions.go
- components/ambient-mcp/client/client.go
- components/ambient-mcp/main.go
- components/ambient-control-plane/internal/reconciler/kube_reconciler.go
- components/ambient-mcp/mention/resolve.go
| scheme := strings.ToLower(parsed.Scheme) | ||
| if scheme != "http" && scheme != "https" { | ||
| return fmt.Errorf("invalid token URL scheme %q (must be http or https)", scheme) | ||
| } |
There was a problem hiding this comment.
Token endpoint validation should not accept plaintext http.
Line 191 currently permits http. Since this endpoint returns the API token, accepting non-TLS transport exposes token theft risk in transit.
Proposed fix
scheme := strings.ToLower(parsed.Scheme)
- if scheme != "http" && scheme != "https" {
- return fmt.Errorf("invalid token URL scheme %q (must be http or https)", scheme)
+ if scheme != "https" {
+ return fmt.Errorf("invalid token URL scheme %q (must be https)", scheme)
}As per coding guidelines, "Flag only errors, security risks, or functionality-breaking problems."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@components/ambient-mcp/tokenexchange/tokenexchange.go` around lines 190 -
193, The token URL validation in tokenexchange.go currently allows plaintext
HTTP by checking parsed.Scheme and permitting "http" or "https"; change the
check to accept only "https" (i.e., require parsed.Scheme == "https"), update
the error message to reflect that only HTTPS is allowed, and ensure this logic
is applied in the function where parsed.Scheme / scheme is used for token
endpoint validation so non‑TLS endpoints are rejected.
There was a problem hiding this comment.
Intentional — this is in-cluster pod-to-pod traffic over Kubernetes service DNS (e.g. http://ambient-control-plane.ns.svc:8080/token). The runner's Python implementation also permits HTTP for the same reason (_grpc_client.py:52-63). Requiring HTTPS would break all in-cluster deployments where TLS termination is handled at the service mesh / ingress layer.
Summary
AMBIENT_TOKENinjection with the same RSA-OAEP token exchange protocol the runner uses — the MCP sidecar now fetches and periodically refreshes its own API token from the CP token servertokenexchangepackage toambient-mcpwith RSA-OAEP encryption, retry/backoff, and background refresh (every 5 min)client.Clientthread-safe (sync.RWMutex+SetToken()) so tokens can be hot-swapped on refreshmention.Resolverto useTokenFuncinstead of a static stringSESSION_IDinto MCP sidecar env vars inbuildMCPSidecar()MCP_API_SERVER_URLtoAMBIENT_API_SERVER_URLand remove the hardcodedruntime-intvalue from the mpp-openshift overlayChanges
< /dev/null | File | What |
|------|------|
|
components/ambient-mcp/tokenexchange/tokenexchange.go| New: RSA-OAEP encrypt + fetch + background refresh ||
components/ambient-mcp/client/client.go| Thread-safe token withSetToken()||
components/ambient-mcp/main.go| Bootstrap via CP token exchange, fallback to staticAMBIENT_TOKEN||
components/ambient-mcp/mention/resolve.go|TokenFuncinstead of static token ||
components/ambient-mcp/tools/sessions.go| Passc.TokenasTokenFunc||
components/ambient-control-plane/internal/reconciler/kube_reconciler.go|buildMCPSidecar(sessionID), injectSESSION_ID, remove static token ||
components/ambient-control-plane/internal/config/config.go|MCPAPIServerURLdefaults toAPIServerURL||
components/manifests/overlays/mpp-openshift/ambient-control-plane.yaml| Remove hardcodedMCP_API_SERVER_URL|Test plan
tokenexchange(18 tests): RSA-OAEP roundtrip, key parsing, URL validation, fetch success/retry/failure, callback, wrong-key rejectionclient(7 tests): construction,SetToken, concurrent access, bearer header propagationmention(9 tests):TokenFunclazy eval, token rotation across requests, resolve by UUID/name, error casesgo build ./...andgo vet ./...pass for bothambient-mcpandambient-control-plane🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests