Add ITA support to CS (google#605)

Author: alexmwu

launcher/container_runner.go

@@ -240,7 +240,7 @@ func NewRunner(ctx context.Context, cdClient *containerd.Client, token oauth2.To
 	asAddr := launchSpec.AttestationServiceAddr
 
 	var verifierClient verifier.Client
-	if launchSpec.ITARegion == "" {
+	if launchSpec.ITAConfig.ITARegion == "" {
 		gcaClient, err := util.NewRESTClient(ctx, asAddr, launchSpec.ProjectID, launchSpec.Region)
 		if err != nil {
 			return nil, fmt.Errorf("failed to create REST verifier client: %v", err)
@@ -582,7 +582,7 @@ func (r *ContainerRunner) Run(ctx context.Context) error {
 	}
 
 	// Only refresh token if agent has a default GCA client (not ITA use case).
-	if r.launchSpec.ITARegion == "" {
+	if r.launchSpec.ITAConfig.ITARegion == "" {
 		if err := r.fetchAndWriteToken(ctx); err != nil {
 			return fmt.Errorf("failed to fetch and write OIDC token: %v", err)
 		}
@@ -591,9 +591,9 @@ func (r *ContainerRunner) Run(ctx context.Context) error {
 	// create and start the TEE server
 	r.logger.Info("EnableOnDemandAttestation is enabled: initializing TEE server.")
 
-	attestClients := &teeserver.AttestClients{}
-	if r.launchSpec.ITARegion != "" {
-		itaClient, err := ita.NewClient(r.launchSpec.ITARegion, r.launchSpec.ITAKey)
+	attestClients := teeserver.AttestClients{}
+	if r.launchSpec.ITAConfig.ITARegion != "" {
+		itaClient, err := ita.NewClient(r.launchSpec.ITAConfig)
 		if err != nil {
 			return fmt.Errorf("failed to create ITA client: %v", err)
 		}

launcher/go.sum

@@ -155,6 +155,7 @@ github.com/caarlos0/ctrlc v1.0.0/go.mod h1:CdXpj4rmq0q/1Eb44M9zi2nKB0QraNKuRGYGr
 github.com/campoy/unique v0.0.0-20180121183637-88950e537e7e/go.mod h1:9IOqJGCPMSc6E5ydlp5NIonxObaeu/Iub/X03EKPVYo=
 github.com/casbin/casbin/v2 v2.1.2/go.mod h1:YcPU1XXisHhLzuxH9coDNf2FbKpjGlbCg3n9yuLkIJQ=
 github.com/cavaliercoder/go-cpio v0.0.0-20180626203310-925f9528c45e/go.mod h1:oDpT4efm8tSYHXV5tHSdRvBet/b/QzxZ+XyyPehvm3A=
+github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4=
 github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM=
 github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
 github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=

launcher/spec/launch_spec.go

@@ -21,6 +21,7 @@ import (
 	"github.com/google/go-tpm-tools/launcher/internal/launchermount"
 	"github.com/google/go-tpm-tools/launcher/internal/logging"
 	"github.com/google/go-tpm-tools/launcher/launcherfile"
+	"github.com/google/go-tpm-tools/verifier"
 	"github.com/google/go-tpm-tools/verifier/util"
 )
 
@@ -124,8 +125,7 @@ type LaunchSpec struct {
 	MonitoringEnabled          MonitoringType
 	LogRedirect                LogRedirectLocation
 	Mounts                     []launchermount.Mount
-	ITARegion                  string
-	ITAKey                     string
+	ITAConfig                  verifier.ITAConfig
 	// DevShmSize is specified in kiB.
 	DevShmSize        int64
 	AddedCapabilities []string
@@ -252,16 +252,14 @@ func (s *LaunchSpec) UnmarshalJSON(b []byte) error {
 		itaRegionVal, itaRegionOK := unmarshaledMap[itaRegion]
 		itaKeyVal, itaKeyOK := unmarshaledMap[itaKey]
 
+		// If key and region are both not in the map, do not set up ITA config.
 		if itaRegionOK != itaKeyOK {
-			return fmt.Errorf("ITA fields %s and %s must both be provided", itaRegion, itaKey)
+			return fmt.Errorf("ITA fields %s and %s must both be provided and non-empty", itaRegion, itaKey)
 		}
 
-		if itaRegionOK {
-			s.ITARegion = itaRegionVal
-		}
-
-		if itaKeyOK {
-			s.ITAKey = itaKeyVal
+		s.ITAConfig = verifier.ITAConfig{
+			ITARegion: itaRegionVal,
+			ITAKey:    itaKeyVal,
 		}
 	}
 
@@ -290,7 +288,7 @@ func (s *LaunchSpec) UnmarshalJSON(b []byte) error {
 // LogFriendly creates a copy of the spec that is safe to log by censoring
 func (s *LaunchSpec) LogFriendly() LaunchSpec {
 	safeSpec := *s
-	safeSpec.ITAKey = strings.Repeat("*", len(s.ITAKey))
+	safeSpec.ITAConfig.ITAKey = strings.Repeat("*", len(s.ITAConfig.ITAKey))
 
 	return safeSpec
 }

launcher/spec/launch_spec_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-tpm-tools/launcher/internal/experiments"
 	"github.com/google/go-tpm-tools/launcher/internal/launchermount"
+	"github.com/google/go-tpm-tools/verifier"
 )
 
 func TestLaunchSpecUnmarshalJSONHappyCases(t *testing.T) {
@@ -64,8 +65,10 @@ func TestLaunchSpecUnmarshalJSONHappyCases(t *testing.T) {
 		DevShmSize:                 234234,
 		Mounts: []launchermount.Mount{launchermount.TmpfsMount{Destination: "/tmpmount", Size: 0},
 			launchermount.TmpfsMount{Destination: "/sized", Size: 222}},
-		ITARegion: "US",
-		ITAKey:    "test-api-key",
+		ITAConfig: verifier.ITAConfig{
+			ITARegion: "US",
+			ITAKey:    "test-api-key",
+		},
 		Experiments: experiments.Experiments{
 			EnableItaVerifier: true,
 		},

launcher/teeserver/tee_server.go

@@ -14,11 +14,15 @@ import (
 	"github.com/google/go-tpm-tools/launcher/spec"
 	"github.com/google/go-tpm-tools/verifier"
 	"github.com/google/go-tpm-tools/verifier/models"
-	"github.com/google/go-tpm-tools/verifier/util"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
 
+const (
+	gcaEndpoint = "/v1/token"
+	itaEndpoint = "/v1/intel/token"
+)
+
 var clientErrorCodes = map[codes.Code]struct{}{
 	codes.InvalidArgument:    {},
 	codes.FailedPrecondition: {},
@@ -43,7 +47,7 @@ type attestHandler struct {
 	// defaultTokenFile string
 	logger     logging.Logger
 	launchSpec spec.LaunchSpec
-	clients    *AttestClients
+	clients    AttestClients
 }
 
 // TeeServer is a server that can be called from a container through a unix
@@ -58,7 +62,7 @@ const (
 )
 
 // New takes in a socket and start to listen to it, and create a server
-func New(ctx context.Context, unixSock string, a agent.AttestationAgent, logger logging.Logger, launchSpec spec.LaunchSpec, clients *AttestClients) (*TeeServer, error) {
+func New(ctx context.Context, unixSock string, a agent.AttestationAgent, logger logging.Logger, launchSpec spec.LaunchSpec, clients AttestClients) (*TeeServer, error) {
 	var err error
 	nl, err := net.Listen("unix", unixSock)
 	if err != nil {
@@ -88,8 +92,8 @@ func (a *attestHandler) Handler() http.Handler {
 	// curl -d '{"audience":"<aud>", "nonces":["<nonce1>"]}' -H "Content-Type: application/json" -X POST
 	//   --unix-socket /tmp/container_launcher/teeserver.sock http://localhost/v1/token
 
-	mux.HandleFunc("/v1/token", a.getToken)
-	mux.HandleFunc("/v1/intel/token", a.getITAToken)
+	mux.HandleFunc(gcaEndpoint, a.getToken)
+	mux.HandleFunc(itaEndpoint, a.getITAToken)
 	return mux
 }
 
@@ -105,16 +109,13 @@ func (a *attestHandler) logAndWriteError(errStr string, status int, w http.Respo
 func (a *attestHandler) getToken(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "text/html")
 
-	// If the handler does not have a GCA client, create one.
-	if a.clients.GCA == nil {
-		gcaClient, err := util.NewRESTClient(a.ctx, a.launchSpec.AttestationServiceAddr, a.launchSpec.ProjectID, a.launchSpec.Region)
-		if err != nil {
-			errStr := fmt.Sprintf("failed to create REST verifier client: %v", err)
-			a.logAndWriteError(errStr, http.StatusInternalServerError, w)
-			return
-		}
+	a.logger.Info(fmt.Sprintf("%s called", gcaEndpoint))
 
-		a.clients.GCA = gcaClient
+	// If the handler does not have an GCA client, return error.
+	if a.clients.GCA == nil {
+		errStr := "no GCA verifier client present, please try rebooting your VM"
+		a.logAndWriteError(errStr, http.StatusInternalServerError, w)
+		return
 	}
 
 	a.attest(w, r, a.clients.GCA)
@@ -124,10 +125,12 @@ func (a *attestHandler) getToken(w http.ResponseWriter, r *http.Request) {
 func (a *attestHandler) getITAToken(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "text/html")
 
+	a.logger.Info(fmt.Sprintf("%s called", itaEndpoint))
+
 	// If the handler does not have an ITA client, return error.
 	if a.clients.ITA == nil {
 		errStr := "no ITA verifier client present - ensure ITA Region and Key are defined in metadata"
-		a.logAndWriteError(errStr, http.StatusPreconditionFailed, w)
+		a.logAndWriteError(errStr, http.StatusInternalServerError, w)
 		return
 	}
 
@@ -176,11 +179,10 @@ func (a *attestHandler) attest(w http.ResponseWriter, r *http.Request, client ve
 		}
 
 		// Do not check that TokenTypeOptions matches TokenType in the launcher.
-
-		tok, err := a.attestAgent.AttestWithClient(a.ctx, agent.AttestAgentOpts{
+		opts := agent.AttestAgentOpts{
 			TokenOptions: &tokenOptions,
-		}, client)
-
+		}
+		tok, err := a.attestAgent.AttestWithClient(a.ctx, opts, client)
 		if err != nil {
 			a.handleAttestError(w, err, "failed to retrieve custom attestation service token")
 			return