feat(federation): admission control on worker registration (S10)
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:
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user