Add worker tagging (#862)

This allows workers to have tags assigned to them, which can then be filtered to control which targets they are allowed to serve.

internal/servers/worker/session.go

@@ -67,6 +67,7 @@ func (w *Worker) getSessionTls(hello *tls.ClientHelloInfo) (*tls.Config, error)
 
 	w.logger.Trace("looking up session", "session_id", sessionId)
 	resp, err := conn.LookupSession(timeoutContext, &pbs.LookupSessionRequest{
+		ServerId:  w.conf.RawConfig.Worker.Name,
 		SessionId: sessionId,
 	})
 	if err != nil {

internal/servers/worker/status.go

@@ -45,143 +45,160 @@ func (w *Worker) startStatusTicking(cancelCtx context.Context) {
 				return
 
 			case <-timer.C:
-				// First send info as-is. We'll perform cleanup duties after we
-				// get cancel/job change info back.
-				var activeJobs []*pbs.JobStatus
-				w.sessionInfoMap.Range(func(key, value interface{}) bool {
-					var jobInfo pbs.SessionJobInfo
-					sessionId := key.(string)
-					si := value.(*sessionInfo)
-					si.RLock()
-					status := si.status
-					connections := make([]*pbs.Connection, 0, len(si.connInfoMap))
-					for k, v := range si.connInfoMap {
-						connections = append(connections, &pbs.Connection{
-							ConnectionId: k,
-							Status:       v.status,
-						})
-					}
-					si.RUnlock()
-					jobInfo.SessionId = sessionId
-					activeJobs = append(activeJobs, &pbs.JobStatus{
-						Job: &pbs.Job{
-							Type: pbs.JOBTYPE_JOBTYPE_SESSION,
-							JobInfo: &pbs.Job_SessionInfo{
-								SessionInfo: &pbs.SessionJobInfo{
-									SessionId:   sessionId,
-									Status:      status,
-									Connections: connections,
-								},
-							},
-						},
-					})
-					return true
-				})
-				client := w.controllerStatusConn.Load().(pbs.ServerCoordinationServiceClient)
-				result, err := client.Status(cancelCtx, &pbs.StatusRequest{
-					Jobs: activeJobs,
-					Worker: &servers.Server{
-						PrivateId:   w.conf.RawConfig.Worker.Name,
-						Name:        w.conf.RawConfig.Worker.Name,
-						Type:        resource.Worker.String(),
-						Description: w.conf.RawConfig.Worker.Description,
-						Address:     w.conf.RawConfig.Worker.PublicAddr,
-					},
-				})
-				if err != nil {
-					w.logger.Error("error making status request to controller", "error", err)
-				} else {
-					w.logger.Trace("successfully sent status to controller")
-					addrs := make([]resolver.Address, 0, len(result.Controllers))
-					strAddrs := make([]string, 0, len(result.Controllers))
-					for _, v := range result.Controllers {
-						addrs = append(addrs, resolver.Address{Addr: v.Address})
-						strAddrs = append(strAddrs, v.Address)
-					}
-					w.logger.Trace("found controllers", "addresses", strAddrs)
-					switch len(strAddrs) {
-					case 0:
-						w.logger.Warn("got no controller addresses from controller; possibly prior to first status save, not persisting")
-					default:
-						w.Resolver().UpdateState(resolver.State{Addresses: addrs})
-					}
-					w.lastStatusSuccess.Store(&LastStatusInformation{StatusResponse: result, StatusTime: time.Now()})
-
-					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()
-								siRaw, ok := w.sessionInfoMap.Load(sessionId)
-								if !ok {
-									w.logger.Warn("asked to cancel session but could not find a local information for it", "session_id", sessionId)
-									continue
-								}
-								si := siRaw.(*sessionInfo)
-								si.Lock()
-								si.status = sessInfo.GetStatus()
-								si.Unlock()
-							}
-						}
-					}
-				}
-
-				// Cleanup: Run through current jobs. Cancel connections for any
-				// canceling session or any session that is expired. Clear out
-				// sessions that are canceled or expired with all connections
-				// marked as closed. Close any that aren't marked as such.
-				closeInfo := make(map[string]string)
-				cleanSessionIds := make([]string, 0)
-				w.sessionInfoMap.Range(func(key, value interface{}) bool {
-					si := value.(*sessionInfo)
-					si.Lock()
-					switch {
-					case si.status == pbs.SESSIONSTATUS_SESSIONSTATUS_CANCELING,
-						si.status == pbs.SESSIONSTATUS_SESSIONSTATUS_TERMINATED,
-						time.Until(si.lookupSessionResponse.Expiration.AsTime()) < 0:
-						var toClose int
-						for k, v := range si.connInfoMap {
-							if v.closeTime.IsZero() {
-								toClose++
-								v.connCancel()
-								w.logger.Info("terminated connection due to cancelation or expiration", "session_id", si.id, "connection_id", k)
-								closeInfo[k] = si.id
-							}
-						}
-						// closeTime is marked by closeConnections iff the
-						// status is returned for that connection as closed. If
-						// the session is no longer valid and all connections
-						// are marked closed, clean up the session.
-						if toClose == 0 {
-							cleanSessionIds = append(cleanSessionIds, si.id)
-						}
-					}
-					si.Unlock()
-					return true
-				})
-
-				// Note that we won't clean these from the info map until the
-				// next time we run this function
-				if len(closeInfo) > 0 {
-					if err := w.closeConnections(cancelCtx, closeInfo); err != nil {
-						w.logger.Error("error marking connections closed", "error", err)
-					}
-				}
-
-				// Forget sessions where the session is expired/canceled and all
-				// connections are canceled and marked closed
-				for _, v := range cleanSessionIds {
-					w.sessionInfoMap.Delete(v)
-				}
-
+				w.sendWorkerStatus(cancelCtx)
 				timer.Reset(getRandomInterval())
 			}
 		}
 	}()
 }
 
