package mill
import (
"context"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"path/filepath"
"strings"
"testing"
"time"
"github.com/gorilla/websocket"
"tangled.org/core/notifier"
"tangled.org/core/spindle/db"
"tangled.org/core/spindle/models"
"tangled.org/core/spindle/observability"
millproto "tangled.org/core/spindle/mill/proto"
millv1 "tangled.org/core/spindle/mill/proto/gen"
)
func discardLogger() *slog.Logger {
return slog.New(slog.NewTextHandler(io.Discard, nil))
}
func nopEncoder() scriptedEncoder {
return scriptedEncoder(func(*millproto.Message) error { return nil })
}
func registerTestExecutor(t *testing.T, d *db.DB, name, tokenHash string, labels []string) {
t.Helper()
if err := d.CreateExecutorToken(tokenHash, nil); err != nil {
t.Fatalf("CreateExecutorToken: %v", err)
}
if err := d.RegisterExecutor(name, tokenHash, labels); err != nil {
t.Fatalf("RegisterExecutor: %v", err)
}
}
func TestHashToken(t *testing.T) {
const raw = "super-secret-executor-token"
if HashToken(raw) != HashToken(raw) {
t.Fatal("HashToken is not deterministic; the same token would stop authenticating")
}
if HashToken("token-a") == HashToken("token-b") {
t.Fatal("HashToken collided two distinct tokens")
}
if HashToken(raw) == raw {
t.Fatal("HashToken returned the raw token; a hash leak would expose a usable credential")
}
}
func TestGenerateTokenDistinct(t *testing.T) {
const n = 100
seen := make(map[string]struct{}, n)
for i := range n {
tok, err := GenerateToken()
if err != nil {
t.Fatalf("GenerateToken: %v", err)
}
if tok == "" {
t.Fatalf("GenerateToken returned an empty token on call %d", i)
}
if _, dup := seen[tok]; dup {
t.Fatalf("GenerateToken repeated a token after %d calls: %q", i, tok)
}
seen[tok] = struct{}{}
}
}
func TestAttachSessionRejectsSecondLiveSession(t *testing.T) {
l := discardLogger()
m := New(l, Config{ReconnectGrace: time.Minute})
sessionOf := func(node string) *millSession {
m.mu.Lock()
defer m.mu.Unlock()
return m.sessions[node]
}
sess1 := newSession("node-1", "inc-1", nil, nopEncoder(), l)
if _, ok := m.attachSession(sess1); !ok {
t.Fatal("first attach of a node was rejected; want accept")
}
if sessionOf("node-1") != sess1 {
t.Fatal("first session was not registered as the live session")
}
sess2 := newSession("node-1", "inc-2", nil, nopEncoder(), l)
if _, ok := m.attachSession(sess2); ok {
t.Fatal("second live attach for an already-live node was accepted; a valid token hijacked the executor")
}
if sessionOf("node-1") != sess1 {
t.Fatal("rejected newcomer evicted the incumbent session")
}
m.detachSession(sess1)
sess3 := newSession("node-1", "inc-3", nil, nopEncoder(), l)
if _, ok := m.attachSession(sess3); !ok {
t.Fatal("attach during the incumbent's reconnect grace was rejected; want adopt")
}
if sessionOf("node-1") != sess3 {
t.Fatal("adopted session was not installed as the live session")
}
}
func TestOnAttemptResultIgnoresForeignLease(t *testing.T) {
ctx := context.Background()
l := discardLogger()
bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db"))
if err != nil {
t.Fatalf("db.Make: %v", err)
}
t.Cleanup(func() { bdb.Close() })
n := notifier.New()
m := New(l, Config{ReconnectGrace: time.Minute})
m.Attach(bdb, &n, testQuotaManager(t, bdb))
foreign := newLease("lease-foreign", "node-a", "inc-a", "dummy")
m.mu.Lock()
m.leases[foreign.id] = foreign
m.mu.Unlock()
sessB := newSession("node-b", "inc-b", nil, nopEncoder(), l)
m.attachSession(sessB)
_ = m.onEventBatch(sessB, &millv1.EventBatch{
Epoch: sessB.epoch,
Events: []*millv1.Event{
{
Seqno: 1,
LeaseId: foreign.id,
Payload: &millv1.Event_AttemptResult{
AttemptResult: &millv1.AttemptResult{
Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS,
},
},
},
},
})
if _, ok := pollTerminal(foreign); ok {
t.Fatal("attempt-result on a foreign lease delivered a terminal; an executor forged another node's job result")
}
if foreign.getState() == leaseDone {
t.Fatal("attempt-result on a foreign lease sealed the lease")
}
}
func TestOnAttemptResultIgnoresAbsentLease(t *testing.T) {
ctx := context.Background()
l := discardLogger()
bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db"))
if err != nil {
t.Fatalf("db.Make: %v", err)
}
t.Cleanup(func() { bdb.Close() })
n := notifier.New()
m := New(l, Config{ReconnectGrace: time.Minute})
m.Attach(bdb, &n, testQuotaManager(t, bdb))
// the reporting node owns this bystander lease, proving an absent-lease
// stream does not spill onto another lease
bystander := newLease("lease-bystander", "node-b", "inc-b", "dummy")
m.mu.Lock()
m.leases[bystander.id] = bystander
m.mu.Unlock()
sessB := newSession("node-b", "inc-b", nil, nopEncoder(), l)
m.attachSession(sessB)
_ = m.onEventBatch(sessB, &millv1.EventBatch{
Epoch: sessB.epoch,
Events: []*millv1.Event{
{
Seqno: 1,
LeaseId: "lease-nonexistent",
Payload: &millv1.Event_AttemptResult{
AttemptResult: &millv1.AttemptResult{
Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS,
},
},
},
},
})
if _, ok := pollTerminal(bystander); ok {
t.Fatal("attempt-result for an absent lease delivered a terminal to a bystander lease")
}
if bystander.getState() == leaseDone {
t.Fatal("attempt-result for an absent lease sealed a bystander lease")
}
}
// owned-lease path proves ignore tests above are not passing merely because
// delivery is broken. correctly owned terminal is delivered
func TestOnAttemptResultDeliversOwnedLease(t *testing.T) {
ctx := context.Background()
l := discardLogger()
bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db"))
if err != nil {
t.Fatalf("db.Make: %v", err)
}
t.Cleanup(func() { bdb.Close() })
n := notifier.New()
m := New(l, Config{ReconnectGrace: time.Minute})
m.Attach(bdb, &n, testQuotaManager(t, bdb))
metrics := observability.NewMetrics()
m.RegisterMetrics(metrics)
owned := newLease("lease-owned", "node-b", "inc-b", "dummy")
m.mu.Lock()
m.leases[owned.id] = owned
m.mu.Unlock()
sessB := newSession("node-b", "inc-b", nil, nopEncoder(), l)
m.attachSession(sessB)
_ = m.onEventBatch(sessB, &millv1.EventBatch{
Epoch: sessB.epoch,
Events: []*millv1.Event{
{
Seqno: 1,
LeaseId: owned.id,
Payload: &millv1.Event_AttemptResult{
AttemptResult: &millv1.AttemptResult{
Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS,
},
},
},
},
})
res, ok := pollTerminal(owned)
if !ok {
t.Fatal("attempt-result on an owned lease was not delivered")
}
if got := res.GetStatus(); got != millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS {
t.Fatalf("delivered terminal status = %v, want %v", got, millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS)
}
if owned.getState() != leaseDone {
t.Fatal("owned lease was not sealed after its terminal was delivered")
}
families, err := metrics.Registry().Gather()
if err != nil {
t.Fatal(err)
}
for _, family := range families {
if family.GetName() == "spindle_workflows_total" {
t.Fatal("mill counted a terminal whose executor retained metric authority")
}
}
}
func TestOnStatusEventOwnership(t *testing.T) {
ctx := context.Background()
l := discardLogger()
bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db"))
if err != nil {
t.Fatalf("db.Make: %v", err)
}
t.Cleanup(func() { bdb.Close() })
n := notifier.New()
m := New(l, Config{ReconnectGrace: time.Minute})
m.Attach(bdb, &n, testQuotaManager(t, bdb))
foreign := newLease("lease-foreign", "node-x", "inc-x", "dummy")
foreign.wid = models.WorkflowId{PipelineId: models.PipelineId("foreign"), Name: "build"}
owned := newLease("lease-owned", "node-z", "inc-z", "dummy")
owned.wid = models.WorkflowId{PipelineId: models.PipelineId("owned"), Name: "build"}
m.mu.Lock()
m.leases[foreign.id] = foreign
m.leases[owned.id] = owned
m.mu.Unlock()
sessY := newSession("node-y", "inc-y", nil, nopEncoder(), l)
m.attachSession(sessY)
sessZ := newSession("node-z", "inc-z", nil, nopEncoder(), l)
m.attachSession(sessZ)
_ = m.onEventBatch(sessY, &millv1.EventBatch{
Epoch: sessY.epoch,
Events: []*millv1.Event{
{
Seqno: 1,
LeaseId: foreign.id,
Payload: &millv1.Event_StatusEvent{
StatusEvent: &millv1.StatusEvent{
Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING,
},
},
},
},
})
if st, err := bdb.GetStatus(foreign.wid); err != nil || st != "" {
t.Fatalf("status stream for a foreign lease authored status %q (err %v); an executor forged another pipeline's status", st, err)
}
_ = m.onEventBatch(sessZ, &millv1.EventBatch{
Epoch: sessZ.epoch,
Events: []*millv1.Event{
{
Seqno: 1,
LeaseId: owned.id,
Payload: &millv1.Event_StatusEvent{
StatusEvent: &millv1.StatusEvent{
Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING,
},
},
},
},
})
st, err := bdb.GetStatus(owned.wid)
if err != nil {
t.Fatalf("owned status stream did not author a status row: %v", err)
}
if st != "running" {
t.Fatalf("owned status = %q, want %q", st, "running")
}
}
func setupTestServer(t *testing.T, authorizedLabels []string) (*Mill, *db.DB, *httptest.Server, string) {
ctx := context.Background()
bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db"))
if err != nil {
t.Fatalf("db.Make: %v", err)
}
t.Cleanup(func() { bdb.Close() })
n := notifier.New()
m := New(discardLogger(), Config{
ReconnectGrace: time.Minute,
})
m.Attach(bdb, &n, testQuotaManager(t, bdb))
const secret = "test-secret"
registerTestExecutor(t, bdb, "dev-node", HashToken(secret), authorizedLabels)
server := httptest.NewServer(http.HandlerFunc(m.HandleExecutorConn))
t.Cleanup(server.Close)
return m, bdb, server, secret
}
func TestAuthLabelEscalation(t *testing.T) {
_, _, server, secret := setupTestServer(t, []string{"linux", "amd64"})
wsUrl := "ws" + strings.TrimPrefix(server.URL, "http")
{
header := http.Header{}
header.Set("Authorization", "Bearer bad-token")
_, resp, err := websocket.DefaultDialer.Dial(wsUrl, header)
if err == nil {
t.Fatal("expected connection with invalid token to fail")
}
if resp != nil && resp.StatusCode != http.StatusUnauthorized {
t.Fatalf("expected 401 Unauthorized, got %d", resp.StatusCode)
}
}
{
header := http.Header{}
header.Set("Authorization", "Bearer "+secret)
conn, _, err := websocket.DefaultDialer.Dial(wsUrl, header)
if err != nil {
t.Fatalf("dial failed: %v", err)
}
defer conn.Close()
stream := millproto.NewWSStream(conn)
enc := millproto.NewEncoder(stream)
dec := millproto.NewDecoder(stream)
hello := &millproto.Message{Hello: &millv1.Hello{ProtocolVersion: millproto.ProtocolVersion,
Arch: "amd64",
Labels: []string{"linux", "gpu"},
Epoch: "inc-1",
CacheNamespace: "tar-zstd-v1/linux/amd64",
}}
if err := enc.Encode(hello); err != nil {
t.Fatalf("encode hello: %v", err)
}
_, err = dec.Decode()
if err == nil {
t.Fatal("expected server to close connection for unauthorized label, but got a message")
}
}
{
header := http.Header{}
header.Set("Authorization", "Bearer "+secret)
conn, _, err := websocket.DefaultDialer.Dial(wsUrl, header)
if err != nil {
t.Fatalf("dial failed: %v", err)
}
defer conn.Close()
stream := millproto.NewWSStream(conn)
enc := millproto.NewEncoder(stream)
dec := millproto.NewDecoder(stream)
hello := &millproto.Message{Hello: &millv1.Hello{ProtocolVersion: millproto.ProtocolVersion,
Arch: "amd64",
Labels: []string{"linux"},
Epoch: "inc-1",
CacheNamespace: "tar-zstd-v1/linux/amd64",
}}
if err := enc.Encode(hello); err != nil {
t.Fatalf("encode hello: %v", err)
}
msg, err := dec.Decode()
if err != nil {
t.Fatalf("expected resume message, got error: %v", err)
}
res := msg.GetResume()
if res == nil {
t.Fatal("expected Resume message, got nil")
}
if res.GetEpoch() != "inc-1" {
t.Fatalf("expected epoch inc-1, got %q", res.GetEpoch())
}
}
}
func TestProtocolVersionMismatchClosesConnection(t *testing.T) {
_, _, server, secret := setupTestServer(t, []string{"linux"})
header := http.Header{}
header.Set("Authorization", "Bearer "+secret)
conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), header)
if err != nil {
t.Fatalf("dial failed: %v", err)
}
defer conn.Close()
stream := millproto.NewWSStream(conn)
enc := millproto.NewEncoder(stream)
if err := enc.Encode(&millproto.Message{Hello: &millv1.Hello{
ProtocolVersion: millproto.ProtocolVersion - 1,
Arch: "amd64",
Labels: []string{"linux"},
Epoch: "inc-1",
CacheNamespace: "tar-zstd-v1/linux/amd64",
}}); err != nil {
t.Fatalf("encode hello: %v", err)
}
if _, err := millproto.NewDecoder(stream).Decode(); err == nil {
t.Fatal("expected server to close a stale protocol connection")
}
}
func TestHandshakeTimeoutAndConcurrency(t *testing.T) {
_, _, server, secret := setupTestServer(t, []string{"linux"})
wsUrl := "ws" + strings.TrimPrefix(server.URL, "http")
// executor that never sends Hello is dropped after the 5s pre-hello deadline
{
header := http.Header{}
header.Set("Authorization", "Bearer "+secret)
conn, _, err := websocket.DefaultDialer.Dial(wsUrl, header)
if err != nil {
t.Fatalf("dial failed: %v", err)
}
defer conn.Close()
time.Sleep(6 * time.Second)
stream := millproto.NewWSStream(conn)
enc := millproto.NewEncoder(stream)
hello := &millproto.Message{Hello: &millv1.Hello{ProtocolVersion: millproto.ProtocolVersion,
Arch: "amd64",
Labels: []string{"linux"},
Epoch: "inc-1",
CacheNamespace: "tar-zstd-v1/linux/amd64",
}}
err = enc.Encode(hello)
dec := millproto.NewDecoder(stream)
_, readErr := dec.Decode()
if readErr == nil {
t.Fatal("expected server to have closed connection due to handshake timeout")
}
}
// second live session for one identity is rejected with 409 before the ws upgrade
{
header := http.Header{}
header.Set("Authorization", "Bearer "+secret)
conn1, _, err := websocket.DefaultDialer.Dial(wsUrl, header)
if err != nil {
t.Fatalf("dial 1 failed: %v", err)
}
defer conn1.Close()
stream1 := millproto.NewWSStream(conn1)
enc1 := millproto.NewEncoder(stream1)
dec1 := millproto.NewDecoder(stream1)
hello1 := &millproto.Message{Hello: &millv1.Hello{ProtocolVersion: millproto.ProtocolVersion,
Arch: "amd64",
Labels: []string{"linux"},
Epoch: "inc-1",
CacheNamespace: "tar-zstd-v1/linux/amd64",
}}
if err := enc1.Encode(hello1); err != nil {
t.Fatalf("encode hello 1: %v", err)
}
_, err = dec1.Decode()
if err != nil {
t.Fatalf("first connection handshake failed: %v", err)
}
_, resp, err := websocket.DefaultDialer.Dial(wsUrl, header)
if err == nil {
t.Fatal("expected second connection for same live identity to be rejected")
}
if resp != nil && resp.StatusCode != http.StatusConflict {
t.Fatalf("expected 409 Conflict for duplicate session, got %d", resp.StatusCode)
}
}
}