Enforce max response metadata size (flyteorg#139)

eapolinario · Nov 13, 2020 · 15545c3 · 15545c3
1 parent 0fcb7cb
commit 15545c3
Show file tree

Hide file tree

Showing 7 changed files with 57 additions and 19 deletions.
diff --git a/cmd/entrypoints/serve.go b/cmd/entrypoints/serve.go
@@ -212,7 +212,8 @@ func serveGatewayInsecure(ctx context.Context, cfg *config.ServerConfig) error {
 	}()
 
 	logger.Infof(ctx, "Starting HTTP/1 Gateway server on %s", cfg.GetHostAddress())
-	httpServer, err := newHTTPServer(ctx, cfg, authContext, cfg.GetGrpcHostAddress(), grpc.WithInsecure())
+	httpServer, err := newHTTPServer(ctx, cfg, authContext, cfg.GetGrpcHostAddress(), grpc.WithInsecure(),
+		grpc.WithMaxHeaderListSize(common.MaxResponseStatusBytes))
 	if err != nil {
 		return err
 	}

diff --git a/pkg/common/constants.go b/pkg/common/constants.go
@@ -9,3 +9,5 @@ const (
 	AuditFieldsContextKey contextutils.Key = "audit_fields"
 	PrincipalContextKey   contextutils.Key = "principal"
 )
+
+const MaxResponseStatusBytes = 32000
diff --git a/pkg/rpc/adminservice/tests/execution_test.go b/pkg/rpc/adminservice/tests/execution_test.go
@@ -6,6 +6,9 @@ import (
 	"testing"
 	"time"
 
+	flyteAdminErrors "github.com/lyft/flyteadmin/pkg/errors"
+	"google.golang.org/grpc/codes"
+
 	"github.com/lyft/flyteidl/gen/pb-go/flyteidl/core"
 
 	"github.com/golang/protobuf/proto"
@@ -178,7 +181,8 @@ func TestCreateWorkflowEventErr(t *testing.T) {
 			Phase:       core.WorkflowExecution_RUNNING,
 		},
 	})
-	assert.EqualError(t, err, "rpc error: code = Internal desc = expected error")
+	assert.EqualError(t, err, "expected error")
+	assert.Equal(t, codes.Internal, err.(flyteAdminErrors.FlyteAdminError).Code())
 	assert.Nil(t, resp)
 }
 
@@ -220,7 +224,7 @@ func TestGetExecutionError(t *testing.T) {
 	actualResponse, err := mockServer.GetExecution(context.Background(), &admin.WorkflowExecutionGetRequest{
 		Id: &workflowExecutionIdentifier,
 	})
-	assert.EqualError(t, err, "rpc error: code = Internal desc = expected error")
+	assert.EqualError(t, err, "expected error")
 	assert.Nil(t, actualResponse)
 }
 
@@ -273,7 +277,8 @@ func TestListExecutionsError(t *testing.T) {
 		},
 		Limit: 1,
 	})
-	assert.EqualError(t, err, "rpc error: code = Internal desc = expected error")
+	assert.EqualError(t, err, "expected error")
+	assert.Equal(t, codes.Internal, err.(flyteAdminErrors.FlyteAdminError).Code())
 	assert.Nil(t, response)
 }
 
@@ -320,6 +325,7 @@ func TestTerminateExecution_Error(t *testing.T) {
 		Id:    &identifier,
 		Cause: abortCause,
 	})
-	assert.EqualError(t, err, "rpc error: code = Internal desc = expected error")
+	assert.EqualError(t, err, "expected error")
+	assert.Equal(t, codes.Internal, err.(flyteAdminErrors.FlyteAdminError).Code())
 	assert.Nil(t, response)
 }
diff --git a/pkg/rpc/adminservice/tests/node_execution_test.go b/pkg/rpc/adminservice/tests/node_execution_test.go
@@ -5,6 +5,9 @@ import (
 	"errors"
 	"testing"
 
+	flyteAdminErrors "github.com/lyft/flyteadmin/pkg/errors"
+	"google.golang.org/grpc/codes"
+
 	"github.com/golang/protobuf/proto"
 
 	"github.com/lyft/flyteidl/gen/pb-go/flyteidl/core"
@@ -68,7 +71,8 @@ func TestCreateNodeEventErr(t *testing.T) {
 			Phase: core.NodeExecution_SKIPPED,
 		},
 	})
-	assert.EqualError(t, err, "rpc error: code = Internal desc = expected error")
+	assert.EqualError(t, err, "expected error")
+	assert.Equal(t, codes.Internal, err.(flyteAdminErrors.FlyteAdminError).Code())
 	assert.Nil(t, resp)
 }
 
@@ -111,7 +115,8 @@ func TestGetNodeExecutionError(t *testing.T) {
 	actualResponse, err := mockServer.GetNodeExecution(context.Background(), &admin.NodeExecutionGetRequest{
 		Id: &nodeExecutionID,
 	})
-	assert.EqualError(t, err, "rpc error: code = Internal desc = expected error")
+	assert.EqualError(t, err, "expected error")
+	assert.Equal(t, codes.Internal, err.(flyteAdminErrors.FlyteAdminError).Code())
 	assert.Nil(t, actualResponse)
 }
 