+// LastStatusSuccess is used in tests (it's exported for tests in other
+// packages) to verify the last time we sent status
 func (w *Worker) LastStatusSuccess() *LastStatusInformation {
 	return w.lastStatusSuccess.Load().(*LastStatusInformation)
 }
+
+func (w *Worker) sendWorkerStatus(cancelCtx context.Context) {
+	// First send info as-is. We'll perform cleanup duties after we
+	// get cancel/job change info back.
+	var activeJobs []*pbs.JobStatus
+
+	// Range over known sessions and collect info
+	w.sessionInfoMap.Range(func(key, value interface{}) bool {
+		var jobInfo pbs.SessionJobInfo
+		sessionId := key.(string)
+		si := value.(*sessionInfo)
+		si.RLock()
+		status := si.status
+		connections := make([]*pbs.Connection, 0, len(si.connInfoMap))
+		for k, v := range si.connInfoMap {
+			connections = append(connections, &pbs.Connection{
+				ConnectionId: k,
+				Status:       v.status,
+			})
+		}
+		si.RUnlock()
+		jobInfo.SessionId = sessionId
+		activeJobs = append(activeJobs, &pbs.JobStatus{
+			Job: &pbs.Job{
+				Type: pbs.JOBTYPE_JOBTYPE_SESSION,
+				JobInfo: &pbs.Job_SessionInfo{
+					SessionInfo: &pbs.SessionJobInfo{
+						SessionId:   sessionId,
+						Status:      status,
+						Connections: connections,
+					},
+				},
+			},
+		})
+		return true
+	})
+
+	// Send status information
+	client := w.controllerStatusConn.Load().(pbs.ServerCoordinationServiceClient)
+	var tags map[string]*servers.TagValues
+	// If we're not going to request a tag update, no reason to have these
+	// marshaled on every status call.
+	if w.updateTags.Load() {
+		tags = w.tags.Load().(map[string]*servers.TagValues)
+	}
+	result, err := client.Status(cancelCtx, &pbs.StatusRequest{
+		Jobs: activeJobs,
+		Worker: &servers.Server{
+			PrivateId:   w.conf.RawConfig.Worker.Name,
+			Type:        resource.Worker.String(),
+			Description: w.conf.RawConfig.Worker.Description,
+			Address:     w.conf.RawConfig.Worker.PublicAddr,
+			Tags:        tags,
+		},
+		UpdateTags: w.updateTags.Load(),
+	})
+	if err != nil {
+		w.logger.Error("error making status request to controller", "error", err)
+	} else {
+		w.logger.Trace("successfully sent status to controller")
+		w.updateTags.Store(false)
+		addrs := make([]resolver.Address, 0, len(result.Controllers))
+		strAddrs := make([]string, 0, len(result.Controllers))
+		for _, v := range result.Controllers {
+			addrs = append(addrs, resolver.Address{Addr: v.Address})
+			strAddrs = append(strAddrs, v.Address)
+		}
+		w.logger.Trace("found controllers", "addresses", strAddrs)
+		switch len(strAddrs) {
+		case 0:
+			w.logger.Warn("got no controller addresses from controller; possibly prior to first status save, not persisting")
+		default:
+			w.Resolver().UpdateState(resolver.State{Addresses: addrs})
+		}
+		w.lastStatusSuccess.Store(&LastStatusInformation{StatusResponse: result, StatusTime: time.Now()})
+
+		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()
+					siRaw, ok := w.sessionInfoMap.Load(sessionId)
+					if !ok {
+						w.logger.Warn("asked to cancel session but could not find a local information for it", "session_id", sessionId)
+						continue
+					}
+					si := siRaw.(*sessionInfo)
+					si.Lock()
+					si.status = sessInfo.GetStatus()
+					si.Unlock()
+				}
+			}
+		}
+	}
+
+	// Cleanup: Run through current jobs. Cancel connections for any
+	// canceling session or any session that is expired. Clear out
+	// sessions that are canceled or expired with all connections
+	// marked as closed. Close any that aren't marked as such.
+	closeInfo := make(map[string]string)
+	cleanSessionIds := make([]string, 0)
+	w.sessionInfoMap.Range(func(key, value interface{}) bool {
+		si := value.(*sessionInfo)
+		si.Lock()
+		switch {
+		case si.status == pbs.SESSIONSTATUS_SESSIONSTATUS_CANCELING,
+			si.status == pbs.SESSIONSTATUS_SESSIONSTATUS_TERMINATED,
+			time.Until(si.lookupSessionResponse.Expiration.AsTime()) < 0:
+			var toClose int
+			for k, v := range si.connInfoMap {
+				if v.closeTime.IsZero() {
+					toClose++
+					v.connCancel()
+					w.logger.Info("terminated connection due to cancelation or expiration", "session_id", si.id, "connection_id", k)
+					closeInfo[k] = si.id
+				}
+			}
+			// closeTime is marked by closeConnections iff the
+			// status is returned for that connection as closed. If
+			// the session is no longer valid and all connections
+			// are marked closed, clean up the session.
+			if toClose == 0 {
+				cleanSessionIds = append(cleanSessionIds, si.id)
+			}
+		}
+		si.Unlock()
+		return true
+	})
+
+	// Note that we won't clean these from the info map until the
+	// next time we run this function
+	if len(closeInfo) > 0 {
+		if err := w.closeConnections(cancelCtx, closeInfo); err != nil {
+			w.logger.Error("error marking connections closed", "error", err)
+		}
+	}
+
+	// Forget sessions where the session is expired/canceled and all
+	// connections are canceled and marked closed
+	for _, v := range cleanSessionIds {
+		w.sessionInfoMap.Delete(v)
+	}
+}

