feat(worker): enforce durable trust state

internal/distworker/runner.go

@@ -2,10 +2,13 @@ package distworker
 
 import (
 	"context"
+	"crypto/ed25519"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"log"
 	"os"
+	"path/filepath"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -27,6 +30,8 @@ const (
 	minQueueCapacity = 16
 
 	malformedTaskEnvelopeError = "malformed_task_envelope"
+	finalDrainLimit            = 32
+	finalDrainTimeout          = 5 * time.Second
 )
 
 // jobPool is the minimal interface the Runner needs from a worker pool. It
@@ -176,6 +181,8 @@ type Runner struct {
 	workerVersion        string
 	capsMu               sync.RWMutex
 	workerCaps           []string
+	identityMu           sync.RWMutex
+	state                workerState
 	serverID             atomic.Int64
 	metricGeneration     atomic.Uint64
 	nextMetricGeneration atomic.Uint64
@@ -188,6 +195,8 @@ type Runner struct {
 	controlWriteMu *sync.Mutex
 	reconnectCh    chan struct{}
 	rotationMu     sync.Mutex
+	leasesMu       sync.Mutex
+	leases         map[string]string
 
 	// outbox holds messages removed from a per-connection writer queue that
 	// could not be written before its websocket closed.
@@ -231,6 +240,7 @@ func NewRunner(cfg *Config) *Runner {
 		resultsBuf:       newResultBuffer(),
 		notificationsBuf: newNotificationBuffer(),
 		peerCache:        newPeerCache(),
+		leases:           make(map[string]string),
 	}
 }
 
@@ -239,7 +249,7 @@ func (r *Runner) Start() error {
 	log.Println("worker: starting...")
 	r.lifecycleMu.Lock()
 
-	if r.config.URL == "" || r.config.Token == "" {
+	if r.config.URL == "" || (r.config.Token == "" && r.config.BootstrapToken == "") {
 		r.lifecycleMu.Unlock()
 		return fmt.Errorf("RSMON_URL and RSMON_TOKEN must be set")
 	}
@@ -253,6 +263,34 @@ func (r *Runner) Start() error {
 		return fmt.Errorf("worker: runner stopped")
 	}
 
+	if r.config.BootstrapToken != "" && (r.config.StateFile == "" || !filepath.IsAbs(r.config.StateFile)) {
+		r.lifecycleMu.Unlock()
+		return fmt.Errorf("bootstrap requires an absolute RSMON_STATE_FILE")
+	}
+	state, err := loadWorkerState(r.config.StateFile)
+	if err != nil {
+		r.lifecycleMu.Unlock()
+		return fmt.Errorf("worker state: %w", err)
+	}
+	if state.Token != "" {
+		r.config.Token = state.Token
+	}
+	if r.config.Token == "" {
+		client := NewClient(r.config.URL, "")
+		boot, err := client.Bootstrap(context.Background(), r.config.WorkerID, r.config.BootstrapToken)
+		if err != nil {
+			r.lifecycleMu.Unlock()
+			return fmt.Errorf("worker bootstrap: %w", err)
+		}
+		state = workerState{Token: boot.AuthToken, WorkerID: boot.WorkerID, VerificationKey: boot.ConfigVerificationKey, SigningKeyID: boot.SigningKeyID}
+		if err := saveWorkerState(r.config.StateFile, state); err != nil {
+			r.lifecycleMu.Unlock()
+			return fmt.Errorf("worker state: %w", err)
+		}
+		r.config.Token = boot.AuthToken
+		r.config.BootstrapToken = ""
+	}
+	r.state = state
 	r.clientMu.Lock()
 	r.client = NewClient(r.config.URL, r.config.Token)
 	r.clientMu.Unlock()
@@ -344,6 +382,7 @@ func (r *Runner) Start() error {
 // Stop gracefully stops the worker
 func (r *Runner) Stop() {
 	r.lifecycleMu.Lock()
+	r.drainFinalResults()
 	select {
 	case <-r.stopCh:
 		// already closed
@@ -681,7 +720,21 @@ func (r *Runner) runWebsocket() error {
 			return err
 		}
 		if (msg.Kind == "init" || msg.Kind == "config") && msg.Init != nil {
-			r.applyInit(msg.Init)
+			if err := r.applyInit(msg.Init); err != nil {
+				log.Printf("worker: rejected config: %v", err)
+				continue
+			}
+			if msg.Init.RotateToken != "" && msg.Init.RotationID != "" {
+				r.applyRotation(conn, &writeMu, msg.Init.RotateToken, msg.Init.RotationID)
+			}
+			continue
+		}
+		if msg.Kind == "stale_lease_ack" && msg.StaleLeaseAck != nil {
+			r.leasesMu.Lock()
+			if r.leases[msg.StaleLeaseAck.JobID] == msg.StaleLeaseAck.LeaseToken {
+				delete(r.leases, msg.StaleLeaseAck.JobID)
+			}
+			r.leasesMu.Unlock()
 			continue
 		}
 		if msg.Kind != "task" {
@@ -709,6 +762,9 @@ func (r *Runner) enqueueTaskMessage(msg wire.WorkerMessage) bool { //nolint:gocr
 		if msg.NotificationTask.JobID == "" || msg.NotificationTask.LeaseToken == "" {
 			return true
 		}
+		if !r.taskInScope(msg.NotificationTask.AccountID) {
+			return r.enqueueFailedNotification(*msg.NotificationTask, "account_scope_mismatch")
+		}
 		log.Printf("worker: received websocket notification task %s method=%s", msg.NotificationTask.JobID, msg.NotificationTask.Method)
 		return r.EnqueueNotification(*msg.NotificationTask)
 	case msg.Task != nil:
@@ -716,6 +772,9 @@ func (r *Runner) enqueueTaskMessage(msg wire.WorkerMessage) bool { //nolint:gocr
 			return true
 		}
 		log.Printf("worker: received websocket task %s", msg.Task.JobID)
+		if !r.taskInScope(msg.Task.AccountID) {
+			return r.enqueueFailedCheck(*msg.Task, "account_scope_mismatch")
+		}
 		if !checkexec.SupportsKind(msg.Task.Kind) {
 			return r.enqueueUnsupportedCheck(*msg.Task)
 		}
@@ -733,12 +792,16 @@ func (r *Runner) enqueueTaskEnvelope(envelope *wire.TaskEnvelope) bool {
 		if !matchingEnvelopeJobID(envelope.JobID, envelope.Job.JobID) || envelope.Job.LeaseToken == "" {
 			return true
 		}
+		if !r.taskInScope(envelope.Job.AccountID) {
+			return r.enqueueFailedCheck(*envelope.Job, "account_scope_mismatch")
+		}
 		if envelope.Type != wire.TaskTypeCheck {
 			return r.enqueueFailedCheck(*envelope.Job, malformedTaskEnvelopeError)
 		}
 		if !checkexec.SupportsKind(envelope.Job.Kind) {
 			return r.enqueueUnsupportedCheck(*envelope.Job)
 		}
+		r.rememberLease(envelope.Job.JobID, envelope.Job.LeaseToken)
 		return r.Enqueue(*envelope.Job)
 	}
 	if !matchingEnvelopeJobID(envelope.JobID, envelope.Notify.JobID) || envelope.Notify.LeaseToken == "" {
@@ -747,9 +810,97 @@ func (r *Runner) enqueueTaskEnvelope(envelope *wire.TaskEnvelope) bool {
 	if envelope.Type != wire.TaskTypeNotification {
 		return r.enqueueFailedNotification(*envelope.Notify, malformedTaskEnvelopeError)
 	}
+	if !r.taskInScope(envelope.Notify.AccountID) {
+		return r.enqueueFailedNotification(*envelope.Notify, "account_scope_mismatch")
+	}
+	r.rememberLease(envelope.Notify.JobID, envelope.Notify.LeaseToken)
 	return r.EnqueueNotification(*envelope.Notify)
 }
 
+func (r *Runner) rememberLease(jobID, token string) {
+	if jobID == "" || token == "" {
+		return
+	}
+	r.leasesMu.Lock()
+	if r.leases == nil {
+		r.leases = make(map[string]string)
+	}
+	r.leases[jobID] = token
+	r.leasesMu.Unlock()
+}
+
+func (r *Runner) forgetLease(jobID string) {
+	r.leasesMu.Lock()
+	delete(r.leases, jobID)
+	r.leasesMu.Unlock()
+}
+
+func (r *Runner) drainFinalResults() {
+	// Shutdown keeps the existing connection just long enough to report a
+	// bounded set of terminal frames and leases that cannot complete.
+	r.clientMu.Lock()
+	conn, writeMu := r.controlConn, r.controlWriteMu
+	r.clientMu.Unlock()
+	if conn == nil || writeMu == nil {
+		return
+	}
+	deadline := time.Now().Add(finalDrainTimeout)
+	count := 0
+	r.leasesMu.Lock()
+	leases := make(map[string]string, len(r.leases))
+	for id, token := range r.leases {
+		leases[id] = token
+	}
+	r.leasesMu.Unlock()
+	for count < finalDrainLimit && time.Now().Before(deadline) {
+		select {
+		case env := <-r.results:
+			for i := range env.reports {
+				env.reports[i].LeaseToken = env.job.LeaseToken
+				if r.writeControlMessage(conn, writeMu, wire.WorkerMessage{Kind: "result", Result: &env.reports[i]}) != nil {
+					return
+				}
+				r.forgetLease(env.job.JobID)
+				count++
+				if count >= finalDrainLimit {
+					return
+				}
+			}
+		case env := <-r.notifyResults:
+			if r.writeControlMessage(conn, writeMu, wire.WorkerMessage{Kind: "result", NotificationResult: &env.report}) != nil {
+				return
+			}
+			r.forgetLease(env.task.JobID)
+			count++
+		default:
+			for id, token := range leases {
+				if count >= finalDrainLimit || time.Now().After(deadline) {
+					return
+				}
+				r.leasesMu.Lock()
+				_, pending := r.leases[id]
+				r.leasesMu.Unlock()
+				if pending {
+					if r.writeControlMessage(conn, writeMu, wire.WorkerMessage{Kind: "stale_lease", StaleLease: &wire.StaleLeaseReport{JobID: id, LeaseToken: token}}) != nil {
+						return
+					}
+					count++
+				}
+			}
+			return
+		}
+	}
+}
+
+func (r *Runner) taskInScope(accountID int64) bool {
+	if accountID == 0 {
+		return true
+	} // legacy task payload
+	r.identityMu.RLock()
+	defer r.identityMu.RUnlock()
+	return r.state.AccountID == 0 || r.state.AccountID == accountID
+}
+
 func matchingEnvelopeJobID(outer, inner string) bool {
 	return outer != "" && outer == inner
 }
@@ -888,6 +1039,7 @@ func (r *Runner) writer(conn *websocket.Conn, writeMu *sync.Mutex, done <-chan s
 					report.JobID, report.CheckID, env.job.Kind, report.State,
 				)
 				r.resultsBuf.add(resultRowFromReport(&env, report, time.Now().UTC()))
+				r.forgetLease(env.job.JobID)
 			}
 		case env, ok := <-r.notifyResults:
 			if !ok {
@@ -907,6 +1059,7 @@ func (r *Runner) writer(conn *websocket.Conn, writeMu *sync.Mutex, done <-chan s
 				"worker: completed notification job=%s method=%s status=%s message=%d duration_ms=%d",
 				env.report.JobID, env.task.Method, env.report.Status, env.task.MessageID, env.report.DurationMs,
 			)
+			r.forgetLease(env.task.JobID)
 		case metric := <-r.metricResults:
 			if metric.generation != metricGeneration {
 				continue
@@ -961,7 +1114,10 @@ func (r *Runner) writeControlMessage(conn *websocket.Conn, writeMu *sync.Mutex,
 	return conn.WriteJSON(message)
 }
 
-func (r *Runner) applyInit(init *wire.WorkerInit) {
+func (r *Runner) applyInit(init *wire.WorkerInit) error {
+	if err := r.verifyInit(init); err != nil {
+		return err
+	}
 	if init.Concurrency > 0 && init.Concurrency != r.Concurrency() {
 		size := init.Concurrency
 		if size > r.maxConcurrency {
@@ -1053,6 +1209,79 @@ func (r *Runner) applyInit(init *wire.WorkerInit) {
 		"config", init.WorkerID, init.RegionCode, init.Version, init.Capabilities,
 		init.Concurrency, len(init.LLMs), smtpCount, tgCount, len(init.SystemContacts), len(init.Peers),
 	)
+	return nil
+}
+
+func (r *Runner) verifyInit(init *wire.WorkerInit) error {
+	// Legacy control planes do not sign config. Do not turn existing workers
+	// into an outage; once bootstrap pins a key, signed config is mandatory.
+	r.identityMu.Lock()
+	defer r.identityMu.Unlock()
+	if r.state.VerificationKey == "" {
+		return nil
+	}
+	if init.Signature == "" || init.SigningKeyID == "" || init.SigningKeyID != r.state.SigningKeyID || init.AccountID == nil || init.WorkerID != r.state.WorkerID {
+		return fmt.Errorf("unsigned or mismatched config")
+	}
+	if r.state.AccountID != 0 && *init.AccountID != r.state.AccountID {
+		return fmt.Errorf("account scope changed")
+	}
+	if init.ConfigVersion < r.state.ConfigVersion {
+		return fmt.Errorf("config version downgrade")
+	}
+	expires, err := time.Parse(time.RFC3339Nano, init.ExpiresAt)
+	if err != nil || !expires.After(time.Now()) {
+		return fmt.Errorf("config expired")
+	}
+	key, err := base64.StdEncoding.DecodeString(r.state.VerificationKey)
+	if err != nil || len(key) != ed25519.PublicKeySize {
+		return fmt.Errorf("invalid pinned verification key")
+	}
+	sig, err := base64.StdEncoding.DecodeString(init.Signature)
+	if err != nil {
+		return fmt.Errorf("invalid config signature")
+	}
+	signature := init.Signature
+	init.Signature = ""
+	body, err := json.Marshal(init)
+	init.Signature = signature
+	if err != nil || !ed25519.Verify(ed25519.PublicKey(key), body, sig) {
+		return fmt.Errorf("invalid config signature")
+	}
+	// Issued/expiry timestamps are signed freshness metadata, regenerated on
+	// reconnect; they do not change immutable configuration content.
+	canonical := *init
+	canonical.Signature, canonical.IssuedAt, canonical.ExpiresAt = "", "", ""
+	canonicalBody, err := json.Marshal(&canonical)
+	if err != nil {
+		return err
+	}
+	payload := base64.StdEncoding.EncodeToString(canonicalBody)
+	if init.ConfigVersion == r.state.ConfigVersion && r.state.ConfigPayload != "" && r.state.ConfigPayload != payload {
+		return fmt.Errorf("config version replay content mismatch")
+	}
+	r.state.AccountID, r.state.ConfigVersion, r.state.ConfigPayload = *init.AccountID, init.ConfigVersion, payload
+	if err := saveWorkerState(r.config.StateFile, r.state); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (r *Runner) applyRotation(conn *websocket.Conn, writeMu *sync.Mutex, token, rotationID string) {
+	if token == "" || token == r.config.Token {
+		return
+	}
+	r.clientMu.Lock()
+	r.config.Token = token
+	r.client = NewClient(r.config.URL, token)
+	r.state.Token = token
+	r.clientMu.Unlock()
+	if err := saveWorkerState(r.config.StateFile, r.state); err != nil {
+		log.Printf("worker: rotation state: %v", err)
+		return
+	}
+	_ = r.writeControlMessage(conn, writeMu, wire.WorkerMessage{Kind: "rotation_ack", RotationAck: &wire.RotationAck{RotationID: rotationID}})
+	_ = conn.Close()
 }
 
 // Credentials returns a snapshot of the current notification credentials.