Add option to expose workload token to task (#15755)

Add `identity` jobspec block to expose workload identity tokens to tasks.

---------

Co-authored-by: Anders <mail@anars.dk>
Co-authored-by: Tim Gross <tgross@hashicorp.com>
Co-authored-by: Michael Schurter <mschurter@hashicorp.com>

client/allocrunner/taskrunner/identity_hook.go

@@ -2,20 +2,33 @@ package taskrunner
 
 import (
 	"context"
+	"fmt"
+	"path/filepath"
 	"sync"
 
 	log "github.com/hashicorp/go-hclog"
 
 	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
+	"github.com/hashicorp/nomad/helper/users"
 )
 
 // identityHook sets the task runner's Nomad workload identity token
 // based on the signed identity stored on the Allocation
+
+const (
+	// wiTokenFile is the name of the file holding the Nomad token inside the
+	// task's secret directory
+	wiTokenFile = "nomad_token"
+)
+
 type identityHook struct {
 	tr       *TaskRunner
 	logger   log.Logger
 	taskName string
 	lock     sync.Mutex
+
+	// tokenPath is the path in which to read and write the token
+	tokenPath string
 }
 
 func newIdentityHook(tr *TaskRunner, logger log.Logger) *identityHook {
@@ -34,21 +47,43 @@ func (*identityHook) Name() string {
 func (h *identityHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
 	h.lock.Lock()
 	defer h.lock.Unlock()
+	h.tokenPath = filepath.Join(req.TaskDir.SecretsDir, wiTokenFile)
 
-	token := h.tr.alloc.SignedIdentities[h.taskName]
-	if token != "" {
-		h.tr.setNomadToken(token)
-	}
-	return nil
+	return h.setToken()
 }
 
 func (h *identityHook) Update(_ context.Context, req *interfaces.TaskUpdateRequest, _ *interfaces.TaskUpdateResponse) error {
 	h.lock.Lock()
 	defer h.lock.Unlock()
 
+	return h.setToken()
+}
+
+// setToken adds the Nomad token to the task's environment and writes it to a
+// file if requested by the jobsepc.
+func (h *identityHook) setToken() error {
 	token := h.tr.alloc.SignedIdentities[h.taskName]
-	if token != "" {
-		h.tr.setNomadToken(token)
+	if token == "" {
+		return nil
 	}
+
+	h.tr.setNomadToken(token)
+
+	if id := h.tr.task.Identity; id != nil && id.File {
+		if err := h.writeToken(token); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// writeToken writes the given token to disk
+func (h *identityHook) writeToken(token string) error {
+	// Write token as owner readable only
+	if err := users.WriteFileFor(h.tokenPath, []byte(token), h.tr.task.User); err != nil {
+		return fmt.Errorf("failed to write nomad token: %w", err)
+	}
+
 	return nil
 }

client/allocrunner/taskrunner/identity_hook_test.go

@@ -0,0 +1,8 @@
+package taskrunner
+
+import "github.com/hashicorp/nomad/client/allocrunner/interfaces"
+
+var _ interfaces.TaskPrestartHook = (*identityHook)(nil)
+var _ interfaces.TaskUpdateHook = (*identityHook)(nil)
+
+// See task_runner_test.go:TestTaskRunner_IdentityHook

client/allocrunner/taskrunner/task_runner_getters.go

@@ -86,6 +86,14 @@ func (tr *TaskRunner) setNomadToken(token string) {
 	tr.nomadTokenLock.Lock()
 	defer tr.nomadTokenLock.Unlock()
 	tr.nomadToken = token
+
+	if id := tr.task.Identity; id != nil {
+		tr.envBuilder.SetWorkloadToken(token, id.Env)
+	} else {
+		// Default to *not* injecting the workload token into the task's
+		// environment.
+		tr.envBuilder.SetWorkloadToken(token, false)
+	}
 }
 
 // getDriverHandle returns a driver handle.

client/allocrunner/taskrunner/task_runner_test.go

@@ -40,6 +40,7 @@ import (
 	"github.com/hashicorp/nomad/plugins/drivers"
 	"github.com/hashicorp/nomad/testutil"
 	"github.com/kr/pretty"
+	"github.com/shoenig/test/must"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -2522,3 +2523,63 @@ func TestTaskRunner_BaseLabels(t *testing.T) {
 	require.Equal(alloc.ID, labels["alloc_id"])
 	require.Equal(alloc.Namespace, labels["namespace"])
 }
+
+// TestTaskRunner_IdentityHook_Enabled asserts that the identity hook exposes a
+// workload identity to a task.
+func TestTaskRunner_IdentityHook_Enabled(t *testing.T) {
+	ci.Parallel(t)
+
+	alloc := mock.BatchAlloc()
+	task := alloc.Job.TaskGroups[0].Tasks[0]
+
+	// Fake an identity and expose it to the task
+	alloc.SignedIdentities = map[string]string{
+		task.Name: "foo",
+	}
+	task.Identity = &structs.WorkloadIdentity{
+		Env:  true,
+		File: true,
+	}
+
+	tr, _, cleanup := runTestTaskRunner(t, alloc, task.Name)
+	defer cleanup()
+
+	testWaitForTaskToDie(t, tr)
+
+	// Assert the token was written to the filesystem
+	tokenBytes, err := os.ReadFile(filepath.Join(tr.taskDir.SecretsDir, "nomad_token"))
+	must.NoError(t, err)
+	must.Eq(t, "foo", string(tokenBytes))
+
+	// Assert the token is built into the task env
+	taskEnv := tr.envBuilder.Build()
+	must.Eq(t, "foo", taskEnv.EnvMap["NOMAD_TOKEN"])
+}
+
+// TestTaskRunner_IdentityHook_Disabled asserts that the identity hook does not
+// expose a workload identity to a task by default.
+func TestTaskRunner_IdentityHook_Disabled(t *testing.T) {
+	ci.Parallel(t)
+
+	alloc := mock.BatchAlloc()
+	task := alloc.Job.TaskGroups[0].Tasks[0]
+
+	// Fake an identity but don't expose it to the task
+	alloc.SignedIdentities = map[string]string{
+		task.Name: "foo",
+	}
+	task.Identity = nil
+
+	tr, _, cleanup := runTestTaskRunner(t, alloc, task.Name)
+	defer cleanup()
+
+	testWaitForTaskToDie(t, tr)
+
+	// Assert the token was written to the filesystem
+	_, err := os.ReadFile(filepath.Join(tr.taskDir.SecretsDir, "nomad_token"))
+	must.Error(t, err)
+
+	// Assert the token is built into the task env
+	taskEnv := tr.envBuilder.Build()
+	must.MapNotContainsKey(t, taskEnv.EnvMap, "NOMAD_TOKEN")
+}