Add experiment-gated migration to VerifyConfidentialSpace (google#586)

Author: jessieqliu

launcher/agent/agent.go

@@ -251,16 +251,24 @@ func (a *agent) AttestWithClient(ctx context.Context, opts AttestAgentOpts, clie
 		a.logger.Info("Found container image signatures: %v\n", signatures)
 	}
 
-	resp, err := client.VerifyAttestation(ctx, req)
+	resp, err := a.verify(ctx, req, client)
 	if err != nil {
 		return nil, err
 	}
+
 	if len(resp.PartialErrs) > 0 {
 		a.logger.Error(fmt.Sprintf("Partial errors from VerifyAttestation: %v", resp.PartialErrs))
 	}
 	return resp.ClaimsToken, nil
 }
 
+func (a *agent) verify(ctx context.Context, req verifier.VerifyAttestationRequest, client verifier.Client) (*verifier.VerifyAttestationResponse, error) {
+	if a.launchSpec.Experiments.EnableVerifyCS {
+		return client.VerifyConfidentialSpace(ctx, req)
+	}
+	return client.VerifyAttestation(ctx, req)
+}
+
 func convertOCIToContainerSignature(ociSig oci.Signature) (*verifier.ContainerSignature, error) {
 	payload, err := ociSig.Payload()
 	if err != nil {

launcher/agent/agent_test.go

@@ -5,6 +5,7 @@ import (
 	"crypto/rand"
 	"crypto/rsa"
 	"encoding/base64"
+	"errors"
 	"fmt"
 	"math"
 	"runtime"
@@ -18,6 +19,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/internal/experiments"
 	"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"
@@ -638,3 +640,77 @@ func measureFakeEvents(attestAgent AttestationAgent) error {
 	}
 	return nil
 }
+
+type testClient struct {
+	verifyCSResp  *verifier.VerifyAttestationResponse
+	verifyAttResp *verifier.VerifyAttestationResponse
+}
+
+func (t *testClient) CreateChallenge(_ context.Context) (*verifier.Challenge, error) {
+	return nil, errors.New("unimplemented")
+}
+
+func (t *testClient) VerifyAttestation(_ context.Context, _ verifier.VerifyAttestationRequest) (*verifier.VerifyAttestationResponse, error) {
+	return t.verifyAttResp, nil
+}
+
+func (t *testClient) VerifyConfidentialSpace(_ context.Context, _ verifier.VerifyAttestationRequest) (*verifier.VerifyAttestationResponse, error) {
+	return t.verifyCSResp, nil
+}
+
+func TestVerify(t *testing.T) {
+	expectedCSResp := &verifier.VerifyAttestationResponse{
+		ClaimsToken: []byte("verify-cs-token"),
+	}
+	expectedAttResp := &verifier.VerifyAttestationResponse{
+		ClaimsToken: []byte("verify-att-token"),
+	}
+
+	vClient := &testClient{
+		verifyCSResp:  expectedCSResp,
+		verifyAttResp: expectedAttResp,
+	}
+
+	testcases := []struct {
+		name         string
+		opts         AttestAgentOpts
+		exps         experiments.Experiments
+		expectedResp *verifier.VerifyAttestationResponse
+	}{
+		{
+			name: "VerifyCS in experiment",
+			exps: experiments.Experiments{
+				EnableVerifyCS: true,
+			},
+			expectedResp: expectedCSResp,
+		},
+		{
+			name: "VerifyAtt in experiment",
+			exps: experiments.Experiments{
+				EnableVerifyCS: false,
+			},
+			expectedResp: expectedAttResp,
+		},
+	}
+
+	ctx := context.Background()
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			attAgent := agent{
+				launchSpec: spec.LaunchSpec{
+					Experiments: tc.exps,
+				},
+			}
+
+			resp, err := attAgent.verify(ctx, verifier.VerifyAttestationRequest{}, vClient)
+			if err != nil {
+				t.Fatalf("verify() returned failure: %v", err)
+			}
+
+			if diff := cmp.Diff(resp, tc.expectedResp); diff != "" {
+				t.Errorf("verify() did not return expected response (-got, +want): %v", diff)
+			}
+		})
+	}
+}

