orchestra/internal/federation/federation.go

package federation

import (
	"crypto/sha256"
	"errors"
	"fmt"
	"sync"
	"time"
)

var ErrUnknownWorker = errors.New("unknown worker")
var ErrUnauthorized = errors.New("worker authentication failed")

type Worker struct {
	ID       string    `json:"id"`
	Address  string    `json:"address"`
	Capacity int       `json:"capacity"`
	LastSeen time.Time `json:"last_seen"`
	Online   bool      `json:"online"`
	Token    string    `json:"-"`
}

// Capture is published by a worker that owns the pane. The coordinator never
// reads a remote herdr socket; this is the worker-pulled counterpart.
type Capture struct {
	TaskID   string    `json:"task_id"`
	PaneID   string    `json:"pane_id"`
	Text     string    `json:"text"`
	Revision uint64    `json:"revision"`
	At       time.Time `json:"at"`
}
type Command struct {
	ID              string    `json:"id"`
	TaskID          string    `json:"task_id"`
	Kind            string    `json:"kind"`
	PaneID          string    `json:"pane_id"`
	CaptureRevision uint64    `json:"capture_revision"`
	CreatedAt       time.Time `json:"created_at"`
	Status          string    `json:"status"`
	Error           string    `json:"error,omitempty"`
}

type Registry struct {
	mu sync.Mutex
	// AdmitToken, if set, is a pre-shared secret every registration must
	// present (S10: registration previously accepted a self-declared id and
	// self-chosen token from any caller — admission-control-free). Leave
	// empty only for a deliberately open deployment.
	AdmitToken string
	workers    map[string]Worker
	TTL        time.Duration
	OnOffline  func(Worker)
	cursors    map[string]uint64
	captures   map[string]Capture // worker/task
	commands   map[string][]Command
}

func (r *Registry) init() {
	if r.TTL <= 0 {
		r.TTL = 90 * time.Second
	}
	if r.workers == nil {
		r.workers = map[string]Worker{}
	}
	if r.cursors == nil {
		r.cursors = map[string]uint64{}
	}
	if r.captures == nil {
		r.captures = map[string]Capture{}
	}
	if r.commands == nil {
		r.commands = map[string][]Command{}
	}
}

func captureKey(worker, task string) string { return worker + "\x00" + task }
func (r *Registry) PutCapture(worker string, c Capture) (Capture, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if _, ok := r.workers[worker]; !ok {
		return Capture{}, ErrUnknownWorker
	}
	if c.TaskID == "" || c.PaneID == "" {
		return Capture{}, errors.New("task_id and pane_id required")
	}
	k := captureKey(worker, c.TaskID)
	old := r.captures[k]
	if old.Text != c.Text || old.PaneID != c.PaneID {
		c.Revision = old.Revision + 1
	}
	if c.Revision == 0 {
		c.Revision = 1
	}
	c.At = time.Now().UTC()
	r.captures[k] = c
	return c, nil
}
func (r *Registry) Capture(worker, task string) (Capture, bool) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	c, ok := r.captures[captureKey(worker, task)]
	return c, ok
}
func (r *Registry) Queue(worker string, c Command) (Command, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if _, ok := r.workers[worker]; !ok {
		return Command{}, ErrUnknownWorker
	}
	if c.TaskID == "" || c.PaneID == "" || c.CaptureRevision == 0 || (c.Kind != "grant_approval" && c.Kind != "deny_approval") {
		return Command{}, errors.New("invalid control command")
	}
	c.ID = fmt.Sprintf("cmd-%x", sha256.Sum256([]byte(fmt.Sprintf("%s/%s/%s/%d/%d", worker, c.TaskID, c.Kind, c.CaptureRevision, time.Now().UnixNano()))))[:20]
	c.CreatedAt = time.Now().UTC()
	c.Status = "pending"
	r.commands[worker] = append(r.commands[worker], c)
	return c, nil
}
func (r *Registry) Commands(worker string) ([]Command, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if _, ok := r.workers[worker]; !ok {
		return nil, ErrUnknownWorker
	}
	var out []Command
	for _, c := range r.commands[worker] {
		if c.Status == "pending" {
			out = append(out, c)
		}
	}
	return out, nil
}
func (r *Registry) Command(worker, id string) (Command, bool) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	for _, c := range r.commands[worker] {
		if c.ID == id {
			return c, true
		}
	}
	return Command{}, false
}
func (r *Registry) CompleteCommand(worker, id, status, message string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	for i := range r.commands[worker] {
		if r.commands[worker][i].ID == id {
			if r.commands[worker][i].Status != "pending" {
				return errors.New("command already resolved")
			}
			r.commands[worker][i].Status = status
			r.commands[worker][i].Error = message
			return nil
		}
	}
	return errors.New("command not found")
}

// Register admits a worker. admitToken must match r.AdmitToken whenever one
// is configured. Re-registering an ID that's already claimed requires that
// worker's own current token, so a caller can't self-declare someone else's
// id and hijack an existing worker's identity/capacity.
func (r *Registry) Register(w Worker, admitToken string) error {
	if w.ID == "" {
		return errors.New("worker id required")
	}
	if w.Token == "" {
		return errors.New("worker token required")
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if r.AdmitToken != "" && admitToken != r.AdmitToken {
		return ErrUnauthorized
	}
	if existing, ok := r.workers[w.ID]; ok && existing.Token != w.Token {
		return ErrUnauthorized
	}
	w.LastSeen = time.Now().UTC()
	w.Online = true
	r.workers[w.ID] = w
	if _, ok := r.cursors[w.ID]; !ok {
		r.cursors[w.ID] = 0
	}
	return nil
}
func (r *Registry) Authenticate(id, token string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	w, ok := r.workers[id]
	if !ok {
		return ErrUnknownWorker
	}
	if w.Token == "" || token == "" || w.Token != token {
		return ErrUnauthorized
	}
	return nil
}
func (r *Registry) Cursor(id string) (uint64, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if _, ok := r.workers[id]; !ok {
		return 0, ErrUnknownWorker
	}
	return r.cursors[id], nil
}
func (r *Registry) Ack(id string, cursor uint64) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if _, ok := r.workers[id]; !ok {
		return ErrUnknownWorker
	}
	if cursor < r.cursors[id] {
		return errors.New("cursor moved backwards")
	}
	r.cursors[id] = cursor
	return nil
}
func (r *Registry) Heartbeat(id string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	w, ok := r.workers[id]
	if !ok {
		return ErrUnknownWorker
	}
	w.LastSeen = time.Now().UTC()
	w.Online = true
	r.workers[id] = w
	return nil
}

// Available refreshes TTL state and reports whether a registered worker owns
// this harness id. Router admission uses it so a reachable TCP bridge alone
// can never make an offline worker eligible for a lease.
func (r *Registry) Available(id string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	w, ok := r.workers[id]
	if !ok {
		return false
	}
	w.Online = time.Since(w.LastSeen) <= r.TTL
	r.workers[id] = w
	return w.Online
}
func (r *Registry) Snapshot() []Worker {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	now := time.Now().UTC()
	out := make([]Worker, 0, len(r.workers))
	for id, w := range r.workers {
		wasOnline := w.Online
		w.Online = now.Sub(w.LastSeen) <= r.TTL
		r.workers[id] = w
		if wasOnline && !w.Online && r.OnOffline != nil {
			go r.OnOffline(w)
		}
		out = append(out, w)
	}
	return out
}