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/controller/handlers/targets/target_service.go

@@ -28,6 +28,7 @@ import (
 	"github.com/hashicorp/boundary/internal/types/resource"
 	"github.com/hashicorp/boundary/internal/types/scope"
 	"github.com/hashicorp/boundary/sdk/strutil"
+	"github.com/hashicorp/go-bexpr"
 	"github.com/mr-tron/base58"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/protobuf/proto"
@@ -418,6 +419,7 @@ HostSetIterationLoop:
 		Endpoint:        endpointUrl.String(),
 		ExpirationTime:  &timestamp.Timestamp{Timestamp: expTime},
 		ConnectionLimit: t.GetSessionConnectionLimit(),
+		WorkerFilter:    t.GetWorkerFilter(),
 	}
 
 	sess, err := session.New(sessionComposition)
@@ -433,15 +435,77 @@ HostSetIterationLoop:
 		return nil, err
 	}
 
+	// WorkerInfo only contains the address; worker IDs below is used to contain
+	// their IDs in the same order. This is used to fetch tags for filtering.
+	// But we avoid allocation unless we actually need it.
 	var workers []*pb.WorkerInfo
+	var workerIds []string
+	hasWorkerFilter := len(t.GetWorkerFilter()) > 0
 	servers, err := serversRepo.ListServers(ctx, servers.ServerTypeWorker)
 	if err != nil {
 		return nil, err
 	}
 	for _, v := range servers {
+		if hasWorkerFilter {
+			workerIds = append(workerIds, v.GetPrivateId())
+		}
 		workers = append(workers, &pb.WorkerInfo{Address: v.Address})
 	}
 
+	if hasWorkerFilter && len(workerIds) > 0 {
+		finalWorkers := make([]*pb.WorkerInfo, 0, len(workers))
+		// Fetch the tags for the given worker IDs
+		tags, err := serversRepo.ListTagsForServers(ctx, workerIds)
+		if err != nil {
+			return nil, err
+		}
+		// Build the map for filtering. This is similar to the filter map we
+		// built from the worker config, but with one extra level: a map of the
+		// worker's ID to its filter map.
+		tagMap := make(map[string]map[string][]string)
+		for _, tag := range tags {
+			currWorkerMap := tagMap[tag.ServerId]
+			if currWorkerMap == nil {
+				currWorkerMap = make(map[string][]string)
+				tagMap[tag.ServerId] = currWorkerMap
+			}
+			currWorkerMap[tag.Key] = append(currWorkerMap[tag.Key], tag.Value)
+			// We don't need to reinsert after the fact because maps are
+			// reference types, so we don't need to re-insert into tagMap
+		}
+
+		// Create the evaluator
+		eval, err := bexpr.CreateEvaluator(t.GetWorkerFilter())
+		if err != nil {
+			return nil, err
+		}
+
+		// Iterate through the known worker IDs, and evaluate. If evaluation
+		// returns true, add to the final worker slice, which is assigned back
+		// to workers after this.
+		for i, worker := range workerIds {
+			filterInput := map[string]interface{}{
+				"name": worker,
+				"tags": tagMap[worker],
+			}
+			ok, err := eval.Evaluate(filterInput)
+			if err != nil {
+				return nil, handlers.ApiErrorWithCodeAndMessage(
+					codes.FailedPrecondition,
+					fmt.Sprintf("Worker filter expression evaluation resulted in error: %s", err))
+			}
+			if ok {
+				finalWorkers = append(finalWorkers, workers[i])
+			}
+		}
+		workers = finalWorkers
+	}
+	if len(workers) == 0 {
+		return nil, handlers.ApiErrorWithCodeAndMessage(
+			codes.FailedPrecondition,
+			"No workers are available to handle this session, or all have been filtered.")
+	}
+
 	sad := &pb.SessionAuthorizationData{
 		SessionId:       sess.PublicId,
 		TargetId:        t.GetPublicId(),
@@ -505,6 +569,9 @@ func (s Service) createInRepo(ctx context.Context, item *pb.Target) (*pb.Target,
 	if item.GetSessionConnectionLimit() != nil {
 		opts = append(opts, target.WithSessionConnectionLimit(item.GetSessionConnectionLimit().GetValue()))
 	}
+	if item.GetWorkerFilter() != nil {
+		opts = append(opts, target.WithWorkerFilter(item.GetWorkerFilter().GetValue()))
+	}
 	tcpAttrs := &pb.TcpTargetAttributes{}
 	if err := handlers.StructToProto(item.GetAttributes(), tcpAttrs); err != nil {
 		return nil, handlers.ApiErrorWithCodeAndMessage(codes.InvalidArgument, "Provided attributes don't match expected format.")
@@ -544,6 +611,9 @@ func (s Service) updateInRepo(ctx context.Context, scopeId, id string, mask []st
 	if item.GetSessionConnectionLimit() != nil {
 		opts = append(opts, target.WithSessionConnectionLimit(item.GetSessionConnectionLimit().GetValue()))
 	}
+	if filter := item.GetWorkerFilter(); filter != nil {
+		opts = append(opts, target.WithWorkerFilter(item.GetWorkerFilter().GetValue()))
+	}
 	tcpAttrs := &pb.TcpTargetAttributes{}
 	if err := handlers.StructToProto(item.GetAttributes(), tcpAttrs); err != nil {
 		return nil, handlers.ApiErrorWithCodeAndMessage(codes.InvalidArgument, "Provided attributes don't match expected format.")
@@ -737,6 +807,9 @@ func toProto(in target.Target, m []*target.TargetSet) (*pb.Target, error) {
 	if in.GetName() != "" {
 		out.Name = wrapperspb.String(in.GetName())
 	}
+	if in.GetWorkerFilter() != "" {
+		out.WorkerFilter = wrapperspb.String(in.GetWorkerFilter())
+	}
 	attrs := &pb.TcpTargetAttributes{}
 	if in.GetDefaultPort() > 0 {
 		attrs.DefaultPort = &wrappers.UInt32Value{Value: in.GetDefaultPort()}
@@ -803,6 +876,11 @@ func validateCreateRequest(req *pbs.CreateTargetRequest) error {
 		default:
 			badFields["type"] = "Unknown type provided."
 		}
+		if filter := req.GetItem().GetWorkerFilter(); filter != nil {
+			if _, err := bexpr.CreateEvaluator(filter.GetValue()); err != nil {
+				badFields["worker_filter"] = "Unable to successfully parse filter expression."
+			}
+		}
 		return badFields
 	})
 }
@@ -838,6 +916,11 @@ func validateUpdateRequest(req *pbs.UpdateTargetRequest) error {
 				badFields["attributes.default_port"] = "This optional field cannot be set to 0."
 			}
 		}
+		if filter := req.GetItem().GetWorkerFilter(); filter != nil {
+			if _, err := bexpr.CreateEvaluator(filter.GetValue()); err != nil {
+				badFields["worker_filter"] = "Unable to successfully parse filter expression."
+			}
+		}
 		return badFields
 	})
 }

