Add device ROT manager

Author: yawangwang

launcher/agent/agent.go

@@ -34,6 +34,7 @@ import (
 	"github.com/google/go-tpm-tools/cel"
 	"github.com/google/go-tpm-tools/client"
 	"github.com/google/go-tpm-tools/internal"
+	"github.com/google/go-tpm-tools/launcher/device"
 	"github.com/google/go-tpm-tools/launcher/internal/logging"
 	"github.com/google/go-tpm-tools/launcher/internal/signaturediscovery"
 	"github.com/google/go-tpm-tools/launcher/spec"
@@ -72,28 +73,13 @@ type attestRoot interface {
 	Attest(nonce []byte) (any, error)
 	// ComputeNonce hashes the challenge and extraData using the algorithm preferred by the attestation root.
 	ComputeNonce(challenge []byte, extraData []byte) []byte
-	// AddDeviceROTs adds detected device RoTs(root of trust).
-	AddDeviceROTs([]DeviceROT)
-	// AttestDeviceROTs fetches a list of runtime device attestation report.
-	AttestDeviceROTs(nonce []byte) ([]any, error)
-}
-
-// DeviceROT defines an interface for all attached devices to collect attestation.
-type DeviceROT interface {
-	// Attest fetches an attestation from the attached device detected by launcher.
-	Attest(nonce []byte) (any, error)
 }
 
 // AttestAgentOpts contains user generated options when calling the
 // VerifyAttestation API
 type AttestAgentOpts struct {
-	TokenOptions *models.TokenOptions
-	*DeviceReportOpts
-}
-
-// DeviceReportOpts contains options for runtime device attestations.
-type DeviceReportOpts struct {
-	EnableRuntimeGPUAttestation bool
+	TokenOptions     *models.TokenOptions
+	DeviceReportOpts device.ReportOpts
 }
 
 type agent struct {
@@ -106,6 +92,7 @@ type agent struct {
 	launchSpec       spec.LaunchSpec
 	logger           logging.Logger
 	sigsCache        *sigsCache
+	deviceROTManager *device.ROTManager
 }
 
 // CreateAttestationAgent returns an agent capable of performing remote
@@ -115,7 +102,7 @@ type agent struct {
 // - principalFetcher is a func to fetch GCE principal tokens for a given audience.
 // - signaturesFetcher is a func to fetch container image signatures associated with the running workload.
 // - logger will log any partial errors returned by VerifyAttestation.
-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) {
+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) {
 	// Fetched the AK and save it, so the agent doesn't need to create a new key everytime
 	ak, err := akFetcher(tpm)
 	if err != nil {
@@ -178,8 +165,8 @@ func CreateAttestationAgent(tpm io.ReadWriteCloser, akFetcher util.TpmKeyFetcher
 		attestAgent.avRot = tpmAR
 	}
 
-	// Add deviceRoTs to the CPU attestation root.
-	attestAgent.avRot.AddDeviceROTs(deviceROTs)
+	// Register device ROT manager to the attestAgent.
+	attestAgent.deviceROTManager = device.NewROTManager(deviceROTs)
 	return attestAgent, nil
 }
 
@@ -360,10 +347,7 @@ func (a *agent) AttestationEvidence(_ context.Context, challenge []byte, extraDa
 }
 
 func (a *agent) attestDeviceROTs(nonce []byte, opts AttestAgentOpts) ([]*attestationpb.DeviceAttestationReport, error) {
-	if opts.DeviceReportOpts == nil {
-		return nil, nil
-	}
-	deviceROTs, err := a.avRot.AttestDeviceROTs(nonce)
+	deviceROTs, err := a.deviceROTManager.AttestDeviceROTs(nonce, opts.DeviceReportOpts)
 	if err != nil {
 		return nil, err
 	}
@@ -411,12 +395,11 @@ func convertOCIToContainerSignature(ociSig oci.Signature) (*verifier.ContainerSi
 }
 
 type tpmAttestRoot struct {
-	tpmMu      sync.Mutex
-	fetchedAK  *client.Key
-	tpm        io.ReadWriteCloser
-	cosCel     gecel.CEL
-	hashAlgos  []crypto.Hash
-	deviceROTs []DeviceROT
+	tpmMu            sync.Mutex
+	fetchedAK        *client.Key
+	tpm              io.ReadWriteCloser
+	cosCel           gecel.CEL
+	hashAlgos        []crypto.Hash
 }
 
 func (t *tpmAttestRoot) GetCEL() gecel.CEL {
@@ -457,22 +440,10 @@ func (t *tpmAttestRoot) ComputeNonce(challenge []byte, extraData []byte) []byte
 	return finalNonce[:]
 }
 
-func (t *tpmAttestRoot) AddDeviceROTs(deviceROTs []DeviceROT) {
-	t.deviceROTs = append(t.deviceROTs, deviceROTs...)
-}
-
-func (t *tpmAttestRoot) AttestDeviceROTs(nonce []byte) ([]any, error) {
-	t.tpmMu.Lock()
-	defer t.tpmMu.Unlock()
-
-	return doAttestDeviceROTs(t.deviceROTs, nonce)
-}
-
 type tdxAttestRoot struct {
-	tdxMu      sync.Mutex
-	qp         *tg.LinuxConfigFsQuoteProvider
-	cosCel     gecel.CEL
-	deviceROTs []DeviceROT
+	tdxMu  sync.Mutex
+	qp     *tg.LinuxConfigFsQuoteProvider
+	cosCel gecel.CEL
 }
 
 func (t *tdxAttestRoot) GetCEL() gecel.CEL {
@@ -513,13 +484,6 @@ func (t *tdxAttestRoot) Attest(nonce []byte) (any, error) {
 	}, nil
 }
 
-func (t *tdxAttestRoot) AttestDeviceROTs(nonce []byte) ([]any, error) {
-	t.tdxMu.Lock()
-	defer t.tdxMu.Unlock()
-
-	return doAttestDeviceROTs(t.deviceROTs, nonce)
-}
-
 func (t *tdxAttestRoot) ComputeNonce(challenge []byte, extraData []byte) []byte {
 	challengeData := challenge
 	if extraData != nil {
@@ -531,10 +495,6 @@ func (t *tdxAttestRoot) ComputeNonce(challenge []byte, extraData []byte) []byte
 	return finalNonce[:]
 }
 
-func (t *tdxAttestRoot) AddDeviceROTs(deviceROTs []DeviceROT) {
-	t.deviceROTs = append(t.deviceROTs, deviceROTs...)
-}
-
 // Refresh refreshes the internal state of the attestation agent.
 // It will reset the container image signatures for now.
 func (a *agent) Refresh(ctx context.Context) error {
@@ -633,15 +593,3 @@ func convertToTPMQuote(v *pb.Attestation) *attestationpb.TpmQuote {
 		},
 	}
 }
-
-func doAttestDeviceROTs(deviceROTs []DeviceROT, nonce []byte) ([]any, error) {
-	var deviceReports []any
-	for _, deviceROT := range deviceROTs {
-		deviceReport, err := deviceROT.Attest(nonce)
-		if err != nil {
-			return nil, err
-		}
-		deviceReports = append(deviceReports, deviceReport)
-	}
-	return deviceReports, nil
-}

launcher/agent/agent_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/google/go-tpm-tools/cel"
 	"github.com/google/go-tpm-tools/client"
 	"github.com/google/go-tpm-tools/internal/test"
+	"github.com/google/go-tpm-tools/launcher/device"
 	"github.com/google/go-tpm-tools/launcher/internal/experiments"
 	"github.com/google/go-tpm-tools/launcher/internal/logging"
 	"github.com/google/go-tpm-tools/launcher/internal/signaturediscovery"
@@ -659,7 +660,6 @@ type fakeTdxAttestRoot struct {
 	cel           gecel.CEL
 	receivedNonce []byte
 	tdxQuote      []byte
-	deviceROTs    []DeviceROT
 }
 
 func (f *fakeTdxAttestRoot) Extend(c gecel.Content) error {
@@ -690,30 +690,9 @@ func (f *fakeTdxAttestRoot) ComputeNonce(challenge []byte, extraData []byte) []b
 	return finalNonce[:]
 }
 
-func (f *fakeTdxAttestRoot) AttestDeviceROTs(nonce []byte) ([]any, error) {
-	var deviceReports []any
-	for _, deviceRoT := range f.deviceROTs {
-		att, err := deviceRoT.Attest(nonce)
-		if err != nil {
-			return nil, err
-		}
-		switch v := att.(type) {
-		case *attestationpb.NvidiaAttestationReport:
-			deviceReports = append(deviceReports, v)
-		default:
-			return nil, fmt.Errorf("unknown device attestation type: %T", v)
-		}
-	}
-	return deviceReports, nil
-}
-
 //go:embed testdata/cel.b64
 var celB64 string
 
-func (f *fakeTdxAttestRoot) AddDeviceROTs(deviceRoTS []DeviceROT) {
-	f.deviceROTs = append(f.deviceROTs, deviceRoTS...)
-}
-
 type fakeGPURoT struct{}
 
 func (f *fakeGPURoT) Attest(nonce []byte) (any, error) {
@@ -728,56 +707,9 @@ func (f *fakeGPURoT) Attest(nonce []byte) (any, error) {
 		},
 	}, nil
 }
-func TestTDXAttestDeviceROTs(t *testing.T) {
-	testCases := []struct {
-		name          string
-		tdxAttestRoot *fakeTdxAttestRoot
-		nonce         []byte
-		wantGPU       bool
-		wantPass      bool
-	}{
-		{
-			name:          "success tdxAttestRoot w/o GPU device",
-			tdxAttestRoot: &fakeTdxAttestRoot{},
-			nonce:         []byte("test-nonce"),
-			wantPass:      true,
-		},
-		{
-			name: "success tdxAttestRoot w/ GPU device",
-			tdxAttestRoot: &fakeTdxAttestRoot{
-				deviceROTs: []DeviceROT{&fakeGPURoT{}},
-			},
-			nonce:    []byte("test-nonce"),
-			wantGPU:  true,
-			wantPass: true,
-		},
-		{
-			name: "failed tdxAttestRoot w/ GPU device",
-			tdxAttestRoot: &fakeTdxAttestRoot{
-				deviceROTs: []DeviceROT{&fakeGPURoT{}},
-			},
-			nonce:    []byte(""),
-			wantPass: false,
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
-			deviceReports, err := tc.tdxAttestRoot.AttestDeviceROTs(tc.nonce)
-			if gotPass := err == nil; gotPass != tc.wantPass {
-				t.Errorf("tdxAttestRoot.AttestDeviceROTs() did not return expected attestation result, got %v, want %v", gotPass, tc.wantPass)
-			}
-			if tc.wantGPU {
-				if len(deviceReports) == 0 {
-					t.Fatalf("tdxAttestRoot.AttestDeviceROTs() didn't return any device reports")
-				}
-				if att := deviceReports[0].(*attestationpb.NvidiaAttestationReport); att == nil {
-					t.Errorf("tdxAttestRoot.AttestDeviceROTs() didn't return expected device report type, want %v, but got nil", &attestationpb.NvidiaAttestationReport{})
-				}
-			}
-		})
-	}
 
+func (f *fakeGPURoT) Vendor() device.Vendor {
+	return device.NvidiaGPU
 }
 
 func TestAttestationEvidence_TDX_Success(t *testing.T) {
@@ -816,8 +748,8 @@ func TestAttestationEvidence_TDX_Success(t *testing.T) {
 				EnableAttestationEvidence: true,
 			},
 		},
+		deviceROTManager: device.NewROTManager([]device.ROT{&fakeGPURoT{}}),
 	}
-	attestAgent.avRot.AddDeviceROTs([]DeviceROT{&fakeGPURoT{}})
 
 	if err := measureFakeEvents(attestAgent); err != nil {
 		t.Fatalf("failed to measure events: %v", err)
@@ -838,7 +770,7 @@ func TestAttestationEvidence_TDX_Success(t *testing.T) {
 		{
 			name: "TDX attestation + runtime GPU attestation",
 			opts: AttestAgentOpts{
-				DeviceReportOpts: &DeviceReportOpts{
+				DeviceReportOpts: device.ReportOpts{
 					EnableRuntimeGPUAttestation: true,
 				},
 			},
@@ -1146,6 +1078,7 @@ func TestAttestationEvidence_TDX_NilExtraData(t *testing.T) {
 				EnableAttestationEvidence: true,
 			},
 		},
+		deviceROTManager: device.NewROTManager(nil),
 	}
 
 	challenge := []byte("test-challenge")

launcher/container_runner.go

@@ -31,6 +31,7 @@ import (
 	"github.com/google/go-tpm-tools/client"
 	workloadservice "github.com/google/go-tpm-tools/keymanager/workload_service"
 	"github.com/google/go-tpm-tools/launcher/agent"
+	"github.com/google/go-tpm-tools/launcher/device"
 	"github.com/google/go-tpm-tools/launcher/internal/gpu"
 	"github.com/google/go-tpm-tools/launcher/internal/healthmonitoring/nodeproblemdetector"
 	"github.com/google/go-tpm-tools/launcher/internal/logging"
@@ -198,7 +199,7 @@ func NewRunner(ctx context.Context, cdClient *containerd.Client, token oauth2.To
 	}
 	specOpts = append(specOpts, cgroupOpts...)
 
-	var deviceROTs []agent.DeviceROT
+	var deviceROTs []device.ROT
 	nvidiaAttester := gpu.NewNvidiaAttester(launchSpec.InstallGpuDriver)
 	if launchSpec.InstallGpuDriver {
 		gpuMounts := []specs.Mount{

launcher/container_runner_test.go

@@ -25,6 +25,7 @@ import (
 	gecel "github.com/google/go-eventlog/cel"
 	"github.com/google/go-tpm-tools/cel"
 	"github.com/google/go-tpm-tools/launcher/agent"
+	"github.com/google/go-tpm-tools/launcher/device"
 	"github.com/google/go-tpm-tools/launcher/internal/gpu"
 	"github.com/google/go-tpm-tools/launcher/internal/logging"
 	"github.com/google/go-tpm-tools/launcher/launcherfile"
@@ -105,6 +106,10 @@ func (f *fakeGPUAttester) EnableReadyState() error {
 	return nil
 }
 
+func (f *fakeGPUAttester) Vendor() device.Vendor {
+	return device.NvidiaGPU
+}
+
 type fakeClaims struct {
 	jwt.RegisteredClaims
 	Signatures []string

launcher/device/device.go

@@ -0,0 +1,63 @@
+package device
+
+import (
+	"sync"
+
+	"go.uber.org/multierr"
+)
+
+// Vendor defines the type for device Root of Trust (ROT) vendors.
+type Vendor int
+
+// Define constants for supported ROT vendors.
+const (
+	Unspecified Vendor = iota
+	NvidiaGPU
+)
+
+// ROT defines an interface for all attached devices to collect attestation.
+type ROT interface {
+	// Attest fetches an attestation from the attached device detected by launcher.
+	Attest(nonce []byte) (any, error)
+	// Vendor returns the device ROT vendor type.
+	Vendor() Vendor
+}
+
+// ROTManager manages the attestation of all attached device ROTs.
+type ROTManager struct {
+	deviceMu sync.Mutex
+	rots     []ROT
+}
+
+// ReportOpts defines the options for device attestation report generation.
+type ReportOpts struct {
+	// EnableRuntimeGPUAttestation indicates whether to include runtime GPU attestation in the device reports.
+	EnableRuntimeGPUAttestation bool
+}
+
+// NewROTManager creates a new ROTManager.
+func NewROTManager(rots []ROT) *ROTManager {
+	return &ROTManager{
+		rots: rots,
+	}
+}
+
+// AttestDeviceROTs fetches attestation reports from all detected device ROTs based on the provided options.
+func (m *ROTManager) AttestDeviceROTs(nonce []byte, opts ReportOpts) ([]any, error) {
+	m.deviceMu.Lock()
+	defer m.deviceMu.Unlock()
+
+	var deviceReports []any
+	var err error
+	for _, deviceROT := range m.rots {
+		if opts.EnableRuntimeGPUAttestation && deviceROT.Vendor() == NvidiaGPU {
+			deviceReport, e := deviceROT.Attest(nonce)
+			if e != nil {
+				err = multierr.Append(err, e)
+			} else {
+				deviceReports = append(deviceReports, deviceReport)
+			}
+		}
+	}
+	return deviceReports, err
+}