From 861ffc1097685f9ecf13adaa381aca5fdf7ef0b4 Mon Sep 17 00:00:00 2001
From: Samuel Karp <samuelkarp@google.com>
Date: Mon, 8 Jun 2026 00:58:16 +0000
Subject: [PATCH] cri: filter CDI annotations on checkpoint restore

Filter out any annotations on the checkpointed container matching
`cdi.k8s.io/` or exactly `cdi.k8s.io` during restore to prevent
unauthorized device restoration. When an annotation is denied, a warning
log is generated.

Tested by:
* Unit tests for exact matching, prefix boundaries, and metadata merging
* Complete CRI integration and checkpoint restore suite

Assisted-by: Antigravity
Signed-off-by: Samuel Karp <samuelkarp@google.com>
---
 .../checkpoint/checkpoint-restore-cri-test.sh |  22 +++-
 .../cri/server/container_checkpoint_linux.go  |  56 ++++++---
 .../server/container_checkpoint_linux_test.go | 107 ++++++++++++++++++
 3 files changed, 167 insertions(+), 18 deletions(-)
 create mode 100644 internal/cri/server/container_checkpoint_linux_test.go

diff --git a/contrib/checkpoint/checkpoint-restore-cri-test.sh b/contrib/checkpoint/checkpoint-restore-cri-test.sh
index 54735db14..82197fa76 100755
--- a/contrib/checkpoint/checkpoint-restore-cri-test.sh
+++ b/contrib/checkpoint/checkpoint-restore-cri-test.sh
@@ -56,6 +56,15 @@ TESTDATA=testdata
 # shellcheck disable=SC2034
 export CONTAINERD_ADDRESS="$TESTDIR/c.sock"
 export CONTAINER_RUNTIME_ENDPOINT="unix:///${CONTAINERD_ADDRESS}"
