cli: add SNP ID block annotations to Pods based on CPU requirements

cli/cmd/assets/snp-id-blocks.json

@@ -0,0 +1 @@
+"THIS FILE IS REPLACED DURING RELEASE BUILD TO INCLUDE SNP ID BLOCKS"

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"
@@ -42,6 +43,22 @@ const (
 //go:embed assets/image-replacements.txt
 var ReleaseImageReplacements []byte
 
+// SNPIDBlockData 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 SNPIDBlockData []byte
+
+// SNPIDBlocks maps runtime -> cpu_count -> product_line -> [SNPIDBlock].
+type SNPIDBlocks map[string]map[string]map[string]SNPIDBlock
+
+// 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.
@@ -514,6 +522,10 @@ func patchTargets(fileMap map[string][]*unstructured.Unstructured, imageReplacem
 			return nil, err
 		}
 
+		if err := patchIDBlockAnnotation(res); err != nil {
+			return nil, fmt.Errorf("injecting ID block annotations: %w", err)
+		}
+
 		return res, nil
 	})
 }
@@ -703,6 +715,80 @@ func patchRuntimeClassName(defaultRuntimeHandler string) func(*applycorev1.PodSp
 	}
 }
 
+func patchIDBlockAnnotation(res any) error {
+	var snpIDBlocks SNPIDBlocks
+	if err := json.Unmarshal(SNPIDBlockData, &snpIDBlocks); err != nil {
+		return fmt.Errorf("unmarshal SNP ID blocks: %w", err)
+	}
+
+	mapFunc := func(meta *applymetav1.ObjectMetaApplyConfiguration, spec *applycorev1.PodSpecApplyConfiguration) (*applymetav1.ObjectMetaApplyConfiguration, *applycorev1.PodSpecApplyConfiguration, error) {
+		if spec == nil || spec.RuntimeClassName == nil {
+			return meta, spec, nil
+		}
+
+		targetPlatform, err := platforms.FromRuntimeClassString(*spec.RuntimeClassName)
+		if err != nil {
+			return meta, spec, fmt.Errorf("could not determine platform for runtime class %q: %w", *spec.RuntimeClassName, err)
+		}
+		if !platforms.IsSNP(targetPlatform) {
+			return meta, spec, nil
+		}
+
+		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 == "" {
+			return meta, spec, fmt.Errorf("missing ID block configuration for runtime %s with %s CPUs", platform, cpuCount)
+		}
+
+		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, nil
+	}
+
+	_, err := kuberesource.MapPodSpecWithMetaAndErrors(res, mapFunc)
+	return err
+}
+
+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/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

