Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483package 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"
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 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)
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_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)
// bystander lease the reporting node owns, 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_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 deliveredfunc 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)
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_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_SUCCESS { t.Fatalf("delivered terminal status = %v, want %v", got, millv1.TerminalStatus_SUCCESS) } if owned.getState() != leaseDone { t.Fatal("owned lease was not sealed after its terminal was delivered") }}
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)
foreign := newLease("lease-foreign", "node-x", "inc-x", "dummy") foreign.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "foreign"}, Name: "build"} owned := newLease("lease-owned", "node-z", "inc-z", "dummy") owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "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_RUNNING, }, }, }, }, }) if _, err := bdb.GetStatus(foreign.wid); err == nil { t.Fatal("status stream for a foreign lease authored a status row; an executor forged another pipeline's status") }
_ = 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_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.Status != "running" { t.Fatalf("owned status = %q, want %q", st.Status, "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)
const secret = "test-secret" if err := bdb.AddExecutorToken("dev-node", HashToken(secret), nil, authorizedLabels); err != nil { t.Fatalf("AddExecutorToken: %v", err) }
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", }} 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", }} 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 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", }} 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", }} 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) } }}