internal/servers/controller/handlers/workers/worker_service.go

@@ -2,15 +2,19 @@ package workers
 
 import (
 	"context"
+	"fmt"
 	"sync"
 	"time"
 
 	"github.com/hashicorp/boundary/internal/gen/controller/api/resources/targets"
 	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
 	"github.com/hashicorp/boundary/internal/kms"
+	"github.com/hashicorp/boundary/internal/servers"
 	"github.com/hashicorp/boundary/internal/servers/controller/common"
+	"github.com/hashicorp/boundary/internal/servers/controller/handlers"
 	"github.com/hashicorp/boundary/internal/session"
 	"github.com/hashicorp/boundary/internal/types/resource"
+	"github.com/hashicorp/go-bexpr"
 	"github.com/hashicorp/go-hclog"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -48,15 +52,15 @@ var (
 )
 
 func (ws *workerServiceServer) Status(ctx context.Context, req *pbs.StatusRequest) (*pbs.StatusResponse, error) {
-	ws.logger.Trace("got status request from worker", "name", req.Worker.Name, "address", req.Worker.Address, "jobs", req.GetJobs())
-	ws.updateTimes.Store(req.Worker.Name, time.Now())
+	ws.logger.Trace("got status request from worker", "name", req.Worker.PrivateId, "address", req.Worker.Address, "jobs", req.GetJobs())
+	ws.updateTimes.Store(req.Worker.PrivateId, time.Now())
 	repo, err := ws.serversRepoFn()
 	if err != nil {
 		ws.logger.Error("error getting servers repo", "error", err)
 		return &pbs.StatusResponse{}, status.Errorf(codes.Internal, "Error aqcuiring repo to store worker status: %v", err)
 	}
 	req.Worker.Type = resource.Worker.String()
-	controllers, _, err := repo.UpsertServer(ctx, req.Worker)
+	controllers, _, err := repo.UpsertServer(ctx, req.Worker, servers.WithUpdateTags(req.GetUpdateTags()))
 	if err != nil {
 		ws.logger.Error("error storing worker status", "error", err)
 		return &pbs.StatusResponse{}, status.Errorf(codes.Internal, "Error storing worker status: %v", err)
@@ -149,6 +153,51 @@ func (ws *workerServiceServer) LookupSession(ctx context.Context, req *pbs.Looku
 		return nil, status.Error(codes.Internal, "Empty session states during lookup.")
 	}
 
+	if sessionInfo.WorkerFilter != "" {
+		if req.ServerId == "" {
+			ws.logger.Error("worker filter enabled for session but got no server ID from worker")
+			return &pbs.LookupSessionResponse{}, status.Errorf(codes.Internal, "Did not receive server ID when looking up session but filtering is enabled: %v", err)
+		}
+		serversRepo, err := ws.serversRepoFn()
+		if err != nil {
+			ws.logger.Error("error getting servers repo", "error", err)
+			return &pbs.LookupSessionResponse{}, status.Errorf(codes.Internal, "Error acquiring server repo when looking up session: %v", err)
+		}
+		tags, err := serversRepo.ListTagsForServers(ctx, []string{req.ServerId})
+		if err != nil {
+			ws.logger.Error("error looking up tags for server", "error", err, "server_id", req.ServerId)
+			return &pbs.LookupSessionResponse{}, status.Errorf(codes.Internal, "Error looking up tags for server: %v", err)
+		}
+		// Build the map for filtering.
+		tagMap := make(map[string][]string)
+		for _, tag := range tags {
+			tagMap[tag.Key] = append(tagMap[tag.Key], tag.Value)
+			// We don't need to reinsert after the fact because maps are
+			// reference types, so we don't need to re-insert into tagMap
+		}
+
+		// Create the evaluator
+		eval, err := bexpr.CreateEvaluator(sessionInfo.WorkerFilter)
+		if err != nil {
+			ws.logger.Error("error creating worker filter evaluator", "error", err, "server_id", req.ServerId)
+			return &pbs.LookupSessionResponse{}, status.Errorf(codes.Internal, "Error creating worker filter evaluator: %v", err)
+		}
+		filterInput := map[string]interface{}{
+			"name": req.ServerId,
+			"tags": tagMap,
+		}
+		ok, err := eval.Evaluate(filterInput)
+		if err != nil {
+			return &pbs.LookupSessionResponse{}, status.Errorf(codes.Internal,
+				fmt.Sprintf("Worker filter expression evaluation resulted in error: %s", err))
+		}
+		if !ok {
+			return nil, handlers.ApiErrorWithCodeAndMessage(
+				codes.FailedPrecondition,
+				"Worker filter expression precludes this worker from serving this session")
+		}
+	}
+
 	resp := &pbs.LookupSessionResponse{
 		Authorization: &targets.SessionAuthorizationData{
 			SessionId:   sessionInfo.GetPublicId(),

internal/servers/controller/testing.go

@@ -305,6 +305,9 @@ type TestControllerOpts struct {
 	// generated, unless provided in a non-nil Config
 	Name string
 
+	// The suffix to use for initial resources
+	InitialResourcesSuffix string
+
 	// The logger to use, or one will be created
 	Logger hclog.Logger
 }
@@ -383,6 +386,18 @@ func NewTestController(t *testing.T, opts *TestControllerOpts) *TestController {
 	}
 	tc.name = opts.Config.Controller.Name
 
+	if opts.InitialResourcesSuffix != "" {
+		suffix := opts.InitialResourcesSuffix
+		tc.b.DevAuthMethodId = "ampw_" + suffix
+		tc.b.DevHostCatalogId = "hcst_" + suffix
+		tc.b.DevHostId = "hst_" + suffix
+		tc.b.DevHostSetId = "hsst_" + suffix
+		tc.b.DevOrgId = "o_" + suffix
+		tc.b.DevProjectId = "p_" + suffix
+		tc.b.DevTargetId = "ttcp_" + suffix
+		tc.b.DevUserId = "u_" + suffix
+	}
+
 	// Set up KMSes
 	switch {
 	case opts.RootKms != nil && opts.WorkerAuthKms != nil: