Add cos_gpu_installer digest verification

Author: meetrajvala

launcher/container_runner.go

@@ -169,7 +169,7 @@ func NewRunner(ctx context.Context, cdClient *containerd.Client, token oauth2.To
 		specOpts = append(specOpts, oci.WithDevShmSize(launchSpec.DevShmSize))
 	}
 
-	if launchSpec.Experiments.EnableGpuDriverInstallation && launchSpec.InstallGpuDriver {
+	if launchSpec.Experiments.EnableConfidentialGPUSupport && launchSpec.InstallGpuDriver {
 		gpuMounts := []specs.Mount{
 			{
 				Type:        "volume",
@@ -341,6 +341,12 @@ func (r *ContainerRunner) measureCELEvents(ctx context.Context) error {
 		return fmt.Errorf("failed to measure memory monitoring state: %v", err)
 	}
 
+	if r.launchSpec.Experiments.EnableConfidentialGPUSupport && r.launchSpec.InstallGpuDriver {
+		if err := r.measureGPUCCMode(); err != nil {
+			return fmt.Errorf("failed to measure gpu cc mode : %v", err)
+		}
+	}
+
 	separator := cel.CosTlv{
 		EventType:    cel.LaunchSeparatorType,
 		EventContent: nil, // Success
@@ -418,6 +424,17 @@ func (r *ContainerRunner) measureMemoryMonitor() error {
 	return nil
 }
 
+func (r *ContainerRunner) measureGPUCCMode() error {
+	ccMode, err := gpu.GetGPUCCMode()
+	if err != nil {
+		return err
+	}
+	if err := r.attestAgent.MeasureEvent(cel.CosTlv{EventType: cel.GpuCCModeType, EventContent: []byte(ccMode)}); err != nil {
+		return err
+	}
+	return nil
+}
+
 // Retrieves the default OIDC token from the attestation service, and returns how long
 // to wait before attemping to refresh it.
 // The token file will be written to a tmp file and then renamed.

launcher/image/preload.sh

@@ -3,6 +3,9 @@
 readonly OEM_PATH='/usr/share/oem'
 readonly CS_PATH="${OEM_PATH}/confidential_space"
 readonly EXPERIMENTS_BINARY="confidential_space_experiments"
+readonly GPU_REF_VALUES_PATH="${CS_PATH}/gpu"
+readonly COS_GPU_INSTALLER_IMAGE_REF="${GPU_REF_VALUES_PATH}/cos_gpu_installer_image_ref"
+readonly COS_GPU_INSTALLER_IMAGE_DIGEST="${GPU_REF_VALUES_PATH}/cos_gpu_installer_image_digest"
 
 copy_launcher() {
   cp launcher "${CS_PATH}/cs_container_launcher"
@@ -100,6 +103,41 @@ configure_systemd_units_for_hardened() {
   disable_unit "var-lib-toolbox.mount"
 }
 
+get_cos_gpu_installer_image_digest() {
+  local image_ref="${1}"
+  local registry
+  local repo_with_image_name
+  local tag
+  local manifest_url
+  local image_digest
+
+  if [[ "$image_ref" =~ ^([^/]+)/([^:]+):([^:]+)$ ]]; then
+    registry="${BASH_REMATCH[1]}"
+    repo_with_image_name="${BASH_REMATCH[2]}"
+    tag="${BASH_REMATCH[3]}"
+  else
+    echo "Error: Invalid image reference format: $image_ref" >&2
+    return 1
+  fi
+
+  manifest_url="https://${registry}/v2/${repo_name}/manifests/${tag}"
+  image_digest=$(curl -s --head ${manifest_url} | grep -i Docker-Content-Digest | cut -d' ' -f2)
+  echo "${image_digest}"
+}
+
+
+set_gpu_driver_ref_values() {
+  local cos_gpu_installer_image_ref
+  local cos_gpu_installer_image_digest
+
+  mkdir ${GPU_REF_VALUES_PATH}
+  cos_gpu_installer_image_ref=$(cos-extensions list -- --gpu-installer)
+  cos_gpu_installer_image_digest=$(get_cos_gpu_installer_image_digest ${cos_gpu_installer_image_ref})
+  
+  echo ${cos_gpu_installer_image_ref} >> ${COS_GPU_INSTALLER_IMAGE_REF}
+  echo ${cos_gpu_installer_image_digest} >> ${COS_GPU_INSTALLER_IMAGE_DIGEST}
+}
+
 main() {
   mount -o remount,rw ${OEM_PATH}
   mkdir ${CS_PATH}
@@ -110,6 +148,7 @@ main() {
   copy_experiment_client
   # Install container launcher.
   copy_launcher
+  set_gpu_driver_ref_values
   setup_launcher_systemd_unit
   # Minimum required COS version for 'e': cos-dev-105-17222-0-0.
   # Minimum required COS version for 'm': cos-dev-113-18203-0-0.

launcher/internal/experiments/experiments.go

@@ -11,9 +11,10 @@ import (
 // Failure to unmarshal the experiment JSON data will result in an empty object being returned
 // to treat experiment flags as their default value. The error should still be checked.
 type Experiments struct {
-	EnableTestFeatureForImage   bool
-	EnableTempFSMount           bool
-	EnableGpuDriverInstallation bool
+	EnableTestFeatureForImage    bool
+	EnableTempFSMount            bool
+	EnableGpuDriverInstallation  bool
+	EnableConfidentialGPUSupport bool
 }
 
 // New takes a filepath, opens the file, and calls ReadJsonInput with the contents

launcher/internal/gpu/config.go

@@ -5,4 +5,8 @@ const (
 	InstallationHostDir = "/var/lib/nvidia"
 	// InstallationContainerDir is the directory where gpu drivers will be available on the workload container.
 	InstallationContainerDir = "/usr/local/nvidia"
+	// InstallerImageRefFile is a filename which has the container image reference of cos_gpu_installer.
+	InstallerImageRefFile = "/usr/share/oem/confidential_space/gpu/cos_gpu_installer_image_ref"
+	// InstallerImageDigestFile is a filename which has the container image digest of cos_gpu_installer.
+	InstallerImageDigestFile = "/usr/share/oem/confidential_space/gpu/cos_gpu_installer_image_digest"
 )

launcher/internal/gpu/driverinstaller.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
+	"regexp"
 	"strings"
 
 	"cos.googlesource.com/cos/tools.git/src/cmd/cos_gpu_installer/deviceinfo"
@@ -18,19 +19,29 @@ import (
 	"github.com/opencontainers/runtime-spec/specs-go"
 )
 
+// CCMode enums
 const (
-	installerContainerID = "tee-gpu-driver-installer-container"
-	installerSnapshotID  = "tee-gpu-driver-installer-snapshot"
+	CCModeON             CCMode = "ON"
+	CCModeOFF            CCMode = "OFF"
+	installerContainerID        = "tee-gpu-driver-installer-container"
+	installerSnapshotID         = "tee-gpu-driver-installer-snapshot"
 )
 
-var supportedGpuTypes = []deviceinfo.GPUType{
-	deviceinfo.L4,
-	deviceinfo.T4,
-	deviceinfo.A100_40GB,
-	deviceinfo.A100_80GB,
+var supportedCGPUTypes = []deviceinfo.GPUType{
 	deviceinfo.H100,
 }
 
+// CCMode represents the status confidential computing mode of the GPU.
+type CCMode string
+
+func (ccm CCMode) isValid() error {
+	switch ccm {
+	case CCModeOFF, CCModeON:
+		return nil
+	}
+	return fmt.Errorf("invalid gpu cc mode: %s", ccm)
+}
+
 // DriverInstaller contains information about the GPU driver installer settings
 type DriverInstaller struct {
 	cdClient   *containerd.Client
@@ -63,7 +74,7 @@ func (di *DriverInstaller) InstallGPUDrivers(ctx context.Context) error {
 	}
 
 	if !gpuType.OpenSupported() {
-		return fmt.Errorf("unsupported GPU type %s, please retry with one of the supported GPU types: %v", gpuType.String(), supportedGpuTypes)
+		return fmt.Errorf("unsupported GPU type %s, please retry with one of the supported confidential GPU types: %v", gpuType.String(), supportedCGPUTypes)
 	}
 
 	ctx = namespaces.WithNamespace(ctx, namespaces.Default)
@@ -79,6 +90,16 @@ func (di *DriverInstaller) InstallGPUDrivers(ctx context.Context) error {
 		return fmt.Errorf("failed to pull installer image: %v", err)
 	}
 
+	installerDigest := image.Target().Digest.String()
+	expectedInstallerDigest, err := os.ReadFile(InstallerImageDigestFile)
+	if err != nil {
+		return fmt.Errorf("failed to read reference image digest from file %s : %v", InstallerImageDigestFile, err)
+	}
+
+	if installerDigest != string(expectedInstallerDigest) {
+		return fmt.Errorf("cos_gpu_installer image digest verification failed - expected : %s, actual : %s", expectedInstallerDigest, installerDigest)
+	}
+
 	mounts := []specs.Mount{
 		{
 			Type:        "volume",
@@ -153,7 +174,7 @@ func (di *DriverInstaller) InstallGPUDrivers(ctx context.Context) error {
 		return fmt.Errorf("failed to verify GPU driver installation: %v", err)
 	}
 
-	ccEnabled, err := isGPUCCModeEnabled(di.logger, gpuType)
+	ccEnabled, err := isGPUCCModeEnabled()
 	if err != nil {
 		return fmt.Errorf("failed to check confidential compute mode status: %v", err)
 	}
@@ -169,11 +190,11 @@ func (di *DriverInstaller) InstallGPUDrivers(ctx context.Context) error {
 }
 
 func getInstallerImageReference() (string, error) {
-	installerImageRefBytes, err := exec.Command("cos-extensions", "list", "--", "--gpu-installer").Output()
+	imageRefBytes, err := os.ReadFile(InstallerImageRefFile)
 	if err != nil {
 		return "", fmt.Errorf("failed to get the cos-gpu-installer version: %v", err)
 	}
-	installerImageRef := strings.TrimSpace(string(installerImageRefBytes))
+	installerImageRef := strings.TrimSpace(string(imageRefBytes))
 	return installerImageRef, nil
 }
 
@@ -208,20 +229,42 @@ func setGPUStateToReady() error {
 	return nil
 }
 
-func isGPUCCModeEnabled(logger logging.Logger, gpuType deviceinfo.GPUType) (bool, error) {
-	// Run nvidia-smi conf-compute command to check if confidential compute mode is ON.
+func isGPUCCModeEnabled() (bool, error) {
+	ccMode, err := GetGPUCCMode()
+	if err != nil {
+		return false, err
+	}
+	return ccMode == CCModeON, nil
+}
+
+// GetGPUCCMode executes nvidia-smi to determine the current Confidential Computing (CC) mode status of the GPU.
+// It returns the CC mode ("ON" or "OFF") and an error if the command fails or if the output cannot be parsed.
+func GetGPUCCMode() (CCMode, error) {
+	// Run nvidia-smi conf-compute command to get the confidential computing mode status.
 	nvidiaSmiCmd := fmt.Sprintf("%s/bin/nvidia-smi", InstallationHostDir)
 	ccModeOutput, err := exec.Command(nvidiaSmiCmd, "conf-compute", "-f").Output()
-	// The nvidia-smi conf-compute command fails for GPU which doesn't support confidential computing.
-	// This check would bypass nvidia-smi conf-compute command for GPU not having confidential compute support.
-	if strings.Contains(string(ccModeOutput), "No CC capable devices found") {
-		logger.Info(fmt.Sprintf("Confidential Computing is not supported for GPU type : %s", gpuType.String()))
-		return false, nil
+	if err != nil {
+		return "", err
 	}
+	ccMode, err := parseCCStatus(string(ccModeOutput))
 	if err != nil {
-		return false, err
+		return "", err
+	}
+	return CCMode(ccMode), nil
+}
+
+func parseCCStatus(output string) (CCMode, error) {
+	re := regexp.MustCompile(`CC status:\s*(ON|OFF)`)
+	match := re.FindStringSubmatch(output)
+
+	if len(match) < 2 {
+		return "", fmt.Errorf("CC status not found in output: %s", output)
+	}
+	ccMode := CCMode(match[1])
+	if err := ccMode.isValid(); err != nil {
+		return "", err
 	}
-	return strings.Contains(string(ccModeOutput), "CC status: ON"), nil
+	return ccMode, nil
 }
 
 func launchNvidiaPersistencedProcess(logger logging.Logger) error {

launcher/internal/gpu/driverinstaller_test.go

@@ -0,0 +1,94 @@
+package gpu
+
+import (
+	"testing"
+)
+
+func TestParseCCStatus(t *testing.T) {
+	tests := []struct {
+		name    string
+		output  string
+		want    CCMode
+		wantErr bool
+	}{
+		{
+			name:    "CC ON",
+			output:  "some text CC status: ON more text",
+			want:    CCModeON,
+			wantErr: false,
+		},
+		{
+			name:    "CC OFF",
+			output:  "another line CC status: OFF at the end",
+			want:    CCModeOFF,
+			wantErr: false,
+		},
+		{
+			name:    "CC ON at the beginning",
+			output:  "CC status: ON some other info",
+			want:    CCModeON,
+			wantErr: false,
+		},
+		{
+			name:    "CC OFF only",
+			output:  "CC status: OFF",
+			want:    CCModeOFF,
+			wantErr: false,
+		},
+		{
+			name:    "CC status not found",
+			output:  "No CC information here",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "CC status misspelled",
+			output:  "CC state: ON",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "CC value missing",
+			output:  "CC status:",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "Invalid CC value",
+			output:  "CC status: ENABLED",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "Multiple CC status lines - picks the first",
+			output:  "CC status: ON\nSome other info\nCC status: OFF",
+			want:    CCModeON,
+			wantErr: false,
+		},
+		{
+			name:    "Case insensitive match",
+			output:  "CC status: on",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "Whitespace around CC value",
+			output:  "CC status:  ON  ",
+			want:    CCModeON,
+			wantErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := parseCCStatus(tt.output)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("parseCCStatus() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("parseCCStatus() got = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

launcher/launcher/main.go

@@ -191,15 +191,15 @@ func startLauncher(launchSpec spec.LaunchSpec, serialConsole *os.File) error {
 
 	ctx := namespaces.WithNamespace(context.Background(), namespaces.Default)
 	if launchSpec.InstallGpuDriver {
-		if launchSpec.Experiments.EnableGpuDriverInstallation {
+		if launchSpec.Experiments.EnableConfidentialGPUSupport {
 			installer := gpu.NewDriverInstaller(containerdClient, launchSpec, logger)
 			err = installer.InstallGPUDrivers(ctx)
 			if err != nil {
 				return fmt.Errorf("failed to install gpu drivers: %v", err)
 			}
 		} else {
-			logger.Info("GPU installation experiment flag is not enabled for this project. Ensure that it is enabled when tee-install-gpu-driver is set to true")
-			return fmt.Errorf("gpu installation experiment flag is not enabled")
+			logger.Info("Confidential GPU support experiment flag is not enabled for this project. Ensure that it is enabled when tee-install-gpu-driver is set to true")
+			return fmt.Errorf("confidential gpu support experiment flag is not enabled")
 		}
 	}