fix: Skip very large gRPC messages, log when it happens (#421)

hermanschaaf · web-flow · commit 0874d585d2fc · 2022-11-22T10:41:18.000Z
Context: cloudquery/cloudquery#4834 This change: - bumps the gRPC max message size from 50MiB to 100MiB - logs the name of the table to Sentry when we see a message over 50MiB - logs to the user if we drop a message because it exceeds 100MiB
diff --git a/clients/constants.go b/clients/constants.go
@@ -1,5 +1,5 @@
 package clients
 
 const (
-	maxMsgSize = 50 * 1024 * 1024 // 50 MiB
+	maxMsgSize = 100 * 1024 * 1024 // 100 MiB
 )
diff --git a/internal/servers/constants.go b/internal/servers/constants.go
@@ -0,0 +1,3 @@
+package servers
+
+const MaxMsgSize = 100 * 1024 * 1024 // 100 MiB
diff --git a/internal/servers/source.go b/internal/servers/source.go
@@ -4,15 +4,18 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 
 	"github.com/cloudquery/plugin-sdk/internal/pb"
 	"github.com/cloudquery/plugin-sdk/plugins"
 	"github.com/cloudquery/plugin-sdk/schema"
 	"github.com/cloudquery/plugin-sdk/specs"
+	"github.com/getsentry/sentry-go"
 	"github.com/rs/zerolog"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/proto"
 )
 
 type SourceServer struct {
@@ -84,9 +87,17 @@ func (s *SourceServer) Sync2(req *pb.Sync2_Request, stream pb.Source_Sync2Server
 			return status.Errorf(codes.Internal, "failed to marshal resource: %v", err)
 		}
 
-		if err := stream.Send(&pb.Sync2_Response{
+		msg := &pb.Sync2_Response{
 			Resource: b,
-		}); err != nil {
+		}
+		err = checkMessageSize(msg, resource)
+		if err != nil {
+			s.Logger.Warn().Str("table", resource.Table.Name).
+				Int("bytes", len(msg.String())).
+				Msg("Row exceeding max bytes ignored")
+			continue
+		}
+		if err := stream.Send(msg); err != nil {
 			return status.Errorf(codes.Internal, "failed to send resource: %v", err)
 		}
 	}
@@ -116,3 +127,19 @@ func (s *SourceServer) GetMetrics(context.Context, *pb.GetSourceMetrics_Request)
 		Metrics: b,
 	}, nil
 }
+
+func checkMessageSize(msg proto.Message, resource *schema.Resource) error {
+	size := proto.Size(msg)
+	// log error to Sentry if row exceeds half of the max size
+	if size > MaxMsgSize/2 {
+		sentry.WithScope(func(scope *sentry.Scope) {
+			scope.SetTag("table", resource.Table.Name)
+			scope.SetExtra("bytes", size)
+			sentry.CurrentHub().CaptureMessage("Large message detected")
+		})
+	}
+	if size > MaxMsgSize {
+		return errors.New("message exceeds max size")
+	}
+	return nil
+}
diff --git a/serve/constants.go b/serve/constants.go
@@ -8,5 +8,4 @@ const (
 	// bufSize used for unit testing grpc server and client
 	testBufSize  = 1024 * 1024
 	flushTimeout = 5 * time.Second
-	maxMsgSize   = 50 * 1024 * 1024 // 50 MiB
 )
diff --git a/serve/destination.go b/serve/destination.go
@@ -107,8 +107,8 @@ func newCmdDestinationServe(destination *destinationServe) *cobra.Command {
 				grpc.ChainStreamInterceptor(
 					logging.StreamServerInterceptor(grpczerolog.InterceptorLogger(logger)),
 				),
-				grpc.MaxRecvMsgSize(maxMsgSize),
-				grpc.MaxSendMsgSize(maxMsgSize),
+				grpc.MaxRecvMsgSize(servers.MaxMsgSize),
+				grpc.MaxSendMsgSize(servers.MaxMsgSize),
 			)
 			pb.RegisterDestinationServer(s, &servers.DestinationServer{
 				Plugin: destination.plugin,
diff --git a/serve/source.go b/serve/source.go
@@ -113,8 +113,8 @@ func newCmdSourceServe(source *sourceServe) *cobra.Command {
 				grpc.ChainStreamInterceptor(
 					logging.StreamServerInterceptor(grpczerolog.InterceptorLogger(logger)),
 				),
-				grpc.MaxRecvMsgSize(maxMsgSize),
-				grpc.MaxSendMsgSize(maxMsgSize),
+				grpc.MaxRecvMsgSize(servers.MaxMsgSize),
+				grpc.MaxSendMsgSize(servers.MaxMsgSize),
 			)
 			source.plugin.SetLogger(logger)
 			pb.RegisterSourceServer(s, &servers.SourceServer{

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`	`1`	`package clients`
`2`	`2`
`3`	`3`	`const (`
`4`		`- maxMsgSize = 50 * 1024 * 1024 // 50 MiB`
	`4`	`+ maxMsgSize = 100 * 1024 * 1024 // 100 MiB`
`5`	`5`	`)`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+package servers`
	`2`	`+`
	`3`	`+const MaxMsgSize = 100 * 1024 * 1024 // 100 MiB`
Original file line number	Diff line number	Diff line change
`@@ -8,5 +8,4 @@ const (`
`8`	`8`	`// bufSize used for unit testing grpc server and client`
`9`	`9`	`testBufSize = 1024 * 1024`
`10`	`10`	`flushTimeout = 5 * time.Second`
`11`		`- maxMsgSize = 50 * 1024 * 1024 // 50 MiB`
`12`	`11`	`)`