launcher/teeserver/tee_server_test.go

@@ -507,3 +507,72 @@ func TestCustomTokenDataParsedSuccessfully(t *testing.T) {
 		}
 	}
 }
+
+func TestCustomHandleAttestError(t *testing.T) {
+	body := `{
+				"audience": "audience",
+				"nonces": ["thisIsAcustomNonce"],
+				"token_type": "OIDC"
+			}`
+
+	testcases := []struct {
+		name           string
+		err            error
+		wantStatusCode int
+	}{
+		{
+			name:           "FailedPrecondition error",
+			err:            status.New(codes.FailedPrecondition, "bad state").Err(),
+			wantStatusCode: http.StatusBadRequest,
+		},
+		{
+			name:           "PermissionDenied error",
+			err:            status.New(codes.PermissionDenied, "denied").Err(),
+			wantStatusCode: http.StatusBadRequest,
+		},
+		{
+			name:           "Internal error",
+			err:            status.New(codes.Internal, "internal server error").Err(),
+			wantStatusCode: http.StatusInternalServerError,
+		},
+		{
+			name:           "Unavailable error",
+			err:            status.New(codes.Unavailable, "service unavailable").Err(),
+			wantStatusCode: http.StatusInternalServerError,
+		},
+		{
+			name:           "non-gRPC error",
+			err:            errors.New("a generic error"),
+			wantStatusCode: http.StatusInternalServerError,
+		},
+	}
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			ah := attestHandler{
+				logger: logging.SimpleLogger(),
+				clients: AttestClients{
+					GCA: &fakeVerifierClient{},
+				},
+				attestAgent: fakeAttestationAgent{
+					attestWithClientFunc: func(context.Context, agent.AttestAgentOpts, verifier.Client) ([]byte, error) {
+						return nil, tc.err
+					},
+				},
+			}
+
+			req := httptest.NewRequest(http.MethodPost, "/v1/token", strings.NewReader(body))
+			w := httptest.NewRecorder()
+
+			ah.getToken(w, req)
+
+			if w.Code != tc.wantStatusCode {
+				t.Errorf("got status code %d, want %d", w.Code, tc.wantStatusCode)
+			}
+
+			_, err := io.ReadAll(w.Result().Body)
+			if err != nil {
+				t.Errorf("failed to read response body: %v", err)
+			}
+		})
+	}
+}

verifier/rest/rest.go

@@ -14,7 +14,6 @@ import (
 	tpb "github.com/google/go-tdx-guest/proto/tdx"
 	"github.com/google/go-tpm-tools/verifier"
 	"github.com/google/go-tpm-tools/verifier/models"
-	"github.com/google/go-tpm-tools/verifier/oci"
 	"github.com/googleapis/gax-go/v2"
 
 	v1 "cloud.google.com/go/confidentialcomputing/apiv1"
@@ -301,25 +300,6 @@ func convertResponseFromREST(resp *ccpb.VerifyAttestationResponse) (*verifier.Ve
 	}, nil
 }
 
-func convertOCISignatureToREST(signature oci.Signature) (*ccpb.ContainerImageSignature, error) {
-	payload, err := signature.Payload()
-	if err != nil {
-		return nil, err
-	}
-	b64Sig, err := signature.Base64Encoded()
-	if err != nil {
-		return nil, err
-	}
-	sigBytes, err := encoding.DecodeString(b64Sig)
-	if err != nil {
-		return nil, err
-	}
-	return &ccpb.ContainerImageSignature{
-		Payload:   payload,
-		Signature: sigBytes,
-	}, nil
-}
-
 func convertSEVSNPProtoToREST(att *spb.Attestation) (*ccpb.VerifyAttestationRequest_SevSnpAttestation, error) {
 	auxBlob := sabi.CertsFromProto(att.GetCertificateChain()).Marshal()
 	rawReport, err := sabi.ReportToAbiBytes(att.GetReport())