Skip to content

Commit e9eb291

Browse files
yawangwangyawangwang
authored
migrate models.NvidiaAttestation to attestation proto library (google#691)
1 parent 262ec22 commit e9eb291

File tree

7 files changed

+84
-76
lines changed

7 files changed

+84
-76
lines changed

‎go.work.sum‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2153,6 +2153,7 @@ github.com/syndtr/gocapability v0.0.0-20200815063812-42c35b437635/go.mod h1:hkRG
21532153
github.com/tchap/go-patricia v2.2.6+incompatible/go.mod h1:bmLyhP68RS6kStMGxByiQ23RP/odRBOTVjwp2cDyi6I=
21542154
github.com/tchap/go-patricia/v2 v2.3.1 h1:6rQp39lgIYZ+MHmdEq4xzuk1t7OdC35z/xm0BGhTkes=
21552155
github.com/tchap/go-patricia/v2 v2.3.1/go.mod h1:VZRHKAb53DLaG+nA9EaYYiaEx6YztwDlLElMsnSHD4k=
2156+
github.com/tink-crypto/tink-go/v2 v2.2.1-0.20241120130117-c41ea0ed393b/go.mod h1:8qt2du2JzY6pUCRZ4cVz/f+gEmznKkvzd1KScaN5Zqk=
21562157
github.com/tj/assert v0.0.0-20171129193455-018094318fb0 h1:Rw8kxzWo1mr6FSaYXjQELRe88y2KdfynXdnK72rdjtA=
21572158
github.com/tj/go-elastic v0.0.0-20171221160941-36157cbbebc2 h1:eGaGNxrtoZf/mBURsnNQKDR7u50Klgcf2eFDQEnc8Bc=
21582159
github.com/tj/go-kinesis v0.0.0-20171128231115-08b17f58cb1b h1:m74UWYy+HBs+jMFR9mdZU6shPewugMyH5+GV6LNgW8w=

‎launcher/agent/agent.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -460,14 +460,14 @@ func (t *tdxAttestRoot) Attest(nonce []byte) (any, error) {
460460
return nil, err
461461
}
462462

463-
var nvAtt *models.NvidiaAttestation
463+
var nvAtt *attestationpb.NvidiaAttestationReport
464464
for _, deviceRoT := range t.deviceROTs {
465465
att, err := deviceRoT.Attest(nonce)
466466
if err != nil {
467467
return nil, err
468468
}
469469
switch v := att.(type) {
470-
case *models.NvidiaAttestation:
470+
case *attestationpb.NvidiaAttestationReport:
471471
nvAtt = v
472472
default:
473473
return nil, fmt.Errorf("unknown device attestation type: %T", v)

‎launcher/agent/agent_test.go‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,6 @@ import (
3333
tpmpb "github.com/google/go-tpm-tools/proto/tpm"
3434
"github.com/google/go-tpm-tools/verifier"
3535
"github.com/google/go-tpm-tools/verifier/fake"
36-
"github.com/google/go-tpm-tools/verifier/models"
3736
"github.com/google/go-tpm-tools/verifier/oci"
3837
"github.com/google/go-tpm-tools/verifier/oci/cosign"
3938
"google.golang.org/protobuf/encoding/protojson"
@@ -668,14 +667,14 @@ func (f *fakeTdxAttestRoot) GetCEL() gecel.CEL {
668667

669668
func (f *fakeTdxAttestRoot) Attest(nonce []byte) (any, error) {
670669
f.receivedNonce = nonce
671-
var nvAtt *models.NvidiaAttestation
670+
var nvAtt *attestationpb.NvidiaAttestationReport
672671
for _, deviceRoT := range f.deviceRoTS {
673672
att, err := deviceRoT.Attest(nonce)
674673
if err != nil {
675674
return nil, err
676675
}
677676
switch v := att.(type) {
678-
case *models.NvidiaAttestation:
677+
case *attestationpb.NvidiaAttestationReport:
679678
nvAtt = v
680679
default:
681680
return nil, fmt.Errorf("unknown device attestation type: %T", v)
@@ -709,9 +708,11 @@ func (f *fakeGPURoT) Attest(nonce []byte) (any, error) {
709708
if len(nonce) == 0 {
710709
return nil, fmt.Errorf("fake GPU attestation failed")
711710
}
712-
return &models.NvidiaAttestation{
713-
CCFeature: &models.NvidiaSinglePassthroughAttestation{
714-
GPUInfo: models.GPUInfo{UUID: "fake-gpu-uuid"},
711+
return &attestationpb.NvidiaAttestationReport{
712+
CcFeature: &attestationpb.NvidiaAttestationReport_Spt{
713+
Spt: &attestationpb.NvidiaAttestationReport_SinglePassthroughAttestation{
714+
GpuQuote: &attestationpb.GpuInfo{Uuid: "fake-gpu-uuid"},
715+
},
715716
},
716717
}, nil
717718
}

‎launcher/internal/gpu/attestation.go‎

Lines changed: 30 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,8 @@ import (
99
"cos.googlesource.com/cos/tools.git/src/cmd/cos_gpu_installer/deviceinfo"
1010
"github.com/NVIDIA/go-nvml/pkg/nvml"
1111
"github.com/confidentsecurity/go-nvtrust/pkg/gonvtrust/gpu"
12-
"github.com/google/go-tpm-tools/verifier/models"
12+
13+
attestationpb "github.com/GoogleCloudPlatform/confidential-space/server/proto/gen/attestation"
1314
)
1415

1516
type attestationType int
@@ -40,7 +41,7 @@ func (a *NvidiaAttester) Attest(nonce []byte) (any, error) {
4041

4142
// collectAttestationEvidence assumes CC GPU devices are in place w/ driver support
4243
// and will try to collect raw attestation evidence and convert it to known data models.
43-
func (a *NvidiaAttester) collectAttestationEvidence(handler gpu.NvmlHandler, nonce []byte) (*models.NvidiaAttestation, error) {
44+
func (a *NvidiaAttester) collectAttestationEvidence(handler gpu.NvmlHandler, nonce []byte) (*attestationpb.NvidiaAttestationReport, error) {
4445
gpuAdmin, err := gpu.NewNvmlGPUAdmin(handler)
4546
if err != nil {
4647
return nil, fmt.Errorf("failed to create GPU admin: %v", err)
@@ -53,7 +54,7 @@ func (a *NvidiaAttester) collectAttestationEvidence(handler gpu.NvmlHandler, non
5354
return nil, fmt.Errorf("failed to collect GPU evidence: %v", err)
5455
}
5556

56-
var gpuInfos []models.GPUInfo
57+
var gpuInfos []*attestationpb.GpuInfo
5758
for i, deviceInfo := range deviceInfos {
5859
device, ret := handler.DeviceGetHandleByIndex(i)
5960
if ret != nvml.SUCCESS {
@@ -84,11 +85,11 @@ func (a *NvidiaAttester) collectAttestationEvidence(handler gpu.NvmlHandler, non
8485
return nil, fmt.Errorf("failed to decode GPU certificate chain: %v", err)
8586
}
8687

87-
gpuInfo := models.GPUInfo{
88-
UUID: uuid,
88+
gpuInfo := &attestationpb.GpuInfo{
89+
Uuid: uuid,
8990
DriverVersion: driverVersion,
90-
VBIOSVersion: vbiosVersion,
91-
GPUArchitectureType: deviceInfo.Arch(),
91+
VbiosVersion: vbiosVersion,
92+
GpuArchitectureType: convertGPUArchToPB(deviceInfo.Arch()),
9293
AttestationReport: deviceInfo.AttestationReport(),
9394
AttestationCertificateChain: attestationCertChainData,
9495
}
@@ -97,15 +98,19 @@ func (a *NvidiaAttester) collectAttestationEvidence(handler gpu.NvmlHandler, non
9798

9899
switch determineAttestationType(gpuInfos) {
99100
case SPT:
100-
return &models.NvidiaAttestation{
101-
CCFeature: &models.NvidiaSinglePassthroughAttestation{
102-
GPUInfo: gpuInfos[0],
101+
return &attestationpb.NvidiaAttestationReport{
102+
CcFeature: &attestationpb.NvidiaAttestationReport_Spt{
103+
Spt: &attestationpb.NvidiaAttestationReport_SinglePassthroughAttestation{
104+
GpuQuote: gpuInfos[0],
105+
},
103106
},
104107
}, nil
105108
case MPT:
106-
return &models.NvidiaAttestation{
107-
CCFeature: &models.NvidiaMultiGpuSecurePassthroughAttestation{
108-
GPUInfos: gpuInfos,
109+
return &attestationpb.NvidiaAttestationReport{
110+
CcFeature: &attestationpb.NvidiaAttestationReport_Mpt{
111+
Mpt: &attestationpb.NvidiaAttestationReport_MultiGpuSecurePassthroughAttestation{
112+
GpuQuotes: gpuInfos,
113+
},
109114
},
110115
}, nil
111116
default:
@@ -116,7 +121,7 @@ func (a *NvidiaAttester) collectAttestationEvidence(handler gpu.NvmlHandler, non
116121
// determineAttesationType auto-detects the GPU attestation type.
117122
// The current implementations "guess" the attestation type.
118123
// Further improvement should be made to parse GPU attesation report to get the actual attestation type.
119-
func determineAttestationType(gpuInfos []models.GPUInfo) attestationType {
124+
func determineAttestationType(gpuInfos []*attestationpb.GpuInfo) attestationType {
120125
gpuType, _ := getGpuTypeInfo()
121126
if gpuType != deviceinfo.H100 && gpuType != deviceinfo.B200 {
122127
return UNSUPPORTED
@@ -126,3 +131,14 @@ func determineAttestationType(gpuInfos []models.GPUInfo) attestationType {
126131
}
127132
return SPT
128133
}
134+
135+
func convertGPUArchToPB(arch string) attestationpb.GpuArchitectureType {
136+
switch arch {
137+
case "HOPPER":
138+
return attestationpb.GpuArchitectureType_GPU_ARCHITECTURE_TYPE_HOPPER
139+
case "BLACKWELL":
140+
return attestationpb.GpuArchitectureType_GPU_ARCHITECTURE_TYPE_BLACKWELL
141+
default:
142+
return attestationpb.GpuArchitectureType_GPU_ARCHITECTURE_TYPE_UNSPECIFIED
143+
}
144+
}

‎launcher/internal/gpu/attestation_test.go‎

Lines changed: 42 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ import (
55

66
"cos.googlesource.com/cos/tools.git/src/cmd/cos_gpu_installer/deviceinfo"
77
"github.com/confidentsecurity/go-nvtrust/pkg/gonvtrust/gpu"
8-
"github.com/google/go-tpm-tools/verifier/models"
8+
9+
attestationpb "github.com/GoogleCloudPlatform/confidential-space/server/proto/gen/attestation"
910
)
1011

1112
func TestCollectAttestationEvidence(t *testing.T) {
@@ -59,8 +60,8 @@ func TestCollectAttestationEvidence(t *testing.T) {
5960
}
6061
if tc.wantPass {
6162
if tc.wantSPT {
62-
if _, ok := attesation.CCFeature.(*models.NvidiaSinglePassthroughAttestation); !ok {
63-
t.Errorf("CollectAttestationEvidence() = %v, want %v", attesation.CCFeature, &models.NvidiaSinglePassthroughAttestation{})
63+
if _, ok := attesation.CcFeature.(*attestationpb.NvidiaAttestationReport_Spt); !ok {
64+
t.Errorf("CollectAttestationEvidence() = %v, want %v", attesation.CcFeature, &attestationpb.NvidiaAttestationReport_Spt{})
6465
}
6566
}
6667
}
@@ -71,39 +72,39 @@ func TestCollectAttestationEvidence(t *testing.T) {
7172
func TestDetermineAttestationType(t *testing.T) {
7273
testCases := []struct {
7374
name string
74-
gpuInfos []models.GPUInfo
75+
gpuInfos []*attestationpb.GpuInfo
7576
gpuType deviceinfo.GPUType
7677
want attestationType
7778
}{
7879
{
7980
name: "UNSUPPORTED GPU type",
80-
gpuInfos: []models.GPUInfo{
81-
{UUID: "gpu-0"},
81+
gpuInfos: []*attestationpb.GpuInfo{
82+
{Uuid: "gpu-0"},
8283
},
8384
gpuType: deviceinfo.Others,
8485
want: UNSUPPORTED,
8586
},
8687
{
8788
name: "SPT attestation type (H100)",
88-
gpuInfos: []models.GPUInfo{
89-
{UUID: "gpu-0"},
89+
gpuInfos: []*attestationpb.GpuInfo{
90+
{Uuid: "gpu-0"},
9091
},
9192
gpuType: deviceinfo.H100,
9293
want: SPT,
9394
},
9495
{
9596
name: "SPT attestation type (B200 with single GPU)",
96-
gpuInfos: []models.GPUInfo{
97-
{UUID: "gpu-0"},
97+
gpuInfos: []*attestationpb.GpuInfo{
98+
{Uuid: "gpu-0"},
9899
},
99100
gpuType: deviceinfo.B200,
100101
want: SPT,
101102
},
102103
{
103104
name: "MPT attestation type (B200 with multiple GPUs)",
104-
gpuInfos: []models.GPUInfo{
105-
{UUID: "gpu-0"},
106-
{UUID: "gpu-1"},
105+
gpuInfos: []*attestationpb.GpuInfo{
106+
{Uuid: "gpu-0"},
107+
{Uuid: "gpu-1"},
107108
},
108109
gpuType: deviceinfo.B200,
109110
want: MPT,
@@ -124,5 +125,33 @@ func TestDetermineAttestationType(t *testing.T) {
124125
}
125126
})
126127
}
128+
}
129+
130+
func TestConvertGPUArchToPB(t *testing.T) {
131+
testCases := []struct {
132+
arch string
133+
wantArch attestationpb.GpuArchitectureType
134+
}{
135+
{
136+
arch: "HOPPER",
137+
wantArch: attestationpb.GpuArchitectureType_GPU_ARCHITECTURE_TYPE_HOPPER,
138+
},
139+
{
140+
arch: "BLACKWELL",
141+
wantArch: attestationpb.GpuArchitectureType_GPU_ARCHITECTURE_TYPE_BLACKWELL,
142+
},
143+
{
144+
arch: "UNSPECIFIED",
145+
wantArch: attestationpb.GpuArchitectureType_GPU_ARCHITECTURE_TYPE_UNSPECIFIED,
146+
},
147+
}
148+
149+
for _, tc := range testCases {
150+
t.Run(tc.arch, func(t *testing.T) {
151+
if got := convertGPUArchToPB(tc.arch); got != tc.wantArch {
152+
t.Errorf("convertGPUArchToPB() = %v, want %v", got, tc.wantArch)
153+
}
154+
})
155+
}
127156

128157
}

‎verifier/client.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ package verifier
55
import (
66
"context"
77

8+
csattestpb "github.com/GoogleCloudPlatform/confidential-space/server/proto/gen/attestation"
89
attestpb "github.com/google/go-tpm-tools/proto/attest"
910
"github.com/google/go-tpm-tools/verifier/models"
1011
"google.golang.org/genproto/googleapis/rpc/status"
@@ -58,7 +59,7 @@ type TDCCELAttestation struct {
5859
// still needs following two for GCE info
5960
AkCert []byte
6061
IntermediateCerts [][]byte
61-
NvidiaAttestation *models.NvidiaAttestation
62+
NvidiaAttestation *csattestpb.NvidiaAttestationReport
6263
}
6364

6465
// VerifyAttestationResponse is the response from a successful

‎verifier/models/nvidia_attestations.go‎

Lines changed: 0 additions & 40 deletions
This file was deleted.

0 commit comments

Comments
 (0)