Skip to content

Commit d700614

Browse files
committed
add nvidia driver runfile hash verification
1 parent 786f494 commit d700614

File tree

5 files changed

+212
-0
lines changed

5 files changed

+212
-0
lines changed

go.work.sum

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -706,6 +706,7 @@ github.com/andreyvit/diff v0.0.0-20170406064948-c7f18ee00883/go.mod h1:rCTlJbsFo
706706
github.com/andygrunwald/go-gerrit v0.0.0-20201231163137-46815e48bfe0/go.mod h1:soxaYLbAFToS0OelBriItCts/mtUZOuLBkCk1Xv4ZSo=
707707
github.com/apache/arrow/go/v15 v15.0.2/go.mod h1:DGXsR3ajT524njufqf95822i+KTh+yea1jass9YXgjA=
708708
github.com/aws/aws-sdk-go v1.43.16/go.mod h1:y4AeaBuwd2Lk+GepC1E9v0qOiTws0MIWAX4oIKwKHZo=
709+
github.com/beevik/etree v1.1.0/go.mod h1:r8Aw8JqVegEf0w2fDnATrX9VpkMcyFeM0FhwO62wh+A=
709710
github.com/blang/semver/v4 v4.0.0/go.mod h1:IbckMUScFkM3pff0VJDNKRiT6TG/YpiHIM2yvyW5YoQ=
710711
github.com/bshuster-repo/logrus-logstash-hook v1.0.0/go.mod h1:zsTqEiSzDgAa/8GZR7E1qaXrhYNDKBYy5/dWPTIflbk=
711712
github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
@@ -954,6 +955,7 @@ github.com/safchain/ethtool v0.0.0-20210803160452-9aa261dae9b1/go.mod h1:Z0q5wiB
954955
github.com/sclevine/agouti v3.0.0+incompatible/go.mod h1:b4WX9W9L1sfQKXeJf1mUTLZKJ48R1S7H23Ji7oFO5Bw=
955956
github.com/seccomp/libseccomp-golang v0.9.2-0.20210429002308-3879420cc921/go.mod h1:JA8cRccbGaA1s33RQf7Y1+q9gHmZX1yB/z9WDN1C6fg=
956957
github.com/sirupsen/logrus v1.9.0/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
958+
github.com/smartystreets/assertions v1.1.1/go.mod h1:tcbTF8ujkAEcZ8TElKY+i30BzYlVhC/LOxJk7iOWnoo=
957959
github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
958960
github.com/spdx/tools-golang v0.5.3/go.mod h1:/ETOahiAo96Ob0/RAIBmFZw6XN0yTnyr/uFZm2NTMhI=
959961
github.com/spf13/afero v1.10.0/go.mod h1:UBogFpq8E9Hx+xc5CNTTEpTnuHVmXDwZcZcE1eb/UhQ=
@@ -984,6 +986,7 @@ github.com/yashtewari/glob-intersection v0.1.0/go.mod h1:LK7pIC3piUjovexikBbJ26Y
984986
github.com/yuin/goldmark v1.4.1/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
985987
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
986988
github.com/zeebo/xxh3 v1.0.2/go.mod h1:5NWz9Sef7zIDm2JHfFlcQvNekmcEl9ekUZQQKCYaDcA=
989+
go.chromium.org/luci v0.0.0-20200722211809-bab0c30be68b/go.mod h1:MIQewVTLvOvc0UioV0JNqTNO/RspKFS0XEeoKrOxsdM=
987990
go.einride.tech/aip v0.67.1/go.mod h1:ZGX4/zKw8dcgzdLsrvpOOGxfxI2QSk12SlP7d6c0/XI=
988991
go.einride.tech/aip v0.68.0/go.mod h1:7y9FF8VtPWqpxuAxl0KQWqaULxW4zFIesD6zF5RIHHg=
989992
go.etcd.io/bbolt v1.3.7/go.mod h1:N9Mkw9X8x5fupy0IKsmuqVtoGDyxsaDlbk4Rd05IAQw=

