diff --git a/client/allocrunner/taskrunner/task_runner_hooks.go b/client/allocrunner/taskrunner/task_runner_hooks.go index 806f429cf..193757127 100644 --- a/client/allocrunner/taskrunner/task_runner_hooks.go +++ b/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, diff --git a/client/allocrunner/taskrunner/task_runner_test.go b/client/allocrunner/taskrunner/task_runner_test.go index 482bf7f62..cb2fe0553 100644 --- a/client/allocrunner/taskrunner/task_runner_test.go +++ b/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) {