client: use WI-issued consul tokens in the template_hook (#18752)

ref https://github.com/hashicorp/team-nomad/issues/404

client/allocrunner/taskrunner/template/template.go

@@ -95,6 +95,10 @@ type TaskTemplateManagerConfig struct {
 	// ConsulNamespace is the Consul namespace for the task
 	ConsulNamespace string
 
+	// ConsulToken is the Consul ACL token fetched by consul_hook using
+	// workload identity
+	ConsulToken string
+
 	// VaultToken is the Vault token for the task.
 	VaultToken string
 
@@ -812,7 +816,15 @@ func newRunnerConfig(config *TaskTemplateManagerConfig,
 	// Set up the Consul config
 	if cc.ConsulConfig != nil {
 		conf.Consul.Address = &cc.ConsulConfig.Addr
-		conf.Consul.Token = &cc.ConsulConfig.Token
+
+		// if we're using WI, use the token from consul_hook
+		// NOTE: from Nomad 1.9 on, WI will be the only supported way of
+		// getting Consul tokens
+		if config.ConsulToken != "" {
+			conf.Consul.Token = &config.ConsulToken
+		} else {
+			conf.Consul.Token = &cc.ConsulConfig.Token
+		}
 
 		// Get the Consul namespace from agent config. This is the lower level
 		// of precedence (beyond default).

client/allocrunner/taskrunner/template/template_test.go

@@ -5,7 +5,6 @@ package template
 
 import (
 	"bytes"
-	"context"
 	"fmt"
 	"io"
 	"os"
@@ -16,7 +15,6 @@ import (
 	"sort"
 	"strconv"
 	"strings"
-	"sync"
 	"syscall"
 	"testing"
 	"time"
@@ -25,6 +23,7 @@ import (
 	ctestutil "github.com/hashicorp/consul/sdk/testutil"
 	"github.com/hashicorp/nomad/ci"
 	"github.com/hashicorp/nomad/client/allocdir"
+	trtesting "github.com/hashicorp/nomad/client/allocrunner/taskrunner/testing"
 	"github.com/hashicorp/nomad/client/config"
 	"github.com/hashicorp/nomad/client/taskenv"
 	clienttestutil "github.com/hashicorp/nomad/client/testutil"
@@ -48,84 +47,6 @@ const (
 	TestTaskName = "test-task"
 )
 
-// MockTaskHooks is a mock of the TaskHooks interface useful for testing
-type MockTaskHooks struct {
-	Restarts  int
-	RestartCh chan struct{}
-
-	Signals    []string
-	SignalCh   chan struct{}
-	signalLock sync.Mutex
-
-	// SignalError is returned when Signal is called on the mock hook
-	SignalError error
-
-	UnblockCh chan struct{}
-
-	KillEvent *structs.TaskEvent
-	KillCh    chan *structs.TaskEvent
-
-	Events      []*structs.TaskEvent
-	EmitEventCh chan *structs.TaskEvent
-
-	// hasHandle can be set to simulate restoring a task after client restart
-	hasHandle bool
-}
-
-func NewMockTaskHooks() *MockTaskHooks {
-	return &MockTaskHooks{
-		UnblockCh:   make(chan struct{}, 1),
-		RestartCh:   make(chan struct{}, 1),
-		SignalCh:    make(chan struct{}, 1),
-		KillCh:      make(chan *structs.TaskEvent, 1),
-		EmitEventCh: make(chan *structs.TaskEvent, 1),
-	}
-}
-func (m *MockTaskHooks) Restart(ctx context.Context, event *structs.TaskEvent, failure bool) error {
-	m.Restarts++
-	select {
-	case m.RestartCh <- struct{}{}:
-	default:
-	}
-	return nil
-}
-
-func (m *MockTaskHooks) Signal(event *structs.TaskEvent, s string) error {
-	m.signalLock.Lock()
-	m.Signals = append(m.Signals, s)
-	m.signalLock.Unlock()
-	select {
-	case m.SignalCh <- struct{}{}:
-	default:
-	}
-
-	return m.SignalError
-}
-
-func (m *MockTaskHooks) Kill(ctx context.Context, event *structs.TaskEvent) error {
-	m.KillEvent = event
-	select {
-	case m.KillCh <- event:
-	default:
-	}
-	return nil
-}
-
-func (m *MockTaskHooks) IsRunning() bool {
-	return m.hasHandle
-}
-
-func (m *MockTaskHooks) EmitEvent(event *structs.TaskEvent) {
-	m.Events = append(m.Events, event)
-	select {
-	case m.EmitEventCh <- event:
-	case <-m.EmitEventCh:
-		m.EmitEventCh <- event
-	}
-}
-
-func (m *MockTaskHooks) SetState(state string, event *structs.TaskEvent) {}
-
 // mockExecutor implements script executor interface
 type mockExecutor struct {
 	DesiredExit int
@@ -140,7 +61,7 @@ func (m *mockExecutor) Exec(timeout time.Duration, cmd string, args []string) ([
 // Consul/Vault as needed
 type testHarness struct {
 	manager        *TaskTemplateManager
-	mockHooks      *MockTaskHooks
+	mockHooks      *trtesting.MockTaskHooks
 	templates      []*structs.Template
 	envBuilder     *taskenv.Builder
 	node           *structs.Node
@@ -160,7 +81,7 @@ func newTestHarness(t *testing.T, templates []*structs.Template, consul, vault b
 	mockNode := mock.Node()
 
 	harness := &testHarness{
-		mockHooks: NewMockTaskHooks(),
+		mockHooks: trtesting.NewMockTaskHooks(),
 		templates: templates,
 		node:      mockNode,
 		config: &config.Config{
@@ -251,7 +172,7 @@ func (h *testHarness) stop() {
 
 func TestTaskTemplateManager_InvalidConfig(t *testing.T) {
 	ci.Parallel(t)
-	hooks := NewMockTaskHooks()
+	hooks := trtesting.NewMockTaskHooks()
 	clientConfig := &config.Config{Region: "global"}
 	taskDir := "foo"
 	a := mock.Alloc()
@@ -846,7 +767,7 @@ func TestTaskTemplateManager_FirstRender_Restored(t *testing.T) {
 	require.Equal(content, string(raw), "Unexpected template data; got %s, want %q", raw, content)
 
 	// task is now running
-	harness.mockHooks.hasHandle = true
+	harness.mockHooks.HasHandle = true
 
 	// simulate a client restart
 	harness.manager.Stop()
@@ -1017,7 +938,7 @@ func TestTaskTemplateManager_Rerender_Signal(t *testing.T) {
 		t.Fatalf("Task unblock should have been called")
 	}
 
-	if len(harness.mockHooks.Signals) != 0 {
+	if len(harness.mockHooks.Signals()) != 0 {
 		t.Fatalf("Should not have received any signals: %+v", harness.mockHooks)
 	}
 
@@ -1033,9 +954,7 @@ OUTER:
 		case <-harness.mockHooks.RestartCh:
 			t.Fatalf("Restart with signal policy: %+v", harness.mockHooks)
 		case <-harness.mockHooks.SignalCh:
-			harness.mockHooks.signalLock.Lock()
-			s := harness.mockHooks.Signals
-			harness.mockHooks.signalLock.Unlock()
+			s := harness.mockHooks.Signals()
 			if len(s) != 2 {
 				continue
 			}

client/allocrunner/taskrunner/template_hook.go

@@ -13,6 +13,7 @@ import (
 	ti "github.com/hashicorp/nomad/client/allocrunner/taskrunner/interfaces"
 	"github.com/hashicorp/nomad/client/allocrunner/taskrunner/template"
 	"github.com/hashicorp/nomad/client/config"
+	cstructs "github.com/hashicorp/nomad/client/structs"
 	"github.com/hashicorp/nomad/client/taskenv"
 	"github.com/hashicorp/nomad/nomad/structs"
 )
@@ -48,6 +49,9 @@ type templateHookConfig struct {
 
 	// renderOnTaskRestart is flag to explicitly render templates on task restart
 	renderOnTaskRestart bool
+
+	// hookResources are used to fetch Consul tokens
+	hookResources *cstructs.AllocHookResources
 }
 
 type templateHook struct {
@@ -78,6 +82,10 @@ type templateHook struct {
 	// nomadToken is the current Nomad token
 	nomadToken string
 
+	// consulToken is the Consul ACL token obtained from consul_hook via
+	// workload identity
+	consulToken string
+
 	// taskDir is the task directory
 	taskDir string
 }
@@ -113,6 +121,27 @@ func (h *templateHook) Prestart(ctx context.Context, req *interfaces.TaskPrestar
 	h.vaultToken = req.VaultToken
 	h.nomadToken = req.NomadToken
 
+	// Set the consul token if the task uses WI
+	if req.Task.Consul != nil {
+		consulTokens := h.config.hookResources.GetConsulTokens()
+
+		var found bool
+		if _, found = consulTokens[req.Task.Consul.Cluster]; !found {
+			return fmt.Errorf(
+				"consul tokens for cluster %s requested by task %s not found",
+				req.Task.Consul.Cluster, req.Task.Name,
+			)
+		}
+
+		h.consulToken, found = consulTokens[req.Task.Consul.Cluster][req.Task.Consul.IdentityName()]
+		if !found {
+			return fmt.Errorf(
+				"consul tokens for cluster %s and identity %s requested by task %s not found",
+				req.Task.Consul.Cluster, req.Task.Consul.IdentityName(), req.Task.Name,
+			)
+		}
+	}
+
 	// Set vault namespace if specified
 	if req.Task.Vault != nil {
 		h.vaultNamespace = req.Task.Vault.Namespace
@@ -162,6 +191,7 @@ func (h *templateHook) newManager() (unblock chan struct{}, err error) {
 		Templates:            h.config.templates,
 		ClientConfig:         h.config.clientConfig,
 		ConsulNamespace:      h.config.consulNamespace,
+		ConsulToken:          h.consulToken,
 		VaultToken:           h.vaultToken,
 		VaultNamespace:       h.vaultNamespace,
 		TaskDir:              h.taskDir,

client/allocrunner/taskrunner/template_hook_test.go

@@ -0,0 +1,121 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package taskrunner
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+
+	"github.com/hashicorp/nomad/ci"
+	"github.com/hashicorp/nomad/client/allocdir"
+	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
+	trtesting "github.com/hashicorp/nomad/client/allocrunner/taskrunner/testing"
+	"github.com/hashicorp/nomad/client/config"
+	cstructs "github.com/hashicorp/nomad/client/structs"
+	"github.com/hashicorp/nomad/client/taskenv"
+	"github.com/hashicorp/nomad/helper/testlog"
+	"github.com/hashicorp/nomad/helper/uuid"
+	"github.com/hashicorp/nomad/nomad/mock"
+	"github.com/hashicorp/nomad/nomad/structs"
+	"github.com/shoenig/test/must"
+)
+
+func Test_templateHook_Prestart_ConsulWI(t *testing.T) {
+	ci.Parallel(t)
+	logger := testlog.HCLogger(t)
+
+	// mock some consul tokens
+	hr := cstructs.NewAllocHookResources()
+	hr.SetConsulTokens(
+		map[string]map[string]string{
+			structs.ConsulDefaultCluster: {
+				fmt.Sprintf("consul_%s", structs.ConsulDefaultCluster): uuid.Generate(),
+			},
+			"test": {
+				"consul_test": uuid.Generate(),
+			},
+		},
+	)
+
+	a := mock.Alloc()
+	clientConfig := &config.Config{Region: "global"}
+	envBuilder := taskenv.NewBuilder(mock.Node(), a, a.Job.TaskGroups[0].Tasks[0], clientConfig.Region)
+	taskHooks := trtesting.NewMockTaskHooks()
+
+	conf := &templateHookConfig{
+		logger:        logger,
+		lifecycle:     taskHooks,
+		events:        &mockEmitter{},
+		clientConfig:  clientConfig,
+		envBuilder:    envBuilder,
+		hookResources: hr,
+	}
+
+	tests := []struct {
+		name        string
+		req         *interfaces.TaskPrestartRequest
+		wantErr     bool
+		wantErrMsg  string
+		consulToken string
+	}{
+		{
+			"task with no Consul WI",
+			&interfaces.TaskPrestartRequest{
+				Task:    &structs.Task{},
+				TaskDir: &allocdir.TaskDir{Dir: "foo"},
+			},
+			false,
+			"",
+			"",
+		},
+		{
+			"task with Consul WI but no corresponding identity",
+			&interfaces.TaskPrestartRequest{
+				Task: &structs.Task{
+					Name:   "foo",
+					Consul: &structs.Consul{Cluster: "bar"},
+				},
+				TaskDir: &allocdir.TaskDir{Dir: "foo"},
+			},
+			true,
+			"consul tokens for cluster bar requested by task foo not found",
+			"",
+		},
+		{
+			"task with Consul WI",
+			&interfaces.TaskPrestartRequest{
+				Task: &structs.Task{
+					Name:   "foo",
+					Consul: &structs.Consul{Cluster: "default"},
+				},
+				TaskDir: &allocdir.TaskDir{Dir: "foo"},
+			},
+			false,
+			"",
+			hr.GetConsulTokens()[structs.ConsulDefaultCluster]["consul_default"],
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			h := &templateHook{
+				config:       conf,
+				logger:       logger,
+				managerLock:  sync.Mutex{},
+				driverHandle: nil,
+			}
+
+			err := h.Prestart(context.Background(), tt.req, nil)
+			if tt.wantErr {
+				must.NotNil(t, err)
+				must.Eq(t, tt.wantErrMsg, err.Error())
+			} else {
+				must.Nil(t, err)
+			}
+
+			must.Eq(t, tt.consulToken, h.consulToken)
+		})
+	}
+}

client/allocrunner/taskrunner/testing/testing.go

@@ -0,0 +1,113 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package testing
+
+import (
+	"context"
+	"sync"
+
+	"github.com/hashicorp/nomad/nomad/structs"
+)
+
+// MockEmitter is a mock of the EventEmitter interface.
+type MockEmitter struct {
+	lock   sync.Mutex
+	events []*structs.TaskEvent
+}
+
+func (m *MockEmitter) EmitEvent(ev *structs.TaskEvent) {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+	m.events = append(m.events, ev)
+}
+
+func (m *MockEmitter) Events() []*structs.TaskEvent {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+	return m.events
+}
+
+// MockTaskHooks is a mock of the TaskHooks interface useful for testing
+type MockTaskHooks struct {
+	Restarts  int
+	RestartCh chan struct{}
+
+	SignalCh   chan struct{}
+	signals    []string
+	signalLock sync.Mutex
+
+	// SignalError is returned when Signal is called on the mock hook
+	SignalError error
+
+	UnblockCh chan struct{}
+
+	KillEvent *structs.TaskEvent
+	KillCh    chan *structs.TaskEvent
+
+	Events      []*structs.TaskEvent
+	EmitEventCh chan *structs.TaskEvent
+
+	// HasHandle can be set to simulate restoring a task after client restart
+	HasHandle bool
+}
+
+func NewMockTaskHooks() *MockTaskHooks {
+	return &MockTaskHooks{
+		UnblockCh:   make(chan struct{}, 1),
+		RestartCh:   make(chan struct{}, 1),
+		SignalCh:    make(chan struct{}, 1),
+		KillCh:      make(chan *structs.TaskEvent, 1),
+		EmitEventCh: make(chan *structs.TaskEvent, 1),
+	}
+}
+func (m *MockTaskHooks) Restart(ctx context.Context, event *structs.TaskEvent, failure bool) error {
+	m.Restarts++
+	select {
+	case m.RestartCh <- struct{}{}:
+	default:
+	}
+	return nil
+}
+
+func (m *MockTaskHooks) Signal(event *structs.TaskEvent, s string) error {
+	m.signalLock.Lock()
+	m.signals = append(m.signals, s)
+	m.signalLock.Unlock()
+	select {
+	case m.SignalCh <- struct{}{}:
+	default:
+	}
+
+	return m.SignalError
+}
+
+func (m *MockTaskHooks) Signals() []string {
+	m.signalLock.Lock()
+	defer m.signalLock.Unlock()
+	return m.signals
+}
+
+func (m *MockTaskHooks) Kill(ctx context.Context, event *structs.TaskEvent) error {
+	m.KillEvent = event
+	select {
+	case m.KillCh <- event:
+	default:
+	}
+	return nil
+}
+
+func (m *MockTaskHooks) IsRunning() bool {
+	return m.HasHandle
+}
+
+func (m *MockTaskHooks) EmitEvent(event *structs.TaskEvent) {
+	m.Events = append(m.Events, event)
+	select {
+	case m.EmitEventCh <- event:
+	case <-m.EmitEventCh:
+		m.EmitEventCh <- event
+	}
+}
+
+func (m *MockTaskHooks) SetState(state string, event *structs.TaskEvent) {}