From 0df838596d4ebf3843d08bede2f6a6e7fdf087da Mon Sep 17 00:00:00 2001
From: dobarx <111326505+dobarx@users.noreply.github.com>
Date: Wed, 28 Feb 2024 19:09:23 +0200
Subject: [PATCH] plugin/pluginapi/v1: increase gRPC call message size (#114)

---
 plugin/pluginapi/v1/client.go |  7 +++++++
 plugin/pluginapi/v1/plugin.go | 14 ++++++++++++--
 plugin/pluginapi/v1/server.go |  6 +++++-
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/plugin/pluginapi/v1/client.go b/plugin/pluginapi/v1/client.go
index 610844b9..2b4c19a5 100644
--- a/plugin/pluginapi/v1/client.go
+++ b/plugin/pluginapi/v1/client.go
@@ -8,6 +8,7 @@ import (
 	"strings"
 
 	goplugin "github.com/hashicorp/go-plugin"
+	"google.golang.org/grpc"
 
 	"github.com/blackstork-io/fabric/pkg/sloghclog"
 	"github.com/blackstork-io/fabric/plugin"
@@ -53,6 +54,12 @@ func NewClient(loc string) (p *plugin.Schema, closefn func() error, err error) {
 			sloghclog.AddSource(false),
 			sloghclog.Level(slog.LevelInfo), // debug is too noisy for plugins
 		),
+		GRPCDialOptions: []grpc.DialOption{
+			grpc.WithDefaultCallOptions(
+				grpc.MaxCallRecvMsgSize(defaultMsgSize),
+				grpc.MaxCallSendMsgSize(defaultMsgSize),
+			),
+		},
 	})
 	rpcClient, err := client.Client()
 	if err != nil {
diff --git a/plugin/pluginapi/v1/plugin.go b/plugin/pluginapi/v1/plugin.go
index de048b45..cfe49690 100644
--- a/plugin/pluginapi/v1/plugin.go
+++ b/plugin/pluginapi/v1/plugin.go
@@ -11,6 +11,8 @@ import (
 	"github.com/blackstork-io/fabric/plugin"
 )
 
+var defaultMsgSize = 1024 * 1024 * 20
+
 var handshake = goplugin.HandshakeConfig{
 	ProtocolVersion:  1,
 	MagicCookieKey:   "PLUGINS_FOR",
@@ -31,6 +33,7 @@ func (p *grpcPlugin) GRPCServer(broker *goplugin.GRPCBroker, s *grpc.Server) err
 
 func (p *grpcPlugin) GRPCClient(ctx context.Context, broker *goplugin.GRPCBroker, c *grpc.ClientConn) (interface{}, error) {
 	client := NewPluginServiceClient(c)
+
 	res, err := client.GetSchema(ctx, &GetSchemaRequest{})
 	if err != nil {
 		return nil, err
@@ -54,6 +57,13 @@ func (p *grpcPlugin) GRPCClient(ctx context.Context, broker *goplugin.GRPCBroker
 	return schema, nil
 }
 
+func (p *grpcPlugin) callOptions() []grpc.CallOption {
+	return []grpc.CallOption{
+		grpc.MaxCallRecvMsgSize(defaultMsgSize),
+		grpc.MaxCallSendMsgSize(defaultMsgSize),
+	}
+}
+
 func (p *grpcPlugin) clientGenerateFunc(name string, client PluginServiceClient) plugin.ProvideContentFunc {
 	return func(ctx context.Context, params *plugin.ProvideContentParams) (*plugin.Content, hcl.Diagnostics) {
 		if params == nil {
@@ -84,7 +94,7 @@ func (p *grpcPlugin) clientGenerateFunc(name string, client PluginServiceClient)
 			Config:      cfgEncoded,
 			Args:        argsEncoded,
 			DataContext: encodeMapData(params.DataContext),
-		})
+		}, p.callOptions()...)
 		if err != nil {
 			return nil, hcl.Diagnostics{{
 				Severity: hcl.DiagError,
@@ -128,7 +138,7 @@ func (p *grpcPlugin) clientDataFunc(name string, client PluginServiceClient) plu
 			Source: name,
 			Config: cfgEncoded,
 			Args:   argsEncoded,
-		})
+		}, p.callOptions()...)
 		if err != nil {
 			return nil, hcl.Diagnostics{{
 				Severity: hcl.DiagError,
diff --git a/plugin/pluginapi/v1/server.go b/plugin/pluginapi/v1/server.go
index 7828a67a..aeef2048 100644
--- a/plugin/pluginapi/v1/server.go
+++ b/plugin/pluginapi/v1/server.go
@@ -4,6 +4,7 @@ import (
 	context "context"
 
 	goplugin "github.com/hashicorp/go-plugin"
+	"google.golang.org/grpc"
 	codes "google.golang.org/grpc/codes"
 	status "google.golang.org/grpc/status"
 
@@ -16,7 +17,10 @@ func Serve(schema *plugin.Schema) {
 		Plugins: map[string]goplugin.Plugin{
 			schema.Name: &grpcPlugin{schema: schema},
 		},
-		GRPCServer: goplugin.DefaultGRPCServer,
+		GRPCServer: func(opts []grpc.ServerOption) *grpc.Server {
+			opts = append(opts, grpc.MaxRecvMsgSize(defaultMsgSize))
+			return grpc.NewServer(opts...)
+		},
 	})
 }