generate: automatic replacement of runtimeClassName

cli/cmd/generate.go

@@ -24,7 +24,10 @@ import (
 	"strings"
 
 	"github.com/edgelesssys/contrast/internal/embedbin"
+	"github.com/edgelesssys/contrast/internal/kuberesource"
 	"github.com/edgelesssys/contrast/internal/manifest"
+	applycorev1 "k8s.io/client-go/applyconfigurations/core/v1"
+
 	"github.com/spf13/cobra"
 )
 
@@ -82,6 +85,7 @@ func runGenerate(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
+	patchTargets(paths, log)
 	if err := generatePolicies(cmd.Context(), flags.policyPath, flags.settingsPath, paths, log); err != nil {
 		return fmt.Errorf("failed to generate policies: %w", err)
 	}
@@ -228,6 +232,62 @@ func generatePolicies(ctx context.Context, regoRulesPath, policySettingsPath str
 	return nil
 }
 
+func patchTargets(paths []string, logger *slog.Logger) {
+	for _, path := range paths {
+		if err := patchTarget(path, logger); err != nil {
+			logger.Warn("Failed to patch resources at", "path", path, "err", err)
+		}
+	}
+}
+
+// patchTarget performs preprocessing steps on kubernetes resources in the given file.
+// It replaces the "runtimeClassName" of pod specs to a versioned one (contrast-cc-<VERSIONHASH>).
+func patchTarget(path string, logger *slog.Logger) error {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return fmt.Errorf("failed to read %s: %w", path, err)
+	}
+	kubeObjs, err := kuberesource.UnmarshalApplyConfigurations(data)
+	if err != nil {
+		return fmt.Errorf("failed to unmarshal %s: %w", path, err)
+	}
+
+	var changed bool
+	replaceRuntimeClassName := runtimeClassNamePatcher(&changed)
+	for i := range kubeObjs {
+		kubeObjs[i] = kuberesource.MapPodSpec(kubeObjs[i], replaceRuntimeClassName)
+	}
+
+	if !changed {
+		logger.Debug("No changes needed for yaml file", "path", path)
+		return nil
+	}
+	logger.Debug("Updating resources in yaml file", "path", path)
+	resource, err := kuberesource.EncodeResources(kubeObjs...)
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(path, resource, os.ModePerm)
+}
+
+func runtimeClassNamePatcher(modified *bool) func(*applycorev1.PodSpecApplyConfiguration) *applycorev1.PodSpecApplyConfiguration {
+	return func(spec *applycorev1.PodSpecApplyConfiguration) *applycorev1.PodSpecApplyConfiguration {
+		switch {
+		case spec.RuntimeClassName == nil:
+			return spec
+		case *spec.RuntimeClassName == runtimeHandler:
+			return spec
+		case strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc"): // patch
+		case *spec.RuntimeClassName == "kata-cc-isolation": // patch
+		default:
+			return spec
+		}
+		*modified = true
+		spec.RuntimeClassName = &runtimeHandler
+		return spec
+	}
+}
+
 func addWorkloadOwnerKeyToManifest(manifst *manifest.Manifest, keyPath string) error {
 	keyData, err := os.ReadFile(keyPath)
 	if err != nil {