bootstrapper: make sure InitServer is only shut down after Init has returned (#1347)

Signed-off-by: Daniel Weiße <dw@edgeless.systems>

bootstrapper/internal/initserver/initserver.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"net"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/edgelesssys/constellation/v2/bootstrapper/initproto"
@@ -51,13 +52,14 @@ import (
 // The server handles initialization calls from the CLI and initializes the
 // Kubernetes cluster.
 type Server struct {
-	nodeLock    locker
-	initializer ClusterInitializer
-	disk        encryptedDisk
-	fileHandler file.Handler
-	grpcServer  serveStopper
-	cleaner     cleaner
-	issuer      atls.Issuer
+	nodeLock     locker
+	initializer  ClusterInitializer
+	disk         encryptedDisk
+	fileHandler  file.Handler
+	grpcServer   serveStopper
+	cleaner      cleaner
+	issuer       atls.Issuer
+	shutdownLock sync.RWMutex
 
 	initSecretHash []byte
 
@@ -113,6 +115,10 @@ func (s *Server) Serve(ip, port string, cleaner cleaner) error {
 
 // Init initializes the cluster.
 func (s *Server) Init(ctx context.Context, req *initproto.InitRequest) (*initproto.InitResponse, error) {
+	// Acquire lock to prevent shutdown while Init is still running
+	s.shutdownLock.RLock()
+	defer s.shutdownLock.RUnlock()
+
 	log := s.log.With(zap.String("peer", grpclog.PeerAddrFromContext(ctx)))
 	log.Infof("Init called")
 
@@ -198,6 +204,9 @@ func (s *Server) Init(ctx context.Context, req *initproto.InitRequest) (*initpro
 func (s *Server) Stop() {
 	s.log.Infof("Stopping")
 
+	// Make sure to only stop the server if no Init calls are running
+	s.shutdownLock.Lock()
+	defer s.shutdownLock.Unlock()
 	s.grpcServer.GracefulStop()
 
 	s.log.Infof("Stopped")