Add option to expose workload token to task

.changelog/15755.txt

@@ -0,0 +1,3 @@
+```release-note:improvement
+identity: Add identity jobspec block for exposing workload identity to tasks
+```

api/tasks.go

@@ -570,6 +570,7 @@ func (g *TaskGroup) Canonicalize(job *Job) {
 	for _, s := range g.Services {
 		s.Canonicalize(nil, g, job)
 	}
+
 }
 
 // These needs to be in sync with DefaultServiceJobRestartPolicy in
@@ -703,6 +704,7 @@ type Task struct {
 	KillSignal      string                 `mapstructure:"kill_signal" hcl:"kill_signal,optional"`
 	Kind            string                 `hcl:"kind,optional"`
 	ScalingPolicies []*ScalingPolicy       `hcl:"scaling,block"`
+	Identity        *WorkloadIdentity      `hcl:"identity,block"`
 }
 
 func (t *Task) Canonicalize(tg *TaskGroup, job *Job) {
@@ -1110,3 +1112,10 @@ func (t *TaskCSIPluginConfig) Canonicalize() {
 		t.HealthTimeout = 30 * time.Second
 	}
 }
+
+// WorkloadIdentity is the jobspec block which determines if and how a workload
+// identity is exposed to tasks.
+type WorkloadIdentity struct {
+	Env  bool `hcl:"env,optional"`
+	File bool `hcl:"file,optional"`
+}

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")
+}

client/taskenv/env.go

@@ -122,6 +122,9 @@ const (
 
 	// VaultNamespace is the environment variable for passing the Vault namespace, if applicable
 	VaultNamespace = "VAULT_NAMESPACE"
+
+	// WorkloadToken is the environment variable for passing the Nomad Workload Identity token
+	WorkloadToken = "NOMAD_TOKEN"
 )
 
 // The node values that can be interpreted.
