Implement mTLS between Router and PicoD

[MregXN]

Background

AgentCube doc docs/design/auth-proposal.md states an intent to bring PicoD into the same zero-trust mTLS boundary as Router and WorkloadManager, defines a PicoD SPIFFE ID (spiffe://<trust-domain>/sa/agentcube-sandbox), and sketches a server/client mTLS code skeleton. It also flags handshake latency (~20–50ms per new connection) as a concern for short-lived sandboxes and mentions keeping a JWT mode as an alternative.

That intent has not yet been implemented on the Router → PicoD hop, and the design proposal leaves several choices open. Before writing code it would be useful to align on those choices, because they lead to visibly different implementations and different follow-up work.

Current state on this hop

Solid arrows describe today's request path; dashed arrows describe optional or planned mechanisms.

flowchart LR
    SDK[SDK / User] -- HTTPS + OIDC --> Router
    Router -- "HTTP + Sandbox JWT (Authorization: Bearer)" --> PicoD[PicoD in sandbox pod]
    Router -. planned: mTLS (transport identity) .-> PicoD
    WM[WorkloadManager] -- "creates sandbox pod<br/>injects PICOD_AUTH_PUBLIC_KEY" --> PicoD
    Router -- "publishes public key<br/>via Secret / ConfigMap" --> WM

Loading

Some relevant current behavior:

• SDK clients reach PicoD only through the Router by default; the SDK data-plane client is bound to ROUTER_URL and the Router reverse-proxies to the sandbox. Direct SDK-to-PicoD is possible via a base_url override but is not the default path.
• PicoD authentication today is application-layer only: a Router-signed JWT verified with the public key delivered via PICOD_AUTH_PUBLIC_KEY. There is no mutual workload identity on the wire between Router and PicoD.
• All PicoD instances share a single logical identity, both today (one Router key pair verifies all sandboxes) and under the proposed SPIFFE ID (sa/agentcube-sandbox). That means transport-layer identity, if added, cannot by itself distinguish which sandbox or which user is on the other end — tenant isolation still has to come from an application-layer signal.
• PicoD is effectively scoped to the CodeInterpreter path: only CodeInterpreter provisioning injects PICOD_AUTH_PUBLIC_KEY and defaults to AuthMode: picod. AgentRuntime uses user-defined pod specs and is not on this hop.

Implementation shape

The following sketches the concrete change surface, taking the Alignment before implementation recommendations below as the working assumption (mTLS coexists with the existing JWT and is selectable; certificates are provisioned by SPIRE via a spiffe-helper sidecar, matching the existing WorkloadManager pattern). It is included so reviewers can gauge scope; the actual choices depend on the alignment discussion.

End-to-end sequence under the working assumption above — provisioning through spiffe-helper (top), Router → PicoD handshake with the existing Sandbox JWT layered on top (middle), and continuous SVID rotation (bottom):

sequenceDiagram
    autonumber
    participant SDK as SDK / User
    participant Router as Router<br/>(SVID: agentcube-router)
    participant WM as WorkloadManager
    participant K8s as K8s API
    participant SPIRE as SPIRE Agent<br/>(node-local)
    participant Helper as spiffe-helper<br/>(sandbox sidecar)
    participant PicoD as PicoD<br/>(SVID: agentcube-sandbox)

    Note over Router,PicoD: 1. Sandbox provisioning (SPIRE-driven cert delivery)
    SDK->>Router: create CodeInterpreter session
    Router->>WM: request sandbox (mTLS, existing hop)
    WM->>K8s: create ns-local ServiceAccount<br/>"agentcube-sandbox"
    WM->>K8s: create sandbox Pod<br/>(app: picod, SA: agentcube-sandbox,<br/>spiffe-helper sidecar, shared emptyDir)
    K8s-->>Helper: sidecar starts
    Helper->>SPIRE: request SVID via Workload API socket
    SPIRE->>SPIRE: attest pod<br/>(label + SA match PicoD ClusterSPIFFEID CR)
    SPIRE-->>Helper: SVID + trust bundle
    Helper->>PicoD: write cert.pem / key.pem / bundle.pem<br/>into shared emptyDir
    PicoD->>PicoD: WaitForCertificateFiles → ListenAndServeTLS

    Note over Router,PicoD: 2. Request path (mTLS + JWT stacked)
    SDK->>Router: invoke (HTTPS + OIDC)
    Router->>PicoD: TLS handshake<br/>(Router SVID ↔ PicoD SVID, both verify URI SAN)
    Router->>PicoD: HTTP request<br/>Authorization: Bearer <Sandbox JWT><br/>(session_id, user_sub)
    PicoD->>PicoD: JWT middleware verifies signature<br/>(Router pubkey, application layer)
    PicoD-->>Router: response
    Router-->>SDK: response

    Note over SPIRE,PicoD: 3. Continuous rotation (no restart)
    SPIRE-->>Helper: push new SVID before TTL
    Helper->>PicoD: atomically rewrite PEM files
    PicoD->>PicoD: CertWatcher fsnotify → hot-reload cert & CA

Loading

1. pkg/mtls — add PicoD SPIFFE ID

pkg/mtls/spiffeid.go currently exposes RouterSPIFFEID and WorkloadManagerSPIFFEID, both built by a componentSPIFFEID(td, ns, sa) helper that assumes the /ns/<ns>/sa/<sa> layout. PicoD's ID from auth-proposal.md does not have a namespace segment, so a second helper (or a dedicated constant) is needed:

// spiffe://<trust-domain>/sa/agentcube-sandbox
PicoDSPIFFEID = fmt.Sprintf("spiffe://%s/sa/agentcube-sandbox", configuredTrustDomain())

No changes required in loader.go, watcher.go, or wait.go — they are already component-agnostic and hot-reload cert/key/CA via CertWatcher.

2. pkg/picod — TLS listener + mode switch

pkg/picod/server.go currently calls httpServer.ListenAndServe() unconditionally. Extend the server config to accept a mode and cert paths:

• Add flags in cmd/picod: --picod-auth-mode=mtls|jwt (default TBD by the latency discussion above), --tls-cert, --tls-key, --tls-ca.
• When mode is mtls: call mtls.WaitForCertificateFiles(cfg, DefaultCertificateFileWaitTimeout) (matches the pattern used in cmd/router/main.go and cmd/workload-manager/main.go), then build tls.Config via mtls.LoadServerConfig(cfg), and use httpServer.ListenAndServeTLS. LoadServerConfig already sets RequireAnyClientCert + VerifyPeerCertificate to verify the client chain against the dynamic CA pool.
• When mode is jwt: current behavior unchanged — plain HTTP, PICOD_AUTH_PUBLIC_KEY env var, JWT verification middleware.
• The AuthMiddleware on /api/* can remain in either mode (it is the application-layer JWT check; keeping it under mTLS is what makes the two layers stack).

3. pkg/router — client transport for the sandbox reverse proxy

Router already has an mTLS client transport for its upstream WM connection (session_manager.go, built via mtls.LoadClientConfig(cfg, WorkloadManagerSPIFFEID), with http:// rewritten to https:// on the endpoint). The sandbox reverse-proxy path (server.go / handlers.go's forwardToSandbox) uses a separate plain http.Transport.

The change is to give the sandbox transport the same treatment: build its TLSClientConfig via mtls.LoadClientConfig(cfg, PicoDSPIFFEID) and force https:// on the resolved sandbox endpoint. One transport is enough for all sandboxes because they share a single SPIFFE identity.

Router flags: reuse the existing --mtls-cert/--mtls-key/--mtls-ca — Router presents the same SVID to any peer. A parallel --picod-mtls-* set is only worth adding if operators want to separate the two trust boundaries; not recommended by default.

The existing Sandbox JWT signing in generateSandboxJWT / configureProxyDirector is untouched and continues to run inside the mTLS connection — that's what makes the two layers stack.

4. pkg/workloadmanager — deliver certificates into the sandbox pod

The current public-key injection lives in:

• pkg/workloadmanager/codeinterpreter_controller.go (convertToPodTemplate, around the PICOD_AUTH_PUBLIC_KEY env var branch)
• pkg/workloadmanager/workload_builder.go (buildCodeInterpreterEnvVars, gated on AuthMode == AuthModePicoD)

Add a new AuthModeType: mtls alongside the existing picod and none. When it is selected, WM constructs the PicoD pod to receive its SVID from SPIRE, mirroring what manifests/charts/base/templates/workloadmanager.yaml already does for WorkloadManager itself:

• Inject a spiffe-helper sidecar container that fetches the SVID + trust bundle from the local SPIRE Agent socket and writes them into a shared emptyDir volume.
• Mount that volume read-only into the PicoD container at the same path the --tls-cert / --tls-key / --tls-ca flags in step 2 point at.
• Add a hostPath mount for the SPIRE Agent socket (/run/spire/sockets).
• Set the PicoD pod's serviceAccountName to a dedicated agentcube-sandbox ServiceAccount and add the app: picod label; both are required by the PicoD ClusterSPIFFEID selectors (see step 5).
• Fail-closed analogue of the current IsPublicKeyCached check: refuse to create the sandbox if SPIRE prerequisites (namespace ServiceAccount, SPIRE Agent socket path availability at the node level) are not in place.

5. Deployment / manifests

• agentcube-router.yaml: no structural change if Router flags are reused; may need a PicoD SPIFFE ID / trust-domain env var.
• CodeInterpreter Helm values: expose the auth-mode selector and the sandbox spiffe-helper sidecar config, modelled on the existing spire.spiffeHelper block in values.yaml.
• New sandbox-side SPIRE wiring under manifests/charts/base/templates/spire/:
  • A PicoD ClusterSPIFFEID CR (the CRD itself is already installed by spire-controller-manager; only a new CR is added). Existing Router / WM CRs cannot be reused — their selectors match different labels/namespaces and their template produces spiffe://<td>/ns/<ns>/sa/<sa>, whereas PicoD's identity is namespace-less (spiffe://<td>/sa/agentcube-sandbox). The new CR is namespace-agnostic and uses k8s:pod-label:app:picod + k8s:sa:agentcube-sandbox as selectors.
  • A dedicated agentcube-sandbox ServiceAccount, created by WorkloadManager in the target namespace during sandbox provisioning — the SA gate is what prevents arbitrary workloads from claiming the PicoD SPIFFE ID.
  • Optional ValidatingAdmissionPolicy restricting the app: picod label to pods created by the WorkloadManager ServiceAccount (defense in depth, per auth-proposal.md).

Alignment before implementation

This issue is not intended to land an implementation. The goal is to align on the following points before opening focused follow-up issues. Each point lists a suggested direction as a starting position.

• mTLS: replacement or additional layer? All sandboxes share one SPIFFE identity, so mTLS alone can't tell which sandbox or user is on the other end — that's the JWT's job.
  Suggested: stack, don't replace. mTLS for machine identity, Sandbox JWT for session/user identity. Selectable via --picod-auth-mode=mtls|jwt.

• Certificate lifecycle: which source? File-based Secret, cert-manager, and SPIRE are all viable via pkg/mtls. SPIRE + spiffe-helper is already in the chart for the Router → WM hop.
  Suggested: go directly with SPIRE. Reuse the existing operational pattern; avoid introducing a second, weaker cert lifecycle just for the sandbox hop.

• Latency posture: default on or opt-in? New mTLS connections add ~20–50ms of handshake — negligible with warm pools, per-invocation for cold starts.
  Suggested: single global switch (--picod-auth-mode), no per-workload branching. Default is deployment-dependent; the concrete value is decided in the follow-up implementation issue.

• Direct SDK → PicoD override? The SDK's base_url override bypasses the Router; adding mTLS on this hop leaves that path unprotected.
  Suggested: declare the override unsupported / test-only for this change. If it needs to remain a production path, that's a separate design conversation and should not block this work.

Related: docs/design/auth-proposal.md, docs/design/PicoD-Plain-Authentication-Design.md.

---

/cc @hzxuzhonghu @kevin-wangzefeng @tjucoder @VanderChen @YaoZengzeng

Looking for review on the current-state description and the four Suggested directions above. Happy to open focused follow-up PRs (PicoD SPIFFE ID → PicoD TLS listener → Router client transport → sandbox-side SPIRE wiring → WM AuthModeType: mtls) once consensus lands.

[ranxi2001]

I checked the proposed flow against current main at 58422c. The current-state description is accurate, but four contracts need to be resolved before implementation:

• Key isolation: the shared cert volume makes the SVID private key readable inside PicoD. The image runs as root and /api/execute launches arbitrary commands without changing credentials, so user code can export it. Because every PicoD uses one SPIFFE URI, the certificate represents the whole class until expiry. Should TLS terminate in a separate sidecar, or is command isolation a prerequisite?
• Peer/session binding: mtls.LoadServerConfig accepts any client signed by the trusted CA; it does not require the Router SPIFFE ID. PicoD's JWT middleware verifies signature/time but not session_id or user_sub. The proposed layers therefore do not yet bind Router to a specific sandbox/user as the sequence suggests.
• Attestation: Kubernetes documents that workload creators can use any ServiceAccount in the namespace, so agentcube-sandbox is not an anti-spoofing boundary by itself. Admission would be mandatory, and the actual Pod creator is the agent-sandbox controller, not WorkloadManager.
• Policy ownership: the issue proposes both per-CodeInterpreter AuthMode=mtls and a global switch. If workloads can opt out, that becomes a downgrade of operator transport policy.

Could this first become a focused child of #441 and explicitly reconcile the PicoD scope decision in #293 plus the session-key work in #352? That would let follow-up PRs split cleanly without freezing the wrong identity model.