client: split identity_hook across allocrunner and taskrunner

client/allocrunner/alloc_runner.go

@@ -206,8 +206,11 @@ type allocRunner struct {
 	// partitions is an interface for managing cpuset partitions
 	partitions cinterfaces.CPUPartitions
 
-	// widmgr fetches workload identities
-	widmgr *widmgr.WIDMgr
+	// widsigner signs workload identities
+	widsigner widmgr.IdentitySigner
+
+	// widmgr manages workload identity signatures
+	widmgr widmgr.IdentityManager
 }
 
 // NewAllocRunner returns a new allocation runner.
@@ -251,7 +254,7 @@ func NewAllocRunner(config *config.AllocRunnerConfig) (interfaces.AllocRunner, e
 		wranglers:                config.Wranglers,
 		partitions:               config.Partitions,
 		hookResources:            cstructs.NewAllocHookResources(),
-		widmgr:                   config.WIDMgr,
+		widsigner:                config.WIDSigner,
 	}
 
 	// Create the logger based on the allocation ID
@@ -269,6 +272,10 @@ func NewAllocRunner(config *config.AllocRunnerConfig) (interfaces.AllocRunner, e
 	ar.shutdownDelayCtx = shutdownDelayCtx
 	ar.shutdownDelayCancelFn = shutdownDelayCancel
 
+	// initialize the workload identity manager
+	widmgr := widmgr.NewWIDMgr(ar.widsigner, alloc, ar.logger)
+	ar.widmgr = widmgr
+
 	// Initialize the runners hooks.
 	if err := ar.initRunnerHooks(config.ClientConfig); err != nil {
 		return nil, err

client/allocrunner/alloc_runner_hooks.go

@@ -118,6 +118,7 @@ func (ar *allocRunner) initRunnerHooks(config *clientconfig.Config) error {
 	// directory path exists for other hooks.
 	alloc := ar.Alloc()
 	ar.runnerHooks = []interfaces.RunnerHook{
+		newIdentityHook(hookLogger, ar.widmgr),
 		newAllocDirHook(hookLogger, ar.allocDir),
 		newUpstreamAllocsHook(hookLogger, ar.prevAllocWatcher),
 		newDiskMigrationHook(hookLogger, ar.prevAllocMigrator, ar.allocDir),

client/allocrunner/identity_hook.go

@@ -0,0 +1,49 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package allocrunner
+
+import (
+	"context"
+
+	log "github.com/hashicorp/go-hclog"
+	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
+	"github.com/hashicorp/nomad/client/widmgr"
+)
+
+type identityHook struct {
+	widmgr widmgr.IdentityManager
+	logger log.Logger
+}
+
+func newIdentityHook(logger log.Logger, widmgr widmgr.IdentityManager) *identityHook {
+	h := &identityHook{
+		widmgr: widmgr,
+	}
+	h.logger = logger.Named(h.Name())
+	return h
+}
+
+func (*identityHook) Name() string {
+	return "identity"
+}
+
+func (h *identityHook) Prerun() error {
+	// run the renewal
+	if err := h.widmgr.Run(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Stop implements interfaces.TaskStopHook
+func (h *identityHook) Stop(context.Context, *interfaces.TaskStopRequest, *interfaces.TaskStopResponse) error {
+	h.widmgr.Shutdown()
+	return nil
+}
+
+// Shutdown implements interfaces.ShutdownHook
+func (h *identityHook) Shutdown() {
+	h.widmgr.Shutdown()
+}

client/allocrunner/identity_hook_test.go

@@ -0,0 +1,88 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package allocrunner
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/shoenig/test/must"
+
+	"github.com/hashicorp/nomad/ci"
+	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
+	cstructs "github.com/hashicorp/nomad/client/structs"
+	"github.com/hashicorp/nomad/client/widmgr"
+	"github.com/hashicorp/nomad/helper/testlog"
+	"github.com/hashicorp/nomad/nomad/mock"
+	"github.com/hashicorp/nomad/nomad/structs"
+)
+
+// statically assert network hook implements the expected interfaces
+var _ interfaces.RunnerPrerunHook = (*identityHook)(nil)
+var _ interfaces.ShutdownHook = (*identityHook)(nil)
+var _ interfaces.TaskStopHook = (*identityHook)(nil)
+
+func TestIdentityHook_Prerun(t *testing.T) {
+	ci.Parallel(t)
+
+	ttl := 30 * time.Second
+
+	wid := &structs.WorkloadIdentity{
+		Name:     "testing",
+		Audience: []string{"consul.io"},
+		Env:      true,
+		File:     true,
+		TTL:      ttl,
+	}
+
+	alloc := mock.Alloc()
+	task := alloc.LookupTask("web")
+	task.Identity = wid
+	task.Identities = []*structs.WorkloadIdentity{wid}
+
+	allocrunner, stopAR := TestAllocRunnerFromAlloc(t, alloc)
+	defer stopAR()
+
+	logger := testlog.HCLogger(t)
+
+	// setup mock signer and WIDMgr
+	mockSigner := widmgr.NewMockWIDSigner(task.Identities)
+	mockWIDMgr := widmgr.NewWIDMgr(mockSigner, alloc, logger)
+	allocrunner.widmgr = mockWIDMgr
+	allocrunner.widsigner = mockSigner
+
+	// do the initial signing
+	_, err := mockSigner.SignIdentities(1, []*structs.WorkloadIdentityRequest{
+		{
+			AllocID:      alloc.ID,
+			TaskName:     task.Name,
+			IdentityName: task.Identities[0].Name,
+		},
+	})
+	must.NoError(t, err)
+
+	start := time.Now()
+	hook := newIdentityHook(logger, mockWIDMgr)
+	must.Eq(t, hook.Name(), "identity")
+	must.NoError(t, hook.Prerun())
+
+	time.Sleep(time.Second) // give goroutines a moment to run
+	sid, err := hook.widmgr.Get(cstructs.TaskIdentity{
+		TaskName:     task.Name,
+		IdentityName: task.Identities[0].Name},
+	)
+	must.Nil(t, err)
+	must.Eq(t, sid.IdentityName, task.Identity.Name)
+	must.NotEq(t, sid.JWT, "")
+
+	// pad expiry time with a second to be safe
+	must.Between(t,
+		start.Add(ttl).Add(-1*time.Second).Unix(),
+		sid.Expiration.Unix(),
+		start.Add(ttl).Add(1*time.Second).Unix(),
+	)
+
+	must.NoError(t, hook.Stop(context.Background(), nil, nil))
+}