diff --git a/docs/runner-protocol-roadmap.md b/docs/runner-protocol-roadmap.md index 90b3ffb..52f0899 100644 --- a/docs/runner-protocol-roadmap.md +++ b/docs/runner-protocol-roadmap.md @@ -22,6 +22,12 @@ controller 使用单一 Go 二进制；默认在同一进程启用 `scheduler` `vm-worker`，也可通过 `--components` 只启用其中一部分。组件是独立应用服务边界， 共享进程不意味着共享后端状态或把 assignment 降级为内存 channel。 +executor 直接运行固定版本的官方 Gitea Runner 二进制，不 fork workflow 执行引擎。 +controller 暴露兼容 RunnerService 的 facade：`FetchTask` 只返回已分配 assignment， +`UpdateTask` 与 `UpdateLog` 转发真实 Gitea。facade 同时验证逻辑 SPIFFE ID、assignment +ID，以及由 controller 密钥确定性生成的 assignment HMAC capability；该 capability +只绑定执行实例，不参与 Zot/OpenBao 等业务授权。 + 这与“收到 webhook 后临时注册另一个 act_runner”不同。`FetchTask` 已经完成任务分配， 不能再期待 Gitea 把同一个 task 分配给随后启动的 runner。协议调度器必须让 executor 执行已经领取的 task，并继续完成日志、状态、心跳、取消和最终结果上报。 @@ -58,6 +64,8 @@ controller 使用单一 Go 二进制；默认在同一进程启用 `scheduler` Pod UID 等短暂未就绪状态以及临时后端错误使用延迟 NAK。 - assignment ACK 后的运行、结果回报和清理由 backend reconciler 根据 Kubernetes、 OpenSandbox 与 Gitea 的事实状态驱动，不继续占用 JetStream delivery。 +- consumer 在 executor 使用上述 facade 成功 claim task 后确认 assignment；无需把完整 + task 写入 Pod annotation、OpenSandbox metadata 或环境变量。 - Pod 与 VM 共享 task/executor 协议，只有环境创建和销毁实现不同。 ## 实现顺序 diff --git a/internal/assignmentqueue/jetstream.go b/internal/assignmentqueue/jetstream.go index 0d22884..cb2e651 100644 --- a/internal/assignmentqueue/jetstream.go +++ b/internal/assignmentqueue/jetstream.go @@ -52,6 +52,11 @@ type Accepter interface { Accept(context.Context, taskassignment.Assignment) (bool, error) } +type Claims interface { + Offer(taskassignment.Assignment) (<-chan struct{}, error) + WaitClaimed(context.Context, string) error +} + // Message is the subset of jetstream.Msg needed by one reconciliation. type Message interface { Data() []byte @@ -62,19 +67,24 @@ type Message interface { // Processor maps one delivery to one idempotent worker reconciliation. type Processor struct { - TrustDomain string - Accepter Accepter - RetryDelay time.Duration + TrustDomain string + Accepter Accepter + Claims Claims + RetryDelay time.Duration + ClaimTimeout time.Duration } func (p Processor) Process(ctx context.Context, message Message) error { - if p.Accepter == nil { - return errors.New("assignment accepter is required") + if p.Accepter == nil || p.Claims == nil { + return errors.New("assignment accepter and claim registry are required") } assignment, err := taskassignment.Unmarshal(message.Data(), p.TrustDomain) if err != nil { return errors.Join(err, message.TermWithReason("invalid assignment")) } + if _, err := p.Claims.Offer(assignment); err != nil { + return errors.Join(err, message.TermWithReason("conflicting assignment")) + } accepted, err := p.Accepter.Accept(ctx, assignment) if err != nil { delay := p.RetryDelay @@ -84,6 +94,20 @@ func (p Processor) Process(ctx context.Context, message Message) error { return errors.Join(err, message.NakWithDelay(delay)) } if accepted { + timeout := p.ClaimTimeout + if timeout <= 0 { + timeout = 4 * time.Minute + } + claimContext, cancel := context.WithTimeout(ctx, timeout) + err := p.Claims.WaitClaimed(claimContext, assignment.ID) + cancel() + if err != nil { + delay := p.RetryDelay + if delay <= 0 { + delay = 2 * time.Second + } + return errors.Join(err, message.NakWithDelay(delay)) + } if err := message.DoubleAck(ctx); err != nil { return fmt.Errorf("ack assignment %s: %w", assignment.ID, err) } diff --git a/internal/assignmentqueue/jetstream_test.go b/internal/assignmentqueue/jetstream_test.go index df58f46..5ea2cec 100644 --- a/internal/assignmentqueue/jetstream_test.go +++ b/internal/assignmentqueue/jetstream_test.go @@ -57,6 +57,26 @@ type fakeAccepter struct { err error } +type fakeClaims struct { + claimed bool +} + +func (c *fakeClaims) Offer(taskassignment.Assignment) (<-chan struct{}, error) { + ready := make(chan struct{}) + if c.claimed { + close(ready) + } + return ready, nil +} + +func (c *fakeClaims) WaitClaimed(ctx context.Context, _ string) error { + if c.claimed { + return nil + } + <-ctx.Done() + return ctx.Err() +} + func (a *fakeAccepter) Accept(context.Context, taskassignment.Assignment) (bool, error) { return a.accepted, a.err } @@ -84,7 +104,7 @@ func encodedAssignment(t *testing.T) []byte { func TestProcessorAcknowledgesPersistedHandoff(t *testing.T) { message := &fakeMessage{data: encodedAssignment(t)} - processor := Processor{TrustDomain: "ddupan.top", Accepter: &fakeAccepter{accepted: true}} + processor := Processor{TrustDomain: "ddupan.top", Accepter: &fakeAccepter{accepted: true}, Claims: &fakeClaims{claimed: true}} if err := processor.Process(context.Background(), message); err != nil { t.Fatal(err) } @@ -95,7 +115,7 @@ func TestProcessorAcknowledgesPersistedHandoff(t *testing.T) { func TestProcessorRetriesUntilBackendHandoffIsDurable(t *testing.T) { message := &fakeMessage{data: encodedAssignment(t)} - processor := Processor{TrustDomain: "ddupan.top", Accepter: &fakeAccepter{}, RetryDelay: 2 * time.Second} + processor := Processor{TrustDomain: "ddupan.top", Accepter: &fakeAccepter{}, Claims: &fakeClaims{}, RetryDelay: 2 * time.Second} if err := processor.Process(context.Background(), message); err != nil { t.Fatal(err) } @@ -109,6 +129,7 @@ func TestProcessorRetriesBackendFailureAndTerminatesPoisonMessage(t *testing.T) processor := Processor{ TrustDomain: "ddupan.top", Accepter: &fakeAccepter{err: errors.New("backend unavailable")}, + Claims: &fakeClaims{}, RetryDelay: time.Minute, } if err := processor.Process(context.Background(), retry); err == nil { diff --git a/internal/giteaactions/client.go b/internal/giteaactions/client.go index 299bdae..6353434 100644 --- a/internal/giteaactions/client.go +++ b/internal/giteaactions/client.go @@ -55,6 +55,16 @@ func (c *Client) FetchTask(ctx context.Context, tasksVersion int64) (*runnerv1.F return response.Msg, nil } +// UpdateTask forwards executor state through the scheduler runner identity. +func (c *Client) UpdateTask(ctx context.Context, request *connect.Request[runnerv1.UpdateTaskRequest]) (*connect.Response[runnerv1.UpdateTaskResponse], error) { + return c.runner.UpdateTask(ctx, request) +} + +// UpdateLog forwards executor log rows through the scheduler runner identity. +func (c *Client) UpdateLog(ctx context.Context, request *connect.Request[runnerv1.UpdateLogRequest]) (*connect.Response[runnerv1.UpdateLogResponse], error) { + return c.runner.UpdateLog(ctx, request) +} + // DefaultHTTPClient is suitable for the scheduler's long-lived connection. func DefaultHTTPClient() *http.Client { return &http.Client{Transport: http.DefaultTransport} diff --git a/internal/giteaactions/client_test.go b/internal/giteaactions/client_test.go index db6a9fa..e4301fe 100644 --- a/internal/giteaactions/client_test.go +++ b/internal/giteaactions/client_test.go @@ -18,6 +18,8 @@ type runnerService struct { t *testing.T declaredLabels []string fetchedVersion int64 + updatedTask int64 + updatedLog int64 expectedUUID string expectedToken string } @@ -47,6 +49,18 @@ func (s *runnerService) FetchTask(_ context.Context, request *connect.Request[ru }), nil } +func (s *runnerService) UpdateTask(_ context.Context, request *connect.Request[runnerv1.UpdateTaskRequest]) (*connect.Response[runnerv1.UpdateTaskResponse], error) { + s.checkAuth(request) + s.updatedTask = request.Msg.GetState().GetId() + return connect.NewResponse(&runnerv1.UpdateTaskResponse{State: request.Msg.State}), nil +} + +func (s *runnerService) UpdateLog(_ context.Context, request *connect.Request[runnerv1.UpdateLogRequest]) (*connect.Response[runnerv1.UpdateLogResponse], error) { + s.checkAuth(request) + s.updatedLog = request.Msg.GetTaskId() + return connect.NewResponse(&runnerv1.UpdateLogResponse{}), nil +} + func TestClientUsesOfficialRunnerProtocol(t *testing.T) { service := &runnerService{ t: t, @@ -75,4 +89,13 @@ func TestClientUsesOfficialRunnerProtocol(t *testing.T) { if service.fetchedVersion != 7 || response.GetTasksVersion() != 8 || response.GetTask().GetId() != 42 { t.Fatalf("unexpected FetchTask exchange: request=%d response=%v", service.fetchedVersion, response) } + if _, err := client.UpdateTask(context.Background(), connect.NewRequest(&runnerv1.UpdateTaskRequest{State: &runnerv1.TaskState{Id: 42}})); err != nil { + t.Fatal(err) + } + if _, err := client.UpdateLog(context.Background(), connect.NewRequest(&runnerv1.UpdateLogRequest{TaskId: 42})); err != nil { + t.Fatal(err) + } + if service.updatedTask != 42 || service.updatedLog != 42 { + t.Fatalf("updated task=%d log=%d", service.updatedTask, service.updatedLog) + } } diff --git a/internal/runnerfacade/capability.go b/internal/runnerfacade/capability.go new file mode 100644 index 0000000..e6e1bfb --- /dev/null +++ b/internal/runnerfacade/capability.go @@ -0,0 +1,36 @@ +package runnerfacade + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/base64" + "errors" +) + +// Capabilities are deterministic per assignment so controller restarts do not +// require a per-task token database. +type Capabilities struct{ key []byte } + +func NewCapabilities(key []byte) (Capabilities, error) { + if len(key) < 32 { + return Capabilities{}, errors.New("runner facade capability key must be at least 32 bytes") + } + return Capabilities{key: append([]byte(nil), key...)}, nil +} + +func (c Capabilities) Issue(assignmentID string) string { + if len(c.key) < 32 || assignmentID == "" { + return "" + } + mac := hmac.New(sha256.New, c.key) + _, _ = mac.Write([]byte("gitea-runner-assignment\x00" + assignmentID)) + return base64.RawURLEncoding.EncodeToString(mac.Sum(nil)) +} + +func (c Capabilities) Verify(assignmentID, token string) bool { + want := c.Issue(assignmentID) + if want == "" || token == "" { + return false + } + return hmac.Equal([]byte(want), []byte(token)) +} diff --git a/internal/runnerfacade/facade.go b/internal/runnerfacade/facade.go new file mode 100644 index 0000000..d0560ce --- /dev/null +++ b/internal/runnerfacade/facade.go @@ -0,0 +1,134 @@ +// Package runnerfacade presents pre-assigned tasks to unmodified Gitea Runner binaries. +package runnerfacade + +import ( + "context" + "errors" + "net/http" + + "connectrpc.com/connect" + "gitea.dev/actionslib/pkg/protocol" + runnerv1 "gitea.dev/actionslib/runner/v1" + "gitea.dev/actionslib/runner/v1/runnerv1connect" + + "git.ddupan.top/panxiao81/gitea-dynamic-runner/internal/taskassignment" +) + +type identityKey struct{} + +func WithSPIFFEID(ctx context.Context, id string) context.Context { + return context.WithValue(ctx, identityKey{}, id) +} + +// SPIFFEMiddleware extracts the authenticated workload identity from the mTLS +// peer certificate. TLS verification itself is configured by the server with +// go-spiffe; this layer only passes the verified ID into Connect handlers. +func SPIFFEMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + if request.TLS == nil || len(request.TLS.PeerCertificates) == 0 { + http.Error(response, "client SPIFFE identity required", http.StatusUnauthorized) + return + } + var spiffeID string + for _, uri := range request.TLS.PeerCertificates[0].URIs { + if uri.Scheme == "spiffe" { + if spiffeID != "" { + http.Error(response, "multiple client SPIFFE identities", http.StatusUnauthorized) + return + } + spiffeID = uri.String() + } + } + if spiffeID == "" { + http.Error(response, "client SPIFFE identity required", http.StatusUnauthorized) + return + } + next.ServeHTTP(response, request.WithContext(WithSPIFFEID(request.Context(), spiffeID))) + }) +} + +type Upstream interface { + UpdateTask(context.Context, *connect.Request[runnerv1.UpdateTaskRequest]) (*connect.Response[runnerv1.UpdateTaskResponse], error) + UpdateLog(context.Context, *connect.Request[runnerv1.UpdateLogRequest]) (*connect.Response[runnerv1.UpdateLogResponse], error) +} + +type Facade struct { + runnerv1connect.UnimplementedRunnerServiceHandler + Registry *Registry + Capabilities Capabilities + Upstream Upstream +} + +func (f *Facade) Handler() (string, http.Handler) { + return runnerv1connect.NewRunnerServiceHandler(f) +} + +func (f *Facade) Register(context.Context, *connect.Request[runnerv1.RegisterRequest]) (*connect.Response[runnerv1.RegisterResponse], error) { + return nil, connect.NewError(connect.CodePermissionDenied, errors.New("executor registration is disabled")) +} + +func (f *Facade) Declare(ctx context.Context, request *connect.Request[runnerv1.DeclareRequest]) (*connect.Response[runnerv1.DeclareResponse], error) { + if _, _, err := f.authenticate(ctx, request); err != nil { + return nil, err + } + return connect.NewResponse(&runnerv1.DeclareResponse{}), nil +} + +func (f *Facade) FetchTask(ctx context.Context, request *connect.Request[runnerv1.FetchTaskRequest]) (*connect.Response[runnerv1.FetchTaskResponse], error) { + assignmentID, spiffeID, err := f.authenticate(ctx, request) + if err != nil { + return nil, err + } + assignment, err := f.Registry.Claim(assignmentID, spiffeID) + if err != nil { + return nil, connect.NewError(connect.CodeFailedPrecondition, err) + } + return connect.NewResponse(&runnerv1.FetchTaskResponse{Task: assignment.Task}), nil +} + +func (f *Facade) UpdateTask(ctx context.Context, request *connect.Request[runnerv1.UpdateTaskRequest]) (*connect.Response[runnerv1.UpdateTaskResponse], error) { + assignment, err := f.authorizeClaimed(ctx, request) + if err != nil { + return nil, err + } + if request.Msg.GetState().GetId() != assignment.Task.GetId() { + return nil, connect.NewError(connect.CodePermissionDenied, errors.New("task update does not match assignment")) + } + return f.Upstream.UpdateTask(ctx, connect.NewRequest(request.Msg)) +} + +func (f *Facade) UpdateLog(ctx context.Context, request *connect.Request[runnerv1.UpdateLogRequest]) (*connect.Response[runnerv1.UpdateLogResponse], error) { + assignment, err := f.authorizeClaimed(ctx, request) + if err != nil { + return nil, err + } + if request.Msg.GetTaskId() != assignment.Task.GetId() { + return nil, connect.NewError(connect.CodePermissionDenied, errors.New("log update does not match assignment")) + } + return f.Upstream.UpdateLog(ctx, connect.NewRequest(request.Msg)) +} + +func (f *Facade) authorizeClaimed(ctx context.Context, request connect.AnyRequest) (taskassignment.Assignment, error) { + assignmentID, spiffeID, err := f.authenticate(ctx, request) + if err != nil { + return taskassignment.Assignment{}, err + } + assignment, err := f.Registry.Resolve(assignmentID, spiffeID) + if err != nil { + return taskassignment.Assignment{}, connect.NewError(connect.CodePermissionDenied, err) + } + return assignment, nil +} + +func (f *Facade) authenticate(ctx context.Context, request connect.AnyRequest) (string, string, error) { + if f.Registry == nil || f.Upstream == nil { + return "", "", connect.NewError(connect.CodeInternal, errors.New("runner facade is not configured")) + } + assignmentID := request.Header().Get(protocol.UUIDHeader) + token := request.Header().Get(protocol.TokenHeader) + spiffeID, _ := ctx.Value(identityKey{}).(string) + if assignmentID == "" || spiffeID == "" || !f.Capabilities.Verify(assignmentID, token) { + return "", "", connect.NewError(connect.CodeUnauthenticated, errors.New("invalid executor identity or capability")) + } + return assignmentID, spiffeID, nil +} diff --git a/internal/runnerfacade/facade_test.go b/internal/runnerfacade/facade_test.go new file mode 100644 index 0000000..16c5d78 --- /dev/null +++ b/internal/runnerfacade/facade_test.go @@ -0,0 +1,126 @@ +package runnerfacade + +import ( + "context" + "testing" + + "connectrpc.com/connect" + "gitea.dev/actionslib/pkg/protocol" + runnerv1 "gitea.dev/actionslib/runner/v1" + "google.golang.org/protobuf/types/known/structpb" + + "git.ddupan.top/panxiao81/gitea-dynamic-runner/internal/taskassignment" +) + +type fakeUpstream struct { + taskUpdates int + logUpdates int +} + +func (u *fakeUpstream) UpdateTask(_ context.Context, request *connect.Request[runnerv1.UpdateTaskRequest]) (*connect.Response[runnerv1.UpdateTaskResponse], error) { + u.taskUpdates++ + return connect.NewResponse(&runnerv1.UpdateTaskResponse{State: request.Msg.State}), nil +} +func (u *fakeUpstream) UpdateLog(_ context.Context, request *connect.Request[runnerv1.UpdateLogRequest]) (*connect.Response[runnerv1.UpdateLogResponse], error) { + u.logUpdates++ + return connect.NewResponse(&runnerv1.UpdateLogResponse{AckIndex: request.Msg.Index + int64(len(request.Msg.Rows))}), nil +} + +func facadeAssignment(t *testing.T) taskassignment.Assignment { + t.Helper() + fields, err := structpb.NewStruct(map[string]any{"repository": "owner/repo"}) + if err != nil { + t.Fatal(err) + } + assignment, err := taskassignment.New(&runnerv1.Task{ + Id: 42, Context: fields, + WorkflowPayload: []byte("jobs:\n publish:\n runs-on: [self-hosted, pod]\n steps: []\n"), + }, "ddupan.top") + if err != nil { + t.Fatal(err) + } + return assignment +} + +func testFacade(t *testing.T) (*Facade, taskassignment.Assignment, string) { + t.Helper() + capabilities, err := NewCapabilities([]byte("0123456789abcdef0123456789abcdef")) + if err != nil { + t.Fatal(err) + } + assignment := facadeAssignment(t) + registry := NewRegistry() + if _, err := registry.Offer(assignment); err != nil { + t.Fatal(err) + } + return &Facade{Registry: registry, Capabilities: capabilities, Upstream: &fakeUpstream{}}, assignment, capabilities.Issue(assignment.ID) +} + +func authenticatedRequest[T any](message *T, assignmentID, token string) *connect.Request[T] { + request := connect.NewRequest(message) + request.Header().Set(protocol.UUIDHeader, assignmentID) + request.Header().Set(protocol.TokenHeader, token) + return request +} + +func TestFacadeReturnsOnlyPreassignedTaskAndSignalsClaim(t *testing.T) { + facade, assignment, token := testFacade(t) + ctx := WithSPIFFEID(context.Background(), assignment.Identity.SPIFFEID) + response, err := facade.FetchTask(ctx, authenticatedRequest(&runnerv1.FetchTaskRequest{}, assignment.ID, token)) + if err != nil { + t.Fatal(err) + } + if response.Msg.GetTask().GetId() != assignment.Task.GetId() { + t.Fatalf("task = %#v", response.Msg.GetTask()) + } + if err := facade.Registry.WaitClaimed(context.Background(), assignment.ID); err != nil { + t.Fatal(err) + } + if _, err := facade.FetchTask(ctx, authenticatedRequest(&runnerv1.FetchTaskRequest{}, assignment.ID, token)); connect.CodeOf(err) != connect.CodeFailedPrecondition { + t.Fatalf("second FetchTask error = %v", err) + } +} + +func TestFacadeRejectsWrongIdentityOrCapability(t *testing.T) { + facade, assignment, token := testFacade(t) + wrongIdentity := WithSPIFFEID(context.Background(), "spiffe://ddupan.top/ci/owner/repo/other") + if _, err := facade.FetchTask(wrongIdentity, authenticatedRequest(&runnerv1.FetchTaskRequest{}, assignment.ID, token)); connect.CodeOf(err) != connect.CodeFailedPrecondition { + t.Fatalf("wrong identity error = %v", err) + } + ctx := WithSPIFFEID(context.Background(), assignment.Identity.SPIFFEID) + if _, err := facade.FetchTask(ctx, authenticatedRequest(&runnerv1.FetchTaskRequest{}, assignment.ID, "wrong")); connect.CodeOf(err) != connect.CodeUnauthenticated { + t.Fatalf("wrong capability error = %v", err) + } +} + +func TestFacadeForwardsOnlyMatchingTaskAndLogUpdates(t *testing.T) { + facade, assignment, token := testFacade(t) + ctx := WithSPIFFEID(context.Background(), assignment.Identity.SPIFFEID) + if _, err := facade.FetchTask(ctx, authenticatedRequest(&runnerv1.FetchTaskRequest{}, assignment.ID, token)); err != nil { + t.Fatal(err) + } + if _, err := facade.UpdateTask(ctx, authenticatedRequest(&runnerv1.UpdateTaskRequest{State: &runnerv1.TaskState{Id: 42}}, assignment.ID, token)); err != nil { + t.Fatal(err) + } + if _, err := facade.UpdateLog(ctx, authenticatedRequest(&runnerv1.UpdateLogRequest{TaskId: 42}, assignment.ID, token)); err != nil { + t.Fatal(err) + } + upstream := facade.Upstream.(*fakeUpstream) + if upstream.taskUpdates != 1 || upstream.logUpdates != 1 { + t.Fatalf("task updates=%d log updates=%d", upstream.taskUpdates, upstream.logUpdates) + } + if _, err := facade.UpdateLog(ctx, authenticatedRequest(&runnerv1.UpdateLogRequest{TaskId: 99}, assignment.ID, token)); connect.CodeOf(err) != connect.CodePermissionDenied { + t.Fatalf("mismatched log error = %v", err) + } +} + +func TestCapabilitiesAreDeterministicAndAssignmentScoped(t *testing.T) { + capabilities, err := NewCapabilities([]byte("0123456789abcdef0123456789abcdef")) + if err != nil { + t.Fatal(err) + } + token := capabilities.Issue("gitea-task-42") + if !capabilities.Verify("gitea-task-42", token) || capabilities.Verify("gitea-task-43", token) { + t.Fatal("capability scope is invalid") + } +} diff --git a/internal/runnerfacade/registry.go b/internal/runnerfacade/registry.go new file mode 100644 index 0000000..ff00650 --- /dev/null +++ b/internal/runnerfacade/registry.go @@ -0,0 +1,92 @@ +package runnerfacade + +import ( + "context" + "errors" + "fmt" + "sync" + + "git.ddupan.top/panxiao81/gitea-dynamic-runner/internal/taskassignment" +) + +type claim struct { + assignment taskassignment.Assignment + claimed bool + ready chan struct{} +} + +// Registry holds pending task payloads and an active authorization cache. +// Pending entries are rebuilt by JetStream redelivery; active authorization is +// reconstructable from backend metadata and is not an independent state store. +type Registry struct { + mu sync.Mutex + claims map[string]*claim +} + +func NewRegistry() *Registry { return &Registry{claims: make(map[string]*claim)} } + +func (r *Registry) Offer(assignment taskassignment.Assignment) (<-chan struct{}, error) { + r.mu.Lock() + defer r.mu.Unlock() + if existing := r.claims[assignment.ID]; existing != nil { + if existing.assignment.Identity != assignment.Identity || existing.assignment.Task.GetId() != assignment.Task.GetId() { + return nil, fmt.Errorf("assignment %s was offered with different task data", assignment.ID) + } + return existing.ready, nil + } + entry := &claim{assignment: assignment, ready: make(chan struct{})} + r.claims[assignment.ID] = entry + return entry.ready, nil +} + +func (r *Registry) Claim(assignmentID, spiffeID string) (taskassignment.Assignment, error) { + r.mu.Lock() + defer r.mu.Unlock() + entry := r.claims[assignmentID] + if entry == nil { + return taskassignment.Assignment{}, errors.New("assignment is not pending") + } + if entry.assignment.Identity.SPIFFEID != spiffeID { + return taskassignment.Assignment{}, errors.New("executor SPIFFE ID does not match assignment") + } + if entry.claimed { + return taskassignment.Assignment{}, errors.New("assignment was already claimed") + } + entry.claimed = true + close(entry.ready) + return entry.assignment, nil +} + +func (r *Registry) Resolve(assignmentID, spiffeID string) (taskassignment.Assignment, error) { + r.mu.Lock() + defer r.mu.Unlock() + entry := r.claims[assignmentID] + if entry == nil || !entry.claimed { + return taskassignment.Assignment{}, errors.New("assignment is not claimed") + } + if entry.assignment.Identity.SPIFFEID != spiffeID { + return taskassignment.Assignment{}, errors.New("executor SPIFFE ID does not match assignment") + } + return entry.assignment, nil +} + +func (r *Registry) WaitClaimed(ctx context.Context, assignmentID string) error { + r.mu.Lock() + entry := r.claims[assignmentID] + r.mu.Unlock() + if entry == nil { + return errors.New("assignment is not pending") + } + select { + case <-entry.ready: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (r *Registry) Remove(assignmentID string) { + r.mu.Lock() + defer r.mu.Unlock() + delete(r.claims, assignmentID) +}