prefactor: centralize Contrast runtime class matching

msanft · msanft · commit 4fe5f053d6d2 · 2026-04-27T18:15:07.000+02:00
diff --git a/cli/cmd/generate.go b/cli/cmd/generate.go
@@ -284,17 +284,17 @@ func runGenerate(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
-// mapCCWorkloads applies the given function to all workloads with the 'contrast-cc' runtime class.
+// mapContrastWorkloads applies the given function to all workloads with a Contrast runtime class.
 // The callback receives an apply configuration together with the file path and index the unstructured object has in the file map.
 // Changes to the apply configuration are not applied to the original unstructured object.
-func mapCCWorkloads(fileMap map[string][]*unstructured.Unstructured, f func(res any, path string, idx int) (any, error)) error {
+func mapContrastWorkloads(fileMap map[string][]*unstructured.Unstructured, f func(res any, path string, idx int) (any, error)) error {
 	for path, resources := range fileMap {
 		for idx, r := range resources {
 			applyConfig, err := kuberesource.UnstructuredToApplyConfiguration(r)
 			if err != nil {
 				continue
 			}
-			if !isCCWorkload(applyConfig) {
+			if !isContrastWorkload(applyConfig) {
 				continue
 			}
 			changed, err := f(applyConfig, path, idx)
@@ -313,11 +313,9 @@ func mapCCWorkloads(fileMap map[string][]*unstructured.Unstructured, f func(res
 	return nil
 }
 
-func isCCWorkload(resource any) (ret bool) {
+func isContrastWorkload(resource any) (ret bool) {
 	kuberesource.MapPodSpec(resource, func(spec *applycorev1.PodSpecApplyConfiguration) *applycorev1.PodSpecApplyConfiguration {
-		if spec != nil && spec.RuntimeClassName != nil && strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") {
-			ret = true
-		}
+		ret = kuberesource.IsContrastPod(spec)
 		return spec
 	})
 	return ret
@@ -339,7 +337,7 @@ func isCoordinator(resource any) bool {
 func runVerifiers(fileMap map[string][]*unstructured.Unstructured, verifiers []verifier.Verifier) error {
 	var findings error
 	for _, v := range verifiers {
-		_ = mapCCWorkloads(fileMap, func(res any, path string, idx int) (any, error) {
+		_ = mapContrastWorkloads(fileMap, func(res any, path string, idx int) (any, error) {
 			if err := v.Verify(res); err != nil {
 				findings = errors.Join(findings, fmt.Errorf("failed to verify resource %q in file %q: %w", fileMap[path][idx].GetName(), path, err))
 			}
@@ -406,7 +404,7 @@ func extractTargets(paths []string, configFile io.Writer, logger *slog.Logger) (
 			applyConfig, err := kuberesource.UnstructuredToApplyConfiguration(object)
 			if err != nil {
 				logger.Warn("Could not convert resource into ApplyConfiguration", "path", path, "err", err)
-			} else if isCCWorkload(applyConfig) {
+			} else if isContrastWorkload(applyConfig) {
 				containsCC = true
 				if isCoordinator(applyConfig) {
 					r, ok := applyConfig.(*applyappsv1.StatefulSetApplyConfiguration)
@@ -454,7 +452,7 @@ func generatePolicies(ctx context.Context, flags *generateFlags, fileMap map[str
 		}
 	}()
 
-	return mapCCWorkloads(fileMap, func(res any, path string, idx int) (any, error) {
+	return mapContrastWorkloads(fileMap, func(res any, path string, idx int) (any, error) {
 		initdataAnno, err := runner.Run(ctx, res, extraPath, logger)
 		if err != nil {
 			return nil, fmt.Errorf("failed to generate policy for %q in %q: %w", fileMap[path][idx].GetName(), path, err)
@@ -496,7 +494,7 @@ func patchTargets(fileMap map[string][]*unstructured.Unstructured, imageReplacem
 			return fmt.Errorf("parsing release image definitions %s: %w", ReleaseImageReplacements, err)
 		}
 	}
-	return mapCCWorkloads(fileMap, func(res any, _ string, _ int) (any, error) {
+	return mapContrastWorkloads(fileMap, func(res any, _ string, _ int) (any, error) {
 		if flags.insecureEnableDebugShell {
 			if _, err := kuberesource.AddDebugShell(res, kuberesource.DebugShell()); err != nil {
 				return nil, fmt.Errorf("injecting debug shell container: %w", err)
diff --git a/cli/cmd/policies.go b/cli/cmd/policies.go
@@ -18,7 +18,7 @@ import (
 )
 
 func manipulateInitdata(fileMap map[string][]*unstructured.Unstructured, manipulators ...func(*initdata.Initdata) error) error {
-	return mapCCWorkloads(fileMap, func(res any, path string, _ int) (resource any, retErr error) {
+	return mapContrastWorkloads(fileMap, func(res any, path string, _ int) (resource any, retErr error) {
 		return kuberesource.MapPodSpecWithMeta(res, func(meta *applymetav1.ObjectMetaApplyConfiguration, spec *applycorev1.PodSpecApplyConfiguration) (*applymetav1.ObjectMetaApplyConfiguration, *applycorev1.PodSpecApplyConfiguration) {
 			if meta == nil {
 				return meta, spec
@@ -60,7 +60,7 @@ func manipulateInitdata(fileMap map[string][]*unstructured.Unstructured, manipul
 
 func policiesFromKubeResources(fileMap map[string][]*unstructured.Unstructured) ([]deployment, error) {
 	var deployments []deployment
-	if err := mapCCWorkloads(fileMap, func(res any, path string, idx int) (any, error) {
+	if err := mapContrastWorkloads(fileMap, func(res any, path string, idx int) (any, error) {
 		name := fileMap[path][idx].GetName()
 		namespace := orDefault(fileMap[path][idx].GetNamespace(), "default")
 		gvk := fileMap[path][idx].GetObjectKind().GroupVersionKind()
diff --git a/cli/verifier/image_ref_valid.go b/cli/verifier/image_ref_valid.go
@@ -25,7 +25,7 @@ func (v *ImageRefValid) Verify(toVerify any) error {
 	kuberesource.MapPodSpec(toVerify, func(
 		spec *applycorev1.PodSpecApplyConfiguration,
 	) *applycorev1.PodSpecApplyConfiguration {
-		if spec == nil || spec.RuntimeClassName == nil || !strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") {
+		if !kuberesource.IsContrastPod(spec) {
 			return spec
 		}
 
diff --git a/cli/verifier/no_shared_fs_mount.go b/cli/verifier/no_shared_fs_mount.go
@@ -6,7 +6,6 @@ package verifier
 import (
 	"errors"
 	"fmt"
-	"strings"
 
 	"github.com/edgelesssys/contrast/internal/kuberesource"
 
@@ -25,7 +24,7 @@ func (v *NoSharedFSMount) Verify(toVerify any) error {
 	// get all volume mounts that are referenced in containers
 	isNonCC := false
 	kuberesource.MapPodSpec(toVerify, func(spec *applycorev1.PodSpecApplyConfiguration) *applycorev1.PodSpecApplyConfiguration {
-		if spec == nil || spec.RuntimeClassName == nil || !strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") {
+		if !kuberesource.IsContrastPod(spec) {
 			// this isn't a confidential pod so we don't need to check further
 			isNonCC = true
 			return spec
diff --git a/cli/verifier/versions_match.go b/cli/verifier/versions_match.go
@@ -34,7 +34,7 @@ func (v *VersionsMatch) Verify(toVerify any) error {
 		meta *applymetav1.ObjectMetaApplyConfiguration,
 		spec *applycorev1.PodSpecApplyConfiguration,
 	) (*applymetav1.ObjectMetaApplyConfiguration, *applycorev1.PodSpecApplyConfiguration) {
-		if spec == nil || spec.RuntimeClassName == nil || !strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") {
+		if !kuberesource.IsContrastPod(spec) {
 			return meta, spec
 		}
 
diff --git a/internal/kuberesource/mutators.go b/internal/kuberesource/mutators.go
@@ -34,6 +34,19 @@ const (
 	imageStoreSizeAnnotationKey   = "contrast.edgeless.systems/image-store-size"
 )
 
+// contrastRuntimeClassPrefixes lists runtime class prefixes that identify Contrast pods.
+var contrastRuntimeClassPrefixes = []string{"contrast-cc"}
+
+// IsContrastPod reports whether a pod uses a Contrast runtime.
+func IsContrastPod(spec *applycorev1.PodSpecApplyConfiguration) bool {
+	if spec == nil || spec.RuntimeClassName == nil {
+		return false
+	}
+	return slices.ContainsFunc(contrastRuntimeClassPrefixes, func(p string) bool {
+		return strings.HasPrefix(*spec.RuntimeClassName, p)
+	})
+}
+
 // AddInitializer adds an initializer and its shared volume to the resource.
 //
 // If the resource does not contain a PodSpec, this function does nothing.
@@ -46,7 +59,7 @@ func AddInitializer(
 		if meta != nil && meta.Annotations[skipInitializerAnnotationKey] == "true" {
 			return meta, spec
 		}
-		if spec == nil || spec.RuntimeClassName == nil || !strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") {
+		if !IsContrastPod(spec) {
 			return meta, spec
 		}
 		if meta != nil && meta.Annotations[securePVAnnotationKey] != "" {
@@ -173,7 +186,7 @@ func AddServiceMesh(
 	serviceMeshProxy *applycorev1.ContainerApplyConfiguration,
 ) (res any, retErr error) {
 	res = MapPodSpecWithMeta(resource, func(meta *applymetav1.ObjectMetaApplyConfiguration, spec *applycorev1.PodSpecApplyConfiguration) (*applymetav1.ObjectMetaApplyConfiguration, *applycorev1.PodSpecApplyConfiguration) {
-		if spec == nil || spec.RuntimeClassName == nil || !strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") {
+		if !IsContrastPod(spec) {
 			return meta, spec
 		}
 
@@ -230,7 +243,7 @@ func AddDebugShell(
 	debugShell *applycorev1.ContainerApplyConfiguration,
 ) (any, error) {
 	return MapPodSpec(resource, func(spec *applycorev1.PodSpecApplyConfiguration) *applycorev1.PodSpecApplyConfiguration {
-		if spec == nil || spec.RuntimeClassName == nil || !strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") {
+		if !IsContrastPod(spec) {
 			return spec
 		}
 
@@ -319,7 +332,7 @@ func AddDmesg(resources []any) []any {
 			WithPrivileged(true).SecurityContextApplyConfiguration)
 
 	addDmesg := func(spec *applycorev1.PodSpecApplyConfiguration) *applycorev1.PodSpecApplyConfiguration {
-		if spec == nil || spec.RuntimeClassName == nil || !strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") {
+		if !IsContrastPod(spec) {
 			return spec
 		}
 		spec.Containers = append(spec.Containers, *dmesgContainer)
@@ -380,7 +393,7 @@ func AddImageStore(resources []any) []any {
 
 	addPvc := func(meta *applymetav1.ObjectMetaApplyConfiguration, spec *applycorev1.PodSpecApplyConfiguration,
 	) (*applymetav1.ObjectMetaApplyConfiguration, *applycorev1.PodSpecApplyConfiguration) {
-		if spec == nil || spec.RuntimeClassName == nil || !strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") {
+		if !IsContrastPod(spec) {
 			return meta, spec
 		}
 
@@ -733,7 +746,7 @@ func PatchNodeSelector(resources []any) []any {
 	var out []any
 	for _, resource := range resources {
 		out = append(out, MapPodSpec(resource, func(spec *applycorev1.PodSpecApplyConfiguration) *applycorev1.PodSpecApplyConfiguration {
-			if spec == nil || spec.RuntimeClassName == nil || !strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") {
+			if !IsContrastPod(spec) {
 				return spec
 			}
 			spec = spec.WithNodeSelector(map[string]string{

Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@ func (v *ImageRefValid) Verify(toVerify any) error {`
`25`	`25`	`kuberesource.MapPodSpec(toVerify, func(`
`26`	`26`	`spec *applycorev1.PodSpecApplyConfiguration,`
`27`	`27`	`) *applycorev1.PodSpecApplyConfiguration {`
`28`		`- if spec == nil \|\| spec.RuntimeClassName == nil \|\| !strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") {`
	`28`	`+ if !kuberesource.IsContrastPod(spec) {`
`29`	`29`	`return spec`
`30`	`30`	`}`
`31`	`31`
Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,7 @@ func (v *VersionsMatch) Verify(toVerify any) error {`
`34`	`34`	`meta *applymetav1.ObjectMetaApplyConfiguration,`
`35`	`35`	`spec *applycorev1.PodSpecApplyConfiguration,`
`36`	`36`	`) (applymetav1.ObjectMetaApplyConfiguration, applycorev1.PodSpecApplyConfiguration) {`
`37`		`- if spec == nil \|\| spec.RuntimeClassName == nil \|\| !strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") {`
	`37`	`+ if !kuberesource.IsContrastPod(spec) {`
`38`	`38`	`return meta, spec`
`39`	`39`	`}`
`40`	`40`