alice.pds.demo.boltless.dev
This repository has no description
1package mill
2
3import (
4 "context"
5 "io"
6 "log/slog"
7 "net/http"
8 "net/http/httptest"
9 "path/filepath"
10 "strings"
11 "testing"
12 "time"
13
14 "github.com/gorilla/websocket"
15 "tangled.org/core/notifier"
16 "tangled.org/core/spindle/db"
17 "tangled.org/core/spindle/models"
18
19 millproto "tangled.org/core/spindle/mill/proto"
20 millv1 "tangled.org/core/spindle/mill/proto/gen"
21)
22
23func discardLogger() *slog.Logger {
24 return slog.New(slog.NewTextHandler(io.Discard, nil))
25}
26
27func nopEncoder() scriptedEncoder {
28 return scriptedEncoder(func(*millproto.Message) error { return nil })
29}
30
31func TestHashToken(t *testing.T) {
32 const raw = "super-secret-executor-token"
33
34 if HashToken(raw) != HashToken(raw) {
35 t.Fatal("HashToken is not deterministic; the same token would stop authenticating")
36 }
37 if HashToken("token-a") == HashToken("token-b") {
38 t.Fatal("HashToken collided two distinct tokens")
39 }
40 if HashToken(raw) == raw {
41 t.Fatal("HashToken returned the raw token; a hash leak would expose a usable credential")
42 }
43}
44
45func TestGenerateTokenDistinct(t *testing.T) {
46 const n = 100
47 seen := make(map[string]struct{}, n)
48 for i := range n {
49 tok, err := GenerateToken()
50 if err != nil {
51 t.Fatalf("GenerateToken: %v", err)
52 }
53 if tok == "" {
54 t.Fatalf("GenerateToken returned an empty token on call %d", i)
55 }
56 if _, dup := seen[tok]; dup {
57 t.Fatalf("GenerateToken repeated a token after %d calls: %q", i, tok)
58 }
59 seen[tok] = struct{}{}
60 }
61}
62
63func TestAttachSessionRejectsSecondLiveSession(t *testing.T) {
64 l := discardLogger()
65 m := New(l, Config{ReconnectGrace: time.Minute})
66
67 sessionOf := func(node string) *millSession {
68 m.mu.Lock()
69 defer m.mu.Unlock()
70 return m.sessions[node]
71 }
72
73 sess1 := newSession("node-1", "inc-1", nil, nopEncoder(), l)
74 if _, ok := m.attachSession(sess1); !ok {
75 t.Fatal("first attach of a node was rejected; want accept")
76 }
77 if sessionOf("node-1") != sess1 {
78 t.Fatal("first session was not registered as the live session")
79 }
80
81 sess2 := newSession("node-1", "inc-2", nil, nopEncoder(), l)
82 if _, ok := m.attachSession(sess2); ok {
83 t.Fatal("second live attach for an already-live node was accepted; a valid token hijacked the executor")
84 }
85 if sessionOf("node-1") != sess1 {
86 t.Fatal("rejected newcomer evicted the incumbent session")
87 }
88
89 m.detachSession(sess1)
90 sess3 := newSession("node-1", "inc-3", nil, nopEncoder(), l)
91 if _, ok := m.attachSession(sess3); !ok {
92 t.Fatal("attach during the incumbent's reconnect grace was rejected; want adopt")
93 }
94 if sessionOf("node-1") != sess3 {
95 t.Fatal("adopted session was not installed as the live session")
96 }
97}
98
99func TestOnAttemptResultIgnoresForeignLease(t *testing.T) {
100 ctx := context.Background()
101 l := discardLogger()
102 bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db"))
103 if err != nil {
104 t.Fatalf("db.Make: %v", err)
105 }
106 t.Cleanup(func() { bdb.Close() })
107 n := notifier.New()
108
109 m := New(l, Config{ReconnectGrace: time.Minute})
110 m.Attach(bdb, &n)
111
112 foreign := newLease("lease-foreign", "node-a", "inc-a", "dummy")
113 m.mu.Lock()
114 m.leases[foreign.id] = foreign
115 m.mu.Unlock()
116
117 sessB := newSession("node-b", "inc-b", nil, nopEncoder(), l)
118 m.attachSession(sessB)
119 _ = m.onEventBatch(sessB, &millv1.EventBatch{
120 Epoch: sessB.epoch,
121 Events: []*millv1.Event{
122 {
123 Seqno: 1,
124 LeaseId: foreign.id,
125 Payload: &millv1.Event_AttemptResult{
126 AttemptResult: &millv1.AttemptResult{
127 Status: millv1.TerminalStatus_SUCCESS,
128 },
129 },
130 },
131 },
132 })
133
134 if _, ok := pollTerminal(foreign); ok {
135 t.Fatal("attempt-result on a foreign lease delivered a terminal; an executor forged another node's job result")
136 }
137 if foreign.getState() == leaseDone {
138 t.Fatal("attempt-result on a foreign lease sealed the lease")
139 }
140}
141
142func TestOnAttemptResultIgnoresAbsentLease(t *testing.T) {
143 ctx := context.Background()
144 l := discardLogger()
145 bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db"))
146 if err != nil {
147 t.Fatalf("db.Make: %v", err)
148 }
149 t.Cleanup(func() { bdb.Close() })
150 n := notifier.New()
151
152 m := New(l, Config{ReconnectGrace: time.Minute})
153 m.Attach(bdb, &n)
154
155 // bystander lease the reporting node owns, proving an absent-lease stream
156 // does not spill onto another lease
157 bystander := newLease("lease-bystander", "node-b", "inc-b", "dummy")
158 m.mu.Lock()
159 m.leases[bystander.id] = bystander
160 m.mu.Unlock()
161
162 sessB := newSession("node-b", "inc-b", nil, nopEncoder(), l)
163 m.attachSession(sessB)
164
165 _ = m.onEventBatch(sessB, &millv1.EventBatch{
166 Epoch: sessB.epoch,
167 Events: []*millv1.Event{
168 {
169 Seqno: 1,
170 LeaseId: "lease-nonexistent",
171 Payload: &millv1.Event_AttemptResult{
172 AttemptResult: &millv1.AttemptResult{
173 Status: millv1.TerminalStatus_SUCCESS,
174 },
175 },
176 },
177 },
178 })
179
180 if _, ok := pollTerminal(bystander); ok {
181 t.Fatal("attempt-result for an absent lease delivered a terminal to a bystander lease")
182 }
183 if bystander.getState() == leaseDone {
184 t.Fatal("attempt-result for an absent lease sealed a bystander lease")
185 }
186}
187
188// owned-lease path proves ignore tests above are not passing merely because
189// delivery is broken. correctly owned terminal is delivered
190func TestOnAttemptResultDeliversOwnedLease(t *testing.T) {
191 ctx := context.Background()
192 l := discardLogger()
193 bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db"))
194 if err != nil {
195 t.Fatalf("db.Make: %v", err)
196 }
197 t.Cleanup(func() { bdb.Close() })
198 n := notifier.New()
199
200 m := New(l, Config{ReconnectGrace: time.Minute})
201 m.Attach(bdb, &n)
202
203 owned := newLease("lease-owned", "node-b", "inc-b", "dummy")
204 m.mu.Lock()
205 m.leases[owned.id] = owned
206 m.mu.Unlock()
207
208 sessB := newSession("node-b", "inc-b", nil, nopEncoder(), l)
209 m.attachSession(sessB)
210 _ = m.onEventBatch(sessB, &millv1.EventBatch{
211 Epoch: sessB.epoch,
212 Events: []*millv1.Event{
213 {
214 Seqno: 1,
215 LeaseId: owned.id,
216 Payload: &millv1.Event_AttemptResult{
217 AttemptResult: &millv1.AttemptResult{
218 Status: millv1.TerminalStatus_SUCCESS,
219 },
220 },
221 },
222 },
223 })
224
225 res, ok := pollTerminal(owned)
226 if !ok {
227 t.Fatal("attempt-result on an owned lease was not delivered")
228 }
229 if got := res.GetStatus(); got != millv1.TerminalStatus_SUCCESS {
230 t.Fatalf("delivered terminal status = %v, want %v", got, millv1.TerminalStatus_SUCCESS)
231 }
232 if owned.getState() != leaseDone {
233 t.Fatal("owned lease was not sealed after its terminal was delivered")
234 }
235}
236
237func TestOnStatusEventOwnership(t *testing.T) {
238 ctx := context.Background()
239 l := discardLogger()
240
241 bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db"))
242 if err != nil {
243 t.Fatalf("db.Make: %v", err)
244 }
245 t.Cleanup(func() { bdb.Close() })
246 n := notifier.New()
247
248 m := New(l, Config{ReconnectGrace: time.Minute})
249 m.Attach(bdb, &n)
250
251 foreign := newLease("lease-foreign", "node-x", "inc-x", "dummy")
252 foreign.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "foreign"}, Name: "build"}
253 owned := newLease("lease-owned", "node-z", "inc-z", "dummy")
254 owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "owned"}, Name: "build"}
255 m.mu.Lock()
256 m.leases[foreign.id] = foreign
257 m.leases[owned.id] = owned
258 m.mu.Unlock()
259
260 sessY := newSession("node-y", "inc-y", nil, nopEncoder(), l)
261 m.attachSession(sessY)
262 sessZ := newSession("node-z", "inc-z", nil, nopEncoder(), l)
263 m.attachSession(sessZ)
264
265 _ = m.onEventBatch(sessY, &millv1.EventBatch{
266 Epoch: sessY.epoch,
267 Events: []*millv1.Event{
268 {
269 Seqno: 1,
270 LeaseId: foreign.id,
271 Payload: &millv1.Event_StatusEvent{
272 StatusEvent: &millv1.StatusEvent{
273 Status: millv1.NonterminalStatus_RUNNING,
274 },
275 },
276 },
277 },
278 })
279 if _, err := bdb.GetStatus(foreign.wid); err == nil {
280 t.Fatal("status stream for a foreign lease authored a status row; an executor forged another pipeline's status")
281 }
282
283 _ = m.onEventBatch(sessZ, &millv1.EventBatch{
284 Epoch: sessZ.epoch,
285 Events: []*millv1.Event{
286 {
287 Seqno: 1,
288 LeaseId: owned.id,
289 Payload: &millv1.Event_StatusEvent{
290 StatusEvent: &millv1.StatusEvent{
291 Status: millv1.NonterminalStatus_RUNNING,
292 },
293 },
294 },
295 },
296 })
297 st, err := bdb.GetStatus(owned.wid)
298 if err != nil {
299 t.Fatalf("owned status stream did not author a status row: %v", err)
300 }
301 if st.Status != "running" {
302 t.Fatalf("owned status = %q, want %q", st.Status, "running")
303 }
304}
305
306func setupTestServer(t *testing.T, authorizedLabels []string) (*Mill, *db.DB, *httptest.Server, string) {
307 ctx := context.Background()
308 bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db"))
309 if err != nil {
310 t.Fatalf("db.Make: %v", err)
311 }
312 t.Cleanup(func() { bdb.Close() })
313 n := notifier.New()
314
315 m := New(discardLogger(), Config{
316 ReconnectGrace: time.Minute,
317 })
318 m.Attach(bdb, &n)
319
320 const secret = "test-secret"
321 if err := bdb.AddExecutorToken("dev-node", HashToken(secret), nil, authorizedLabels); err != nil {
322 t.Fatalf("AddExecutorToken: %v", err)
323 }
324
325 server := httptest.NewServer(http.HandlerFunc(m.HandleExecutorConn))
326 t.Cleanup(server.Close)
327
328 return m, bdb, server, secret
329}
330
331func TestAuthLabelEscalation(t *testing.T) {
332 _, _, server, secret := setupTestServer(t, []string{"linux", "amd64"})
333
334 wsUrl := "ws" + strings.TrimPrefix(server.URL, "http")
335
336 {
337 header := http.Header{}
338 header.Set("Authorization", "Bearer bad-token")
339 _, resp, err := websocket.DefaultDialer.Dial(wsUrl, header)
340 if err == nil {
341 t.Fatal("expected connection with invalid token to fail")
342 }
343 if resp != nil && resp.StatusCode != http.StatusUnauthorized {
344 t.Fatalf("expected 401 Unauthorized, got %d", resp.StatusCode)
345 }
346 }
347
348 {
349 header := http.Header{}
350 header.Set("Authorization", "Bearer "+secret)
351 conn, _, err := websocket.DefaultDialer.Dial(wsUrl, header)
352 if err != nil {
353 t.Fatalf("dial failed: %v", err)
354 }
355 defer conn.Close()
356
357 stream := millproto.NewWSStream(conn)
358 enc := millproto.NewEncoder(stream)
359 dec := millproto.NewDecoder(stream)
360
361 hello := &millproto.Message{Hello: &millv1.Hello{
362 ProtocolVersion: millproto.ProtocolVersion,
363 Arch: "amd64",
364 Labels: []string{"linux", "gpu"},
365 Epoch: "inc-1",
366 }}
367 if err := enc.Encode(hello); err != nil {
368 t.Fatalf("encode hello: %v", err)
369 }
370
371 _, err = dec.Decode()
372 if err == nil {
373 t.Fatal("expected server to close connection for unauthorized label, but got a message")
374 }
375 }
376
377 {
378 header := http.Header{}
379 header.Set("Authorization", "Bearer "+secret)
380 conn, _, err := websocket.DefaultDialer.Dial(wsUrl, header)
381 if err != nil {
382 t.Fatalf("dial failed: %v", err)
383 }
384 defer conn.Close()
385
386 stream := millproto.NewWSStream(conn)
387 enc := millproto.NewEncoder(stream)
388 dec := millproto.NewDecoder(stream)
389
390 hello := &millproto.Message{Hello: &millv1.Hello{
391 ProtocolVersion: millproto.ProtocolVersion,
392 Arch: "amd64",
393 Labels: []string{"linux"},
394 Epoch: "inc-1",
395 }}
396 if err := enc.Encode(hello); err != nil {
397 t.Fatalf("encode hello: %v", err)
398 }
399
400 msg, err := dec.Decode()
401 if err != nil {
402 t.Fatalf("expected resume message, got error: %v", err)
403 }
404 res := msg.GetResume()
405 if res == nil {
406 t.Fatal("expected Resume message, got nil")
407 }
408 if res.GetEpoch() != "inc-1" {
409 t.Fatalf("expected epoch inc-1, got %q", res.GetEpoch())
410 }
411 }
412}
413
414func TestHandshakeTimeoutAndConcurrency(t *testing.T) {
415 _, _, server, secret := setupTestServer(t, []string{"linux"})
416 wsUrl := "ws" + strings.TrimPrefix(server.URL, "http")
417
418 // executor that never sends Hello is dropped after the 5s pre-hello deadline
419 {
420 header := http.Header{}
421 header.Set("Authorization", "Bearer "+secret)
422 conn, _, err := websocket.DefaultDialer.Dial(wsUrl, header)
423 if err != nil {
424 t.Fatalf("dial failed: %v", err)
425 }
426 defer conn.Close()
427
428 time.Sleep(6 * time.Second)
429
430 stream := millproto.NewWSStream(conn)
431 enc := millproto.NewEncoder(stream)
432 hello := &millproto.Message{Hello: &millv1.Hello{
433 ProtocolVersion: millproto.ProtocolVersion,
434 Arch: "amd64",
435 Labels: []string{"linux"},
436 Epoch: "inc-1",
437 }}
438 err = enc.Encode(hello)
439 dec := millproto.NewDecoder(stream)
440 _, readErr := dec.Decode()
441 if readErr == nil {
442 t.Fatal("expected server to have closed connection due to handshake timeout")
443 }
444 }
445
446 // second live session for one identity is rejected with 409 before the ws upgrade
447 {
448 header := http.Header{}
449 header.Set("Authorization", "Bearer "+secret)
450
451 conn1, _, err := websocket.DefaultDialer.Dial(wsUrl, header)
452 if err != nil {
453 t.Fatalf("dial 1 failed: %v", err)
454 }
455 defer conn1.Close()
456
457 stream1 := millproto.NewWSStream(conn1)
458 enc1 := millproto.NewEncoder(stream1)
459 dec1 := millproto.NewDecoder(stream1)
460 hello1 := &millproto.Message{Hello: &millv1.Hello{
461 ProtocolVersion: millproto.ProtocolVersion,
462 Arch: "amd64",
463 Labels: []string{"linux"},
464 Epoch: "inc-1",
465 }}
466 if err := enc1.Encode(hello1); err != nil {
467 t.Fatalf("encode hello 1: %v", err)
468 }
469 _, err = dec1.Decode()
470 if err != nil {
471 t.Fatalf("first connection handshake failed: %v", err)
472 }
473
474 _, resp, err := websocket.DefaultDialer.Dial(wsUrl, header)
475 if err == nil {
476 t.Fatal("expected second connection for same live identity to be rejected")
477 }
478 if resp != nil && resp.StatusCode != http.StatusConflict {
479 t.Fatalf("expected 409 Conflict for duplicate session, got %d", resp.StatusCode)
480 }
481 }
482}