internal/servers/worker/worker.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/hashicorp/boundary/internal/cmd/config"
+	"github.com/hashicorp/boundary/internal/servers"
 	"github.com/hashicorp/go-hclog"
 	"github.com/hashicorp/vault/sdk/helper/base62"
 	"github.com/hashicorp/vault/sdk/helper/mlock"
@@ -33,6 +34,14 @@ type Worker struct {
 
 	controllerSessionConn *atomic.Value
 	sessionInfoMap        *sync.Map
+
+	// We store the current set in an atomic value so that we can add
+	// reload-on-sighup behavior later
+	tags *atomic.Value
+	// This stores whether or not to send updated tags on the next status
+	// request. It can be set via startup in New below, or (eventually) via
+	// SIGHUP.
+	updateTags ua.Bool
 }
 
 func New(conf *Config) (*Worker, error) {
@@ -44,12 +53,15 @@ func New(conf *Config) (*Worker, error) {
 		controllerResolver:    new(atomic.Value),
 		controllerSessionConn: new(atomic.Value),
 		sessionInfoMap:        new(sync.Map),
+		tags:                  new(atomic.Value),
 	}
 
 	w.lastStatusSuccess.Store((*LastStatusInformation)(nil))
-	w.started.Store(false)
 	w.controllerResolver.Store((*manual.Resolver)(nil))
 
+	w.parseAndStoreTags(conf.RawConfig.Worker.Tags)
+	w.updateTags.Store(true)
+
 	if conf.SecureRandomReader == nil {
 		conf.SecureRandomReader = rand.Reader
 	}
@@ -135,3 +147,17 @@ func (w *Worker) Resolver() *manual.Resolver {
 	}
 	return raw.(*manual.Resolver)
 }
+
+func (w *Worker) parseAndStoreTags(incoming map[string][]string) {
+	if len(incoming) == 0 {
+		w.tags.Store(map[string]*servers.TagValues{})
+		return
+	}
+	tags := make(map[string]*servers.TagValues, len(incoming))
+	for k, v := range incoming {
+		tags[k] = &servers.TagValues{
+			Values: append(make([]string, 0, len(v)), v...),
+		}
+	}
+	w.tags.Store(tags)
+}