require an env interpolation function for service.IdentityHandle

hashicorp · Apr 11, 2024 · 825efce · 825efce
1 parent de3943b
commit 825efce
Show file tree

Hide file tree

Showing 7 changed files with 19 additions and 16 deletions.
diff --git a/client/allocrunner/consul_hook.go b/client/allocrunner/consul_hook.go
@@ -181,8 +181,8 @@ func (h *consulHook) prepareConsulTokensForServices(services []*structs.Service,
 		}
 
 		// Find signed identity workload.
-		identity := taskenv.InterpolateWIHandle(env, *service.IdentityHandle())
-		jwt, err := h.widmgr.Get(identity)
+		handle := *service.IdentityHandle(env.ReplaceEnv)
+		jwt, err := h.widmgr.Get(handle)
 		if err != nil {
 			mErr = multierror.Append(mErr, fmt.Errorf(
 				"error getting signed identity for service %s: %v",
@@ -196,7 +196,7 @@ func (h *consulHook) prepareConsulTokensForServices(services []*structs.Service,
 			JWT:            jwt.JWT,
 			AuthMethodName: consulConfig.ServiceIdentityAuthMethod,
 			Meta: map[string]string{
-				"requested_by": fmt.Sprintf("nomad_service_%s", identity.InterpolatedWorkloadIdentifier),
+				"requested_by": fmt.Sprintf("nomad_service_%s", handle.InterpolatedWorkloadIdentifier),
 			},
 		}
 		token, err := h.getConsulToken(clusterName, req)

diff --git a/client/allocrunner/consul_hook_test.go b/client/allocrunner/consul_hook_test.go
@@ -158,7 +158,7 @@ func Test_consulHook_prepareConsulTokensForServices(t *testing.T) {
 	hashedJWT := make(map[string]string)
 
 	for _, s := range services {
-		widHandle := taskenv.InterpolateWIHandle(env, *s.IdentityHandle())
+		widHandle := *s.IdentityHandle(env.ReplaceEnv)
 		jwt, err := hook.widmgr.Get(widHandle)
 		must.NoError(t, err)
 

diff --git a/client/widmgr/widmgr.go b/client/widmgr/widmgr.go
@@ -62,7 +62,7 @@ func NewWIDMgr(signer IdentitySigner, a *structs.Allocation, db cstate.StateDB,
 
 	for _, service := range tg.Services {
 		if service.Identity != nil {
-			handle := taskenv.InterpolateWIHandle(allocEnv, *service.IdentityHandle())
+			handle := *service.IdentityHandle(allocEnv.ReplaceEnv)
 			widspecs[handle] = service.Identity
 		}
 	}
@@ -77,7 +77,7 @@ func NewWIDMgr(signer IdentitySigner, a *structs.Allocation, db cstate.StateDB,
 		taskEnv := envBuilder.UpdateTask(a, task).Build()
 		for _, service := range task.Services {
 			if service.Identity != nil {
-				handle := taskenv.InterpolateWIHandle(taskEnv, *service.IdentityHandle())
+				handle := *service.IdentityHandle(taskEnv.ReplaceEnv)
 				widspecs[handle] = service.Identity
 			}
 		}

diff --git a/client/widmgr/widmgr_test.go b/client/widmgr/widmgr_test.go
@@ -54,10 +54,7 @@ func TestWIDMgr_Restore(t *testing.T) {
 	widSpecs[2].TTL = time.Second
 	signer.setWIDs(widSpecs)
 
-	wiHandle := service.IdentityHandle()
-	wiHandle.InterpolatedWorkloadIdentifier = envBuilder.Build().ReplaceEnv(
-		wiHandle.WorkloadIdentifier)
-
+	wiHandle := service.IdentityHandle(envBuilder.Build().ReplaceEnv)
 	mgr.widSpecs[*wiHandle].TTL = time.Second
 
 	// force a re-sign to re-populate the lastToken and save to the db

diff --git a/nomad/alloc_endpoint.go b/nomad/alloc_endpoint.go
@@ -637,13 +637,13 @@ func (a *Alloc) signServices(
 	// services can be on the level of task groups or tasks
 	for _, tg := range job.TaskGroups {
 		for _, service := range tg.Services {
-			if service.IdentityHandle().Equal(wid) {
+			if service.IdentityHandle(nil).Equal(wid) {
 				return true, a.signIdentities(alloc, service.Identity, idReq, reply, now)
 			}
 		}
 		for _, task := range tg.Tasks {
 			for _, service := range task.Services {
-				if service.IdentityHandle().Equal(wid) {
+				if service.IdentityHandle(nil).Equal(wid) {
 					return true, a.signIdentities(alloc, service.Identity, idReq, reply, now)
 				}
 			}

diff --git a/nomad/structs/services.go b/nomad/structs/services.go
@@ -802,15 +802,21 @@ func (s *Service) MakeUniqueIdentityName() string {
 	return fmt.Sprintf("%s_%v-%v", prefix, s.Name, s.PortLabel)
 }
 
+type envReplacer func(string) string
+
 // IdentityHandle returns a WorkloadIdentityHandle which is a pair of service
 // identity name and service name.
-func (s *Service) IdentityHandle() *WIHandle {
+func (s *Service) IdentityHandle(replace envReplacer) *WIHandle {
 	if s.Identity != nil {
-		return &WIHandle{
+		wi := &WIHandle{
 			IdentityName:       s.Identity.Name,
 			WorkloadIdentifier: s.Name,
 			WorkloadType:       WorkloadTypeService,
 		}
+		if replace != nil {
+			wi.InterpolatedWorkloadIdentifier = replace(s.Name)
+		}
+		return wi
 	}
 	return nil
 }

diff --git a/nomad/structs/structs_test.go b/nomad/structs/structs_test.go
@@ -8240,7 +8240,7 @@ func TestNewIdentityClaims(t *testing.T) {
 				name:           path,
 				group:          tg.Name,
 				wid:            s.Identity,
-				wiHandle:       s.IdentityHandle(),
+				wiHandle:       s.IdentityHandle(nil),
 				expectedClaims: expectedClaims[path],
 			})
 		}
@@ -8269,7 +8269,7 @@ func TestNewIdentityClaims(t *testing.T) {
 					name:           path,
 					group:          tg.Name,
 					wid:            s.Identity,
-					wiHandle:       s.IdentityHandle(),
+					wiHandle:       s.IdentityHandle(nil),
 					expectedClaims: expectedClaims[path],
 				})
 			}