From a18591ae8f797c9288c0e0e44cec09a7b947b80d Mon Sep 17 00:00:00 2001 From: Moritz Sanft <58110325+msanft@users.noreply.github.com> Date: Mon, 27 Apr 2026 14:54:51 +0200 Subject: [PATCH] prefactor: centralize Contrast runtime class matching --- cli/cmd/generate.go | 20 +++++++++----------- cli/cmd/policies.go | 4 ++-- cli/verifier/image_ref_valid.go | 2 +- cli/verifier/no_shared_fs_mount.go | 3 +-- cli/verifier/versions_match.go | 2 +- internal/kuberesource/mutators.go | 25 +++++++++++++++++++------ 6 files changed, 33 insertions(+), 23 deletions(-) diff --git a/cli/cmd/generate.go b/cli/cmd/generate.go index b3021cda343..927508ae91f 100644 --- 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 index 2c7d03c70be..a559bb6cb53 100644 --- 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 index 85b54289cdc..cb25b388ab7 100644 --- 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 index 2ffb62d6f05..befc92bedd8 100644 --- 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 index 718b225c508..34ed57af3dc 100644 --- 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 index a1e54450e91..854145c6e86 100644 --- 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{