Add test helper for creating a multihop worker setup (#2513)

* Add test helper for creating a multihop worker setup.

internal/daemon/worker/testing.go

@@ -12,12 +12,20 @@ import (
 
 	"github.com/hashicorp/boundary/internal/cmd/base"
 	"github.com/hashicorp/boundary/internal/cmd/config"
+	"github.com/hashicorp/boundary/internal/daemon/controller/common"
 	"github.com/hashicorp/boundary/internal/db"
 	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
 	"github.com/hashicorp/boundary/internal/observability/event"
+	"github.com/hashicorp/boundary/internal/server"
+	"github.com/hashicorp/boundary/internal/server/store"
+	"github.com/hashicorp/boundary/internal/types/scope"
 	"github.com/hashicorp/go-hclog"
 	wrapping "github.com/hashicorp/go-kms-wrapping/v2"
+	"github.com/hashicorp/nodeenrollment/types"
+	"github.com/mr-tron/base58"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
 )
 
 // TestWorker wraps a base.Server and Worker to provide a
@@ -367,3 +375,84 @@ func (tw *TestWorker) AddClusterWorkerMember(t testing.TB, opts *TestWorkerOpts)
 	}
 	return NewTestWorker(t, nextOpts)
 }
+
+// NewTestMultihopWorkers creates a KMS and PKI worker with the controller as an upstream, and a
+// child PKI worker as a downstream of the PKI worker connected to the controller.
+// Tags for the PKI and child PKI worker can be passed in, if desired
+func NewTestMultihopWorkers(t testing.TB, logger hclog.Logger, controllerContext context.Context, clusterAddrs []string,
+	workerAuthKms wrapping.Wrapper, serversRepoFn common.ServersRepoFactory, pkiTags,
+	childPkiTags map[string][]string,
+) (kmsWorker *TestWorker, pkiWorker *TestWorker, childPkiWorker *TestWorker) {
+	require := require.New(t)
+	kmsWorker = NewTestWorker(t, &TestWorkerOpts{
+		WorkerAuthKms:    workerAuthKms,
+		InitialUpstreams: clusterAddrs,
+		Logger:           logger.Named("kmsWorker"),
+	})
+
+	// names should not be set when using pki workers
+	pkiWorkerConf, err := config.DevWorker()
+	require.NoError(err)
+	pkiWorkerConf.Worker.Name = ""
+	if pkiTags != nil {
+		pkiWorkerConf.Worker.Tags = pkiTags
+	}
+	pkiWorkerConf.Worker.InitialUpstreams = clusterAddrs
+	pkiWorker = NewTestWorker(t, &TestWorkerOpts{
+		InitialUpstreams: clusterAddrs,
+		Logger:           logger.Named("pkiWorker"),
+		Config:           pkiWorkerConf,
+	})
+	t.Cleanup(pkiWorker.Shutdown)
+
+	// Get a server repo and worker auth repo
+	serversRepo, err := serversRepoFn()
+	require.NoError(err)
+
+	// Perform initial authentication of worker to controller
+	reqBytes, err := base58.FastBase58Decoding(pkiWorker.Worker().WorkerAuthRegistrationRequest)
+	require.NoError(err)
+
+	// Decode the proto into the request and create the worker
+	pkiWorkerReq := new(types.FetchNodeCredentialsRequest)
+	require.NoError(proto.Unmarshal(reqBytes, pkiWorkerReq))
+	_, err = serversRepo.CreateWorker(controllerContext, &server.Worker{
+		Worker: &store.Worker{
+			ScopeId: scope.Global.String(),
+		},
+	}, server.WithFetchNodeCredentialsRequest(pkiWorkerReq))
+	require.NoError(err)
+
+	childPkiWorkerConf, err := config.DevWorker()
+	require.NoError(err)
+	childPkiWorkerConf.Worker.Name = ""
+	if childPkiTags != nil {
+		childPkiWorkerConf.Worker.Tags = childPkiTags
+	}
+	childPkiWorkerConf.Worker.InitialUpstreams = kmsWorker.ProxyAddrs()
+
+	childPkiWorker = NewTestWorker(t, &TestWorkerOpts{
+		InitialUpstreams: kmsWorker.ProxyAddrs(),
+		Logger:           logger.Named("childPkiWorker"),
+		Config:           childPkiWorkerConf,
+	})
+
+	// Perform initial authentication of worker to controller
+	reqBytes, err = base58.FastBase58Decoding(childPkiWorker.Worker().WorkerAuthRegistrationRequest)
+	require.NoError(err)
+
+	// Decode the proto into the request and create the worker
+	childPkiWorkerReq := new(types.FetchNodeCredentialsRequest)
+	require.NoError(proto.Unmarshal(reqBytes, childPkiWorkerReq))
+	_, err = serversRepo.CreateWorker(controllerContext, &server.Worker{
+		Worker: &store.Worker{
+			ScopeId: scope.Global.String(),
+		},
+	}, server.WithFetchNodeCredentialsRequest(childPkiWorkerReq))
+	require.NoError(err)
+
+	// Sleep so that workers can startup and connect.
+	time.Sleep(10 * time.Second)
+
+	return kmsWorker, pkiWorker, childPkiWorker
+}

