Reorganize logic to allow a poison pill (#2317)

go.mod

@@ -92,7 +92,7 @@ require github.com/hashicorp/go-dbw v0.0.0-20220412153211-c470aec9369f // this i
 require (
 	github.com/DATA-DOG/go-sqlmock v1.5.0
 	github.com/hashicorp/go-kms-wrapping/extras/kms/v2 v2.0.0-20220711120347-32232bae6803
-	github.com/hashicorp/nodeenrollment v0.1.11
+	github.com/hashicorp/nodeenrollment v0.1.13
 )
 
 require (

go.sum

@@ -740,8 +740,8 @@ github.com/hashicorp/golang-lru v0.5.4 h1:YDjusn29QI/Das2iO9M0BHnIbxPeyuCHsjMW+l
 github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
 github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
-github.com/hashicorp/nodeenrollment v0.1.11 h1:8D161n9BqAJ1YVZSasA3oLqCJuh/R4TZNT1cYShobPA=
-github.com/hashicorp/nodeenrollment v0.1.11/go.mod h1:LIPKi+g0g/vl3xhpbzugCalHSxX1PMeqnatkAsxRgyM=
+github.com/hashicorp/nodeenrollment v0.1.13 h1:h830EbaWqt1wyRi3DERDtUHLW3tXl4cLfbVkPeD5Twc=
+github.com/hashicorp/nodeenrollment v0.1.13/go.mod h1:LIPKi+g0g/vl3xhpbzugCalHSxX1PMeqnatkAsxRgyM=
 github.com/hashicorp/vault/api v1.3.1 h1:pkDkcgTh47PRjY1NEFeofqR4W/HkNUi9qIakESO2aRM=
 github.com/hashicorp/vault/api v1.3.1/go.mod h1:QeJoWxMFt+MsuWcYhmwRLwKEXrjwAFFywzhptMsTIUw=
 github.com/hashicorp/vault/sdk v0.1.13/go.mod h1:B+hVj7TpuQY1Y/GPbCpffmgd+tSEwvhkWnjtSYCaS2M=

internal/cmd/base/servers.go

@@ -74,6 +74,8 @@ func init() {
 type Server struct {
 	*Command
 
+	ServerSideShutdownCh chan struct{}
+
 	InfoKeys []string
 	Info     map[string]string
 
@@ -153,6 +155,7 @@ type Server struct {
 func NewServer(cmd *Command) *Server {
 	return &Server{
 		Command:              cmd,
+		ServerSideShutdownCh: make(chan struct{}),
 		InfoKeys:             make([]string, 0, 20),
 		Info:                 make(map[string]string),
 		SecureRandomReader:   rand.Reader,

internal/cmd/commands/dev/dev.go

@@ -352,7 +352,7 @@ func (c *Command) Flags() *base.FlagSets {
 	f.BoolVar(&base.BoolVar{
 		Name:   "worker-auth-storage-skip-cleanup",
 		Target: &c.flagWorkerAuthStorageSkipCleanup,
-		Usage:  "Prevents deletion of temp worker credential storage directory if set",
+		Usage:  "Prevents deletion of temp worker credential storage directory if set.",
 	})
 
 	f.BoolVar(&base.BoolVar{
@@ -775,28 +775,42 @@ func (c *Command) Run(args []string) int {
 	// Wait for shutdown
 	shutdownTriggered := false
 
-	for !shutdownTriggered && !errorEncountered.Load() {
-		select {
-		case <-c.ShutdownCh:
-			c.UI.Output("==> Boundary dev environment shutdown triggered, interrupt again to force")
+	// Add a force-shutdown goroutine to consume more interrupts
+	abortForceShutdownCh := make(chan struct{})
+	defer close(abortForceShutdownCh)
 
-			// Add a force-shutdown goroutine to consume another interrupt
-			abortForceShutdownCh := make(chan struct{})
-			defer close(abortForceShutdownCh)
-			go func() {
-				shutdownCh := make(chan os.Signal, 4)
-				signal.Notify(shutdownCh, os.Interrupt, syscall.SIGTERM)
+	runShutdownLogic := func() {
+		go func() {
+			shutdownCh := make(chan os.Signal, 4)
+			signal.Notify(shutdownCh, os.Interrupt, syscall.SIGTERM)
+			for {
 				select {
 				case <-shutdownCh:
-					c.UI.Error("Second interrupt received, forcing shutdown")
+					c.UI.Error("Forcing shutdown")
 					os.Exit(base.CommandUserError)
 
+				case <-c.ServerSideShutdownCh:
+					// Drain connections in case this is hit more than once
+
 				case <-abortForceShutdownCh:
 					// No-op, we just use this to shut down the goroutine
+					return
 				}
-			}()
+			}
+		}()
 
-			shutdownTriggered = true
+		shutdownTriggered = true
+	}
+
+	for !shutdownTriggered && !errorEncountered.Load() {
+		select {
+		case <-c.ServerSideShutdownCh:
+			c.UI.Output("==> Boundary dev environment self-terminating")
+			runShutdownLogic()
+
+		case <-c.ShutdownCh:
+			c.UI.Output("==> Boundary dev environment shutdown triggered, interrupt again to force")
+			runShutdownLogic()
 
 		case <-c.SigUSR2Ch:
 			buf := make([]byte, 32*1024*1024)

internal/cmd/commands/server/server.go

@@ -705,52 +705,67 @@ func (c *Command) WaitForInterrupt() int {
 	// Wait for shutdown
 	shutdownTriggered := false
 
-	for !shutdownTriggered {
-		select {
-		case <-c.ShutdownCh:
-			c.UI.Output("==> Boundary server shutdown triggered, interrupt again to force")
+	// Add a force-shutdown goroutine to consume another interrupt
+	abortForceShutdownCh := make(chan struct{})
+	defer close(abortForceShutdownCh)
 
-			// Add a force-shutdown goroutine to consume another interrupt
-			abortForceShutdownCh := make(chan struct{})
-			defer close(abortForceShutdownCh)
-			go func() {
-				shutdownCh := make(chan os.Signal, 4)
-				signal.Notify(shutdownCh, os.Interrupt, syscall.SIGTERM)
+	runShutdownLogic := func() {
+		go func() {
+			shutdownCh := make(chan os.Signal, 4)
+			signal.Notify(shutdownCh, os.Interrupt, syscall.SIGTERM)
+			for {
 				select {
 				case <-shutdownCh:
-					c.UI.Error("Second interrupt received, forcing shutdown")
+					c.UI.Error("Forcing shutdown")
 					os.Exit(base.CommandUserError)
+
+				case <-c.ServerSideShutdownCh:
+					// Drain connections in case this is hit more than once
+
 				case <-abortForceShutdownCh:
 					// No-op, we just use this to shut down the goroutine
-				}
-			}()
-
-			if c.Config.Controller != nil && c.opsServer != nil {
-				c.opsServer.WaitIfHealthExists(c.Config.Controller.GracefulShutdownWaitDuration, c.UI)
-			}
-
-			// Do worker shutdown
-			if c.Config.Worker != nil {
-				if err := c.worker.Shutdown(); err != nil {
-					c.UI.Error(fmt.Errorf("Error shutting down worker: %w", err).Error())
+					return
 				}
 			}
+		}()
 
-			// Do controller shutdown
-			if c.Config.Controller != nil {
-				if err := c.controller.Shutdown(); err != nil {
-					c.UI.Error(fmt.Errorf("Error shutting down controller: %w", err).Error())
-				}
+		if c.Config.Controller != nil && c.opsServer != nil {
+			c.opsServer.WaitIfHealthExists(c.Config.Controller.GracefulShutdownWaitDuration, c.UI)
+		}
+
+		// Do worker shutdown
+		if c.Config.Worker != nil {
+			if err := c.worker.Shutdown(); err != nil {
+				c.UI.Error(fmt.Errorf("Error shutting down worker: %w", err).Error())
 			}
+		}
 
-			if c.opsServer != nil {
-				err := c.opsServer.Shutdown()
-				if err != nil {
-					c.UI.Error(fmt.Errorf("Error shutting down ops listeners: %w", err).Error())
-				}
+		// Do controller shutdown
+		if c.Config.Controller != nil {
+			if err := c.controller.Shutdown(); err != nil {
+				c.UI.Error(fmt.Errorf("Error shutting down controller: %w", err).Error())
 			}
+		}
 
-			shutdownTriggered = true
+		if c.opsServer != nil {
+			err := c.opsServer.Shutdown()
+			if err != nil {
+				c.UI.Error(fmt.Errorf("Error shutting down ops listeners: %w", err).Error())
+			}
+		}
+
+		shutdownTriggered = true
+	}
+
+	for !shutdownTriggered {
+		select {
+		case <-c.ServerSideShutdownCh:
+			c.UI.Output("==> Boundary server self-terminating")
+			runShutdownLogic()
+
+		case <-c.ShutdownCh:
+			c.UI.Output("==> Boundary server shutdown triggered, interrupt again to force")
+			runShutdownLogic()
 
 		case <-c.SighupCh:
 			c.UI.Output("==> Boundary server reload triggered")
@@ -847,7 +862,7 @@ func (c *Command) verifyKmsSetup() error {
 	ctx := context.Background()
 	kmsCache, err := kms.New(ctx, rw, rw)
 	if err != nil {
-		return fmt.Errorf("error creating kms: %w", err)
+		return fmt.Errorf("%s: error creating kms: %w", op, err)
 	}
 	if err := kmsCache.VerifyGlobalRoot(ctx); err != nil {
 		return err

internal/cmd/config/config.go

@@ -501,7 +501,7 @@ func Parse(d string) (*Config, error) {
 						var tags []string
 						err = json.Unmarshal([]byte(parsed), &tags)
 						if err != nil {
-							return nil, fmt.Errorf("Error unmarshalling env var/file contents: %w", err)
+							return nil, fmt.Errorf("Error unmarshaling env var/file contents: %w", err)
 						}
 						m[k] = tags
 					}

internal/cmd/config/config_test.go

@@ -753,7 +753,7 @@ func TestWorkerTags(t *testing.T) {
 			`,
 			expWorkerTags: nil,
 			expErr:        true,
-			expErrStr:     "Error unmarshalling env var/file contents: json: cannot unmarshal object into Go value of type []string",
+			expErrStr:     "Error unmarshaling env var/file contents: json: cannot unmarshal object into Go value of type []string",
 		},
 		{
 			name: "no clean mapping to internal structures",

internal/daemon/controller/controller_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"os"
 	"path/filepath"
+	"sync"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -24,10 +25,12 @@ func TestController_New(t *testing.T) {
 		testCtx := context.Background()
 		ctx, cancel := context.WithCancel(context.Background())
 		tc := &TestController{
-			t:      t,
-			ctx:    ctx,
-			cancel: cancel,
-			opts:   nil,
+			t:              t,
+			ctx:            ctx,
+			cancel:         cancel,
+			opts:           nil,
+			shutdownDoneCh: make(chan struct{}),
+			shutdownOnce:   new(sync.Once),
 		}
 
 		// TestControllerConfig(...) will create initial scopes
@@ -207,10 +210,12 @@ func TestControllerNewListenerConfig(t *testing.T) {
 			ctx, cancel := context.WithCancel(context.Background())
 
 			tc := &TestController{
-				t:      t,
-				ctx:    ctx,
-				cancel: cancel,
-				opts:   nil,
+				t:              t,
+				ctx:            ctx,
+				cancel:         cancel,
+				opts:           nil,
+				shutdownDoneCh: make(chan struct{}),
+				shutdownOnce:   new(sync.Once),
 			}
 			conf := TestControllerConfig(t, ctx, tc, nil)
 			conf.Listeners = tt.listeners
@@ -234,10 +239,12 @@ func TestController_NewPluginsConfig(t *testing.T) {
 	testCtx := context.Background()
 	ctx, cancel := context.WithCancel(context.Background())
 	tc := &TestController{
-		t:      t,
-		ctx:    ctx,
-		cancel: cancel,
-		opts:   nil,
+		t:              t,
+		ctx:            ctx,
+		cancel:         cancel,
+		opts:           nil,
+		shutdownDoneCh: make(chan struct{}),
+		shutdownOnce:   new(sync.Once),
 	}
 
 	initialConfig, err := config.DevController()

internal/daemon/controller/listeners.go

@@ -12,13 +12,10 @@ import (
 
 	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	"github.com/hashicorp/boundary/internal/cmd/base"
-	"github.com/hashicorp/boundary/internal/daemon/cluster/handlers"
 	"github.com/hashicorp/boundary/internal/daemon/common"
 	"github.com/hashicorp/boundary/internal/daemon/controller/internal/metric"
 	"github.com/hashicorp/boundary/internal/errors"
-	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
 	"github.com/hashicorp/go-multierror"
-	"github.com/hashicorp/nodeenrollment/multihop"
 	nodeenet "github.com/hashicorp/nodeenrollment/net"
 	"github.com/hashicorp/nodeenrollment/protocol"
 	"google.golang.org/grpc"
@@ -212,20 +209,11 @@ func (c *Controller) configureForCluster(ln *base.ServerListener) (func(), error
 		),
 	)
 
-	workerService := handlers.NewWorkerServiceServer(c.ServersRepoFn, c.SessionRepoFn, c.ConnectionRepoFn,
-		c.workerStatusUpdateTimes, c.kms)
-	pbs.RegisterServerCoordinationServiceServer(workerServer, workerService)
-	pbs.RegisterSessionServiceServer(workerServer, workerService)
-
-	multihopService, err := handlers.NewMultihopServiceServer(
-		workerAuthStorage,
-		true,
-		nil,
-	)
-	if err != nil {
-		return nil, fmt.Errorf("%s: error creating multihop service handler: %w", op, err)
+	for _, fn := range controllerGrpcServiceRegistrationFunctions {
+		if err := fn(c.baseContext, c, workerServer); err != nil {
+			return nil, err
+		}
 	}
-	multihop.RegisterMultihopServiceServer(workerServer, multihopService)
 
 	metric.InitializeClusterCollectors(c.conf.PrometheusRegisterer, workerServer)
 

internal/daemon/controller/listeners_test.go

@@ -16,6 +16,7 @@ import (
 	"net/http"
 	"os"
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 
@@ -235,7 +236,14 @@ func TestStartListeners(t *testing.T) {
 
 			ctx, cancel := context.WithCancel(context.Background())
 
-			tc := &TestController{t: t, ctx: ctx, cancel: cancel, opts: &TestControllerOpts{}}
+			tc := &TestController{
+				t:              t,
+				ctx:            ctx,
+				cancel:         cancel,
+				shutdownDoneCh: make(chan struct{}),
+				shutdownOnce:   new(sync.Once),
+				opts:           new(TestControllerOpts),
+			}
 			t.Cleanup(func() { tc.Shutdown() })
 
 			conf := TestControllerConfig(t, ctx, tc, nil)

internal/daemon/controller/rpc_registration.go

@@ -0,0 +1,87 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/hashicorp/boundary/internal/daemon/cluster/handlers"
+	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
+	"github.com/hashicorp/nodeenrollment"
+	"github.com/hashicorp/nodeenrollment/multihop"
+	"google.golang.org/grpc"
+)
+
+var controllerGrpcServiceRegistrationFunctions []func(context.Context, *Controller, *grpc.Server) error
+
+func init() {
+	controllerGrpcServiceRegistrationFunctions = append(controllerGrpcServiceRegistrationFunctions,
+		registerControllerServerCoordinationService,
+		registerControllerSessionService,
+		registerControllerMultihopService,
+	)
+}
+
+func registerControllerServerCoordinationService(ctx context.Context, c *Controller, server *grpc.Server) error {
+	const op = "controller.registerControllerServerCoordinationService"
+
+	switch {
+	case nodeenrollment.IsNil(ctx):
+		return fmt.Errorf("%s: context is nil", op)
+	case c == nil:
+		return fmt.Errorf("%s: controller is nil", op)
+	case server == nil:
+		return fmt.Errorf("%s: server is nil", op)
+	}
+
+	workerService := handlers.NewWorkerServiceServer(c.ServersRepoFn, c.SessionRepoFn, c.ConnectionRepoFn,
+		c.workerStatusUpdateTimes, c.kms)
+	pbs.RegisterServerCoordinationServiceServer(server, workerService)
+	return nil
+}
+
+func registerControllerSessionService(ctx context.Context, c *Controller, server *grpc.Server) error {
+	const op = "controller.registerControllerSessionService"
+
+	switch {
+	case nodeenrollment.IsNil(ctx):
+		return fmt.Errorf("%s: context is nil", op)
+	case c == nil:
+		return fmt.Errorf("%s: controller is nil", op)
+	case server == nil:
+		return fmt.Errorf("%s: server is nil", op)
+	}
+
+	workerService := handlers.NewWorkerServiceServer(c.ServersRepoFn, c.SessionRepoFn, c.ConnectionRepoFn,
+		c.workerStatusUpdateTimes, c.kms)
+	pbs.RegisterSessionServiceServer(server, workerService)
+	return nil
+}
+
+func registerControllerMultihopService(ctx context.Context, c *Controller, server *grpc.Server) error {
+	const op = "controller.registerControllerMultihopService"
+
+	switch {
+	case nodeenrollment.IsNil(ctx):
+		return fmt.Errorf("%s: context is nil", op)
+	case c == nil:
+		return fmt.Errorf("%s: controller is nil", op)
+	case server == nil:
+		return fmt.Errorf("%s: server is nil", op)
+	}
+
+	workerAuthStorage, err := c.WorkerAuthRepoStorageFn()
+	if err != nil {
+		return fmt.Errorf("%s: error fetching worker auth storage: %w", op, err)
+	}
+
+	multihopService, err := handlers.NewMultihopServiceServer(
+		workerAuthStorage,
+		true,
+		nil,
+	)
+	if err != nil {
+		return fmt.Errorf("%s: error creating multihop service handler: %w", op, err)
+	}
+	multihop.RegisterMultihopServiceServer(server, multihopService)
+	return nil
+}

internal/daemon/controller/testing.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -58,16 +59,18 @@ const (
 // fully-programmatic controller for tests. Error checking (for instance, for
 // valid config) is not stringent at the moment.
 type TestController struct {
-	b            *base.Server
-	c            *Controller
-	t            testing.TB
-	apiAddrs     []string // The address the Controller API is listening on
-	clusterAddrs []string
-	client       *api.Client
-	ctx          context.Context
-	cancel       context.CancelFunc
-	name         string
-	opts         *TestControllerOpts
+	b              *base.Server
+	c              *Controller
+	t              testing.TB
+	apiAddrs       []string // The address the Controller API is listening on
+	clusterAddrs   []string
+	client         *api.Client
+	ctx            context.Context
+	cancel         context.CancelFunc
+	name           string
+	opts           *TestControllerOpts
+	shutdownDoneCh chan struct{}
+	shutdownOnce   *sync.Once
 }
 
 // Server returns the underlying base server
@@ -282,29 +285,33 @@ func (tc *TestController) buildClient() {
 // Shutdown runs any cleanup functions; be sure to run this after your test is
 // done
 func (tc *TestController) Shutdown() {
-	if tc.b != nil {
-		close(tc.b.ShutdownCh)
-	}
-
-	tc.cancel()
-
-	if tc.c != nil {
-		if err := tc.c.Shutdown(); err != nil {
-			tc.t.Error(err)
+	tc.shutdownOnce.Do(func() {
+		if tc.b != nil {
+			close(tc.b.ShutdownCh)
 		}
-	}
-	if tc.b != nil {
-		if err := tc.b.RunShutdownFuncs(); err != nil {
-			tc.t.Error(err)
+
+		tc.cancel()
+
+		if tc.c != nil {
+			if err := tc.c.Shutdown(); err != nil {
+				tc.t.Error(err)
+			}
 		}
-		if !tc.opts.DisableDatabaseDestruction {
-			if tc.b.DestroyDevDatabase(tc.ctx) != nil {
-				if err := tc.b.DestroyDevDatabase(tc.ctx); err != nil {
-					tc.t.Error(err)
+		if tc.b != nil {
+			if err := tc.b.RunShutdownFuncs(); err != nil {
+				tc.t.Error(err)
+			}
+			if !tc.opts.DisableDatabaseDestruction {
+				if tc.b.DestroyDevDatabase(tc.ctx) != nil {
+					if err := tc.b.DestroyDevDatabase(tc.ctx); err != nil {
+						tc.t.Error(err)
+					}
 				}
 			}
 		}
-	}
+
+		close(tc.shutdownDoneCh)
+	})
 }
 
 type TestControllerOpts struct {
@@ -434,10 +441,12 @@ func NewTestController(t testing.TB, opts *TestControllerOpts) *TestController {
 	}
 
 	tc := &TestController{
-		t:      t,
-		ctx:    ctx,
-		cancel: cancel,
-		opts:   opts,
+		t:              t,
+		ctx:            ctx,
+		cancel:         cancel,
+		opts:           opts,
+		shutdownDoneCh: make(chan struct{}),
+		shutdownOnce:   new(sync.Once),
 	}
 
 	conf := TestControllerConfig(t, ctx, tc, opts)
@@ -450,6 +459,21 @@ func NewTestController(t testing.TB, opts *TestControllerOpts) *TestController {
 
 	tc.buildClient()
 
+	// The real server functions will listen for shutdown cues and act so mimic
+	// that here, and ensure that channels get drained
+	go func() {
+		for {
+			select {
+			case <-tc.b.ShutdownCh:
+				tc.Shutdown()
+			case <-tc.b.ServerSideShutdownCh:
+				tc.Shutdown()
+			case <-tc.shutdownDoneCh:
+				return
+			}
+		}
+	}()
+
 	if !opts.DisableAutoStart {
 		if err := tc.c.Start(); err != nil {
 			tc.Shutdown()

internal/daemon/worker/controller_connection.go

@@ -109,9 +109,10 @@ func (w *Worker) controllerDialerFunc(extraAlpnProtos ...string) func(context.Co
 		}
 
 		if err == nil && conn != nil {
-			if !w.everAuthenticated.Load() {
-				w.everAuthenticated.Store(true)
+			if w.everAuthenticated.Load() == authenticationStatusNeverAuthenticated {
+				w.everAuthenticated.Store(authenticationStatusFirstAuthentication)
 			}
+
 			event.WriteSysEvent(ctx, op, "worker has successfully authenticated")
 		}
 
@@ -204,10 +205,11 @@ func (w *Worker) createClientConn(addr string) error {
 	if err != nil {
 		return fmt.Errorf("error dialing controller for worker auth: %w", err)
 	}
-	w.GrpcClientConn = cc
 
+	w.GrpcClientConn = cc
 	w.controllerStatusConn.Store(pbs.NewServerCoordinationServiceClient(cc))
 	w.controllerMultihopConn.Store(multihop.NewMultihopServiceClient(cc))
+
 	return nil
 }
 

internal/daemon/worker/listeners.go

@@ -13,11 +13,9 @@ import (
 	"time"
 
 	"github.com/hashicorp/boundary/internal/cmd/base"
-	"github.com/hashicorp/boundary/internal/daemon/cluster/handlers"
 	"github.com/hashicorp/boundary/internal/daemon/common"
 	"github.com/hashicorp/boundary/internal/daemon/worker/session"
 	"github.com/hashicorp/boundary/internal/errors"
-	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
 	"github.com/hashicorp/boundary/internal/observability/event"
 	"github.com/hashicorp/go-multierror"
 	nodee "github.com/hashicorp/nodeenrollment"
@@ -110,7 +108,6 @@ func (w *Worker) configureForWorker(ln *base.ServerListener, logger *log.Logger,
 		if !ok {
 			return nil, temperror.New(stderrers.New("client could not be understood as a multihop service client"))
 		}
-		// log.Println("proxying fetch node credentials request")
 		return multihopClient.FetchNodeCredentials(ctx, req)
 	}
 
@@ -128,7 +125,6 @@ func (w *Worker) configureForWorker(ln *base.ServerListener, logger *log.Logger,
 		if !ok {
 			return nil, temperror.New(stderrers.New("client could not be understood as a multihop service client"))
 		}
-		// log.Println("proxying generate server cert request")
 		return multihopClient.GenerateServerCertificates(ctx, req)
 	}
 
@@ -179,18 +175,12 @@ func (w *Worker) configureForWorker(ln *base.ServerListener, logger *log.Logger,
 		grpc.MaxRecvMsgSize(math.MaxInt32),
 		grpc.MaxSendMsgSize(math.MaxInt32),
 	)
-	multihopService, err := handlers.NewMultihopServiceServer(
-		w.WorkerAuthStorage,
-		false,
-		w.controllerMultihopConn,
-	)
-	if err != nil {
-		return nil, fmt.Errorf("%s: error creating multihop service handler: %w", op, err)
+
+	for _, fn := range workerGrpcServiceRegistrationFunctions {
+		if err := fn(cancelCtx, w, downstreamServer); err != nil {
+			return nil, err
+		}
 	}
-	multihop.RegisterMultihopServiceServer(downstreamServer, multihopService)
-	statusSessionService := NewWorkerProxyServiceServer(w.GrpcClientConn, w.controllerStatusConn)
-	pbs.RegisterServerCoordinationServiceServer(downstreamServer, statusSessionService)
-	pbs.RegisterSessionServiceServer(downstreamServer, statusSessionService)
 
 	ln.GrpcServer = downstreamServer
 

internal/daemon/worker/rpc_registration.go

@@ -0,0 +1,63 @@
+package worker
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/hashicorp/boundary/internal/daemon/cluster/handlers"
+	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
+	"github.com/hashicorp/nodeenrollment"
+	"github.com/hashicorp/nodeenrollment/multihop"
+	"google.golang.org/grpc"
+)
+
+var workerGrpcServiceRegistrationFunctions []func(context.Context, *Worker, *grpc.Server) error
+
+func init() {
+	workerGrpcServiceRegistrationFunctions = append(workerGrpcServiceRegistrationFunctions,
+		registerWorkerStatusSessionService,
+		registerWorkerMultihopService,
+	)
+}
+
+func registerWorkerStatusSessionService(ctx context.Context, w *Worker, server *grpc.Server) error {
+	const op = "worker.registerWorkerStatusSessionService"
+
+	switch {
+	case nodeenrollment.IsNil(ctx):
+		return fmt.Errorf("%s: context is nil", op)
+	case w == nil:
+		return fmt.Errorf("%s: worker is nil", op)
+	case server == nil:
+		return fmt.Errorf("%s: server is nil", op)
+	}
+
+	statusSessionService := NewWorkerProxyServiceServer(w.GrpcClientConn, w.controllerStatusConn)
+	pbs.RegisterServerCoordinationServiceServer(server, statusSessionService)
+	pbs.RegisterSessionServiceServer(server, statusSessionService)
+	return nil
+}
+
+func registerWorkerMultihopService(ctx context.Context, w *Worker, server *grpc.Server) error {
+	const op = "worker.registerWorkerMultihopService"
+
+	switch {
+	case nodeenrollment.IsNil(ctx):
+		return fmt.Errorf("%s: context is nil", op)
+	case w == nil:
+		return fmt.Errorf("%s: worker is nil", op)
+	case server == nil:
+		return fmt.Errorf("%s: server is nil", op)
+	}
+
+	multihopService, err := handlers.NewMultihopServiceServer(
+		w.WorkerAuthStorage,
+		false,
+		w.controllerMultihopConn,
+	)
+	if err != nil {
+		return fmt.Errorf("%s: error creating multihop service handler: %w", op, err)
+	}
+	multihop.RegisterMultihopServiceServer(server, multihopService)
+	return nil
+}

internal/daemon/worker/status.go

@@ -16,6 +16,8 @@ import (
 	"github.com/hashicorp/go-secure-stdlib/strutil"
 )
 
+var firstStatusCheckPostHooks []func(context.Context, *Worker) error
+
 type LastStatusInformation struct {
 	*pbs.StatusResponse
 	StatusTime              time.Time
@@ -49,7 +51,7 @@ func (w *Worker) startStatusTicking(cancelCtx context.Context, sessionManager se
 			// If we've never managed to successfully authenticate then we won't have
 			// any session information anyways and this will produce a ton of noise in
 			// observability, so skip calling the function and retry in a short duration
-			if !w.everAuthenticated.Load() {
+			if w.everAuthenticated.Load() == authenticationStatusNeverAuthenticated {
 				timer.Reset(10 * time.Millisecond)
 				continue
 			}
@@ -190,68 +192,74 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context, sessionManager sess
 			// Exit out of status function; our work here is done and we don't need to create closeConnection requests
 			return
 		}
-	} else {
-		w.updateTags.Store(false)
-		var addrs []string
-		// This may be empty if we are in a multiple hop scenario
-		if len(result.CalculatedUpstreams) > 0 {
-			addrs = make([]string, 0, len(result.CalculatedUpstreams))
-			for _, v := range result.CalculatedUpstreams {
+
+		// Standard cleanup: Run through current jobs. Cancel connections
+		// for any canceling session or any session that is expired.
+		w.cleanupConnections(cancelCtx, false, sessionManager)
+		return
+	}
+
+	w.updateTags.Store(false)
+	var addrs []string
+	// This may be empty if we are in a multiple hop scenario
+	if len(result.CalculatedUpstreams) > 0 {
+		addrs = make([]string, 0, len(result.CalculatedUpstreams))
+		for _, v := range result.CalculatedUpstreams {
+			addrs = append(addrs, v.Address)
+		}
+	} else if w.conf.RawConfig.HcpbClusterId != "" {
+		// This is a worker that is one hop away from managed workers, so attempt to get that list
+		hcpbWorkersCtx, hcpbWorkersCancel := context.WithTimeout(cancelCtx, common.StatusTimeout)
+		defer hcpbWorkersCancel()
+		workersResp, err := client.ListHcpbWorkers(hcpbWorkersCtx, &pbs.ListHcpbWorkersRequest{})
+		if err != nil {
+			event.WriteError(hcpbWorkersCtx, op, err, event.WithInfoMsg("error fetching managed worker information"))
+		} else {
+			addrs = make([]string, 0, len(workersResp.Workers))
+			for _, v := range workersResp.Workers {
 				addrs = append(addrs, v.Address)
 			}
-
-		} else if w.conf.RawConfig.HcpbClusterId != "" {
-			// This is a worker that is one hop away from managed workers, so attempt to get that list
-			hcpbWorkersCtx, hcpbWorkersCancel := context.WithTimeout(cancelCtx, common.StatusTimeout)
-			defer hcpbWorkersCancel()
-			workersResp, err := client.ListHcpbWorkers(hcpbWorkersCtx, &pbs.ListHcpbWorkersRequest{})
-			if err != nil {
-				event.WriteError(hcpbWorkersCtx, op, err, event.WithInfoMsg("error fetching managed worker information"))
-			} else {
-				addrs = make([]string, 0, len(workersResp.Workers))
-				for _, v := range workersResp.Workers {
-					addrs = append(addrs, v.Address)
-				}
-			}
 		}
-		if len(addrs) > 0 {
-			lastStatus := w.lastStatusSuccess.Load().(*LastStatusInformation)
-			// Compare upstreams; update resolver if there is a difference, and emit an event with old and new addresses
-			if lastStatus != nil && !strutil.EquivalentSlices(lastStatus.LastCalculatedUpstreams, addrs) {
-				upstreamsMessage := fmt.Sprintf("Upstreams has changed; old upstreams were: %s, new upstreams are: %s", lastStatus.LastCalculatedUpstreams, addrs)
-				event.WriteSysEvent(cancelCtx, op, upstreamsMessage)
-				for _, as := range addressReceivers {
-					as.SetAddresses(addrs)
-				}
-			} else if lastStatus == nil {
-				for _, as := range addressReceivers {
-					as.SetAddresses(addrs)
-				}
-				event.WriteSysEvent(cancelCtx, op, fmt.Sprintf("Upstreams after first status set to: %s", addrs))
+	}
+
+	if len(addrs) > 0 {
+		lastStatus := w.lastStatusSuccess.Load().(*LastStatusInformation)
+		// Compare upstreams; update resolver if there is a difference, and emit an event with old and new addresses
+		if lastStatus != nil && !strutil.EquivalentSlices(lastStatus.LastCalculatedUpstreams, addrs) {
+			upstreamsMessage := fmt.Sprintf("Upstreams has changed; old upstreams were: %s, new upstreams are: %s", lastStatus.LastCalculatedUpstreams, addrs)
+			event.WriteSysEvent(cancelCtx, op, upstreamsMessage)
+			for _, as := range addressReceivers {
+				as.SetAddresses(addrs)
 			}
+		} else if lastStatus == nil {
+			for _, as := range addressReceivers {
+				as.SetAddresses(addrs)
+			}
+			event.WriteSysEvent(cancelCtx, op, fmt.Sprintf("Upstreams after first status set to: %s", addrs))
 		}
-		w.lastStatusSuccess.Store(&LastStatusInformation{StatusResponse: result, StatusTime: time.Now(), LastCalculatedUpstreams: addrs})
+	}
 
-		for _, request := range result.GetJobsRequests() {
-			switch request.GetRequestType() {
-			case pbs.CHANGETYPE_CHANGETYPE_UPDATE_STATE:
-				switch request.GetJob().GetType() {
-				case pbs.JOBTYPE_JOBTYPE_SESSION:
-					sessInfo := request.GetJob().GetSessionInfo()
-					sessionId := sessInfo.GetSessionId()
-					si := sessionManager.Get(sessionId)
-					if si == nil {
-						event.WriteError(statusCtx, op, errors.New("session change requested but could not find local information for it"), event.WithInfo("session_id", sessionId))
-						continue
-					}
-					si.ApplyLocalStatus(sessInfo.GetStatus())
+	w.lastStatusSuccess.Store(&LastStatusInformation{StatusResponse: result, StatusTime: time.Now(), LastCalculatedUpstreams: addrs})
 
-					// Update connection state if there are any connections in
-					// the request.
-					for _, conn := range sessInfo.GetConnections() {
-						if err := si.ApplyLocalConnectionStatus(conn.GetConnectionId(), conn.GetStatus()); err != nil {
-							event.WriteError(statusCtx, op, err, event.WithInfo("connection_id", conn.GetConnectionId()))
-						}
+	for _, request := range result.GetJobsRequests() {
+		switch request.GetRequestType() {
+		case pbs.CHANGETYPE_CHANGETYPE_UPDATE_STATE:
+			switch request.GetJob().GetType() {
+			case pbs.JOBTYPE_JOBTYPE_SESSION:
+				sessInfo := request.GetJob().GetSessionInfo()
+				sessionId := sessInfo.GetSessionId()
+				si := sessionManager.Get(sessionId)
+				if si == nil {
+					event.WriteError(statusCtx, op, errors.New("session change requested but could not find local information for it"), event.WithInfo("session_id", sessionId))
+					continue
+				}
+				si.ApplyLocalStatus(sessInfo.GetStatus())
+
+				// Update connection state if there are any connections in
+				// the request.
+				for _, conn := range sessInfo.GetConnections() {
+					if err := si.ApplyLocalConnectionStatus(conn.GetConnectionId(), conn.GetStatus()); err != nil {
+						event.WriteError(statusCtx, op, err, event.WithInfo("connection_id", conn.GetConnectionId()))
 					}
 				}
 			}
@@ -261,6 +269,22 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context, sessionManager sess
 	// Standard cleanup: Run through current jobs. Cancel connections
 	// for any canceling session or any session that is expired.
 	w.cleanupConnections(cancelCtx, false, sessionManager)
+
+	// If we have post hooks for after the first status check, run them now
+	if w.everAuthenticated.CAS(authenticationStatusFirstAuthentication, authenticationStatusFirstStatusRpcSuccessful) {
+		for _, fn := range firstStatusCheckPostHooks {
+			if err := fn(cancelCtx, w); err != nil {
+				// If we can't verify status we can't be expected to behave
+				// properly so error and trigger shutdown
+				event.WriteError(cancelCtx, op, fmt.Errorf("error running first status check post hook: %w", err))
+				// We don't use a non-blocking select here to ensure that it
+				// happens; we should catch blocks in tests but we want to
+				// ensure the signal is being listened to
+				w.conf.ServerSideShutdownCh <- struct{}{}
+				return
+			}
+		}
+	}
 }
 
 // cleanupConnections walks all sessions and shuts down all proxy connections.

internal/daemon/worker/testing.go

@@ -6,6 +6,7 @@ import (
 	"net"
 	"os"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -23,13 +24,15 @@ import (
 // fully-programmatic worker for tests. Error checking (for instance, for
 // valid config) is not stringent at the moment.
 type TestWorker struct {
-	b      *base.Server
-	w      *Worker
-	t      testing.TB
-	addrs  []string // The address the worker proxies are listening on
-	ctx    context.Context
-	cancel context.CancelFunc
-	name   string
+	b              *base.Server
+	w              *Worker
+	t              testing.TB
+	addrs          []string // The address the worker proxies are listening on
+	ctx            context.Context
+	cancel         context.CancelFunc
+	name           string
+	shutdownDoneCh chan struct{}
+	shutdownOnce   *sync.Once
 }
 
 // Worker returns the underlying worker
@@ -139,22 +142,26 @@ func (tw *TestWorker) LookupSession(id string) (TestSessionInfo, bool) {
 // Shutdown runs any cleanup functions; be sure to run this after your test is
 // done
 func (tw *TestWorker) Shutdown() {
-	if tw.b != nil {
-		close(tw.b.ShutdownCh)
-	}
-
-	tw.cancel()
-
-	if tw.w != nil {
-		if err := tw.w.Shutdown(); err != nil {
-			tw.t.Error(err)
+	tw.shutdownOnce.Do(func() {
+		if tw.b != nil {
+			close(tw.b.ShutdownCh)
 		}
-	}
-	if tw.b != nil {
-		if err := tw.b.RunShutdownFuncs(); err != nil {
-			tw.t.Error(err)
+
+		tw.cancel()
+
+		if tw.w != nil {
+			if err := tw.w.Shutdown(); err != nil {
+				tw.t.Error(err)
+			}
 		}
-	}
+		if tw.b != nil {
+			if err := tw.b.RunShutdownFuncs(); err != nil {
+				tw.t.Error(err)
+			}
+		}
+
+		close(tw.shutdownDoneCh)
+	})
 }
 
 type TestWorkerOpts struct {
@@ -204,9 +211,11 @@ func NewTestWorker(t testing.TB, opts *TestWorkerOpts) *TestWorker {
 	ctx, cancel := context.WithCancel(context.Background())
 
 	tw := &TestWorker{
-		t:      t,
-		ctx:    ctx,
-		cancel: cancel,
+		t:              t,
+		ctx:            ctx,
+		cancel:         cancel,
+		shutdownDoneCh: make(chan struct{}),
+		shutdownOnce:   new(sync.Once),
 	}
 
 	if opts == nil {
@@ -308,6 +317,21 @@ func NewTestWorker(t testing.TB, opts *TestWorkerOpts) *TestWorker {
 		tw.w.nonceFn = opts.NonceFn
 	}
 
+	// The real server functions will listen for shutdown cues and act so mimic
+	// that here, and ensure that channels get drained
+	go func() {
+		for {
+			select {
+			case <-tw.b.ShutdownCh:
+				tw.Shutdown()
+			case <-tw.b.ServerSideShutdownCh:
+				tw.Shutdown()
+			case <-tw.shutdownDoneCh:
+				return
+			}
+		}
+	}()
+
 	if !opts.DisableAutoStart {
 		if err := tw.w.Start(); err != nil {
 			tw.Shutdown()

internal/daemon/worker/worker.go

@@ -39,6 +39,12 @@ import (
 
 type randFn func(length int) (string, error)
 
+const (
+	authenticationStatusNeverAuthenticated uint32 = iota
+	authenticationStatusFirstAuthentication
+	authenticationStatusFirstStatusRpcSuccessful
+)
+
 type Worker struct {
 	conf   *Config
 	logger hclog.Logger
@@ -60,7 +66,7 @@ type Worker struct {
 	sessionManager session.Manager
 
 	controllerStatusConn *atomic.Value
-	everAuthenticated    *ua.Bool
+	everAuthenticated    *ua.Uint32
 	lastStatusSuccess    *atomic.Value
 	workerStartTime      time.Time
 
@@ -102,7 +108,7 @@ func New(conf *Config) (*Worker, error) {
 		logger:                 conf.Logger.Named("worker"),
 		started:                ua.NewBool(false),
 		controllerStatusConn:   new(atomic.Value),
-		everAuthenticated:      ua.NewBool(false),
+		everAuthenticated:      ua.NewUint32(authenticationStatusNeverAuthenticated),
 		lastStatusSuccess:      new(atomic.Value),
 		controllerMultihopConn: new(atomic.Value),
 		tags:                   new(atomic.Value),