From 31cd119a7e91451f84d6b434b81e11cfa7a273cf Mon Sep 17 00:00:00 2001
From: Cody Soyland <cody@soyland.org>
Date: Thu, 8 Oct 2020 15:08:57 -0500
Subject: [PATCH] Move grpc.Server creation into initializer

This fixes an issue with the grpc-web middleware using a nil
grpc.Server instance.
---
 server/grpc.go   | 52 +++++++++++++++++++++++++-----------------------
 server/server.go |  5 +++--
 2 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/server/grpc.go b/server/grpc.go
index c13a91955..8a4e74c80 100644
--- a/server/grpc.go
+++ b/server/grpc.go
@@ -1178,9 +1178,9 @@ func fieldDataType(f *pilosa.Field) string {
 
 type grpcServer struct {
 	api        *pilosa.API
-	mu         sync.Mutex
 	grpcServer *grpc.Server
 	ln         net.Listener
+	tlsConfig  *tls.Config
 
 	logger logger.Logger
 	stats  stats.StatsClient
@@ -1202,6 +1202,13 @@ func OptGRPCServerListener(ln net.Listener) grpcServerOption {
 	}
 }
 
+func OptGRPCServerTLSConfig(tlsConfig *tls.Config) grpcServerOption {
+	return func(s *grpcServer) error {
+		s.tlsConfig = tlsConfig
+		return nil
+	}
+}
+
 func OptGRPCServerLogger(logger logger.Logger) grpcServerOption {
 	return func(s *grpcServer) error {
 		s.logger = logger
@@ -1216,26 +1223,9 @@ func OptGRPCServerStats(stats stats.StatsClient) grpcServerOption {
 	}
 }
 
-func (s *grpcServer) Serve(tlsConfig *tls.Config) error {
+func (s *grpcServer) Serve() error {
 	s.logger.Printf("enabled grpc listening on %s", s.ln.Addr())
 
-	opts := make([]grpc.ServerOption, 0)
-	if tlsConfig != nil {
-		creds := credentials.NewTLS(tlsConfig)
-		opts = append(opts, grpc.Creds(creds))
-	}
-
-	// create grpc server
-	s.mu.Lock()
-	s.grpcServer = grpc.NewServer(opts...)
-	grpcHandler := NewGRPCHandler(s.api).WithLogger(s.logger).WithStats(s.stats)
-	pb.RegisterPilosaServer(s.grpcServer, grpcHandler)
-	vdsm_pb.RegisterMoleculaServer(s.grpcServer, NewVDSMGRPCHandler(grpcHandler, s.api).WithLogger(s.logger).WithStats(s.stats))
-
-	// register the server so its services are available to grpc_cli and others
-	reflection.Register(s.grpcServer)
-	s.mu.Unlock()
-
 	// and start...
 	if err := s.grpcServer.Serve(s.ln); err != nil {
 		return errors.Wrap(err, "starting grpc server")
@@ -1243,7 +1233,7 @@ func (s *grpcServer) Serve(tlsConfig *tls.Config) error {
 	return nil
 }
 
-func (s *grpcServer) Middleware(origins []string) func(http.Handler) http.Handler {
+func (s *grpcServer) middleware(origins []string) func(http.Handler) http.Handler {
 	httpOriginFunc := grpcweb.WithOriginFunc(func(origin string) bool {
 		for _, x := range origins {
 			if origin == x {
@@ -1269,11 +1259,7 @@ func (s *grpcServer) Middleware(origins []string) func(http.Handler) http.Handle
 // Stop stops the GRPC server. There's no error because the underlying GRPC
 // stuff doesn't report an error.
 func (s *grpcServer) Stop() {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	if s.grpcServer != nil {
-		s.grpcServer.Stop()
-	}
+	s.grpcServer.Stop()
 }
 
 func NewGRPCServer(opts ...grpcServerOption) (*grpcServer, error) {
@@ -1286,5 +1272,21 @@ func NewGRPCServer(opts ...grpcServerOption) (*grpcServer, error) {
 			return nil, errors.Wrap(err, "applying option")
 		}
 	}
+
+	gopts := make([]grpc.ServerOption, 0)
+	if server.tlsConfig != nil {
+		creds := credentials.NewTLS(server.tlsConfig)
+		gopts = append(gopts, grpc.Creds(creds))
+	}
+
+	// create grpc server
+	server.grpcServer = grpc.NewServer(gopts...)
+	grpcHandler := NewGRPCHandler(server.api).WithLogger(server.logger).WithStats(server.stats)
+	pb.RegisterPilosaServer(server.grpcServer, grpcHandler)
+	vdsm_pb.RegisterMoleculaServer(server.grpcServer, NewVDSMGRPCHandler(grpcHandler, server.api).WithLogger(server.logger).WithStats(server.stats))
+
+	// register the server so its services are available to grpc_cli and others
+	reflection.Register(server.grpcServer)
+
 	return server, nil
 }
diff --git a/server/server.go b/server/server.go
index ff6dcce66..f172f5647 100644
--- a/server/server.go
+++ b/server/server.go
@@ -171,7 +171,7 @@ func (m *Command) Start() (err error) {
 
 	m.logger.Printf("listening as %s\n", m.listenURI)
 	go func() {
-		if err := m.grpcServer.Serve(m.tlsConfig); err != nil {
+		if err := m.grpcServer.Serve(); err != nil {
 			m.logger.Printf("grpc server error: %v", err)
 		}
 	}()
@@ -432,6 +432,7 @@ func (m *Command) SetupServer() error {
 	m.grpcServer, err = NewGRPCServer(
 		OptGRPCServerAPI(m.API),
 		OptGRPCServerListener(m.grpcLn),
+		OptGRPCServerTLSConfig(m.tlsConfig),
 		OptGRPCServerLogger(m.logger),
 		OptGRPCServerStats(statsClient),
 	)
@@ -446,7 +447,7 @@ func (m *Command) SetupServer() error {
 		http.OptHandlerFileSystem(&statik.FileSystem{}),
 		http.OptHandlerListener(m.ln),
 		http.OptHandlerCloseTimeout(m.closeTimeout),
-		http.OptHandlerMiddleware(m.grpcServer.Middleware(m.Config.Handler.AllowedOrigins)),
+		http.OptHandlerMiddleware(m.grpcServer.middleware(m.Config.Handler.AllowedOrigins)),
 	)
 	return errors.Wrap(err, "new handler")
 }