kata-runtime: select correct SNP ID block for platform at runtime

cli/cmd/common.go

@@ -17,6 +17,7 @@ import (
 	"github.com/edgelesssys/contrast/internal/attestation/certcache"
 	"github.com/edgelesssys/contrast/internal/constants"
 	"github.com/edgelesssys/contrast/internal/fsstore"
+	"github.com/google/go-sev-guest/abi"
 	"github.com/spf13/afero"
 	"github.com/spf13/cobra"
 	"golang.org/x/term"
@@ -41,6 +42,19 @@ const (
 //go:embed assets/image-replacements.txt
 var ReleaseImageReplacements []byte
 
+// SNPIDBlocks contains the SNP ID blocks for different vCPU counts and CPU generations
+// as a JSON map of platform -> vCPU count -> CPU generation -> SnpIDBlock.
+//
+//go:embed assets/snp-id-blocks.json
+var SNPIDBlocks []byte
+
+// SnpIDBlock represents the SNP ID block and ID auth used for SEV-SNP guests.
+type SnpIDBlock struct {
+	IDBlock     string        `json:"idBlock"`
+	IDAuth      string        `json:"idAuth"`
+	GuestPolicy abi.SnpPolicy `json:"guestPolicy"`
+}
+
 func commandOut() io.Writer {
 	if term.IsTerminal(int(os.Stdout.Fd())) {
 		return nil // use out writer of parent

cli/cmd/generate.go

@@ -16,6 +16,7 @@ import (
 	"os"
 	"path/filepath"
 	"slices"
+	"strconv"
 	"strings"
 
 	"github.com/edgelesssys/contrast/cli/genpolicy"
@@ -24,6 +25,8 @@ import (
 	"github.com/edgelesssys/contrast/internal/kuberesource"
 	"github.com/edgelesssys/contrast/internal/manifest"
 	"github.com/edgelesssys/contrast/internal/platforms"
+	"github.com/google/go-sev-guest/abi"
+	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	applyappsv1 "k8s.io/client-go/applyconfigurations/apps/v1"
 	applycorev1 "k8s.io/client-go/applyconfigurations/core/v1"
@@ -33,8 +36,13 @@ import (
 )
 
 const (
-	contrastRoleAnnotationKey     = "contrast.edgeless.systems/pod-role"
-	workloadSecretIDAnnotationKey = "contrast.edgeless.systems/workload-secret-id"
+	annotationPrefix              = "contrast.edgeless.systems/"
+	kataAnnotationPrefix          = "io.katacontainers.config.hypervisor."
+	contrastRoleAnnotationKey     = annotationPrefix + "pod-role"
+	workloadSecretIDAnnotationKey = annotationPrefix + "workload-secret-id"
+	idBlockAnnotation             = kataAnnotationPrefix + "snp_id_block_"
+	idAuthAnnotationKey           = kataAnnotationPrefix + "snp_id_auth_"
+	guestPolicyAnnotationKey      = kataAnnotationPrefix + "snp_guest_policy_"
 )
 
 // NewGenerateCmd creates the contrast generate subcommand.
@@ -511,6 +519,10 @@ func patchTargets(logger *slog.Logger, fileMap map[string][]*unstructured.Unstru
 		replaceRuntimeClassName := patchRuntimeClassName(logger, runtimeHandler)
 		kuberesource.MapPodSpec(res, replaceRuntimeClassName)
 
+		if err := patchIDBlockAnnotation(res, logger); err != nil {
+			return nil, fmt.Errorf("injecting ID block annotations: %w", err)
+		}
+
 		return res, nil
 	})
 }
@@ -706,6 +718,84 @@ func patchRuntimeClassName(logger *slog.Logger, defaultRuntimeHandler string) fu
 	}
 }
 
+func patchIDBlockAnnotation(res any, logger *slog.Logger) error {
+	// runtime -> cpu_count -> product_line -> ID block
+	var snpIDBlocks map[string]map[string]map[string]SnpIDBlock
+	if err := json.Unmarshal(SNPIDBlocks, &snpIDBlocks); err != nil {
+		return fmt.Errorf("unmarshal SNP ID blocks: %w", err)
+	}
+
+	var mapErr error
+	mapFunc := func(meta *applymetav1.ObjectMetaApplyConfiguration, spec *applycorev1.PodSpecApplyConfiguration) (*applymetav1.ObjectMetaApplyConfiguration, *applycorev1.PodSpecApplyConfiguration) {
+		if spec == nil || spec.RuntimeClassName == nil {
+			return meta, spec
+		}
+
+		targetPlatform, err := platforms.FromRuntimeClassString(*spec.RuntimeClassName)
+		if err != nil {
+			logger.Error("could not determine platform for runtime class", "runtime-class-name", *spec.RuntimeClassName, "err", err)
+			return meta, spec
+		}
+		if !platforms.IsSNP(targetPlatform) {
+			return meta, spec
+		}
+
+		var regularContainersCPU int64
+		for _, container := range spec.Containers {
+			regularContainersCPU += getCPUCount(container.Resources)
+		}
+		var initContainersCPU int64
+		for _, container := range spec.InitContainers {
+			cpuCount := getCPUCount(container.Resources)
+			initContainersCPU += cpuCount
+			// Sidecar containers remain running alongside the actual application, consuming CPU resources
+			if container.RestartPolicy != nil && *container.RestartPolicy == corev1.ContainerRestartPolicyAlways {
+				regularContainersCPU += cpuCount
+			}
+		}
+		podLevelCPU := getCPUCount(spec.Resources)
+
+		// Convert milliCPUs to number of CPUs (rounding up), and add 1 for hypervisor overhead
+		totalMilliCPUs := max(regularContainersCPU, initContainersCPU, podLevelCPU)
+		cpuCount := strconv.FormatInt((totalMilliCPUs+999)/1000+1, 10)
+
+		platform := strings.ToLower(targetPlatform.String())
+
+		genoa, milan := string(manifest.Genoa), string(manifest.Milan)
+		// Ensure we pre-calculated the required blocks
+		if snpIDBlocks[platform] == nil || snpIDBlocks[platform][cpuCount] == nil ||
+			snpIDBlocks[platform][cpuCount][genoa].IDBlock == "" || snpIDBlocks[platform][cpuCount][milan].IDBlock == "" {
+			mapErr = fmt.Errorf("missing ID block configuration for runtime %s with %s CPUs", platform, cpuCount)
+			return meta, spec
+		}
+
+		if meta == nil {
+			meta = &applymetav1.ObjectMetaApplyConfiguration{}
+		}
+		if meta.Annotations == nil {
+			meta.Annotations = make(map[string]string, 6)
+		}
+		meta.Annotations[idBlockAnnotation+genoa] = snpIDBlocks[platform][cpuCount][genoa].IDBlock
+		meta.Annotations[idBlockAnnotation+milan] = snpIDBlocks[platform][cpuCount][milan].IDBlock
+		meta.Annotations[idAuthAnnotationKey+genoa] = snpIDBlocks[platform][cpuCount][genoa].IDAuth
+		meta.Annotations[idAuthAnnotationKey+milan] = snpIDBlocks[platform][cpuCount][milan].IDAuth
+		meta.Annotations[guestPolicyAnnotationKey+genoa] = strconv.FormatUint(abi.SnpPolicyToBytes(snpIDBlocks[platform][cpuCount][genoa].GuestPolicy), 10)
+		meta.Annotations[guestPolicyAnnotationKey+milan] = strconv.FormatUint(abi.SnpPolicyToBytes(snpIDBlocks[platform][cpuCount][milan].GuestPolicy), 10)
+
+		return meta, spec
+	}
+
+	kuberesource.MapPodSpecWithMeta(res, mapFunc)
+	return mapErr
+}
+
+func getCPUCount(resources *applycorev1.ResourceRequirementsApplyConfiguration) int64 {
+	if resources != nil && resources.Limits != nil {
+		return resources.Limits.Cpu().MilliValue()
+	}
+	return 0
+}
+
 type generateFlags struct {
 	policyPath               string
 	settingsPath             string

cli/main.go

@@ -78,6 +78,7 @@ func buildVersionString() (string, error) {
 		}
 		for _, snp := range values.SNP {
 			fmt.Fprintf(versionsWriter, "\t- product name:\t%s\n", snp.ProductName)
+			fmt.Fprintf(versionsWriter, "\t  vCPUs:\t%d\n", snp.CPUs)
 			fmt.Fprintf(versionsWriter, "\t  launch digest:\t%s\n", snp.TrustedMeasurement.String())
 			fmt.Fprint(versionsWriter, "\t  default SNP TCB:\t\n")
 			printOptionalSVN("bootloader", snp.MinimumTCB.BootloaderVersion)

internal/kuberesource/sets.go

@@ -156,11 +156,15 @@ func MultiCPU() []any {
 							Container().
 								WithName("multi-cpu").
 								WithImage("ghcr.io/edgelesssys/contrast/ubuntu:24.04@sha256:0f9e2b7901aa01cf394f9e1af69387e2fd4ee256fd8a95fb9ce3ae87375a31e6").
-								WithCommand("/usr/bin/bash", "-c", "sleep infinity").
+								WithCommand("/bin/sh", "-c", "sleep inf").
 								WithResources(ResourceRequirements().
-									// Explicitly set a CPU limit to test assignement of CPUs to VMs.
+									WithRequests(corev1.ResourceList{
+										corev1.ResourceMemory: resource.MustParse("200Mi"),
+									}).
+									// Explicitly set a CPU limit to test assignment of CPUs to VMs.
 									WithLimits(corev1.ResourceList{
-										corev1.ResourceCPU: resource.MustParse("1"),
+										corev1.ResourceMemory: resource.MustParse("200Mi"),
+										corev1.ResourceCPU:    resource.MustParse("1"),
 									}),
 								),
 						),

internal/manifest/referencevalues.go

@@ -200,6 +200,10 @@ type SNPReferenceValues struct {
 	PlatformInfo            abi.SnpPlatformInfo
 	MinimumMitigationVector uint64
 	AllowedChipIDs          []HexString
+	// CPUs is the number of vCPUs assigned to the VM.
+	// This field is purely informative as [SNPReferenceValues.TrustedMeasurement]
+	// already implicitly contains the number of vCPUs
+	CPUs uint64
 }
 
 // Validate checks the validity of all fields in the AKS reference values.

nodeinstaller/internal/kataconfig/config.go

@@ -4,17 +4,11 @@
 package kataconfig
 
 import (
-	"bytes"
 	_ "embed"
-	"encoding/json"
 	"fmt"
 	"path/filepath"
-	"strings"
 
 	"github.com/edgelesssys/contrast/internal/platforms"
-	"github.com/google/go-sev-guest/abi"
-	"github.com/google/go-sev-guest/kds"
-	"github.com/google/go-sev-guest/proto/sevsnp"
 	"github.com/pelletier/go-toml/v2"
 )
 
@@ -31,9 +25,6 @@ var (
 	//go:embed configuration-qemu-snp.toml
 	kataBareMetalQEMUSNPBaseConfig string
 
-	//go:embed snp-id-blocks.json
-	snpIDBlocks string
-
 	// RuntimeNamePlaceholder is the placeholder for the per-runtime path (i.e. /opt/edgeless/contrast-cc...) in the target file paths.
 	RuntimeNamePlaceholder = "@@runtimeName@@"
 )
@@ -43,10 +34,10 @@ func KataRuntimeConfig(
 	baseDir string,
 	platform platforms.Platform,
 	qemuExtraKernelParams string,
-	snpIDBlock SnpIDBlock,
 	imagepullerConfigPath string,
 	debug bool,
 ) (*Config, error) {
+	var customContrastAnnotations []string
 	var config Config
 	switch {
 	case platforms.IsTDX(platform):
@@ -59,11 +50,13 @@ func KataRuntimeConfig(
 			return nil, fmt.Errorf("failed to unmarshal kata runtime configuration: %w", err)
 		}
 
+		for _, productLine := range []string{"_Milan", "_Genoa"} {
+			for _, annotationType := range []string{"snp_id_block", "snp_id_auth", "snp_guest_policy"} {
+				customContrastAnnotations = append(customContrastAnnotations, annotationType+productLine)
+			}
+		}
+
 		config.Hypervisor["qemu"]["firmware"] = filepath.Join(baseDir, "snp", "share", "OVMF.fd")
-		// Add SNP ID block to protect against migration attacks.
-		config.Hypervisor["qemu"]["snp_id_block"] = snpIDBlock.IDBlock
-		config.Hypervisor["qemu"]["snp_id_auth"] = snpIDBlock.IDAuth
-		config.Hypervisor["qemu"]["snp_guest_policy"] = abi.SnpPolicyToBytes(snpIDBlock.GuestPolicy)
 	default:
 		return nil, fmt.Errorf("unsupported platform: %s", platform)
 	}
@@ -105,7 +98,7 @@ func KataRuntimeConfig(
 	config.Hypervisor["qemu"]["enable_debug"] = debug
 	// Disable all annotations, as we don't support these. Some will mess up measurements,
 	// others bypass things you can archive via correct resource declaration anyway.
-	config.Hypervisor["qemu"]["enable_annotations"] = []string{"cc_init_data"}
+	config.Hypervisor["qemu"]["enable_annotations"] = append(customContrastAnnotations, "cc_init_data")
 	// Fix and align guest memory calculation.
 	config.Hypervisor["qemu"]["default_memory"] = platforms.DefaultMemoryInMebiBytes(platform)
 	config.Runtime["sandbox_cgroup_only"] = true
@@ -115,36 +108,6 @@ func KataRuntimeConfig(
 	return &config, nil
 }
 
-// SnpIDBlock represents the SNP ID block and ID auth used for SEV-SNP guests.
-type SnpIDBlock struct {
-	IDBlock     string        `json:"idBlock"`
-	IDAuth      string        `json:"idAuth"`
-	GuestPolicy abi.SnpPolicy `json:"guestPolicy"`
-}
-
-// platform -> product -> snpIDBlock.
-type snpIDBlockMap map[string]map[string]SnpIDBlock
-
-// SnpIDBlockForPlatform returns the embedded SNP ID block and ID auth for the given platform and product.
-func SnpIDBlockForPlatform(platform platforms.Platform, productName sevsnp.SevProduct_SevProductName) (SnpIDBlock, error) {
-	var blocks snpIDBlockMap
-	decoder := json.NewDecoder(bytes.NewReader([]byte(snpIDBlocks)))
-	decoder.DisallowUnknownFields()
-	if err := decoder.Decode(&blocks); err != nil {
-		return SnpIDBlock{}, fmt.Errorf("unmarshaling embedded SNP ID blocks: %w", err)
-	}
-	blockForPlatform, ok := blocks[strings.ToLower(platform.String())]
-	if !ok {
-		return SnpIDBlock{}, fmt.Errorf("no SNP ID block found for platform %s", platform)
-	}
-	productLine := kds.ProductLine(&sevsnp.SevProduct{Name: productName})
-	block, ok := blockForPlatform[productLine]
-	if !ok {
-		return SnpIDBlock{}, fmt.Errorf("no SNP ID block found for product %s", productLine)
-	}
-	return block, nil
-}
-
 // Config is the configuration for the Kata runtime.
 // Source: https://github.com/kata-containers/kata-containers/blob/4029d154ba0c26fcf4a8f9371275f802e3ef522c/src/runtime/pkg/katautils/Config.go
 // This is a simplified version of the actual configuration.

nodeinstaller/internal/kataconfig/config_test.go

@@ -51,11 +51,7 @@ func TestKataRuntimeConfig(t *testing.T) {
 			require := require.New(t)
 			assert := assert.New(t)
 
-			snpIDBlock := kataconfig.SnpIDBlock{
-				IDAuth:  "PLACEHOLDER_ID_AUTH",
-				IDBlock: "PLACEHOLDER_ID_BLOCK",
-			}
-			cfg, err := kataconfig.KataRuntimeConfig("/", platform, "", snpIDBlock, "", false)
+			cfg, err := kataconfig.KataRuntimeConfig("/", platform, "", "", false)
 			require.NoError(err)
 
 			configBytes, err := cfg.Marshal()

nodeinstaller/internal/kataconfig/testdata/expected-configuration-qemu-snp-gpu.toml

@@ -20,7 +20,7 @@ disable_image_nvdimm = true
 disable_nesting_checks = true
 disable_selinux = false
 disable_vhost_net = false
-enable_annotations = ['cc_init_data']
+enable_annotations = ['snp_id_block_Milan', 'snp_id_auth_Milan', 'snp_guest_policy_Milan', 'snp_id_block_Genoa', 'snp_id_auth_Genoa', 'snp_guest_policy_Genoa', 'cc_init_data']
 enable_debug = false
 enable_guest_swap = false
 enable_hugepages = false
@@ -57,9 +57,9 @@ rx_rate_limiter_max_rate = 0
 seccompsandbox = ''
 sev_snp_guest = true
 shared_fs = 'none'
-snp_guest_policy = 131072
-snp_id_auth = 'PLACEHOLDER_ID_AUTH'
-snp_id_block = 'PLACEHOLDER_ID_BLOCK'
+snp_guest_policy = 196608
+snp_id_auth = ''
+snp_id_block = ''
 tx_rate_limiter_max_rate = 0
 use_legacy_serial = false
 valid_entropy_sources = ['/dev/urandom', '/dev/random', '']

nodeinstaller/internal/kataconfig/testdata/expected-configuration-qemu-snp.toml

@@ -19,7 +19,7 @@ disable_image_nvdimm = true
 disable_nesting_checks = true
 disable_selinux = false
 disable_vhost_net = false
-enable_annotations = ['cc_init_data']
+enable_annotations = ['snp_id_block_Milan', 'snp_id_auth_Milan', 'snp_guest_policy_Milan', 'snp_id_block_Genoa', 'snp_id_auth_Genoa', 'snp_guest_policy_Genoa', 'cc_init_data']
 enable_debug = false
 enable_guest_swap = false
 enable_hugepages = false
@@ -56,9 +56,9 @@ rx_rate_limiter_max_rate = 0
 seccompsandbox = ''
 sev_snp_guest = true
 shared_fs = 'none'
-snp_guest_policy = 131072
-snp_id_auth = 'PLACEHOLDER_ID_AUTH'
-snp_id_block = 'PLACEHOLDER_ID_BLOCK'
+snp_guest_policy = 196608
+snp_id_auth = ''
+snp_id_block = ''
 tx_rate_limiter_max_rate = 0
 use_legacy_serial = false
 valid_entropy_sources = ['/dev/urandom', '/dev/random', '']

nodeinstaller/internal/kataconfig/update-testdata/main.go

@@ -48,11 +48,6 @@ func main() {
 		},
 	}
 
-	snpIDBlock := kataconfig.SnpIDBlock{
-		IDAuth:  "PLACEHOLDER_ID_AUTH",
-		IDBlock: "PLACEHOLDER_ID_BLOCK",
-	}
-
 	for platform, platformConfig := range platforms {
 		upstreamFile := filepath.Join(tarball, "opt", "kata", "share", "defaults", "kata-containers", fmt.Sprintf("configuration-%s.toml", platformConfig.upstream))
 		configFile := filepath.Join(gitroot, "nodeinstaller", "internal", "kataconfig", fmt.Sprintf("configuration-%s.toml", platformConfig.config))
@@ -80,7 +75,7 @@ func main() {
 			log.Fatalf("failed to write new config: %s", err)
 		}
 
-		cfg, err := kataconfig.KataRuntimeConfig("/", platform, "", snpIDBlock, "", false)
+		cfg, err := kataconfig.KataRuntimeConfig("/", platform, "", "", false)
 		if err != nil {
 			log.Fatalf("failed to create config: %s", err)
 		}