feat(nvsnap): shared-token authentication for the agent API

src/compute-plane-services/nvsnap/cmd/agent/main.go

@@ -90,6 +90,11 @@ func main() {
 	flag.StringVar(&config.FSStorePath, "fsstore-path", os.Getenv("NVSNAP_FSSTORE_PATH"),
 		"Path to a shared filesystem mounted on every node (Lustre/Weka/EFS/Filestore/NFS). When set, captures are published here and the restore cascade copies from this path before peer fanout. Empty disables.")
 	flag.StringVar(&config.ListenAddr, "listen", ":8081", "Listen address")
+	// The token itself is env-only, never a flag: flag values show up in the
+	// pod spec and in `ps`, and this is a credential.
+	var authMode string
+	flag.StringVar(&authMode, "auth-mode", os.Getenv("NVSNAP_AGENT_AUTH_MODE"),
+		"Agent API authentication: disabled (default), permissive (check, log failures, still serve), or required (401). Token comes from NVSNAP_AGENT_TOKEN. See GH #486.")
 	flag.StringVar(&config.CheckpointDir, "checkpoint-dir", "/var/lib/nvsnap/checkpoints", "Checkpoint storage directory (in-agent-container path)")
 	flag.StringVar(&config.CheckpointHostDir, "checkpoint-host-dir", "/var/lib/containerd/nvsnap-checkpoints", "Host path that backs --checkpoint-dir (must match the DaemonSet hostPath mount; used to translate paths for the capture-write writer Job)")
 	flag.StringVar(&config.CRIUPath, "criu-path", "/usr/local/sbin/criu", "Path to CRIU binary (on host filesystem)")
@@ -204,6 +209,19 @@ func main() {
 		"imagePullSecret name for the mount-holder pod (created by operators in the workload namespace). Defaults to nvsnap-agent-pull; set to '-' to disable.")
 
 	flag.Parse()
+
+	// Fail startup on a bad mode rather than falling back to disabled: an
+	// operator who typo'd --auth-mode should hear about it now, not discover
+	// months later that the API was open the whole time.
+	var authErr error
+	if config.AuthMode, authErr = agent.ParseAuthMode(authMode); authErr != nil {
+		logrus.WithError(authErr).Fatal("invalid --auth-mode")
+	}
+	config.AuthToken = os.Getenv("NVSNAP_AGENT_TOKEN")
+	if config.AuthMode != agent.AuthDisabled && config.AuthToken == "" {
+		logrus.Fatalf("--auth-mode=%s requires NVSNAP_AGENT_TOKEN to be set", config.AuthMode)
+	}
+
 	config.RootfsCapture.WarmupDelay = time.Duration(rootfsWarmupSec) * time.Second
 	for _, b := range strings.Split(replicationPeerBuckets, ",") {
 		if b = strings.TrimSpace(b); b != "" {

src/compute-plane-services/nvsnap/cmd/nvsnap-mount-prep/main.go

@@ -40,6 +40,9 @@ limitations under the License.
 //	NVSNAP_POD_UID         (required) downward API: metadata.uid
 //	NVSNAP_RESTORE_HASH    (required) full sha256 of the capture
 //	NVSNAP_AGENT_URL       (required) e.g. http://$(HOST_IP):8081
+//	NVSNAP_AGENT_TOKEN     (optional) bearer token for the agent API (GH #486);
+//	                       empty sends no header, which is correct while the
+//	                       agent still runs with auth disabled
 //	NVSNAP_CAPTURE_NODE    (optional) where capture data lives; empty=this node
 //	NVSNAP_PREP_MOUNTS     (required) JSON-encoded []VolumeMeta from the manifest
 //	NVSNAP_PREP_DEADLINE   (optional) duration; default 15m
@@ -196,6 +199,7 @@ func startWithRetry(agentURL string, req prepRequest) error {
 			return err
 		}
 		httpReq.Header.Set("Content-Type", "application/json")
+		setAgentAuth(httpReq)
 		resp, err := http.DefaultClient.Do(httpReq)
 		if err != nil {
 			lastErr = err
@@ -220,6 +224,7 @@ func startWithRetry(agentURL string, req prepRequest) error {
 
 func getStatus(agentURL, podUID string) (*prepStatus, error) {
 	httpReq, err := http.NewRequestWithContext(context.Background(), http.MethodGet, agentURL+"/v1/restore/prep/"+podUID, http.NoBody)
+	setAgentAuth(httpReq)
 	if err != nil {
 		return nil, err
 	}
@@ -244,3 +249,15 @@ func getStatus(agentURL, podUID string) (*prepStatus, error) {
 	}
 	return &s, nil
 }
+
+// setAgentAuth attaches the agent API bearer token when one is configured.
+// Empty is the normal state until the operator turns auth on, and sending no
+// header is exactly what a disabled or permissive agent expects. See GH #486.
+func setAgentAuth(r *http.Request) {
+	if r == nil {
+		return
+	}
+	if tok := os.Getenv("NVSNAP_AGENT_TOKEN"); tok != "" {
+		r.Header.Set("Authorization", "Bearer "+tok)
+	}
+}

src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-daemonset.yaml

@@ -82,6 +82,12 @@ spec:
           image: {{ include "nvsnap.agent.image" . }}
           imagePullPolicy: {{ .Values.agent.image.pullPolicy }}
           args:
+            {{- if .Values.agent.auth.enabled }}
+            # permissive counts and logs unauthenticated callers but still
+            # serves them; required returns 401. Roll out on permissive until
+            # nvsnap_agent_auth_total{result="missing"} is zero.
+            - --auth-mode={{ .Values.agent.auth.mode }}
+            {{- end }}
             - --cuda-checkpoint-path=/criu-bundle/cuda-checkpoint
             - --criu-path=/criu-bundle/criu
             # Translate in-container --checkpoint-dir to the host path
@@ -168,6 +174,17 @@ spec:
               valueFrom:
                 fieldRef:
                   fieldPath: status.hostIP
+            {{- if .Values.agent.auth.enabled }}
+            # Shared bearer token for the agent API (GH #486). Env rather than
+            # a flag: flag values are visible in the pod spec and in `ps`.
+            # The agent uses it both to verify inbound requests and to sign
+            # its own peer calls.
+            - name: NVSNAP_AGENT_TOKEN
+              valueFrom:
+                secretKeyRef:
+                  name: nvsnap-agent-token
+                  key: token
+            {{- end }}
             {{- if .Values.server.enabled }}
             - name: NVSNAP_CATALOG_URL
               value: "http://nvsnap-server.{{ .Release.Namespace }}.svc.cluster.local:{{ .Values.server.service.port }}"

src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-token-secret.yaml

@@ -0,0 +1,45 @@
+{{- if .Values.agent.auth.enabled }}
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Shared bearer token for the agent HTTP API (GH #486).
+#
+# The agent API is the control surface of a privileged process and the
+# DaemonSet binds it to every node's IP, so it needs authentication in the
+# request path. This Secret holds the token both the agent (to verify) and its
+# callers (to present) read.
+#
+# Generated once and preserved across upgrades: `helm upgrade` re-renders every
+# template, so a freshly random token on each upgrade would rotate the
+# credential out from under running callers and cause a self-inflicted outage
+# mid-rollout. The lookup below reuses the existing value when the Secret is
+# already present. Set agent.auth.token explicitly to manage it yourself (or to
+# rotate deliberately).
+{{- $ns := .Release.Namespace }}
+{{- $name := "nvsnap-agent-token" }}
+{{- $existing := lookup "v1" "Secret" $ns $name }}
+{{- $token := "" }}
+{{- if .Values.agent.auth.token }}
+{{-   $token = .Values.agent.auth.token | b64enc }}
+{{- else if and $existing $existing.data $existing.data.token }}
+{{-   $token = $existing.data.token }}
+{{- else }}
+{{-   $token = randAlphaNum 48 | b64enc }}
+{{- end }}
+apiVersion: v1
+kind: Secret
+metadata:
+  name: {{ $name }}
+  namespace: {{ $ns }}
+  labels:
+    app.kubernetes.io/name: nvsnap
+    app.kubernetes.io/part-of: nvsnap
+  annotations:
+    # helm.sh/resource-policy keeps the Secret if the release is removed with
+    # --keep-history style workflows; without it a delete/reinstall cycle
+    # silently rotates the token.
+    helm.sh/resource-policy: keep
+type: Opaque
+data:
+  token: {{ $token }}
+{{- end }}

src/compute-plane-services/nvsnap/deploy/helm/nvsnap/values.yaml

@@ -183,6 +183,29 @@ agent:
 
   # Host paths the agent bind-mounts. Override if your nodes use
   # non-standard layouts (e.g. K3s on a single laptop).
+
+  # Authentication for the agent HTTP API (GH #486).
+  #
+  # The agent API restores and deletes checkpoints, serves any file inside a
+  # checkpoint, and exposes pprof, on a process running privileged with
+  # /var/lib and the containerd root bind-mounted. The DaemonSet binds it to
+  # every node's IP (hostNetwork + hostPort), and NetworkPolicy cannot fence a
+  # hostNetwork pod, so access control has to live in the request path.
+  #
+  # Rollout: enable with mode=permissive first. The agent then counts and logs
+  # unauthenticated callers via nvsnap_agent_auth_total{result="missing"} but
+  # still serves them, so nothing breaks while callers pick up the token.
+  # Switch to required once that series is flat at zero.
+  auth:
+    # Off by default so an upgrade does not lock out callers that have not
+    # been given the token yet. The agent logs a warning while it is off.
+    enabled: false
+    # permissive | required. Ignored when enabled=false.
+    mode: permissive
+    # Leave empty to have the chart generate one and preserve it across
+    # upgrades. Set explicitly to manage or rotate the credential yourself.
+    token: ""
+
   hostPaths:
     checkpoints: /var/lib/containerd/nvsnap-checkpoints
     containerdSock: /run/containerd/containerd.sock

src/compute-plane-services/nvsnap/internal/agent/BUILD.bazel

@@ -5,6 +5,7 @@ go_library(
     srcs = [
         "agent.go",
         "blob_uploader.go",
+        "auth.go",
         "capture_cascade.go",
         "capture_peer.go",
         "cascade_fetch.go",
@@ -90,6 +91,7 @@ go_test(
         "l2_promote_async_test.go",
         "l2_writer_test.go",
         "nim_backend_test.go",
+        "auth_test.go",
         "pathsafe_test.go",
         "peer_fanout_test.go",
         "peer_load_test.go",

src/compute-plane-services/nvsnap/internal/agent/agent.go

@@ -67,7 +67,15 @@ type Config struct {
 	CRIUPath            string
 	NodeName            string
 	LogLevel            string
-	UseNsenter          bool // Run CRIU/cuda-checkpoint in host mount namespace (for containerized agents)
+
+	// AuthToken is the shared bearer token callers must present on the
+	// agent API. Sourced from a Secret rather than a flag so it does not
+	// land in the pod spec or in `ps` output. Empty disables the check.
+	AuthToken string
+	// AuthMode is disabled (default), permissive, or required. See auth.go
+	// for why the rollout needs a permissive state.
+	AuthMode   AuthMode
+	UseNsenter bool // Run CRIU/cuda-checkpoint in host mount namespace (for containerized agents)
 
 	// Prewarm enables agent-side page-cache prewarm of the rox-backed
 	// overlay lowerdir on restore (--prewarm, default true). Reads the
@@ -437,6 +445,35 @@ func (a *Agent) Run(ctx context.Context) error {
 
 	router := mux.NewRouter()
 
+	// RED metrics for every route: rate and status via APIRequestsTotal,
+	// duration via APIRequestDuration, keyed on the route TEMPLATE rather
+	// than the concrete path so checkpoint IDs never become label values.
+	// Registered outermost so it also observes requests the auth guard
+	// rejects -- a spike of 401s is exactly what the rollout needs to see.
+	router.Use(metrics.InstrumentRoute())
+
+	// Present the token on our own peer calls too. Set unconditionally: an
+	// agent in permissive mode still has peers that may already require it.
+	SetOutboundToken(a.config.AuthToken)
+
+	// Order matters: gorilla/mux runs middleware in registration order, so
+	// auth is registered FIRST. Otherwise pathVarGuard answers a malformed
+	// {id} with 400 before the caller is authenticated, telling an
+	// unauthenticated client which routes exist and how their variables are
+	// shaped. Authenticate, then validate.
+	if guard := tokenGuard(a.config.AuthMode, a.config.AuthToken, a.log); guard != nil {
+		router.Use(guard)
+		a.log.WithField("mode", a.config.AuthMode).Info("Agent API authentication enabled")
+	} else {
+		a.log.Warn("Agent API is UNAUTHENTICATED: set NVSNAP_AGENT_TOKEN and " +
+			"--auth-mode to require a bearer token (GH #486)")
+	}
+
+	// Every {id}/{hash}/{pod-uid} below names a directory under a hostPath
+	// mount. Validate them in one place so a route added later is covered
+	// without remembering to. See pathVarGuard in pathsafe.go.
+	router.Use(pathVarGuard)
+
 	// Metrics endpoint
 	router.Handle("/metrics", metrics.Handler()).Methods("GET")
 
@@ -835,6 +872,11 @@ func (a *Agent) gpuRestoreHandler(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, fmt.Sprintf("invalid request: %v", err), http.StatusBadRequest)
 		return
 	}
+	// Body-supplied, unlike the {id} routes the middleware covers.
+	if err := validPathSegment("checkpointId", req.CheckpointID); err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
 
 	log := a.log.WithFields(logrus.Fields{
 		"checkpointId": req.CheckpointID,