@@ -406,25 +409,27 @@ type Builder struct {
 	// clientTaskSecretsDir is the secrets dir from the client's perspective; eg <client_task_root>/secrets
 	clientTaskSecretsDir string
 
-	cpuCores         string
-	cpuLimit         int64
-	memLimit         int64
-	memMaxLimit      int64
-	taskName         string
-	allocIndex       int
-	datacenter       string
-	cgroupParent     string
-	namespace        string
-	region           string
-	allocId          string
-	allocName        string
-	groupName        string
-	vaultToken       string
-	vaultNamespace   string
-	injectVaultToken bool
-	jobID            string
-	jobName          string
-	jobParentID      string
+	cpuCores            string
+	cpuLimit            int64
+	memLimit            int64
+	memMaxLimit         int64
+	taskName            string
+	allocIndex          int
+	datacenter          string
+	cgroupParent        string
+	namespace           string
+	region              string
+	allocId             string
+	allocName           string
+	groupName           string
+	vaultToken          string
+	vaultNamespace      string
+	injectVaultToken    bool
+	workloadToken       string
+	injectWorkloadToken bool
+	jobID               string
+	jobName             string
+	jobParentID         string
 
 	// otherPorts for tasks in the same alloc
 	otherPorts map[string]string
@@ -567,6 +572,11 @@ func (b *Builder) buildEnv(allocDir, localDir, secretsDir string,
 		envMap[VaultNamespace] = b.vaultNamespace
 	}
 
+	// Build the Nomad Workload Token
+	if b.injectWorkloadToken && b.workloadToken != "" {
+		envMap[WorkloadToken] = b.workloadToken
+	}
+
 	// Copy and interpolate task meta
 	for k, v := range b.taskMeta {
 		envMap[hargs.ReplaceEnv(k, nodeAttrs, envMap)] = hargs.ReplaceEnv(v, nodeAttrs, envMap)
@@ -1018,6 +1028,14 @@ func (b *Builder) SetVaultToken(token, namespace string, inject bool) *Builder {
 	return b
 }
 
+func (b *Builder) SetWorkloadToken(token string, inject bool) *Builder {
+	b.mu.Lock()
+	b.workloadToken = token
+	b.injectWorkloadToken = inject
+	b.mu.Unlock()
+	return b
+}
+
 // addPort keys and values for other tasks to an env var map
 func addPort(m map[string]string, taskName, ip, portLabel string, port int) {
 	key := fmt.Sprintf("%s%s_%s", AddrPrefix, taskName, portLabel)

command/agent/job_endpoint.go

@@ -1151,6 +1151,13 @@ func ApiTaskToStructsTask(job *structs.Job, group *structs.TaskGroup,
 	structsTask.Affinities = ApiAffinitiesToStructs(apiTask.Affinities)
 	structsTask.CSIPluginConfig = ApiCSIPluginConfigToStructsCSIPluginConfig(apiTask.CSIPluginConfig)
 
+	if apiTask.Identity != nil {
+		structsTask.Identity = &structs.WorkloadIdentity{
+			Env:  apiTask.Identity.Env,
+			File: apiTask.Identity.File,
+		}
+	}
+
 	if apiTask.RestartPolicy != nil {
 		structsTask.RestartPolicy = &structs.RestartPolicy{
 			Attempts: *apiTask.RestartPolicy.Attempts,

command/assets/connect.nomad.hcl

@@ -356,6 +356,14 @@ job "countdash" {
       #   max_file_size = 15
       # }
 
+      # The "identity" block instructs Nomad to expose the task's workload
+      # identity token as an environment variable and in the file
+      # secrets/nomad_token.
+      # identity {
+      #   env  = true
+      #   file = true
+      # }
+
       # The "resources" block describes the requirements a task needs to
       # execute. Resource requirements include memory, network, cpu, and more.
       # This ensures the task will execute on a machine that contains enough

command/assets/example-short.nomad.hcl

@@ -16,6 +16,11 @@ job "example" {
         auth_soft_fail = true
       }
 
+      identity {
+        env  = true
+        file = true
+      }
+
       resources {
         cpu    = 500
         memory = 256

command/assets/example.nomad.hcl

@@ -331,7 +331,6 @@ job "example" {
       #   }
       # }
 
-
       # The "logs" block instructs the Nomad client on how many log files and
       # the maximum size of those logs files to retain. Logging is enabled by
       # default, but the "logs" block allows for finer-grained control over
@@ -347,6 +346,14 @@ job "example" {
       #   max_file_size = 15
       # }
 
+      # The "identity" block instructs Nomad to expose the task's workload
+      # identity token as an environment variable and in the file
+      # secrets/nomad_token.
+      identity {
+        env  = true
+        file = true
+      }
+
       # The "resources" block describes the requirements a task needs to
       # execute. Resource requirements include memory, cpu, and more.
       # This ensures the task will execute on a machine that contains enough

e2e/e2eutil/utils.go

@@ -222,10 +222,12 @@ func WaitForAllocsStopped(t *testing.T, nomadClient *api.Client, allocIDs []stri
 	}
 }
 
-func WaitForAllocStopped(t *testing.T, nomadClient *api.Client, allocID string) {
+func WaitForAllocStopped(t *testing.T, nomadClient *api.Client, allocID string) *api.Allocation {
+	var alloc *api.Allocation
+	var err error
 	testutil.WaitForResultRetries(retries, func() (bool, error) {
 		time.Sleep(time.Millisecond * 100)
-		alloc, _, err := nomadClient.Allocations().Info(allocID, nil)
+		alloc, _, err = nomadClient.Allocations().Info(allocID, nil)
 		if err != nil {
 			return false, err
 		}
@@ -243,6 +245,7 @@ func WaitForAllocStopped(t *testing.T, nomadClient *api.Client, allocID string)
 	}, func(err error) {
 		require.NoError(t, err, "failed to wait on alloc")
 	})
+	return alloc
 }
 
 func WaitForAllocStatus(t *testing.T, nomadClient *api.Client, allocID string, status string) {