client: pass alloc hook resources to template hook (#19040)

The task template hook uses the alloc resource to retrieve Consul
tokens, so it must be passed from the allocation.

client/allocrunner/taskrunner/task_runner_hooks.go

@@ -121,6 +121,7 @@ func (tr *TaskRunner) initHooks() {
 			templates:           task.Templates,
 			clientConfig:        tr.clientConfig,
 			envBuilder:          tr.envBuilder,
+			hookResources:       tr.allocHookResources,
 			consulNamespace:     consulNamespace,
 			nomadNamespace:      tr.alloc.Job.Namespace,
 			renderOnTaskRestart: task.RestartPolicy.RenderTemplates,

client/allocrunner/taskrunner/task_runner_test.go

@@ -5,13 +5,16 @@ package taskrunner
 
 import (
 	"context"
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"path"
 	"path/filepath"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -29,6 +32,9 @@ import (
 	"github.com/hashicorp/nomad/client/serviceregistration/wrapper"
 	cstate "github.com/hashicorp/nomad/client/state"
 	cstructs "github.com/hashicorp/nomad/client/structs"
+	"github.com/hashicorp/nomad/helper"
+	structsc "github.com/hashicorp/nomad/nomad/structs/config"
+
 	ctestutil "github.com/hashicorp/nomad/client/testutil"
 	"github.com/hashicorp/nomad/client/vaultclient"
 	"github.com/hashicorp/nomad/client/widmgr"
@@ -154,6 +160,13 @@ func testTaskRunnerConfig(t *testing.T, alloc *structs.Allocation, taskName stri
 		AllocHookResources:    cstructs.NewAllocHookResources(),
 	}
 
+	// The alloc runner identity_hook is responsible for running the WIDMgr, so
+	// run it before returning to simulate that.
+	if err := conf.WIDMgr.Run(); err != nil {
+		trCleanup()
+		t.Fatalf("failed to run WIDMgr: %v", err)
+	}
+
 	return conf, trCleanup
 }
 
@@ -2360,6 +2373,152 @@ func TestTaskRunner_Template_BlockingPreStart(t *testing.T) {
 	}
 }
 
+func TestTaskRunner_TemplateWorkloadIdendity(t *testing.T) {
+	ci.Parallel(t)
+
+	expectedConsulValue := "consul-value"
+	consulKVResp := fmt.Sprintf(`
+[
+    {
+        "LockIndex": 0,
+        "Key": "consul-key",
+        "Flags": 0,
+        "Value": "%s",
+        "CreateIndex": 57,
+        "ModifyIndex": 57
+    }
+]
+`, base64.StdEncoding.EncodeToString([]byte(expectedConsulValue)))
+
+	expectedVaultSecret := "vault-secret"
+	vaultSecretResp := fmt.Sprintf(`
+{
+  "data": {
+    "data": {
+      "secret": "%s"
+    },
+    "metadata": {
+      "created_time": "2023-10-18T15:58:29.65137Z",
+      "custom_metadata": null,
+      "deletion_time": "",
+      "destroyed": false,
+      "version": 1
+    }
+  }
+}`, expectedVaultSecret)
+
+	// Start a test server for Consul and Vault.
+	vaultServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintln(w, vaultSecretResp)
+	}))
+	t.Cleanup(vaultServer.Close)
+
+	firstConsulRequest := &atomic.Bool{}
+	firstConsulRequest.Store(true)
+	consulServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !firstConsulRequest.Load() {
+			// Simulate a blocking query to avoid a tight loop in
+			// consul-template.
+			var wait time.Duration
+			if waitStr := r.FormValue("wait"); waitStr != "" {
+				wait, _ = time.ParseDuration(waitStr)
+			}
+
+			if wait != 0 {
+				doneCh := make(chan struct{})
+				t.Cleanup(func() { close(doneCh) })
+
+				timer, stop := helper.NewSafeTimer(wait)
+				t.Cleanup(stop)
+
+				select {
+				case <-timer.C:
+				case <-doneCh:
+					return
+				}
+			}
+		} else {
+			// Send response immediately if this is the first request.
+			firstConsulRequest.Store(true)
+		}
+
+		// Return an index so consul-template knows that the blocking query
+		// didn't timeout on first run.
+		// https://github.com/hashicorp/consul-template/blob/2d2654ffe96210db43306922aaefbb730a8e07f9/watch/view.go#L267-L272
+		w.Header().Set("X-Consul-Index", "57")
+		fmt.Fprintln(w, consulKVResp)
+	}))
+	t.Cleanup(consulServer.Close)
+
+	// Create allocation with a template that reads from Consul and Vault.
+	alloc := mock.BatchAlloc()
+	alloc.Job.TaskGroups[0].Count = 1
+	task := alloc.Job.TaskGroups[0].Tasks[0]
+	task.Driver = "mock_driver"
+	task.Config = map[string]interface{}{
+		"run_for": "2s",
+	}
+	task.Consul = &structs.Consul{
+		Cluster: structs.ConsulDefaultCluster,
+	}
+	task.Vault = &structs.Vault{
+		Cluster: structs.VaultDefaultCluster,
+	}
+	task.Identities = []*structs.WorkloadIdentity{
+		{Name: task.Consul.IdentityName()},
+		{Name: task.Vault.IdentityName()},
+	}
+	task.Templates = []*structs.Template{
+		{
+			EmbeddedTmpl: `
+{{key "consul-key"}}
+{{with secret "secret/data/vault-key"}}{{.Data.data.secret}}{{end}}
+`,
+			DestPath: "local/out.txt",
+		},
+	}
+
+	// Create task runner with a Consul and Vault cluster configured.
+	conf, cleanup := testTaskRunnerConfig(t, alloc, task.Name, nil)
+	conf.ClientConfig.ConsulConfigs = map[string]*structsc.ConsulConfig{
+		structs.ConsulDefaultCluster: {
+			Addr: consulServer.URL,
+		},
+	}
+	conf.ClientConfig.VaultConfigs = map[string]*structsc.VaultConfig{
+		structs.VaultDefaultCluster: {
+			Enabled: pointer.Of(true),
+			Addr:    vaultServer.URL,
+		},
+	}
+	conf.AllocHookResources.SetConsulTokens(map[string]map[string]string{
+		structs.ConsulDefaultCluster: {
+			task.Consul.IdentityName(): "consul-task-token",
+		},
+	})
+	t.Cleanup(cleanup)
+
+	tr, err := NewTaskRunner(conf)
+	must.NoError(t, err)
+
+	// Run the task runner.
+	go tr.Run()
+	t.Cleanup(func() {
+		tr.Kill(context.Background(), structs.NewTaskEvent("cleanup"))
+	})
+
+	// Wait for task to complete.
+	testWaitForTaskToStart(t, tr)
+
+	// Verify template was rendered with the expected content.
+	data, err := os.ReadFile(path.Join(conf.TaskDir.LocalDir, "out.txt"))
+	must.NoError(t, err)
+
+	renderedTmpl := string(data)
+	must.StrContains(t, renderedTmpl, expectedConsulValue)
+	must.StrContains(t, renderedTmpl, expectedVaultSecret)
+}
+
 // TestTaskRunner_Template_NewVaultToken asserts that a new vault token is
 // created when rendering template and that it is revoked on alloc completion
 func TestTaskRunner_Template_NewVaultToken(t *testing.T) {