diff --git a/server/server.go b/server/server.go index ab7a08bb..dce01f4b 100644 --- a/server/server.go +++ b/server/server.go @@ -43,6 +43,10 @@ type Config struct { GRPCStreamMiddleware []grpc.StreamServerInterceptor `yaml:"-"` HTTPMiddleware []middleware.Interface `yaml:"-"` + GPRCServerMaxRecvMsgSize int `yaml:"grpc_server_max_recv_msg_size"` + GRPCServerMaxSendMsgSize int `yaml:"grpc_server_max_send_msg_size"` + GPRCServerMaxConcurrentStreams uint `yaml:"grpc_server_max_concurrent_streams"` + LogLevel logging.Level `yaml:"log_level"` Log logging.Interface `yaml:"-"` @@ -58,6 +62,9 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) { f.DurationVar(&cfg.HTTPServerReadTimeout, "server.http-read-timeout", 30*time.Second, "Read timeout for HTTP server") f.DurationVar(&cfg.HTTPServerWriteTimeout, "server.http-write-timeout", 30*time.Second, "Write timeout for HTTP server") f.DurationVar(&cfg.HTTPServerIdleTimeout, "server.http-idle-timeout", 120*time.Second, "Idle timeout for HTTP server") + f.IntVar(&cfg.GPRCServerMaxRecvMsgSize, "server.grpc-max-recv-msg-size-bytes", 4*1024*1024, "Limit on the size of a gRPC message this server can receive (bytes).") + f.IntVar(&cfg.GRPCServerMaxSendMsgSize, "server.grpc-max-send-msg-size-bytes", 4*1024*1024, "Limit on the size of a gRPC message this server can send (bytes).") + f.UintVar(&cfg.GPRCServerMaxConcurrentStreams, "server.grpc-max-concurrent-streams", 100, "Limit on the number of concurrent streams for gRPC calls (0 = unlimited)") f.StringVar(&cfg.PathPrefix, "server.path-prefix", "", "Base path to serve all API routes from (e.g. /v1/)") cfg.LogLevel.RegisterFlags(f) } @@ -134,6 +141,9 @@ func New(cfg Config) (*Server, error) { grpc.StreamInterceptor(grpc_middleware.ChainStreamServer( grpcStreamMiddleware..., )), + grpc.MaxRecvMsgSize(cfg.GPRCServerMaxRecvMsgSize), + grpc.MaxSendMsgSize(cfg.GRPCServerMaxSendMsgSize), + grpc.MaxConcurrentStreams(uint32(cfg.GPRCServerMaxConcurrentStreams)), } grpcOptions = append(grpcOptions, cfg.GRPCOptions...) grpcServer := grpc.NewServer(grpcOptions...) diff --git a/server/server_test.go b/server/server_test.go index ffee7257..aedbeaa3 100644 --- a/server/server_test.go +++ b/server/server_test.go @@ -2,6 +2,7 @@ package server import ( "errors" + "flag" "net/http" "strconv" "testing" @@ -33,7 +34,9 @@ func (f FakeServer) Succeed(ctx context.Context, req *google_protobuf.Empty) (*g } func TestErrorInstrumentationMiddleware(t *testing.T) { - cfg := Config{GRPCListenPort: 1234} + var cfg Config + cfg.RegisterFlags(flag.NewFlagSet("", flag.ExitOnError)) + cfg.GRPCListenPort = 1234 server, err := New(cfg) require.NoError(t, err)