github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/client/allocrunner/taskrunner/identity_hook.go (about)

     1  package taskrunner
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	log "github.com/hashicorp/go-hclog"
     8  
     9  	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
    10  )
    11  
    12  // identityHook sets the task runner's Nomad workload identity token
    13  // based on the signed identity stored on the Allocation
    14  type identityHook struct {
    15  	tr       *TaskRunner
    16  	logger   log.Logger
    17  	taskName string
    18  	lock     sync.Mutex
    19  }
    20  
    21  func newIdentityHook(tr *TaskRunner, logger log.Logger) *identityHook {
    22  	h := &identityHook{
    23  		tr:       tr,
    24  		taskName: tr.taskName,
    25  	}
    26  	h.logger = logger.Named(h.Name())
    27  	return h
    28  }
    29  
    30  func (*identityHook) Name() string {
    31  	return "identity"
    32  }
    33  
    34  func (h *identityHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
    35  	h.lock.Lock()
    36  	defer h.lock.Unlock()
    37  
    38  	token := h.tr.alloc.SignedIdentities[h.taskName]
    39  	if token != "" {
    40  		h.tr.setNomadToken(token)
    41  	}
    42  	return nil
    43  }
    44  
    45  func (h *identityHook) Update(_ context.Context, req *interfaces.TaskUpdateRequest, _ *interfaces.TaskUpdateResponse) error {
    46  	h.lock.Lock()
    47  	defer h.lock.Unlock()
    48  
    49  	token := h.tr.alloc.SignedIdentities[h.taskName]
    50  	if token != "" {
    51  		h.tr.setNomadToken(token)
    52  	}
    53  	return nil
    54  }