generate: patch runtime class name

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,9 @@ func runGenerate(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
+	if err := patchTargets(paths, log); err != nil {
+		return fmt.Errorf("failed to patch targets: %w", err)
+	}
 	if err := generatePolicies(cmd.Context(), flags.policyPath, flags.settingsPath, paths, log); err != nil {
 		return fmt.Errorf("failed to generate policies: %w", err)
 	}
@@ -231,6 +237,53 @@ func generatePolicies(ctx context.Context, regoRulesPath, policySettingsPath str
 	return nil
 }
 
+func patchTargets(paths []string, logger *slog.Logger) error {
+	for _, path := range paths {
+		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)
+			continue
+		}
+		logger.Debug("Updating resources in yaml file", "path", path)
+		resource, err := kuberesource.EncodeResources(kubeObjs...)
+		if err != nil {
+			return err
+		}
+		if err := os.WriteFile(path, resource, os.ModePerm); err != nil {
+			return fmt.Errorf("failed to write %s: %w", path, err)
+		}
+	}
+	return nil
+}
+
+func runtimeClassNamePatcher(modified *bool) func(*applycorev1.PodSpecApplyConfiguration) *applycorev1.PodSpecApplyConfiguration {
+	return func(spec *applycorev1.PodSpecApplyConfiguration) *applycorev1.PodSpecApplyConfiguration {
+		if spec.RuntimeClassName == nil || *spec.RuntimeClassName == runtimeHandler {
+			return spec
+		}
+
+		if strings.HasPrefix(*spec.RuntimeClassName, "contrast-cc") || *spec.RuntimeClassName == "kata-cc-isolation" {
+			*modified = true
+			spec.RuntimeClassName = &runtimeHandler
+		}
+		return spec
+	}
+}
+
 func addWorkloadOwnerKeyToManifest(manifst *manifest.Manifest, keyPath string) error {
 	keyData, err := os.ReadFile(keyPath)
 	if err != nil {

e2e/internal/contrasttest/contrasttest.go

@@ -19,8 +19,8 @@ import (
 
 	"github.com/edgelesssys/contrast/cli/cmd"
 	"github.com/edgelesssys/contrast/e2e/internal/kubeclient"
-	"github.com/edgelesssys/contrast/e2e/internal/kuberesource"
 	"github.com/edgelesssys/contrast/internal/kubeapi"
+	"github.com/edgelesssys/contrast/internal/kuberesource"
 	"github.com/stretchr/testify/require"
 )
 

e2e/openssl/openssl_test.go

@@ -18,7 +18,7 @@ import (
 
 	"github.com/edgelesssys/contrast/e2e/internal/contrasttest"
 	"github.com/edgelesssys/contrast/e2e/internal/kubeclient"
-	"github.com/edgelesssys/contrast/e2e/internal/kuberesource"
+	"github.com/edgelesssys/contrast/internal/kuberesource"
 	"github.com/stretchr/testify/require"
 )
 

e2e/servicemesh/servicemesh_test.go

@@ -20,7 +20,7 @@ import (
 
 	"github.com/edgelesssys/contrast/e2e/internal/contrasttest"
 	"github.com/edgelesssys/contrast/e2e/internal/kubeclient"
-	"github.com/edgelesssys/contrast/e2e/internal/kuberesource"
+	"github.com/edgelesssys/contrast/internal/kuberesource"
 	"github.com/stretchr/testify/require"
 )
 

go.mod

@@ -21,6 +21,7 @@ require (
 	k8s.io/apimachinery v0.30.0
 	k8s.io/client-go v0.30.0
 	k8s.io/utils v0.0.0-20240423183400-0849a56e8f22
+	sigs.k8s.io/yaml v1.4.0
 )
 
 require (
@@ -66,5 +67,4 @@ require (
 	k8s.io/kube-openapi v0.0.0-20240228011516-70dd3763d340 // indirect
 	sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect
 	sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect
-	sigs.k8s.io/yaml v1.4.0 // indirect
 )

e2e/internal/kuberesource/mutators.go → internal/kuberesource/mutators.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 
 	applyappsv1 "k8s.io/client-go/applyconfigurations/apps/v1"
+	applybatchv1 "k8s.io/client-go/applyconfigurations/batch/v1"
 	applycorev1 "k8s.io/client-go/applyconfigurations/core/v1"
 )
 
@@ -123,3 +124,29 @@ func AddLogging(resources []any, level string) []any {
 	}
 	return resources
 }
+
+// MapPodSpec applies a function to a PodSpec in a Kubernetes resource.
+func MapPodSpec(resource any, f func(spec *applycorev1.PodSpecApplyConfiguration) *applycorev1.PodSpecApplyConfiguration) any {
+	if resource == nil {
+		return nil
+	}
+	switch r := resource.(type) {
+	case *applybatchv1.CronJobApplyConfiguration:
+		r.Spec.JobTemplate.Spec.Template.Spec = f(r.Spec.JobTemplate.Spec.Template.Spec)
+	case *applyappsv1.DaemonSetApplyConfiguration:
+		r.Spec.Template.Spec = f(r.Spec.Template.Spec)
+	case *applyappsv1.DeploymentApplyConfiguration:
+		r.Spec.Template.Spec = f(r.Spec.Template.Spec)
+	case *applybatchv1.JobApplyConfiguration:
+		r.Spec.Template.Spec = f(r.Spec.Template.Spec)
+	case *applycorev1.PodApplyConfiguration:
+		r.Spec = f(r.Spec)
+	case *applyappsv1.ReplicaSetApplyConfiguration:
+		r.Spec.Template.Spec = f(r.Spec.Template.Spec)
+	case *applycorev1.ReplicationControllerApplyConfiguration:
+		r.Spec.Template.Spec = f(r.Spec.Template.Spec)
+	case *applyappsv1.StatefulSetApplyConfiguration:
+		r.Spec.Template.Spec = f(r.Spec.Template.Spec)
+	}
+	return resource
+}

internal/kuberesource/reader.go

@@ -0,0 +1,32 @@
+// Copyright 2024 Edgeless Systems GmbH
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package kuberesource
+
+import (
+	"fmt"
+
+	"github.com/edgelesssys/contrast/internal/kubeapi"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/client-go/applyconfigurations"
+)
+
+// UnmarshalApplyConfigurations unmarshals a YAML document into a list of ApplyConfigurations.
+func UnmarshalApplyConfigurations(data []byte) ([]any, error) {
+	objs, err := kubeapi.UnmarshalUnstructuredK8SResource(data)
+	if err != nil {
+		return nil, err
+	}
+	var result []any
+	for _, obj := range objs {
+		applyConfig := applyconfigurations.ForKind(obj.GetObjectKind().GroupVersionKind())
+		if applyConfig == nil {
+			return nil, fmt.Errorf("unmarshalling: unsupported resource type %s for %q", obj.GroupVersionKind().String(), obj.GetName())
+		}
+		if err := runtime.DefaultUnstructuredConverter.FromUnstructured(obj.UnstructuredContent(), applyConfig); err != nil {
+			return nil, err
+		}
+		result = append(result, applyConfig)
+	}
+	return result, nil
+}

internal/kuberesource/reader_writer_test.go

@@ -0,0 +1,32 @@
+// Copyright 2024 Edgeless Systems GmbH
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package kuberesource
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestEncodeDecode(t *testing.T) {
+	require := require.New(t)
+
+	fixture := `apiVersion: v1
+kind: Pod
+metadata:
+  name: foo
+spec:
+  containers:
+  - image: image
+    name: bar
+`
+
+	resources, err := UnmarshalApplyConfigurations([]byte(fixture))
+	require.NoError(err)
+
+	got, err := EncodeResources(resources...)
+	require.NoError(err)
+
+	require.Equal(fixture, string(got))
+}

e2e/internal/kuberesource/resourcegen/main.go → internal/kuberesource/resourcegen/main.go

@@ -9,7 +9,7 @@ import (
 	"log"
 	"os"
 
-	"github.com/edgelesssys/contrast/e2e/internal/kuberesource"
+	"github.com/edgelesssys/contrast/internal/kuberesource"
 )
 
 func main() {

e2e/internal/kuberesource/writer.go → internal/kuberesource/writer.go

@@ -6,12 +6,12 @@ package kuberesource
 import (
 	"bytes"
 
-	"gopkg.in/yaml.v3"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/yaml"
 )
 
-// EncodeResources encodes a list of resources into a single YAML document.
+// EncodeResources encodes a list of apply configurations into a single YAML document.
 func EncodeResources(resources ...any) ([]byte, error) {
 	unstructuredResources, err := ResourcesToUnstructured(resources)
 	if err != nil {

packages/by-name/contrast/package.nix

@@ -19,7 +19,8 @@ let
     ldflags = [
       "-s"
       "-X github.com/edgelesssys/contrast/internal/manifest.trustedMeasurement=${launchDigest}"
-      "-X github.com/edgelesssys/contrast/e2e/internal/kuberesource.runtimeHandler=${runtimeHandler}"
+      "-X github.com/edgelesssys/contrast/cli/cmd.runtimeHandler=${runtimeHandler}"
+      "-X github.com/edgelesssys/contrast/internal/kuberesource.runtimeHandler=${runtimeHandler}"
     ];
 
     subPackages = [ "e2e/openssl" "e2e/servicemesh" ];
@@ -62,7 +63,7 @@ buildGoModule rec {
   proxyVendor = true;
   vendorHash = "sha256-i+7DhygotCNhczpaZlI9O7enKVOW7smauOKcGQhOtzI=";
 
-  subPackages = packageOutputs ++ [ "e2e/internal/kuberesource/resourcegen" ];
+  subPackages = packageOutputs ++ [ "internal/kuberesource/resourcegen" ];
 
   prePatch = ''
     install -D ${lib.getExe genpolicy} cli/cmd/assets/genpolicy
@@ -77,7 +78,7 @@ buildGoModule rec {
     "-X main.version=v${version}"
     "-X github.com/edgelesssys/contrast/internal/manifest.trustedMeasurement=${launchDigest}"
     "-X github.com/edgelesssys/contrast/cli/cmd.runtimeHandler=${runtimeHandler}"
-    "-X github.com/edgelesssys/contrast/e2e/internal/kuberesource.runtimeHandler=${runtimeHandler}"
+    "-X github.com/edgelesssys/contrast/internal/kuberesource.runtimeHandler=${runtimeHandler}"
   ];
 
   preCheck = ''