feat(worker): enforce durable trust state

internal/distworker/runner_protocol_test.go

@@ -1,8 +1,13 @@
 package distworker
 
 import (
+	"net/http"
+	"net/http/httptest"
+	"sync"
 	"testing"
+	"time"
 
+	"github.com/gorilla/websocket"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -18,6 +23,44 @@ func TestEnqueueTaskMessagePrefersEnvelopeOverLegacyCheck(t *testing.T) {
 	assert.Empty(t, r.jobQueue, "the sibling legacy field must not run a second check")
 }
 
+func TestDrainFinalResultsBoundsAndReportsStaleLeases(t *testing.T) {
+	upgrader := websocket.Upgrader{}
+	reports := make(chan wire.WorkerMessage, finalDrainLimit+2)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		conn, err := upgrader.Upgrade(w, req, nil)
+		require.NoError(t, err)
+		defer conn.Close()
+		for {
+			var message wire.WorkerMessage
+			if err := conn.ReadJSON(&message); err != nil {
+				return
+			}
+			reports <- message
+		}
+	}))
+	defer server.Close()
+	wsURL := "ws" + server.URL[len("http"):]
+	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	require.NoError(t, err)
+	defer conn.Close()
+	r := &Runner{controlConn: conn, controlWriteMu: &sync.Mutex{}, results: make(chan resultEnvelope, 2), notifyResults: make(chan notifyResultEnvelope, 2), leases: map[string]string{"stale": "lease-stale"}}
+	r.results <- resultEnvelope{job: wire.CheckJob{JobID: "done", LeaseToken: "lease-done"}, reports: []wire.CheckResultReport{{JobID: "done"}}}
+	r.drainFinalResults()
+	found := false
+	deadline := time.After(time.Second)
+	for !found {
+		select {
+		case message := <-reports:
+			if message.StaleLease != nil {
+				assert.Equal(t, "stale", message.StaleLease.JobID)
+				found = true
+			}
+		case <-deadline:
+			t.Fatal("missing stale lease report")
+		}
+	}
+}
+
 func TestEnqueueTaskMessageRejectsMismatchedEnvelopeJobIDs(t *testing.T) {
 	r := &Runner{jobQueue: make(chan wire.CheckJob, 1), results: make(chan resultEnvelope, 1), stopCh: make(chan struct{})}
 	message := wire.WorkerMessage{Kind: "task", TaskEnvelope: &wire.TaskEnvelope{
@@ -170,3 +213,11 @@ func TestEnqueueTaskMessageRejectsEmptyLeaseWithoutSideEffects(t *testing.T) {
 		})
 	}
 }
+
+func TestEnqueueTaskMessageRejectsCrossAccountTask(t *testing.T) {
+	r := &Runner{jobQueue: make(chan wire.CheckJob, 1), results: make(chan resultEnvelope, 1), stopCh: make(chan struct{}), state: workerState{AccountID: 4}}
+	require.True(t, r.enqueueTaskMessage(wire.WorkerMessage{Kind: "task", Task: &wire.CheckJob{JobID: "job", LeaseToken: "lease", AccountID: 5}}))
+	env := <-r.results
+	require.NotNil(t, env.reports[0].Error)
+	assert.Equal(t, "account_scope_mismatch", *env.reports[0].Error)
+}