TEE Server Error code Translation (google#587)

* Attest now will pass in a default audience instead of passing an error

* TestCustomToken now lets empty audience through

* Moved audience check from tee server to agent

* removed unused constant

* tokenOpts is now passed as request

* added in constant for audience

* stashing testing comment

* adding in changes from jessieqliu comments

* changes for tee-server error

* changed test TestRequestFailurePassedToCaller to return correct status message

* removed comments and code from another PR

* refactor based on review comments

* added in fix for using empty map and additional comment and nit changes

---------

Co-authored-by: Sibghat Shah <sibghat@google.com>

Author: Sibcgh

launcher/teeserver/tee_server.go

@@ -15,8 +15,21 @@ import (
 	"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"
 )
 
+var clientErrorCodes = map[codes.Code]struct{}{
+	codes.InvalidArgument:    {},
+	codes.FailedPrecondition: {},
+	codes.PermissionDenied:   {},
+	codes.Unauthenticated:    {},
+	codes.NotFound:           {},
+	codes.Aborted:            {},
+	codes.OutOfRange:         {},
+	codes.Canceled:           {},
+}
+
 // AttestClients contains clients for supported verifier services that can be used to
 // get attestation tokens.
 type AttestClients struct {
@@ -120,15 +133,13 @@ func (a *attestHandler) attest(w http.ResponseWriter, r *http.Request, client ve
 	switch r.Method {
 	case http.MethodGet:
 		if err := a.attestAgent.Refresh(a.ctx); err != nil {
-			errStr := fmt.Sprintf("failed to refresh attestation agent: %v", err)
-			a.logAndWriteError(errStr, http.StatusInternalServerError, w)
+			a.logAndWriteHTTPError(w, http.StatusInternalServerError, fmt.Errorf("failed to refresh attestation agent: %w", err))
 			return
 		}
 
 		token, err := a.attestAgent.AttestWithClient(a.ctx, agent.AttestAgentOpts{}, client)
 		if err != nil {
-			errStr := fmt.Sprintf("failed to retrieve attestation service token: %v", err)
-			a.logAndWriteError(errStr, http.StatusInternalServerError, w)
+			a.handleAttestError(w, err, "failed to retrieve attestation service token")
 			return
 		}
 
@@ -165,7 +176,8 @@ func (a *attestHandler) attest(w http.ResponseWriter, r *http.Request, client ve
 			TokenOptions: &tokenOptions,
 		}, client)
 		if err != nil {
-			a.logAndWriteHTTPError(w, http.StatusBadRequest, err)
+
+			a.handleAttestError(w, err, "failed to retrieve custom attestation service token")
 			return
 		}
 
@@ -203,3 +215,20 @@ func (s *TeeServer) Shutdown(ctx context.Context) error {
 	}
 	return nil
 }
+
+func (a *attestHandler) handleAttestError(w http.ResponseWriter, err error, message string) {
+	st, ok := status.FromError(err)
+	if ok {
+		if _, exists := clientErrorCodes[st.Code()]; exists {
+			// User errors, like invalid arguments. Map user errors to 400 Bad Request.
+			a.logAndWriteHTTPError(w, http.StatusBadRequest, fmt.Errorf("%s: %w", message, err))
+			return
+		}
+		// Server-side or transient errors. Map user errors 500 Internal Server Error.
+		a.logAndWriteHTTPError(w, http.StatusInternalServerError, fmt.Errorf("%s: %w", message, err))
+		return
+	}
+	// If it's not a gRPC error, it's likely an internal error within the launcher.
+	// Map user errors 500 Internal Server Error
+	a.logAndWriteHTTPError(w, http.StatusInternalServerError, fmt.Errorf("%s: %w", message, err))
+}

launcher/teeserver/tee_server_test.go

@@ -16,6 +16,8 @@ import (
 	"github.com/google/go-tpm-tools/launcher/internal/logging"
 	"github.com/google/go-tpm-tools/verifier"
 	"github.com/google/go-tpm-tools/verifier/models"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 // Implements verifier.Client interface so it can be used to initialize test attestHandlers
@@ -58,8 +60,6 @@ func (f fakeAttestationAgent) Close() error {
 func TestGetDefaultToken(t *testing.T) {
 	testTokenContent := "test token"
 
-	// An empty attestHandler is fine for now as it is not being used
-	// in the handler.
 	ah := attestHandler{
 		logger: logging.SimpleLogger(),
 		clients: &AttestClients{
@@ -83,8 +83,40 @@ func TestGetDefaultToken(t *testing.T) {
 	if w.Code != http.StatusOK {
 		t.Errorf("got return code: %d, want: %d", w.Code, http.StatusOK)
 	}
-	if string(data) != testTokenContent {
-		t.Errorf("got content: %v, want: %s", testTokenContent, string(data))
+	if diff := cmp.Diff(testTokenContent, string(data)); diff != "" {
+		t.Errorf("getToken() response body mismatch (-want +got):\n%s", diff)
+	}
+}
+
+func TestGetDefaultTokenServerError(t *testing.T) {
+	// An empty attestHandler is fine for now as it is not being used
+	// in the handler.
+	ah := attestHandler{
+		logger: logging.SimpleLogger(),
+		clients: &AttestClients{
+			GCA: &fakeVerifierClient{},
+		},
+		attestAgent: fakeAttestationAgent{
+			attestWithClientFunc: func(context.Context, agent.AttestAgentOpts, verifier.Client) ([]byte, error) {
+				return nil, errors.New("internal server error from agent")
+			},
+		}}
+
+	req := httptest.NewRequest(http.MethodGet, "/v1/token", nil)
+	w := httptest.NewRecorder()
+
+	ah.getToken(w, req)
+	data, err := io.ReadAll(w.Result().Body)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if w.Code != http.StatusInternalServerError {
+		t.Errorf("got return code: %d, want: %d", w.Code, http.StatusInternalServerError)
+	}
+	expectedError := "failed to retrieve attestation service token: internal server error from agent"
+	if diff := cmp.Diff(expectedError, string(data)); diff != "" {
+		t.Errorf("getToken() response body mismatch (-want +got):\n%s", diff)
 	}
 }
 
@@ -118,7 +150,7 @@ func TestCustomToken(t *testing.T) {
 			attestWithClientFunc: func(context.Context, agent.AttestAgentOpts, verifier.Client) ([]byte, error) {
 				return nil, errors.New("Error")
 			},
-			want: http.StatusBadRequest,
+			want: http.StatusInternalServerError,
 		},
 		{
 			testName: "TestTokenTypeRequired",
@@ -167,8 +199,6 @@ func TestCustomToken(t *testing.T) {
 	}
 
 	for i, test := range tests {
-		// An empty attestHandler is fine for now as it is not being used
-		// in the handler.
 		ah := attestHandler{
 			logger: logging.SimpleLogger(),
 			clients: &AttestClients{
@@ -330,3 +360,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)
+			}
+		})
+	}
+}