@@ -47,6 +47,7 @@ func KataRuntimeConfig(
 	imagepullerConfigPath string,
 	debug bool,
 ) (*Config, error) {
+	var customContrastAnnotations []string
 	var config Config
 	switch {
 	case platforms.IsTDX(platform):
@@ -62,6 +63,12 @@ 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
@@ -108,7 +115,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
@@ -127,8 +134,8 @@ type SnpIDBlock struct {
 	GuestPolicy abi.SnpPolicy `json:"guestPolicy"`
 }
 
-// platform -> product -> snpIDBlock.
-type snpIDBlockMap map[string]map[string]SnpIDBlock
+// platform -> cpu_count -> product -> snpIDBlock.
+type snpIDBlockMap map[string]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) {
@@ -142,8 +149,12 @@ func SnpIDBlockForPlatform(platform platforms.Platform, productName sevsnp.SevPr
 	if !ok {
 		return SnpIDBlock{}, fmt.Errorf("no SNP ID block found for platform %s", platform)
 	}
+	// TODO: Get correct ID block based on requested vCPU count at runtime
+	if blockForPlatform["1"] == nil {
+		return SnpIDBlock{}, fmt.Errorf("no SNP ID blocks found for platform %s", platform)
+	}
 	productLine := kds.ProductLine(&sevsnp.SevProduct{Name: productName})
-	block, ok := blockForPlatform[productLine]
+	block, ok := blockForPlatform["1"][productLine]
 	if !ok {
 		return SnpIDBlock{}, fmt.Errorf("no SNP ID block found for product %s", productLine)
 	}

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

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

packages/by-name/contrast/cli/package.nix

@@ -9,6 +9,7 @@
   installShellFiles,
   contrastPkgsStatic,
   reference-values,
+  snp-id-blocks,
 }:
 
 buildGoModule (finalAttrs: {
@@ -56,6 +57,7 @@ buildGoModule (finalAttrs: {
     install -D ${lib.getExe contrastPkgsStatic.kata.genpolicy} cli/genpolicy/assets/genpolicy-kata
     install -D ${kata.genpolicy.rules}/genpolicy-rules.rego cli/genpolicy/assets/genpolicy-rules-kata.rego
     install -D ${reference-values} internal/manifest/assets/reference-values.json
+    install -D ${snp-id-blocks} cli/cmd/assets/snp-id-blocks.json
   '';
 
   # postPatch will be overwritten by the .#base.contrast.cli-release derivation, prePatch won't.

packages/by-name/contrast/reference-values/package.nix

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BUSL-1.1
 
 {
+  lib,
   kata,
   OVMF-TDX,
   node-installer-image,
@@ -24,23 +25,29 @@ let
         platformInfo = {
           SMTEnabled = true;
         };
-        launch-digest = kata.calculateSnpLaunchDigest {
-          inherit os-image;
-          inherit (node-installer-image) withDebug;
-        };
+        vcpuCounts = lib.range 1 8;
+        products = [
+          "Milan"
+          "Genoa"
+        ];
+
+        generateRefVal =
+          vcpus: product:
+          let
+            launch-digest = kata.calculateSnpLaunchDigest {
+              inherit os-image vcpus;
+              inherit (node-installer-image) withDebug;
+            };
+            filename = "${lib.toLower product}.hex";
+          in
+          {
+            inherit guestPolicy platformInfo;
+            trustedMeasurement = builtins.readFile "${launch-digest}/${filename}";
+            productName = product;
+            cpus = vcpus;
+          };
       in
-      [
-        {
-          inherit guestPolicy platformInfo;
-          trustedMeasurement = builtins.readFile "${launch-digest}/milan.hex";
-          productName = "Milan";
-        }
-        {
-          inherit guestPolicy platformInfo;
-          trustedMeasurement = builtins.readFile "${launch-digest}/genoa.hex";
-          productName = "Genoa";
-        }
-      ];
+      builtins.concatLists (map (vcpus: map (product: generateRefVal vcpus product) products) vcpuCounts);
   };
 
   snpRefVals = snpRefValsWith node-installer-image.os-image;

packages/by-name/contrast/snp-id-blocks/package.nix

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BUSL-1.1
 
 {
+  lib,
   kata,
   calculateSnpIDBlock,
   node-installer-image,
@@ -12,27 +13,40 @@ let
     os-image:
     let
       guestPolicy = builtins.fromJSON (builtins.readFile ../reference-values/snpGuestPolicyQEMU.json);
-      launch-digest = kata.calculateSnpLaunchDigest {
-        inherit os-image;
-        inherit (node-installer-image) withDebug;
-      };
-      idBlocks = calculateSnpIDBlock {
-        snp-launch-digest = launch-digest;
-        snp-guest-policy = ../reference-values/snpGuestPolicyQEMU.json;
-      };
+      idBlocksForVcpus =
+        vcpus:
+        let
+          launch-digest = kata.calculateSnpLaunchDigest {
+            inherit os-image vcpus;
+            inherit (node-installer-image) withDebug;
+          };
+          idBlocks = calculateSnpIDBlock {
+            snp-launch-digest = launch-digest;
+            snp-guest-policy = ../reference-values/snpGuestPolicyQEMU.json;
+          };
+        in
+        {
+          Milan = {
+            idBlock = builtins.readFile "${idBlocks}/id-block-milan.base64";
+            idAuth = builtins.readFile "${idBlocks}/id-auth-milan.base64";
+            inherit guestPolicy;
+          };
+          Genoa = {
+            idBlock = builtins.readFile "${idBlocks}/id-block-genoa.base64";
+            idAuth = builtins.readFile "${idBlocks}/id-auth-genoa.base64";
+            inherit guestPolicy;
+          };
+        };
+
+      vcpuCounts = lib.range 1 8;
+      allVcpuBlocks = builtins.listToAttrs (
+        map (vcpus: {
+          name = toString vcpus;
+          value = idBlocksForVcpus vcpus;
+        }) vcpuCounts
+      );
     in
-    {
-      Milan = {
-        idBlock = builtins.readFile "${idBlocks}/id-block-milan.base64";
-        idAuth = builtins.readFile "${idBlocks}/id-auth-milan.base64";
-        inherit guestPolicy;
-      };
-      Genoa = {
-        idBlock = builtins.readFile "${idBlocks}/id-block-genoa.base64";
-        idAuth = builtins.readFile "${idBlocks}/id-auth-genoa.base64";
-        inherit guestPolicy;
-      };
-    };
+    allVcpuBlocks;
 in
 
 builtins.toFile "snp-id-blocks.json" (