internal/daemon/worker/testing_test.go

@@ -10,10 +10,14 @@ import (
 	"testing"
 	"time"
 
+	"github.com/hashicorp/boundary/internal/cmd/config"
+	"github.com/hashicorp/boundary/internal/daemon/controller"
 	"github.com/hashicorp/boundary/internal/daemon/worker/session"
 	"github.com/hashicorp/boundary/internal/db"
 	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
+	"github.com/hashicorp/boundary/internal/server"
 	"github.com/hashicorp/boundary/sdk/pbs/controller/api/resources/targets"
+	"github.com/hashicorp/go-hclog"
 	wrapping "github.com/hashicorp/go-kms-wrapping/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -122,6 +126,51 @@ func TestTestWorker_WorkerAuthStorageKms(t *testing.T) {
 	}
 }
 
+func TestNewTestMultihopWorkers(t *testing.T) {
+	ctx := context.Background()
+	logger := hclog.New(&hclog.LoggerOptions{
+		Level: hclog.Trace,
+	})
+	conf, err := config.DevController()
+	require.NoError(t, err)
+	c := controller.NewTestController(t, &controller.TestControllerOpts{
+		Config: conf,
+		Logger: logger.Named("controller"),
+	})
+	t.Cleanup(c.Shutdown)
+	pkiTags := map[string][]string{"connected": {"directly"}}
+	childPkiTags := map[string][]string{"connected": {"multihop"}}
+
+	kmsWorker, pkiWorker, childPkiWorker := NewTestMultihopWorkers(t, logger, c.Context(), c.ClusterAddrs(),
+		c.Config().WorkerAuthKms, c.Controller().ServersRepoFn, pkiTags, childPkiTags)
+
+	srvRepo, err := c.Controller().ServersRepoFn()
+	workers, err := srvRepo.ListWorkers(ctx, []string{"global"})
+	assert.Len(t, workers, 3)
+	require.NoError(t, err)
+	var kmsW, pkiW, childPkiW *server.Worker
+	for _, w := range workers {
+		switch w.GetAddress() {
+		case kmsWorker.ProxyAddrs()[0]:
+			kmsW = w
+		case pkiWorker.ProxyAddrs()[0]:
+			pkiW = w
+		case childPkiWorker.ProxyAddrs()[0]:
+			childPkiW = w
+		}
+	}
+	assert.NotNil(t, kmsW)
+	assert.NotNil(t, pkiW)
+	assert.NotNil(t, childPkiW)
+
+	assert.NotZero(t, kmsW.GetLastStatusTime())
+	assert.NotZero(t, pkiW.GetLastStatusTime())
+	assert.NotZero(t, childPkiW.GetLastStatusTime())
+
+	assert.Equal(t, pkiTags, pkiW.GetConfigTags())
+	assert.Equal(t, childPkiTags, childPkiW.GetConfigTags())
+}
+
 func createTestCert(t *testing.T) ([]byte, ed25519.PublicKey, ed25519.PrivateKey) {
 	pub, priv, err := ed25519.GenerateKey(rand.Reader)
 	require.NoError(t, err)