kami/orchestra
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

feat(federation): admission control on worker registration (S10)

Browse Source 
Register() previously trusted a self-declared id and self-chosen token
from any caller, and let a second caller silently hijack an existing
worker id by re-registering it with a different token. Adds an optional
pre-shared AdmitToken (ORCHESTRA_FEDERATION_ADMIT_TOKEN) and requires a
same-id re-registration to present the existing worker's own token.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W1rkJ2hBMybnJctPbcy4tT
This commit is contained in:
kami
2026-07-27 23:34:02 +04:00
parent fab9225a78
commit 8b4955a687
5 changed files with 85 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/orchestra/main.go
+8 -3
View File
@@ -128,7 +128,7 @@ func main() {
}
}
mux := http.NewServeMux()
workers := &federation.Registry{}
workers := &federation.Registry{AdmitToken: os.Getenv("ORCHESTRA_FEDERATION_ADMIT_TOKEN")}
workers.OnOffline = func(w federation.Worker) {
for _, t := range s.Tasks() {
if t.State == domain.StateLeased && t.Lease != nil && t.Lease.HarnessID == w.ID {
@@ -703,8 +703,13 @@ func main() {
http.Error(w, "token required", 400)
return
}
if err := workers.Register(worker); err != nil {
http.Error(w, err.Error(), 400)
admitToken := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
if err := workers.Register(worker, admitToken); err != nil {
status := 400
if err == federation.ErrUnauthorized {
status = 401
}
http.Error(w, err.Error(), status)
return
}
w.WriteHeader(http.StatusCreated)
internal/federation/federation.go
+24 -6
View File
@@ -19,11 +19,16 @@ type Worker struct {
}
type Registry struct {
mu sync.Mutex
workers map[string]Worker
TTL time.Duration
OnOffline func(Worker)
cursors map[string]uint64
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
}
func (r *Registry) init() {
@@ -37,13 +42,26 @@ func (r *Registry) init() {
r.cursors = map[string]uint64{}
}
}
func (r *Registry) Register(w Worker) error {
// 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
internal/federation/federation_test.go
+23 -2
View File
@@ -7,7 +7,7 @@ import (
func TestCursorIsMonotonicAndAuthenticationIsRequired(t *testing.T) {
r := &Registry{}
if err := r.Register(Worker{ID: "workpc", Token: "secret"}); err != nil {
if err := r.Register(Worker{ID: "workpc", Token: "secret"}, ""); err != nil {
t.Fatal(err)
}
if err := r.Authenticate("workpc", "wrong"); err != ErrUnauthorized {
@@ -27,10 +27,31 @@ func TestCursorIsMonotonicAndAuthenticationIsRequired(t *testing.T) {
}
}
func TestRegisterRequiresAdmitTokenAndOwnToken(t *testing.T) {
r := &Registry{AdmitToken: "admit-secret"}
if err := r.Register(Worker{ID: "workpc", Token: "secret"}, "wrong"); err != ErrUnauthorized {
t.Fatalf("wrong admit token: got %v", err)
}
if err := r.Register(Worker{ID: "workpc", Token: "secret"}, "admit-secret"); err != nil {
t.Fatal(err)
}
// Re-registering the same id with a different token is a hijack
// attempt (S10), not a legitimate re-registration, and must be refused
// even with a valid admit token.
if err := r.Register(Worker{ID: "workpc", Token: "different"}, "admit-secret"); err != ErrUnauthorized {
t.Fatalf("hijack with different token: got %v", err)
}
// The same worker re-registering with its own token (e.g. after a
// restart) must still succeed.
if err := r.Register(Worker{ID: "workpc", Token: "secret"}, "admit-secret"); err != nil {
t.Fatalf("legitimate re-registration: %v", err)
}
}
func TestOfflineHookRunsOnceOnTransition(t *testing.T) {
called := make(chan Worker, 1)
r := &Registry{TTL: time.Millisecond, OnOffline: func(w Worker) { called <- w }}
if err := r.Register(Worker{ID: "workpc"}); err != nil {
if err := r.Register(Worker{ID: "workpc", Token: "secret"}, ""); err != nil {
t.Fatal(err)
}
r.mu.Lock()
internal/integration/federation_lease_test.go
+1 -1
View File
@@ -80,7 +80,7 @@ func TestCrossMachineLeaseAnchorAndQuotaArePerHost(t *testing.T) {
// worker (spec §2.1): it must be registered and reachable before the
// router would ever consider leasing to it.
workers := &federation.Registry{}
if err := workers.Register(federation.Worker{ID: "workpc", Address: "workpc.mesh", Token: "secret"}); err != nil {
if err := workers.Register(federation.Worker{ID: "workpc", Address: "workpc.mesh", Token: "secret"}, ""); err != nil {
t.Fatal(err)
}
if err := workers.Heartbeat("workpc"); err != nil {
progress.md
+29 -1
View File
@@ -245,7 +245,35 @@ Fixed so far:
existing pattern rather than introducing a one-off test harness for one
handler.
Not yet started: Codex/opencode Stop-hook-equivalent scripts, S8–S11. See
- **S9** — the coordinator's `Monitor` loop (30s ticker, `Coordinator.expire`)
and `main.go`'s own 1s reclaim ticker both called `Store.ExpireLeases`
independently. AUDIT.md filed this as "harmless — CAS rejects the loser,"
but the real effect was worse: the coordinator's `expire()` is the *only*
place that kills the herdr session/pane for an expired lease, and since the
1s ticker ran 30x more often it almost always expired the lease first,
leaving the coordinator's own `ExpireLeases` call with nothing left to
expire — so its pane-kill path silently never ran, orphaning herdr panes
past their TTL whenever a coordinator was configured. Fixed by having
`main.go`'s ticker skip `ExpireLeases` entirely when `coordinator != nil`
and defer reclaim to the coordinator's loop, keeping only `AssignPending`
as a periodic retry. No coordinator (e.g. no herdrs configured) still uses
the direct `ExpireLeases` path, since nothing else would reclaim leases in
that case.
- **S10** — `federation.Registry.Register` accepted a self-declared `id` and
self-chosen `token` from any caller with no admission control, and
silently let a second caller re-register an existing worker id with a
*different* token, hijacking that worker's identity/capacity out from
under it. Added `Registry.AdmitToken` (a pre-shared secret, wired from
`ORCHESTRA_FEDERATION_ADMIT_TOKEN`, checked against the registration
request's `Authorization: Bearer` header in `main.go`) and a same-ID
re-registration now requires presenting the existing worker's own token.
Covered by `TestRegisterRequiresAdmitTokenAndOwnToken`
(`internal/federation/federation_test.go`): wrong admit token rejected,
correct admit token accepted, same-id-different-token rejected as a
hijack, same-id-same-token (legitimate restart) still succeeds.
Not yet started: Codex/opencode Stop-hook-equivalent scripts, S8, S11. See
`AUDIT.md` for the full plan.
**Phase 0 done (2026-07-27):** this box has live TCP reachability to the real