@@ -160,7 +165,8 @@ func TestListNodeExecutionsError(t *testing.T) {
 		Limit: 1,
 		Token: "20",
 	})
-	assert.EqualError(t, err, "rpc error: code = Internal desc = expected error")
+	assert.EqualError(t, err, "expected error")
+	assert.Equal(t, codes.Internal, err.(flyteAdminErrors.FlyteAdminError).Code())
 	assert.Nil(t, response)
 }
 
@@ -211,7 +217,8 @@ func TestListNodeExecutionsForTaskError(t *testing.T) {
 		Limit: 1,
 		Token: "20",
 	})
-	assert.EqualError(t, err, "rpc error: code = Internal desc = expected error")
+	assert.EqualError(t, err, "expected error")
+	assert.Equal(t, codes.Internal, err.(flyteAdminErrors.FlyteAdminError).Code())
 	assert.Nil(t, response)
 }
 

diff --git a/pkg/rpc/adminservice/tests/task_execution_test.go b/pkg/rpc/adminservice/tests/task_execution_test.go
@@ -5,7 +5,10 @@ import (
 	"errors"
 	"testing"
 
+	"google.golang.org/grpc/codes"
+
 	"github.com/golang/protobuf/proto"
+	flyteAdminErrors "github.com/lyft/flyteadmin/pkg/errors"
 	"github.com/lyft/flyteadmin/pkg/manager/mocks"
 	"github.com/lyft/flyteidl/gen/pb-go/flyteidl/admin"
 	"github.com/lyft/flyteidl/gen/pb-go/flyteidl/core"
@@ -177,7 +180,8 @@ func TestTaskExecution(t *testing.T) {
 				RetryAttempt:    retryAttempt,
 			},
 		})
-		assert.EqualError(t, err, "rpc error: code = Internal desc = expected error")
+		assert.EqualError(t, err, "expected error")
+		assert.Equal(t, codes.Internal, err.(flyteAdminErrors.FlyteAdminError).Code())
 		assert.Nil(t, resp)
 	})
 

diff --git a/pkg/rpc/adminservice/util/transformers.go b/pkg/rpc/adminservice/util/transformers.go
@@ -1,20 +1,25 @@
 package util
 
 import (
+	"github.com/lyft/flyteadmin/pkg/common"
 	"github.com/lyft/flyteadmin/pkg/errors"
 
 	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
 )
 
-// Transforms errors to grpc-compatible error types.
+// Transforms errors to grpc-compatible error types and optionally truncates it if necessary.
 func TransformAndRecordError(err error, metrics *RequestMetrics) error {
-	switch err := err.(type) {
-	case errors.FlyteAdminError:
-		metrics.Record(err.(errors.FlyteAdminError).Code())
-		return err
-	default:
-		metrics.Record(codes.Internal)
-		return status.Error(codes.Internal, err.Error())
+	var errorMessage = err.Error()
+	concatenateErrMessage := false
+	if len(errorMessage) > common.MaxResponseStatusBytes {
+		errorMessage = err.Error()[:common.MaxResponseStatusBytes]
+		concatenateErrMessage = true
 	}
+	if flyteAdminError, ok := err.(errors.FlyteAdminError); !ok {
+		err = errors.NewFlyteAdminError(codes.Internal, errorMessage)
+	} else if concatenateErrMessage {
+		err = errors.NewFlyteAdminError(flyteAdminError.Code(), errorMessage)
+	}
+	metrics.Record(err.(errors.FlyteAdminError).Code())
+	return err
 }
diff --git a/pkg/rpc/adminservice/util/transformers_test.go b/pkg/rpc/adminservice/util/transformers_test.go
@@ -5,6 +5,8 @@ import (
 	"errors"
 	"testing"
 
+	"github.com/lyft/flyteadmin/pkg/common"
+
 	adminErrors "github.com/lyft/flyteadmin/pkg/errors"
 	mockScope "github.com/lyft/flytestdlib/promutils"
 	"github.com/stretchr/testify/assert"
@@ -38,3 +40,14 @@ func TestTransformError_BasicError(t *testing.T) {
 	assert.True(t, ok)
 	assert.Equal(t, codes.Internal, transormerStatus.Code())
 }
+
+func TestTruncateErrorMessage(t *testing.T) {
+	errorMessage := make([]byte, common.MaxResponseStatusBytes+1)
+	for i := 0; i <= common.MaxResponseStatusBytes; i++ {
+		errorMessage[i] = byte('a')
+	}
+
+	err := adminErrors.NewFlyteAdminError(codes.InvalidArgument, string(errorMessage))
+	transformedError := TransformAndRecordError(err, &testRequestMetrics)
+	assert.Len(t, transformedError.Error(), common.MaxResponseStatusBytes)
+}