Disconnect PKI workers from upstreams when they are no longer authorized (#2515)

internal/daemon/controller/controller.go

@@ -14,6 +14,7 @@ import (
 	"github.com/hashicorp/boundary/internal/cmd/config"
 	credstatic "github.com/hashicorp/boundary/internal/credential/static"
 	"github.com/hashicorp/boundary/internal/credential/vault"
+	"github.com/hashicorp/boundary/internal/daemon/cluster"
 	"github.com/hashicorp/boundary/internal/daemon/controller/common"
 	"github.com/hashicorp/boundary/internal/daemon/controller/handlers/health"
 	"github.com/hashicorp/boundary/internal/daemon/controller/internal/metric"
@@ -132,6 +133,8 @@ type Controller struct {
 	// Used to signal the Health Service to start
 	// replying to queries with "503 Service Unavailable".
 	HealthService *health.Service
+
+	pkiConnManager *cluster.DownstreamManager
 }
 
 func New(ctx context.Context, conf *Config) (*Controller, error) {
@@ -146,6 +149,7 @@ func New(ctx context.Context, conf *Config) (*Controller, error) {
 		workerStatusUpdateTimes: new(sync.Map),
 		enabledPlugins:          conf.Server.EnabledPlugins,
 		apiListeners:            make([]*base.ServerListener, 0),
+		pkiConnManager:          cluster.NewDownstreamManager(),
 	}
 
 	if downstreamRouterFactory != nil {
@@ -385,7 +389,7 @@ func New(ctx context.Context, conf *Config) (*Controller, error) {
 
 func (c *Controller) Start() error {
 	const op = "controller.(Controller).Start"
-	if c.started.Load() {
+	if c.started.Swap(true) {
 		event.WriteSysEvent(context.TODO(), op, "already started, skipping")
 		return nil
 	}
@@ -424,10 +428,9 @@ func (c *Controller) Start() error {
 		defer c.tickerWg.Done()
 		c.startCloseExpiredPendingTokens(c.baseContext)
 	}()
-	go func() {
-		defer c.tickerWg.Done()
-		c.started.Store(true)
-	}()
+	if err := c.startWorkerConnectionMaintenanceTicking(c.baseContext, c.tickerWg, c.pkiConnManager); err != nil {
+		return errors.Wrap(c.baseContext, err, op)
+	}
 
 	if c.downstreamRoutes != nil {
 		c.tickerWg.Add(2)

internal/daemon/controller/listeners.go

@@ -12,6 +12,7 @@ import (
 
 	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	"github.com/hashicorp/boundary/internal/cmd/base"
+	"github.com/hashicorp/boundary/internal/daemon/cluster"
 	"github.com/hashicorp/boundary/internal/daemon/common"
 	"github.com/hashicorp/boundary/internal/daemon/controller/internal/metric"
 	"github.com/hashicorp/boundary/internal/errors"
@@ -157,12 +158,20 @@ func (c *Controller) configureForCluster(ln *base.ServerListener) (func(), error
 	if err != nil {
 		return nil, fmt.Errorf("%s: error creating eventing listener: %w", op, err)
 	}
+
+	// This wraps the normal pki worker connections with a listener which adds
+	// the worker key id of the connections to the controller's pkiConnManager.
+	pkiWorkerTrackingListener, err := cluster.NewTrackingListener(c.baseContext, eventingAuthedListener, c.pkiConnManager)
+	if err != nil {
+		return nil, fmt.Errorf("%s: error creating pki worker tracking listener: %w", op, err)
+	}
+
 	// Create a multiplexer to unify connections between PKI and KMS
 	multiplexingAuthedListener, err := nodeenet.NewMultiplexingListener(c.baseContext, nodeeAuthedListener.Addr())
 	if err != nil {
 		return nil, fmt.Errorf("error instantiating authed multiplexed listener: %w", err)
 	}
-	if err := multiplexingAuthedListener.IngressListener(eventingAuthedListener); err != nil {
+	if err := multiplexingAuthedListener.IngressListener(pkiWorkerTrackingListener); err != nil {
 		return nil, fmt.Errorf("error adding authed evented listener to multiplexed listener: %w", err)
 	}
 	// Connections that came in on the cluster listener that are not authed via
@@ -178,12 +187,20 @@ func (c *Controller) configureForCluster(ln *base.ServerListener) (func(), error
 	if err != nil {
 		return nil, fmt.Errorf("error instantiating reverse gprc connection split listener: %w", err)
 	}
+
+	// This wraps the reverse grpc pki worker connections with a listener which
+	// adds the worker key id of the connections to the pkiConnManager.
+	revPkiWorkerTrackingListener, err := cluster.NewTrackingListener(c.baseContext, reverseGrpcListener, c.pkiConnManager)
+	if err != nil {
+		return nil, fmt.Errorf("%s: error creating reverse grpc pki worker tracking listener: %w", op, err)
+	}
+
 	// Create a multiplexer to unify reverse grpc connections between PKI and KMS
 	multiplexingReverseGrpcListener, err := nodeenet.NewMultiplexingListener(c.baseContext, reverseGrpcListener.Addr())
 	if err != nil {
 		return nil, fmt.Errorf("error instantiating reverse grpc multiplexed listener: %w", err)
 	}
-	if err := multiplexingReverseGrpcListener.IngressListener(reverseGrpcListener); err != nil {
+	if err := multiplexingReverseGrpcListener.IngressListener(revPkiWorkerTrackingListener); err != nil {
 		return nil, fmt.Errorf("error adding reverse grpc listener to multiplexed listener: %w", err)
 	}
 

internal/daemon/controller/rpc_registration.go

@@ -33,8 +33,8 @@ func registerControllerServerCoordinationService(ctx context.Context, c *Control
 		return fmt.Errorf("%s: server is nil", op)
 	}
 
-	workerService := handlers.NewWorkerServiceServer(c.ServersRepoFn, c.SessionRepoFn, c.ConnectionRepoFn,
-		c.workerStatusUpdateTimes, c.kms)
+	workerService := handlers.NewWorkerServiceServer(c.ServersRepoFn, c.WorkerAuthRepoStorageFn,
+		c.SessionRepoFn, c.ConnectionRepoFn, c.workerStatusUpdateTimes, c.kms)
 	pbs.RegisterServerCoordinationServiceServer(server, workerService)
 	return nil
 }
@@ -51,8 +51,8 @@ func registerControllerSessionService(ctx context.Context, c *Controller, server
 		return fmt.Errorf("%s: server is nil", op)
 	}
 
-	workerService := handlers.NewWorkerServiceServer(c.ServersRepoFn, c.SessionRepoFn, c.ConnectionRepoFn,
-		c.workerStatusUpdateTimes, c.kms)
+	workerService := handlers.NewWorkerServiceServer(c.ServersRepoFn, c.WorkerAuthRepoStorageFn,
+		c.SessionRepoFn, c.ConnectionRepoFn, c.workerStatusUpdateTimes, c.kms)
 	pbs.RegisterSessionServiceServer(server, workerService)
 	return nil
 }

internal/daemon/controller/tickers.go

@@ -3,8 +3,10 @@ package controller
 import (
 	"context"
 	"math/rand"
+	"sync"
 	"time"
 
+	"github.com/hashicorp/boundary/internal/daemon/cluster"
 	"github.com/hashicorp/boundary/internal/server/store"
 
 	"github.com/hashicorp/boundary/internal/errors"
@@ -13,8 +15,9 @@ import (
 
 // In the future we could make this configurable
 const (
-	statusInterval      = 10 * time.Second
-	terminationInterval = 1 * time.Minute
+	workerConnectionMaintenanceInterval = 3 * time.Second
+	statusInterval                      = 10 * time.Second
+	terminationInterval                 = 1 * time.Minute
 )
 
 // This is exported so it can be tweaked in tests
@@ -155,3 +158,51 @@ func (c *Controller) startCloseExpiredPendingTokens(cancelCtx context.Context) {
 		}
 	}
 }
+
+func (c *Controller) startWorkerConnectionMaintenanceTicking(cancelCtx context.Context, wg *sync.WaitGroup, m *cluster.DownstreamManager) error {
+	const op = "controller.(Controller).startWorkerConnectionMaintenanceTicking"
+	switch {
+	case m == nil:
+		return errors.New(cancelCtx, errors.InvalidParameter, op, "DownstreamManager is nil")
+	case wg == nil:
+		return errors.New(cancelCtx, errors.InvalidParameter, op, "wait group is nil")
+	}
+	go func() {
+		defer wg.Done()
+		r := rand.New(rand.NewSource(time.Now().UnixNano()))
+		getRandomInterval := func() time.Duration {
+			// 0 to 0.5 adjustment to the base
+			f := r.Float64() / 2
+			// Half a chance to be faster, not slower
+			if r.Float32() > 0.5 {
+				f = -1 * f
+			}
+			return workerConnectionMaintenanceInterval + time.Duration(f*float64(workerConnectionMaintenanceInterval))
+		}
+		timer := time.NewTimer(0)
+		for {
+			select {
+			case <-cancelCtx.Done():
+				event.WriteSysEvent(cancelCtx, op, "context done, shutting down")
+				return
+
+			case <-timer.C:
+				repo, err := c.WorkerAuthRepoStorageFn()
+				if err != nil {
+					event.WriteError(cancelCtx, op, err, event.WithInfoMsg("error fetching repository for cluster connection maintenance"))
+					break
+				}
+				wKeyIds := m.Connected()
+				authorized, err := repo.FilterToAuthorizedWorkerKeyIds(cancelCtx, wKeyIds)
+				if err != nil {
+					event.WriteError(cancelCtx, op, err, event.WithInfoMsg("couldn't get authorized workers from repo"))
+					break
+				}
+				cluster.DisconnectUnauthorized(m, wKeyIds, authorized)
+			}
+
+			timer.Reset(getRandomInterval())
+		}
+	}()
+	return nil
+}