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"
@@ -40,6 +41,18 @@ const (
 //go:embed assets/image-replacements.txt
 var ReleaseImageReplacements []byte
 
+// SNPIDBlocks contains the SNP ID blocks for different vCPU counts and CPU generations.
+//
+//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,7 @@ import (
 	"github.com/edgelesssys/contrast/internal/kuberesource"
 	"github.com/edgelesssys/contrast/internal/manifest"
 	"github.com/edgelesssys/contrast/internal/platforms"
+	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"
@@ -35,6 +37,10 @@ import (
 const (
 	contrastRoleAnnotationKey     = "contrast.edgeless.systems/pod-role"
 	workloadSecretIDAnnotationKey = "contrast.edgeless.systems/workload-secret-id"
+	hypervisorCPUCountAnnotation  = "io.katacontainers.config.hypervisor.default_vcpus"
+	idBlockAnnotation             = "contrast.edgeless.systems/snp-id-block/"
+	amdCPUGenerationMilan         = "Milan"
+	amdCPUGenerationGenoa         = "Genoa"
 )
 
 // NewGenerateCmd creates the contrast generate subcommand.
@@ -511,6 +517,10 @@ func patchTargets(logger *slog.Logger, fileMap map[string][]*unstructured.Unstru
 		replaceRuntimeClassName := patchRuntimeClassName(logger, runtimeHandler)
 		kuberesource.MapPodSpec(res, replaceRuntimeClassName)
 
+		if err := patchIDBlockAnnotation(res); err != nil {
+			return nil, fmt.Errorf("injecting ID block annotations: %w", err)
+		}
+
 		return res, nil
 	})
 }
@@ -692,6 +702,84 @@ func patchRuntimeClassName(logger *slog.Logger, defaultRuntimeHandler string) fu
 	}
 }
 
+func patchIDBlockAnnotation(res any) 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 {
+			mapErr = fmt.Errorf("determining platform from runtime class name %s: %w", *spec.RuntimeClassName, err)
+			return meta, spec
+		}
+		if !platforms.IsSNP(targetPlatform) {
+			return meta, spec
+		}
+
+		var regularContainersCPU int64
+		for _, container := range spec.Containers {
+			regularContainersCPU += getNeededCPUCount(container.Resources)
+		}
+		var initContainersCPU int64
+		for _, container := range spec.InitContainers {
+			cpuCount := getNeededCPUCount(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 := getNeededCPUCount(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())
+
+		// Ensure we pre-calculated the required blocks
+		if snpIDBlocks[platform] == nil || snpIDBlocks[platform][amdCPUGenerationGenoa] == nil || snpIDBlocks[platform][amdCPUGenerationMilan] == nil ||
+			snpIDBlocks[platform][amdCPUGenerationGenoa][cpuCount].IDBlock == "" || snpIDBlocks[platform][amdCPUGenerationMilan][cpuCount].IDBlock == "" {
+			mapErr = fmt.Errorf("missing ID block configuration for runtime %s with %s CPUs", platform, cpuCount)
+			return meta, spec
+		}
+
+		if meta.Annotations == nil {
+			meta.Annotations = make(map[string]string, 3)
+		}
+		meta.Annotations[idBlockAnnotation+platform] = snpIDBlocks[platform][amdCPUGenerationGenoa][cpuCount].IDBlock
+		meta.Annotations[idBlockAnnotation+platform] = snpIDBlocks[platform][amdCPUGenerationMilan][cpuCount].IDBlock
+		meta.Annotations[hypervisorCPUCountAnnotation] = cpuCount
+
+		return meta, spec
+	}
+
+	kuberesource.MapPodSpecWithMeta(res, mapFunc)
+	return mapErr
+}
+
+func getNeededCPUCount(resources *applycorev1.ResourceRequirementsApplyConfiguration) int64 {
+	if resources == nil {
+		return 0
+	}
+	var requests, limits int64
+	if resources.Requests != nil {
+		requests = resources.Requests.Cpu().MilliValue()
+	}
+	if resources.Limits != nil {
+		limits = resources.Limits.Cpu().MilliValue()
+	}
+	return max(requests, limits)
+}
+
 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

@@ -122,8 +122,8 @@ type SnpIDBlock struct {
 	GuestPolicy abi.SnpPolicy `json:"guestPolicy"`
 }
 
-// platform -> product -> snpIDBlock.
-type snpIDBlockMap map[string]map[string]SnpIDBlock
+// platform -> product -> cpuCount -> 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) {
@@ -133,7 +133,11 @@ func SnpIDBlockForPlatform(platform platforms.Platform, productName sevsnp.SevPr
 	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())]
+	// TODO: Get correct ID block based on requested vCPU count at runtime
+	if blocks["1"] == nil {
+		return SnpIDBlock{}, fmt.Errorf("no SNP ID blocks found for platform %s", platform)
+	}
+	blockForPlatform, ok := blocks["1"][strings.ToLower(platform.String())]
 	if !ok {
 		return SnpIDBlock{}, fmt.Errorf("no SNP ID block found for platform %s", platform)
 	}

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 release-cli derivation, prePatch won't.

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

@@ -24,23 +24,29 @@ let
         platformInfo = {
           SMTEnabled = true;
         };
-        launch-digest = kata.calculateSnpLaunchDigest {
-          inherit os-image;
-          debug = node-installer-image.debugRuntime;
-        };
+        vcpuCounts = builtins.genList (x: x + 1) 8;
+        products = [
+          "Milan"
+          "Genoa"
+        ];
+
+        generateRefVal =
+          vcpus: product:
+          let
+            launch-digest = kata.calculateSnpLaunchDigest {
+              inherit os-image vcpus;
+              debug = node-installer-image.debugRuntime;
+            };
+            filename = if product == "Milan" then "milan.hex" else "genoa.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

@@ -12,27 +12,40 @@ let
     os-image:
     let
       guestPolicy = builtins.fromJSON (builtins.readFile ../reference-values/snpGuestPolicyQEMU.json);
-      launch-digest = kata.calculateSnpLaunchDigest {
-        inherit os-image;
-        debug = node-installer-image.debugRuntime;
-      };
-      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;
+            debug = node-installer-image.debugRuntime;
+          };
+          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 = builtins.genList (x: x + 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" (

packages/by-name/kata/calculateSnpLaunchDigest/package.nix

@@ -12,6 +12,7 @@
 {
   os-image,
   debug,
+  vcpus,
 }:
 
 let
@@ -40,7 +41,7 @@ stdenvNoCC.mkDerivation {
     ${lib.getExe python3Packages.sev-snp-measure} \
       --mode snp \
       --ovmf ${ovmf-snp} \
-      --vcpus 1 \
+      --vcpus ${toString vcpus} \
       --vcpu-type EPYC-Milan \
       --kernel ${kernel} \
       --initrd ${initrd} \
@@ -49,7 +50,7 @@ stdenvNoCC.mkDerivation {
     ${lib.getExe python3Packages.sev-snp-measure} \
       --mode snp \
       --ovmf ${ovmf-snp} \
-      --vcpus 1 \
+      --vcpus ${toString vcpus} \
       --vcpu-type EPYC-Genoa \
       --kernel ${kernel} \
       --initrd ${initrd} \