Skip to content

Commit 3e8553f

Browse files
yawangwangyawangwang
committed
Add device ROT manager
1 parent c16668d commit 3e8553f

File tree

9 files changed

+251
-145
lines changed

9 files changed

+251
-145
lines changed

‎launcher/agent/agent.go‎

Lines changed: 16 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ import (
3434
"github.com/google/go-tpm-tools/cel"
3535
"github.com/google/go-tpm-tools/client"
3636
"github.com/google/go-tpm-tools/internal"
37+
"github.com/google/go-tpm-tools/launcher/device"
3738
"github.com/google/go-tpm-tools/launcher/internal/logging"
3839
"github.com/google/go-tpm-tools/launcher/internal/signaturediscovery"
3940
"github.com/google/go-tpm-tools/launcher/spec"
@@ -72,28 +73,13 @@ type attestRoot interface {
7273
Attest(nonce []byte) (any, error)
7374
// ComputeNonce hashes the challenge and extraData using the algorithm preferred by the attestation root.
7475
ComputeNonce(challenge []byte, extraData []byte) []byte
75-
// AddDeviceROTs adds detected device RoTs(root of trust).
76-
AddDeviceROTs([]DeviceROT)
77-
// AttestDeviceROTs fetches a list of runtime device attestation report.
78-
AttestDeviceROTs(nonce []byte) ([]any, error)
79-
}
80-
81-
// DeviceROT defines an interface for all attached devices to collect attestation.
82-
type DeviceROT interface {
83-
// Attest fetches an attestation from the attached device detected by launcher.
84-
Attest(nonce []byte) (any, error)
8576
}
8677

8778
// AttestAgentOpts contains user generated options when calling the
8879
// VerifyAttestation API
8980
type AttestAgentOpts struct {
90-
TokenOptions *models.TokenOptions
91-
*DeviceReportOpts
92-
}
93-
94-
// DeviceReportOpts contains options for runtime device attestations.
95-
type DeviceReportOpts struct {
96-
EnableRuntimeGPUAttestation bool
81+
TokenOptions *models.TokenOptions
82+
DeviceReportOpts device.ReportOpts
9783
}
9884

9985
type agent struct {
@@ -106,6 +92,7 @@ type agent struct {
10692
launchSpec spec.LaunchSpec
10793
logger logging.Logger
10894
sigsCache *sigsCache
95+
deviceROTManager *device.ROTManager
10996
}
11097

11198
// CreateAttestationAgent returns an agent capable of performing remote
@@ -115,7 +102,7 @@ type agent struct {
115102
// - principalFetcher is a func to fetch GCE principal tokens for a given audience.
116103
// - signaturesFetcher is a func to fetch container image signatures associated with the running workload.
117104
// - logger will log any partial errors returned by VerifyAttestation.
118-
func CreateAttestationAgent(tpm io.ReadWriteCloser, akFetcher util.TpmKeyFetcher, verifierClient verifier.Client, principalFetcher principalIDTokenFetcher, sigsFetcher signaturediscovery.Fetcher, launchSpec spec.LaunchSpec, logger logging.Logger, deviceROTs []DeviceROT) (AttestationAgent, error) {
105+
func CreateAttestationAgent(tpm io.ReadWriteCloser, akFetcher util.TpmKeyFetcher, verifierClient verifier.Client, principalFetcher principalIDTokenFetcher, sigsFetcher signaturediscovery.Fetcher, launchSpec spec.LaunchSpec, logger logging.Logger, deviceROTs []device.ROT) (AttestationAgent, error) {
119106
// Fetched the AK and save it, so the agent doesn't need to create a new key everytime
120107
ak, err := akFetcher(tpm)
121108
if err != nil {
@@ -178,8 +165,8 @@ func CreateAttestationAgent(tpm io.ReadWriteCloser, akFetcher util.TpmKeyFetcher
178165
attestAgent.avRot = tpmAR
179166
}
180167

181-
// Add deviceRoTs to the CPU attestation root.
182-
attestAgent.avRot.AddDeviceROTs(deviceROTs)
168+
// Register device ROT manager to the attestAgent.
169+
attestAgent.deviceROTManager = device.NewROTManager(deviceROTs)
183170
return attestAgent, nil
184171
}
185172

@@ -360,10 +347,7 @@ func (a *agent) AttestationEvidence(_ context.Context, challenge []byte, extraDa
360347
}
361348

362349
func (a *agent) attestDeviceROTs(nonce []byte, opts AttestAgentOpts) ([]*attestationpb.DeviceAttestationReport, error) {
363-
if opts.DeviceReportOpts == nil {
364-
return nil, nil
365-
}
366-
deviceROTs, err := a.avRot.AttestDeviceROTs(nonce)
350+
deviceROTs, err := a.deviceROTManager.AttestDeviceROTs(nonce, opts.DeviceReportOpts)
367351
if err != nil {
368352
return nil, err
369353
}
@@ -411,12 +395,11 @@ func convertOCIToContainerSignature(ociSig oci.Signature) (*verifier.ContainerSi
411395
}
412396

413397
type tpmAttestRoot struct {
414-
tpmMu sync.Mutex
415-
fetchedAK *client.Key
416-
tpm io.ReadWriteCloser
417-
cosCel gecel.CEL
418-
hashAlgos []crypto.Hash
419-
deviceROTs []DeviceROT
398+
tpmMu sync.Mutex
399+
fetchedAK *client.Key
400+
tpm io.ReadWriteCloser
401+
cosCel gecel.CEL
402+
hashAlgos []crypto.Hash
420403
}
421404

422405
func (t *tpmAttestRoot) GetCEL() gecel.CEL {
@@ -457,22 +440,10 @@ func (t *tpmAttestRoot) ComputeNonce(challenge []byte, extraData []byte) []byte
457440
return finalNonce[:]
458441
}
459442

460-
func (t *tpmAttestRoot) AddDeviceROTs(deviceROTs []DeviceROT) {
461-
t.deviceROTs = append(t.deviceROTs, deviceROTs...)
462-
}
463-
464-
func (t *tpmAttestRoot) AttestDeviceROTs(nonce []byte) ([]any, error) {
465-
t.tpmMu.Lock()
466-
defer t.tpmMu.Unlock()
467-
468-
return doAttestDeviceROTs(t.deviceROTs, nonce)
469-
}
470-
471443
type tdxAttestRoot struct {
472-
tdxMu sync.Mutex
473-
qp *tg.LinuxConfigFsQuoteProvider
474-
cosCel gecel.CEL
475-
deviceROTs []DeviceROT
444+
tdxMu sync.Mutex
445+
qp *tg.LinuxConfigFsQuoteProvider
446+
cosCel gecel.CEL
476447
}
477448

478449
func (t *tdxAttestRoot) GetCEL() gecel.CEL {
@@ -513,13 +484,6 @@ func (t *tdxAttestRoot) Attest(nonce []byte) (any, error) {
513484
}, nil
514485
}
515486

516-
func (t *tdxAttestRoot) AttestDeviceROTs(nonce []byte) ([]any, error) {
517-
t.tdxMu.Lock()
518-
defer t.tdxMu.Unlock()
519-
520-
return doAttestDeviceROTs(t.deviceROTs, nonce)
521-
}
522-
523487
func (t *tdxAttestRoot) ComputeNonce(challenge []byte, extraData []byte) []byte {
524488
challengeData := challenge
525489
if extraData != nil {
@@ -531,10 +495,6 @@ func (t *tdxAttestRoot) ComputeNonce(challenge []byte, extraData []byte) []byte
531495
return finalNonce[:]
532496
}
533497

534-
func (t *tdxAttestRoot) AddDeviceROTs(deviceROTs []DeviceROT) {
535-
t.deviceROTs = append(t.deviceROTs, deviceROTs...)
536-
}
537-
538498
// Refresh refreshes the internal state of the attestation agent.
539499
// It will reset the container image signatures for now.
540500
func (a *agent) Refresh(ctx context.Context) error {
@@ -633,15 +593,3 @@ func convertToTPMQuote(v *pb.Attestation) *attestationpb.TpmQuote {
633593
},
634594
}
635595
}
636-
637-
func doAttestDeviceROTs(deviceROTs []DeviceROT, nonce []byte) ([]any, error) {
638-
var deviceReports []any
639-
for _, deviceROT := range deviceROTs {
640-
deviceReport, err := deviceROT.Attest(nonce)
641-
if err != nil {
642-
return nil, err
643-
}
644-
deviceReports = append(deviceReports, deviceReport)
645-
}
646-
return deviceReports, nil
647-
}

‎launcher/agent/agent_test.go‎

Lines changed: 6 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import (
2929
"github.com/google/go-tpm-tools/cel"
3030
"github.com/google/go-tpm-tools/client"
3131
"github.com/google/go-tpm-tools/internal/test"
32+
"github.com/google/go-tpm-tools/launcher/device"
3233
"github.com/google/go-tpm-tools/launcher/internal/experiments"
3334
"github.com/google/go-tpm-tools/launcher/internal/logging"
3435
"github.com/google/go-tpm-tools/launcher/internal/signaturediscovery"
@@ -659,7 +660,6 @@ type fakeTdxAttestRoot struct {
659660
cel gecel.CEL
660661
receivedNonce []byte
661662
tdxQuote []byte
662-
deviceROTs []DeviceROT
663663
}
664664

665665
func (f *fakeTdxAttestRoot) Extend(c gecel.Content) error {
@@ -690,30 +690,9 @@ func (f *fakeTdxAttestRoot) ComputeNonce(challenge []byte, extraData []byte) []b
690690
return finalNonce[:]
691691
}
692692

693-
func (f *fakeTdxAttestRoot) AttestDeviceROTs(nonce []byte) ([]any, error) {
694-
var deviceReports []any
695-
for _, deviceRoT := range f.deviceROTs {
696-
att, err := deviceRoT.Attest(nonce)
697-
if err != nil {
698-
return nil, err
699-
}
700-
switch v := att.(type) {
701-
case *attestationpb.NvidiaAttestationReport:
702-
deviceReports = append(deviceReports, v)
703-
default:
704-
return nil, fmt.Errorf("unknown device attestation type: %T", v)
705-
}
706-
}
707-
return deviceReports, nil
708-
}
709-
710693
//go:embed testdata/cel.b64
711694
var celB64 string
712695

713-
func (f *fakeTdxAttestRoot) AddDeviceROTs(deviceRoTS []DeviceROT) {
714-
f.deviceROTs = append(f.deviceROTs, deviceRoTS...)
715-
}
716-
717696
type fakeGPURoT struct{}
718697

719698
func (f *fakeGPURoT) Attest(nonce []byte) (any, error) {
@@ -728,56 +707,9 @@ func (f *fakeGPURoT) Attest(nonce []byte) (any, error) {
728707
},
729708
}, nil
730709
}
731-
func TestTDXAttestDeviceROTs(t *testing.T) {
732-
testCases := []struct {
733-
name string
734-
tdxAttestRoot *fakeTdxAttestRoot
735-
nonce []byte
736-
wantGPU bool
737-
wantPass bool
738-
}{
739-
{
740-
name: "success tdxAttestRoot w/o GPU device",
741-
tdxAttestRoot: &fakeTdxAttestRoot{},
742-
nonce: []byte("test-nonce"),
743-
wantPass: true,
744-
},
745-
{
746-
name: "success tdxAttestRoot w/ GPU device",
747-
tdxAttestRoot: &fakeTdxAttestRoot{
748-
deviceROTs: []DeviceROT{&fakeGPURoT{}},
749-
},
750-
nonce: []byte("test-nonce"),
751-
wantGPU: true,
752-
wantPass: true,
753-
},
754-
{
755-
name: "failed tdxAttestRoot w/ GPU device",
756-
tdxAttestRoot: &fakeTdxAttestRoot{
757-
deviceROTs: []DeviceROT{&fakeGPURoT{}},
758-
},
759-
nonce: []byte(""),
760-
wantPass: false,
761-
},
762-
}
763-
764-
for _, tc := range testCases {
765-
t.Run(tc.name, func(t *testing.T) {
766-
deviceReports, err := tc.tdxAttestRoot.AttestDeviceROTs(tc.nonce)
767-
if gotPass := err == nil; gotPass != tc.wantPass {
768-
t.Errorf("tdxAttestRoot.AttestDeviceROTs() did not return expected attestation result, got %v, want %v", gotPass, tc.wantPass)
769-
}
770-
if tc.wantGPU {
771-
if len(deviceReports) == 0 {
772-
t.Fatalf("tdxAttestRoot.AttestDeviceROTs() didn't return any device reports")
773-
}
774-
if att := deviceReports[0].(*attestationpb.NvidiaAttestationReport); att == nil {
775-
t.Errorf("tdxAttestRoot.AttestDeviceROTs() didn't return expected device report type, want %v, but got nil", &attestationpb.NvidiaAttestationReport{})
776-
}
777-
}
778-
})
779-
}
780710

711+
func (f *fakeGPURoT) Vendor() device.Vendor {
712+
return device.NvidiaGPU
781713
}
782714

783715
func TestAttestationEvidence_TDX_Success(t *testing.T) {
@@ -816,8 +748,8 @@ func TestAttestationEvidence_TDX_Success(t *testing.T) {
816748
EnableAttestationEvidence: true,
817749
},
818750
},
751+
deviceROTManager: device.NewROTManager([]device.ROT{&fakeGPURoT{}}),
819752
}
820-
attestAgent.avRot.AddDeviceROTs([]DeviceROT{&fakeGPURoT{}})
821753

822754
if err := measureFakeEvents(attestAgent); err != nil {
823755
t.Fatalf("failed to measure events: %v", err)
@@ -838,7 +770,7 @@ func TestAttestationEvidence_TDX_Success(t *testing.T) {
838770
{
839771
name: "TDX attestation + runtime GPU attestation",
840772
opts: AttestAgentOpts{
841-
DeviceReportOpts: &DeviceReportOpts{
773+
DeviceReportOpts: device.ReportOpts{
842774
EnableRuntimeGPUAttestation: true,
843775
},
844776
},
@@ -1146,6 +1078,7 @@ func TestAttestationEvidence_TDX_NilExtraData(t *testing.T) {
11461078
EnableAttestationEvidence: true,
11471079
},
11481080
},
1081+
deviceROTManager: device.NewROTManager(nil),
11491082
}
11501083

11511084
challenge := []byte("test-challenge")

‎launcher/container_runner.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import (
3131
"github.com/google/go-tpm-tools/client"
3232
workloadservice "github.com/google/go-tpm-tools/keymanager/workload_service"
3333
"github.com/google/go-tpm-tools/launcher/agent"
34+
"github.com/google/go-tpm-tools/launcher/device"
3435
"github.com/google/go-tpm-tools/launcher/internal/gpu"
3536
"github.com/google/go-tpm-tools/launcher/internal/healthmonitoring/nodeproblemdetector"
3637
"github.com/google/go-tpm-tools/launcher/internal/logging"
@@ -198,7 +199,7 @@ func NewRunner(ctx context.Context, cdClient *containerd.Client, token oauth2.To
198199
}
199200
specOpts = append(specOpts, cgroupOpts...)
200201

201-
var deviceROTs []agent.DeviceROT
202+
var deviceROTs []device.ROT
202203
nvidiaAttester := gpu.NewNvidiaAttester(launchSpec.InstallGpuDriver)
203204
if launchSpec.InstallGpuDriver {
204205
gpuMounts := []specs.Mount{

‎launcher/container_runner_test.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import (
2525
gecel "github.com/google/go-eventlog/cel"
2626
"github.com/google/go-tpm-tools/cel"
2727
"github.com/google/go-tpm-tools/launcher/agent"
28+
"github.com/google/go-tpm-tools/launcher/device"
2829
"github.com/google/go-tpm-tools/launcher/internal/gpu"
2930
"github.com/google/go-tpm-tools/launcher/internal/logging"
3031
"github.com/google/go-tpm-tools/launcher/launcherfile"
@@ -105,6 +106,10 @@ func (f *fakeGPUAttester) EnableReadyState() error {
105106
return nil
106107
}
107108

109+
func (f *fakeGPUAttester) Vendor() device.Vendor {
110+
return device.NvidiaGPU
111+
}
112+
108113
type fakeClaims struct {
109114
jwt.RegisteredClaims
110115
Signatures []string

‎launcher/device/device.go‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
// Package device provides an interface and management for device Root of Trust (ROT) attestation in the launcher.
2+
package device
3+
4+
import (
5+
"sync"
6+
7+
"go.uber.org/multierr"
8+
)
9+
10+
// Vendor defines the type for device Root of Trust (ROT) vendors.
11+
type Vendor int
12+
13+
// Define constants for supported ROT vendors.
14+
const (
15+
Unspecified Vendor = iota
16+
NvidiaGPU
17+
)
18+
19+
// ROT defines an interface for all attached devices to collect attestation.
20+
type ROT interface {
21+
// Attest fetches an attestation from the attached device detected by launcher.
22+
Attest(nonce []byte) (any, error)
23+
// Vendor returns the device ROT vendor type.
24+
Vendor() Vendor
25+
}
26+
27+
// ROTManager manages the attestation of all attached device ROTs.
28+
type ROTManager struct {
29+
deviceMu sync.Mutex
30+
rots []ROT
31+
}
32+
33+
// ReportOpts defines the options for device attestation report generation.
34+
type ReportOpts struct {
35+
// EnableRuntimeGPUAttestation indicates whether to include runtime GPU attestation in the device reports.
36+
EnableRuntimeGPUAttestation bool
37+
}
38+
39+
// NewROTManager creates a new ROTManager.
40+
func NewROTManager(rots []ROT) *ROTManager {
41+
return &ROTManager{
42+
rots: rots,
43+
}
44+
}
45+
46+
// AttestDeviceROTs fetches attestation reports from all detected device ROTs based on the provided options.
47+
func (m *ROTManager) AttestDeviceROTs(nonce []byte, opts ReportOpts) ([]any, error) {
48+
m.deviceMu.Lock()
49+
defer m.deviceMu.Unlock()
50+
51+
var deviceReports []any
52+
var err error
53+
for _, deviceROT := range m.rots {
54+
if opts.EnableRuntimeGPUAttestation && deviceROT.Vendor() == NvidiaGPU {
55+
deviceReport, e := deviceROT.Attest(nonce)
56+
if e != nil {
57+
err = multierr.Append(err, e)
58+
} else {
59+
deviceReports = append(deviceReports, deviceReport)
60+
}
61+
}
62+
}
63+
return deviceReports, err
64+
}

0 commit comments

Comments
 (0)