address review comments

Author: meetrajvala

launcher/image/preload.sh

@@ -6,8 +6,8 @@ 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"
-readonly DRIVER_DIGEST_SHA256SUM="${GPU_REF_VALUES_PATH}/driver_digest_sha256sum"
-readonly DRIVERS_GCS_BUCKET="cos-nvidia-gpu-drivers"
+readonly DRIVER_DIGEST="${GPU_REF_VALUES_PATH}/driver_digest"
+readonly DRIVER_GCS_BUCKET="cos-nvidia-gpu-drivers"
 
 copy_launcher() {
   cp launcher "${CS_PATH}/cs_container_launcher"
@@ -128,24 +128,35 @@ get_cos_gpu_installer_image_digest() {
   echo "${image_digest}"
 }
 
-set_reference_driver_digest() {
+validate_sha256_hex() {
+  driver_digest="${1}"
+  # Check for the expected length of the SHA256 digest (64 hex characters)
+  if [ ${#driver_digest} -ne 64 ]; then
+    echo "Error: driver digest has an unexpected length: ${#driver_digest}, Expected 64." >&2
+    return 1
+  fi
+  # Check for valid hexadecimal string
+  if [[ ! ${driver_digest} =~ ^[0-9a-fA-F]+$ ]]; then
+    return "Error: driver digest ${driver_digest} is not a valid hexadecimal string." >&2
+    return 1
+  fi
+}
+
+store_driver_digest() {
+  local gpu_type="${1}"
   local driver_version
   local driver_digest_gcs_url
 
-  # Fetching the default driver version for H100 GPU.
-  driver_version=$(cos-extensions list -- --target-gpu NVIDIA_H100_80GB | grep DEFAULT | cut -d" " -f 1)
-  driver_digest_gcs_url="https://storage.googleapis.com/${DRIVERS_GCS_BUCKET}/sha256/NVIDIA-Linux-x86_64-${driver_version}.run.sha256"
-  if ! curl -sSL ${driver_digest_gcs_url} -o ${DRIVER_DIGEST_SHA256SUM}; then
+  # Fetching the default driver version for the given GPU.
+  driver_version=$(cos-extensions list -- --target-gpu ${gpu_type} | grep DEFAULT | cut -d" " -f 1)
+  driver_digest_gcs_url="https://storage.googleapis.com/${DRIVER_GCS_BUCKET}/sha256/NVIDIA-Linux-x86_64-${driver_version}.run.sha256"
+  if ! curl -sSL ${driver_digest_gcs_url} -o ${DRIVER_DIGEST}; then
     echo "Error: failed to download the driver digest file from ${driver_digest_gcs_url}." >&2
     return 1
   fi
 
-  driver_digest=$(cat ${DRIVER_DIGEST_SHA256SUM} | cut -d " " -f 1)
-  # Check for the expected length of the SHA256 digest (64 hex characters)
-  if [ ${#driver_digest} -ne 64 ]; then
-    echo "Error: driver digest has an unexpected length: ${#driver_digest}, Expected 64." >&2
-    return 1
-  fi
+  driver_digest=$(cat ${DRIVER_DIGEST} | cut -d " " -f 1)
+  validate_sha256_hex ${driver_digest}
 }
 
 
@@ -161,21 +172,17 @@ set_gpu_driver_ref_values() {
   fi
 
   cos_gpu_installer_image_digest=$(get_cos_gpu_installer_image_digest ${cos_gpu_installer_image_ref})
-  if [ -z "${cos_gpu_installer_image_ref}" ]; then
+  if [ -z "${cos_gpu_installer_image_digest}" ]; then
     echo "Error: get_cos_gpu_installer_image_digest returned an empty or invalid digest for: ${cos_gpu_installer_image_ref}." >&2
     return 1
   fi
 
   image_digest_hex_part=$(echo "${cos_gpu_installer_image_digest}" | sed 's/^sha256://' | tr -d '[:space:]')
-  # Check for the expected length of the SHA256 digest (64 hex characters)
-  if [ ${#image_digest_hex_part} -ne 64 ]; then
-    echo "Error: cos_gpu_installer image digest has an unexpected length: ${#image_digest_hex_part}, Expected 64." >&2
-    return 1
-  fi
+  validate_sha256_hex ${image_digest_hex_part}
 
   echo ${cos_gpu_installer_image_ref} > ${COS_GPU_INSTALLER_IMAGE_REF}
   echo ${cos_gpu_installer_image_digest} > ${COS_GPU_INSTALLER_IMAGE_DIGEST}
-  set_reference_driver_digest
+  store_driver_digest "NVIDIA_H100_80GB"
 }
 
 main() {

launcher/internal/gpu/config.go

@@ -10,5 +10,5 @@ const (
 	// 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"
 	// ReferenceDriverDigestFile is a filename which has the reference digest of nvidia driver installer.
-	ReferenceDriverDigestFile = "/usr/share/oem/confidential_space/gpu/driver_digest_sha256sum"
+	ReferenceDriverDigestFile = "/usr/share/oem/confidential_space/gpu/driver_digest"
 )

launcher/internal/gpu/driverinstaller.go

@@ -220,7 +220,7 @@ func verifyDriverDigest(driverFile, referenceHash string) error {
 		return err
 	}
 	if calculatedHash != referenceHash {
-		return fmt.Errorf("gpu driver digest verification failed - expected : %s, got : %s", referenceHash, calculatedHash)
+		return fmt.Errorf("GPU driver digest verification failed - expected : %s, got : %s", referenceHash, calculatedHash)
 	}
 	return nil
 }

launcher/internal/gpu/driverinstaller_test.go

@@ -291,14 +291,14 @@ func TestVerifyDriverDigest(t *testing.T) {
 			driverDigest:    "test-digest",
 			refDriverDigest: "8edf273aa28919d86f9f0ab68b1f267280821a3251c281d19748f940c180d27a",
 			wantErr:         true,
-			errSubstr:       "gpu driver digest verification failed",
+			errSubstr:       "GPU driver digest verification failed",
 		},
 		{
 			name:            "Empty reference driver digest",
 			driverDigest:    "test-digest",
 			refDriverDigest: "",
 			wantErr:         true,
-			errSubstr:       "gpu driver digest verification failed",
+			errSubstr:       "GPU driver digest verification failed",
 		},
 		{
 			name:            "Installed driver file does not exist",