+
+# Generate crictl config file with 30s timeout
+export CRI_CONFIG_FILE="${TESTDIR}/crictl.yaml"
+cat <<EOF > "${CRI_CONFIG_FILE}"
+runtime-endpoint: unix://${CONTAINERD_ADDRESS}
+image-endpoint: unix://${CONTAINERD_ADDRESS}
+timeout: 30
+EOF
+
 TEST_IMAGE=ghcr.io/containerd/alpine
 
 function test_from_archive() {
@@ -69,9 +78,12 @@ function test_from_archive() {
 	echo -n "--> Start pod: "
 	pod_id=$(crictl runp "$POD_JSON")
 	echo "$pod_id"
+	CTR_JSON=$(mktemp)
+	jq '.annotations = {"cdi.k8s.io/device":"gpu","safe.annotation":"true"}' "$TESTDATA"/container_sleep.json >"$CTR_JSON"
 	echo -n "--> Create container: "
-	ctr_id=$(crictl create "$pod_id" "$TESTDATA"/container_sleep.json "$POD_JSON")
+	ctr_id=$(crictl create "$pod_id" "$CTR_JSON" "$POD_JSON")
 	echo "$ctr_id"
+	rm -f "$CTR_JSON"
 	echo -n "--> Start container: "
 	crictl start "$ctr_id"
 	lines_before=$(crictl logs "$ctr_id" | wc -l)
@@ -108,6 +120,12 @@ function test_from_archive() {
 			"should be larger than before checkpointing ($lines_before)"
 		false
 	fi
+	echo "--> Verifying CDI annotation filtering on restore: "
+	actual_annots=$(crictl inspect "$ctr_id" | jq -c '.status.annotations')
+	if jq -e 'has("cdi.k8s.io/device") or (has("safe.annotation") | not)' <<<"$actual_annots" >/dev/null; then
+		echo "error: CDI annotation was not filtered or safe annotation missing: $actual_annots"
+		exit 1
+	fi
 	# Cleanup
 	echo "--> Cleanup images: "
 	crictl rmi "${TEST_IMAGE}" | sed 's/^/----> \t/'
@@ -192,6 +210,8 @@ function test_from_oci() {
 
 cat >"${TESTDIR}/config.toml" <<EOF
 version = 3
+[plugins."io.containerd.cri.v1.runtime"]
+  enable_cdi = false
 [plugins."io.containerd.cri.v1.runtime".containerd]
   default_runtime_name = "test-runtime"
 [plugins.'io.containerd.cri.v1.runtime'.containerd.runtimes.test-runtime]
diff --git a/internal/cri/server/container_checkpoint_linux.go b/internal/cri/server/container_checkpoint_linux.go
index 01eab2fed..3dd714765 100644
--- a/internal/cri/server/container_checkpoint_linux.go
+++ b/internal/cri/server/container_checkpoint_linux.go
@@ -293,23 +293,11 @@ func (c *criService) CRImportCheckpoint(
 		}
 	}
 
-	if createAnnotations != nil {
-		// The hash also needs to be update or Kubernetes thinks the container needs to be restarted
-		_, ok1 := createAnnotations["io.kubernetes.container.hash"]
-		_, ok2 := originalAnnotations["io.kubernetes.container.hash"]
-
-		if ok1 && ok2 {
-			originalAnnotations["io.kubernetes.container.hash"] = createAnnotations["io.kubernetes.container.hash"]
-		}
-
-		// The restart count also needs to be correctly updated
-		_, ok1 = createAnnotations["io.kubernetes.container.restartCount"]
-		_, ok2 = originalAnnotations["io.kubernetes.container.restartCount"]
-
-		if ok1 && ok2 {
-			originalAnnotations["io.kubernetes.container.restartCount"] = createAnnotations["io.kubernetes.container.restartCount"]
-		}
-	}
+	originalAnnotations = filterAndMergeAnnotations(
+		ctx,
+		originalAnnotations,
+		createAnnotations,
+	)
 
 	var containerdImage client.Image
 
@@ -757,3 +745,37 @@ func writeSpecDumpFile(ctx context.Context, store content.Store, desc v1.Descrip
 
 	return nil
 }
+
+func filterAndMergeAnnotations(
+	ctx context.Context,
+	checkpointAnnotations map[string]string,
+	createAnnotations map[string]string,
+) map[string]string {
+	result := make(map[string]string)
+
+	for k, v := range checkpointAnnotations {
+		if strings.HasPrefix(k, "cdi.k8s.io/") || k == "cdi.k8s.io" {
+			log.G(ctx).Warnf("Denying annotation %q in checkpoint restore", k)
+			continue
+		}
+		result[k] = v
+	}
+
+	// The hash also needs to be update or Kubernetes thinks the container needs to be restarted
+	_, ok1 := createAnnotations["io.kubernetes.container.hash"]
+	_, ok2 := result["io.kubernetes.container.hash"]
+
+	if ok1 && ok2 {
+		result["io.kubernetes.container.hash"] = createAnnotations["io.kubernetes.container.hash"]
+	}
+
+	// The restart count also needs to be correctly updated
+	_, ok1 = createAnnotations["io.kubernetes.container.restartCount"]
+	_, ok2 = result["io.kubernetes.container.restartCount"]
+
+	if ok1 && ok2 {
+		result["io.kubernetes.container.restartCount"] = createAnnotations["io.kubernetes.container.restartCount"]
+	}
+
+	return result
+}
diff --git a/internal/cri/server/container_checkpoint_linux_test.go b/internal/cri/server/container_checkpoint_linux_test.go
new file mode 100644
index 000000000..3010e5233
--- /dev/null
+++ b/internal/cri/server/container_checkpoint_linux_test.go
@@ -0,0 +1,107 @@
+//go:build linux
+
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package server
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/containerd/log"
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+)
+
+type testLogHook struct {
+	mu      sync.Mutex
+	entries []string
+}
+
+func (h *testLogHook) Levels() []logrus.Level {
+	return []logrus.Level{logrus.WarnLevel}
+}
+
+func (h *testLogHook) Fire(entry *logrus.Entry) error {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	h.entries = append(h.entries, entry.Message)
+	return nil
+}
+
+func TestFilterAndMergeAnnotations(t *testing.T) {
+	for desc, tc := range map[string]struct {
+		checkpointAnnotations map[string]string
+		createAnnotations     map[string]string
+		expectedAnnotations   map[string]string
+		expectedWarnings      []string
+	}{
+		"cdi denied prefix boundaries": {
+			checkpointAnnotations: map[string]string{
+				"cdi.k8s.io/device":   "gpu",
+				"cdi.k8s.io/":         "true",
+				"cdi.k8s.io":          "true",
+				"safe.org/cdi.k8s.io": "ignored",
+				"other":               "val",
+			},
+			expectedAnnotations: map[string]string{
+				"safe.org/cdi.k8s.io": "ignored",
+				"other":               "val",
+			},
+			expectedWarnings: []string{
+				`Denying annotation "cdi.k8s.io/device" in checkpoint restore`,
+				`Denying annotation "cdi.k8s.io/" in checkpoint restore`,
+				`Denying annotation "cdi.k8s.io" in checkpoint restore`,
+			},
+		},
+
+		"createAnnotations update kubernetes metadata if present in both": {
+			checkpointAnnotations: map[string]string{
+				"io.kubernetes.container.hash":         "old-hash",
+				"io.kubernetes.container.restartCount": "1",
+				"safe.annotation":                      "2",
+			},
+			createAnnotations: map[string]string{
+				"io.kubernetes.container.hash":         "new-hash",
+				"io.kubernetes.container.restartCount": "2",
+			},
+			expectedAnnotations: map[string]string{
+				"io.kubernetes.container.hash":         "new-hash",
+				"io.kubernetes.container.restartCount": "2",
+				"safe.annotation":                      "2",
+			},
+		},
+	} {
+		t.Run(desc, func(t *testing.T) {
+			logger := logrus.New()
+			logger.SetLevel(logrus.WarnLevel)
+			hook := &testLogHook{}
+			logger.AddHook(hook)
+			ctx := log.WithLogger(context.Background(), logrus.NewEntry(logger))
+
+			res := filterAndMergeAnnotations(
+				ctx,
+				tc.checkpointAnnotations,
+				tc.createAnnotations,
+			)
+
+			assert.Equal(t, tc.expectedAnnotations, res)
+			assert.ElementsMatch(t, tc.expectedWarnings, hook.entries)
+		})
+	}
+}