launcher/image/preload.sh

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ readonly EXPERIMENTS_BINARY="confidential_space_experiments"
66
readonly GPU_REF_VALUES_PATH="${CS_PATH}/gpu"
77
readonly COS_GPU_INSTALLER_IMAGE_REF="${GPU_REF_VALUES_PATH}/cos_gpu_installer_image_ref"
88
readonly COS_GPU_INSTALLER_IMAGE_DIGEST="${GPU_REF_VALUES_PATH}/cos_gpu_installer_image_digest"
9+
readonly DRIVER_DIGEST_SHA256SUM="${GPU_REF_VALUES_PATH}/driver_digest_sha256sum"
10+
readonly DRIVERS_GCS_BUCKET="cos-nvidia-gpu-drivers"
911

1012
copy_launcher() {
1113
cp launcher "${CS_PATH}/cs_container_launcher"
@@ -126,6 +128,26 @@ get_cos_gpu_installer_image_digest() {
126128
echo "${image_digest}"
127129
}
128130

131+
set_reference_driver_digest() {
132+
local driver_version
133+
local driver_digest_gcs_url
134+
135+
# Fetching the default driver version for H100 GPU.
136+
driver_version=$(cos-extensions list -- --target-gpu NVIDIA_H100_80GB | grep DEFAULT | cut -d" " -f 1)
137+
driver_digest_gcs_url="https://storage.googleapis.com/${DRIVERS_GCS_BUCKET}/sha256/NVIDIA-Linux-x86_64-${driver_version}.run.sha256"
138+
if ! curl -sSL ${driver_digest_gcs_url} -o ${DRIVER_DIGEST_SHA256SUM}; then
139+
echo "Error: failed to download the driver digest file from ${driver_digest_gcs_url}." >&2
140+
return 1
141+
fi
142+
143+
driver_digest=$(cat ${DRIVER_DIGEST_SHA256SUM} | cut -d " " -f 1)
144+
# Check for the expected length of the SHA256 digest (64 hex characters)
145+
if [ ${#driver_digest} -ne 64 ]; then
146+
echo "Error: driver digest has an unexpected length: ${#driver_digest}, Expected 64." >&2
147+
return 1
148+
fi
149+
}
150+
129151

130152
set_gpu_driver_ref_values() {
131153
local cos_gpu_installer_image_ref
@@ -153,6 +175,7 @@ set_gpu_driver_ref_values() {
153175

154176
echo ${cos_gpu_installer_image_ref} > ${COS_GPU_INSTALLER_IMAGE_REF}
155177
echo ${cos_gpu_installer_image_digest} > ${COS_GPU_INSTALLER_IMAGE_DIGEST}
178+
set_reference_driver_digest
156179
}
157180

158181
main() {

launcher/internal/gpu/config.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,6 @@ const (
99
InstallerImageRefFile = "/usr/share/oem/confidential_space/gpu/cos_gpu_installer_image_ref"
1010
// InstallerImageDigestFile is a filename which has the container image digest of cos_gpu_installer.
1111
InstallerImageDigestFile = "/usr/share/oem/confidential_space/gpu/cos_gpu_installer_image_digest"
12+
// ReferenceDriverDigestFile is a filename which has the reference digest of nvidia driver installer.
13+
ReferenceDriverDigestFile = "/usr/share/oem/confidential_space/gpu/driver_digest_sha256sum"
1214
)

launcher/internal/gpu/driverinstaller.go

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ package gpu
33

44
import (
55
"context"
6+
"crypto/sha256"
7+
"encoding/hex"
68
"fmt"
79
"os"
810
"os/exec"
@@ -150,6 +152,15 @@ func (di *DriverInstaller) InstallGPUDrivers(ctx context.Context) error {
150152
return fmt.Errorf("GPU driver installation task ended with non-zero status code %d", code)
151153
}
152154

155+
referenceDigest, runFilename, err := parseDriverDigestFile(ReferenceDriverDigestFile)
156+
if err != nil {
157+
return fmt.Errorf("failed to get driver filename, got error: %v", err)
158+
}
159+
runFile := fmt.Sprintf("%s/%s", InstallationHostDir, runFilename)
160+
if err = verifyDriverDigest(runFile, referenceDigest); err != nil {
161+
return fmt.Errorf("failed to verify GPU driver digest: %v", err)
162+
}
163+
153164
if err = launchNvidiaPersistencedProcess(di.logger); err != nil {
154165
return fmt.Errorf("failed to start nvidia-persistenced process: %v", err)
155166
}
@@ -203,6 +214,17 @@ func verifyInstallerImageDigest(image containerd.Image, referenceDigestFile stri
203214
return nil
204215
}
205216

217+
func verifyDriverDigest(driverFile, referenceHash string) error {
218+
calculatedHash, err := calculateSHA256Hash(driverFile)
219+
if err != nil {
220+
return err
221+
}
222+
if calculatedHash != referenceHash {
223+
return fmt.Errorf("gpu driver digest verification failed - expected : %s, got : %s", referenceHash, calculatedHash)
224+
}
225+
return nil
226+
}
227+
206228
func remountAsExecutable(dir string) error {
207229
if err := os.MkdirAll(dir, 0755); err != nil {
208230
return fmt.Errorf("failed to create dir %q: %v", dir, err)
@@ -288,3 +310,29 @@ func nvidiaSmiRunFunc(args ...string) nvidiaSmiCmdRun {
288310
cmd := fmt.Sprintf("%s/bin/nvidia-smi", InstallationHostDir)
289311
return func() error { return exec.Command(cmd, args...).Run() }
290312
}
313+
314+
func calculateSHA256Hash(filePath string) (string, error) {
315+
contentBytes, err := os.ReadFile(filePath)
316+
if err != nil {
317+
return "", fmt.Errorf("failed to read the file %s, got error %v", filePath, err)
318+
}
319+
hashBytes := sha256.Sum256(contentBytes)
320+
return hex.EncodeToString(hashBytes[:]), nil
321+
}
322+
323+
// Reference driver digest file contains driver digest along with driver .Run filename.
324+
// parseDriverDigestFile() gets reference digest file and returns the run file's digest and the run filename.
325+
// Sample reference driver digest file content:
326+
//
327+
// 65fe3e2236c1ddab26eaf8e1b3f3b2b0951b8824d7c4a5022552579288ff7fea NVIDIA-Linux-aarch64-570.124.06.run
328+
func parseDriverDigestFile(digestFile string) (string, string, error) {
329+
contentBytes, err := os.ReadFile(digestFile)
330+
if err != nil {
331+
return "", "", fmt.Errorf("failed to read the file %s, got error %v", digestFile, err)
332+
}
333+
fields := strings.Fields(string(contentBytes))
334+
if len(fields) != 2 {
335+
return "", "", fmt.Errorf("unexpected content length in reference file %s", digestFile)
336+
}
337+
return fields[0], fields[1], nil
338+
}

launcher/internal/gpu/driverinstaller_test.go

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -271,3 +271,139 @@ func TestIsConfidentialComputeSupported(t *testing.T) {
271271
})
272272
}
273273
}
274+
275+
func TestVerifyDriverDigest(t *testing.T) {
276+
tests := []struct {
277+
name string
278+
driverDigest string
279+
refDriverDigest string
280+
wantErr bool
281+
errSubstr string
282+
}{
283+
{
284+
name: "Driver digest matches",
285+
driverDigest: "test-digest",
286+
refDriverDigest: "8edf273aa28919d86f9f0ab68b1f267280821a3251c281d19748f940c180d27f",
287+
wantErr: false,
288+
},
289+
{
290+
name: "Driver digest mismatch",
291+
driverDigest: "test-digest",
292+
refDriverDigest: "8edf273aa28919d86f9f0ab68b1f267280821a3251c281d19748f940c180d27a",
293+
wantErr: true,
294+
errSubstr: "gpu driver digest verification failed",
295+
},
296+
{
297+
name: "Empty reference driver digest",
298+
driverDigest: "test-digest",
299+
refDriverDigest: "",
300+
wantErr: true,
301+
errSubstr: "gpu driver digest verification failed",
302+
},
303+
{
304+
name: "Installed driver file does not exist",
305+
driverDigest: "",
306+
refDriverDigest: "8edf273aa28919d86f9f0ab68b1f267280821a3251c281d19748f940c180d27f",
307+
wantErr: true,
308+
errSubstr: "failed to read the file",
309+
},
310+
}
311+
312+
for _, tt := range tests {
313+
t.Run(tt.name, func(t *testing.T) {
314+
315+
tempDir := t.TempDir()
316+
installedFile := fmt.Sprintf("%s/installed_driver_digest", tempDir)
317+
318+
if tt.name != "Installed driver file does not exist" {
319+
err := os.WriteFile(installedFile, []byte(tt.driverDigest), 0644)
320+
if err != nil {
321+
t.Fatalf("failed to write to the driver digest testfile %s: %v", installedFile, err)
322+
}
323+
}
324+
err := verifyDriverDigest(installedFile, tt.refDriverDigest)
325+
if (err != nil) != tt.wantErr {
326+
t.Errorf("VerifyDriverDigest() error = %v, wantErr %v", err, tt.wantErr)
327+
}
328+
if tt.wantErr && !strings.Contains(err.Error(), tt.errSubstr) {
329+
t.Errorf("VerifyDriverDigest() error message %s is expected to contain %s", err.Error(), tt.errSubstr)
330+
}
331+
})
332+
}
333+
}
334+
335+
func TestParseDriverDigestFile(t *testing.T) {
336+
tests := []struct {
337+
name string
338+
refFileContent string
339+
expectedFilename string
340+
expectedDigest string
341+
wantErr bool
342+
errSubstr string
343+
}{
344+
{
345+
name: "Valid content in reference digest file",
346+
refFileContent: "8edf273aa28919d86f9f0ab68b1f267280821a3251c281d19748f940c180d27f test-driver-file",
347+
expectedDigest: "8edf273aa28919d86f9f0ab68b1f267280821a3251c281d19748f940c180d27f",
348+
expectedFilename: "test-driver-file",
349+
wantErr: false,
350+
},
351+
{
352+
name: "Valid content in reference digest file extra whitespaces",
353+
refFileContent: "8edf273aa28919d86f9f0ab68b1f267280821a3251c281d19748f940c180d27f test-driver-file",
354+
expectedDigest: "8edf273aa28919d86f9f0ab68b1f267280821a3251c281d19748f940c180d27f",
355+
expectedFilename: "test-driver-file",
356+
wantErr: false,
357+
},
358+
{
359+
name: "Malformed content in reference digest file",
360+
refFileContent: "8edf273aa28919d86f9f0ab68b1f267280821a3251c281d19748f940c180d27a test-driver-file some-more-data",
361+
expectedFilename: "",
362+
wantErr: true,
363+
errSubstr: "unexpected content length in reference file",
364+
},
365+
{
366+
name: "Empty reference digest file",
367+
refFileContent: "",
368+
expectedFilename: "",
369+
wantErr: true,
370+
errSubstr: "unexpected content length in reference file",
371+
},
372+
{
373+
name: "Reference file does not exist",
374+
refFileContent: "",
375+
expectedFilename: "",
376+
wantErr: true,
377+
errSubstr: "failed to read the file",
378+
},
379+
}
380+
381+
for _, tt := range tests {
382+
t.Run(tt.name, func(t *testing.T) {
383+
384+
tempDir := t.TempDir()
385+
refFile := fmt.Sprintf("%s/reference_driver_digest", tempDir)
386+
387+
if tt.name != "Reference file does not exist" {
388+
err := os.WriteFile(refFile, []byte(tt.refFileContent), 0644)
389+
if err != nil {
390+
t.Fatalf("failed to write to the reference digest testfile %s: %v", refFile, err)
391+
}
392+
}
393+
394+
digest, filename, err := parseDriverDigestFile(refFile)
395+
if (err != nil) != tt.wantErr {
396+
t.Errorf("parseDriverDigestFile() error = %v, wantErr %v", err, tt.wantErr)
397+
}
398+
if tt.wantErr && !strings.Contains(err.Error(), tt.errSubstr) {
399+
t.Errorf("parseDriverDigestFile() error message %s is expected to contain %s", err.Error(), tt.errSubstr)
400+
}
401+
if filename != tt.expectedFilename {
402+
t.Errorf("parseDriverDigestFile() returned unexpected filename, got = %v, want %v", filename, tt.expectedFilename)
403+
}
404+
if digest != tt.expectedDigest {
405+
t.Errorf("parseDriverDigestFile() returned unexpected digest, got = %v, want %v", digest, tt.expectedDigest)
406+
}
407+
})
408+
}
409+
}

0 commit comments

Comments
 (0)