alice.pds.demo.boltless.dev
This repository has no description
1package microvm
2
3import (
4 "context"
5 "errors"
6 "fmt"
7 "io"
8 "log/slog"
9 "net"
10 "sync"
11 "time"
12
13 "github.com/mdlayher/vsock"
14
15 "tangled.org/core/spindle/agentproto"
16 agentv1 "tangled.org/core/spindle/agentproto/gen"
17)
18
19const guestWorkflowUser = "spindle-workflow"
20
21var errGuestTimedOut = errors.New("guest reported step timed out")
22
23type agentHub struct {
24 l *slog.Logger
25 ln *vsock.Listener
26 pending map[uint32]chan net.Conn
27 mu sync.Mutex
28}
29
30func newAgentHub(port uint32, l *slog.Logger) (*agentHub, error) {
31 ln, err := vsock.Listen(port, nil)
32 if err != nil {
33 return nil, fmt.Errorf("listen for agent on vsock port %d: %w", port, err)
34 }
35 h := &agentHub{
36 l: l,
37 ln: ln,
38 pending: make(map[uint32]chan net.Conn),
39 }
40 go h.acceptLoop()
41 return h, nil
42}
43
44func (h *agentHub) expect(cid uint32) (<-chan net.Conn, func(), error) {
45 h.mu.Lock()
46 defer h.mu.Unlock()
47 if _, exists := h.pending[cid]; exists {
48 return nil, nil, fmt.Errorf("already waiting for agent cid %d", cid)
49 }
50 ch := make(chan net.Conn, 1)
51 h.pending[cid] = ch
52 unregister := func() {
53 h.mu.Lock()
54 delete(h.pending, cid)
55 h.mu.Unlock()
56 close(ch)
57 for conn := range ch {
58 if conn != nil {
59 _ = conn.Close()
60 }
61 }
62 }
63 return ch, unregister, nil
64}
65
66func (h *agentHub) acceptLoop() {
67 for {
68 conn, err := h.ln.Accept()
69 if err != nil {
70 h.l.Error("agent vsock accept failed", "error", err)
71 return
72 }
73
74 addr, ok := conn.RemoteAddr().(*vsock.Addr)
75 if !ok {
76 h.l.Warn("agent connection has unexpected remote address", "remote", conn.RemoteAddr())
77 _ = conn.Close()
78 continue
79 }
80
81 h.mu.Lock()
82 ch, ok := h.pending[addr.ContextID]
83 if ok {
84 delete(h.pending, addr.ContextID)
85 }
86 h.mu.Unlock()
87
88 // todo: if / when we add agent recovery (reconnect) we should add a
89 // boot-initialized session credential to prevent random connections...
90 // checking cid here works to ensure for now since we dont attempt to
91 // reconnect, so we block anything else thats not expected (and agent
92 // runs first in the boot sequence always).
93 if !ok {
94 h.l.Warn("dropping agent connection for unknown cid", "cid", addr.ContextID)
95 _ = conn.Close()
96 continue
97 }
98
99 select {
100 case ch <- conn:
101 default:
102 _ = conn.Close()
103 }
104 }
105}
106
107type AgentExec struct {
108 *agentv1.ExecStart
109 ID string
110 Stdin io.Reader
111 Stdout io.Writer
112 Stderr io.Writer
113}
114
115type AgentSession struct {
116 conn net.Conn
117 cid uint32
118 enc *agentproto.Encoder
119 dec *agentproto.Decoder
120 l *slog.Logger
121 mu sync.Mutex
122}
123
124func NewAgentSession(conn net.Conn, cid uint32, l *slog.Logger) *AgentSession {
125 return &AgentSession{
126 conn: conn,
127 cid: cid,
128 enc: agentproto.NewEncoder(conn),
129 dec: agentproto.NewDecoder(conn),
130 l: l,
131 }
132}
133
134func (s *AgentSession) Init(ctx context.Context, init *agentv1.Init) error {
135 s.mu.Lock()
136 defer s.mu.Unlock()
137
138 hello, err := s.decode(ctx)
139 if err != nil {
140 return fmt.Errorf("read agent hello: %w", err)
141 }
142 helloPayload := hello.Hello
143 if helloPayload == nil {
144 return fmt.Errorf("expected agent hello, got nil")
145 }
146 if helloPayload.ProtocolVersion != agentproto.ProtocolVersion {
147 return fmt.Errorf("agent protocol version %d, want %d (stale guest image?)", helloPayload.ProtocolVersion, agentproto.ProtocolVersion)
148 }
149 s.l.Info("agent connected", "protocol", helloPayload.ProtocolVersion, "version", helloPayload.AgentVersion, "boot", helloPayload.BootId, "nix", helloPayload.NixVersion)
150
151 if err := s.enc.Encode(&agentproto.Message{
152 Id: "init",
153 Init: init,
154 }); err != nil {
155 return fmt.Errorf("send agent init: %w", err)
156 }
157 return nil
158}
159
160func (s *AgentSession) Exec(ctx context.Context, exec AgentExec) (int, error) {
161 s.mu.Lock()
162 defer s.mu.Unlock()
163
164 if exec.ID == "" {
165 return 0, fmt.Errorf("empty ID passed to Exec")
166 }
167
168 if exec.ExecStart.TimeoutSeconds == 0 {
169 exec.ExecStart.TimeoutSeconds = timeoutSeconds(ctx, guestTimeoutGrace)
170 }
171
172 ln, port, err := listenRandomVsockPort(ctx)
173 if err != nil {
174 return 0, fmt.Errorf("listen for exec stdio: %w", err)
175 }
176 defer ln.Close()
177 exec.ExecStart.StdioVsockPort = port
178
179 if err := s.enc.Encode(&agentproto.Message{
180 Id: exec.ID,
181 ExecStart: exec.ExecStart,
182 }); err != nil {
183 return 0, fmt.Errorf("send exec_start: %w", err)
184 }
185
186 filtered := &cidFilteredVsockListener{Listener: ln, cid: s.cid, logger: s.l}
187 stdioDone := make(chan error, 1)
188 go func() {
189 stdioDone <- pumpStdio(ctx, filtered, exec.Stdin, exec.Stdout)
190 }()
191
192 for {
193 msg, err := s.decode(ctx)
194 if err != nil {
195 ln.Close()
196 <-stdioDone
197 return 0, err
198 }
199 if msg.BuiltPaths == nil && msg.Id != exec.ID {
200 continue
201 }
202
203 if p := msg.BuiltPaths; p != nil {
204 // s.l.Debug("guest built paths", "reason", p.Reason, "count", len(p.Paths))
205 } else if p := msg.ExecStderr; p != nil {
206 _, _ = exec.Stderr.Write(p.Data)
207 } else if p := msg.ExecExit; p != nil {
208 var err error
209 if p.Error != "" {
210 s.l.Warn("guest exec error", "id", msg.Id, "error", p.Error)
211 err = fmt.Errorf("guest exec error: %s", p.Error)
212 }
213 ln.Close() // wake Accept if the guest never dialed
214 if err := <-stdioDone; err != nil {
215 return 0, err
216 }
217 if p.TimedOut {
218 return int(p.ExitCode), errGuestTimedOut
219 }
220 return int(p.ExitCode), err
221 }
222 }
223}
224
225func pumpStdio(ctx context.Context, ln net.Listener, stdin io.Reader, stdout io.Writer) error {
226 conn, err := ln.Accept()
227 if err != nil {
228 if errors.Is(err, net.ErrClosed) {
229 return nil // exec ended before dialing
230 }
231 return fmt.Errorf("accept guest stdio connection: %w", err)
232 }
233 vsockConn, ok := conn.(*vsock.Conn)
234 if !ok {
235 conn.Close()
236 return fmt.Errorf("guest connection is not a vsock connection")
237 }
238 defer vsockConn.Close()
239 stop := context.AfterFunc(ctx, func() { vsockConn.Close() })
240 defer stop()
241
242 stdinDone := make(chan error, 1)
243 go func() {
244 stdinDone <- writeStdin(vsockConn, stdin)
245 }()
246
247 if _, err := io.Copy(stdout, conn); err != nil {
248 vsockConn.Close()
249 <-stdinDone
250 return fmt.Errorf("read guest stdout: %w", err)
251 }
252 if err := <-stdinDone; err != nil {
253 return fmt.Errorf("write guest stdin: %w", err)
254 }
255 return nil
256}
257
258// send stdin EOF without closing stdout
259func writeStdin(conn *vsock.Conn, r io.Reader) error {
260 if r != nil {
261 if _, err := io.Copy(conn, r); err != nil {
262 return err
263 }
264 }
265 return conn.CloseWrite()
266}
267
268func (s *AgentSession) ActivateConfig(ctx context.Context, id string, req *agentv1.ActivateConfig, out io.Writer) (*agentv1.ActivateConfigResult, error) {
269 s.mu.Lock()
270 defer s.mu.Unlock()
271
272 if id == "" {
273 return nil, fmt.Errorf("empty ID passed to ActivateConfig")
274 }
275 if req.TimeoutSeconds == 0 {
276 req.TimeoutSeconds = timeoutSeconds(ctx, guestTimeoutGrace)
277 }
278 if err := s.enc.Encode(&agentproto.Message{
279 Id: id,
280 ActivateConfig: req,
281 }); err != nil {
282 return nil, fmt.Errorf("send activate_config: %w", err)
283 }
284
285 for {
286 msg, err := s.decode(ctx)
287 if err != nil {
288 return nil, err
289 }
290 if msg.BuiltPaths == nil && msg.Id != id {
291 continue
292 }
293
294 if p := msg.BuiltPaths; p != nil {
295 // s.l.Debug("guest built paths", "reason", p.Reason, "count", len(p.Paths))
296 } else if p := msg.ExecStderr; p != nil {
297 if out != nil {
298 _, _ = out.Write(p.Data)
299 }
300 } else if p := msg.ActivateConfigResult; p != nil {
301 if p.Error != "" {
302 return nil, errors.New(p.Error)
303 }
304 if p.Toplevel == "" {
305 return nil, fmt.Errorf("activate config returned empty toplevel")
306 }
307 return p, nil
308 }
309 }
310}
311
312func (s *AgentSession) Poweroff(ctx context.Context) error {
313 s.mu.Lock()
314 defer s.mu.Unlock()
315
316 id := "poweroff"
317 if err := s.enc.Encode(&agentproto.Message{
318 Id: id,
319 Poweroff: &agentv1.Poweroff{},
320 }); err != nil {
321 return fmt.Errorf("send poweroff: %w", err)
322 }
323
324 for {
325 msg, err := s.decode(ctx)
326 if err != nil {
327 return err
328 }
329 if msg.Id != id {
330 continue
331 }
332 p := msg.PoweroffResult
333 if p == nil {
334 continue
335 }
336 if p.Error != "" {
337 return fmt.Errorf("guest poweroff failed: %s", p.Error)
338 }
339 return nil
340 }
341}
342
343func (s *AgentSession) Drain(ctx context.Context) (uint32, error) {
344 s.mu.Lock()
345 defer s.mu.Unlock()
346
347 drainID := "cache-drain"
348 if err := s.enc.Encode(&agentproto.Message{
349 Id: drainID,
350 CacheDrain: &agentv1.CacheDrain{
351 TimeoutSeconds: timeoutSeconds(ctx, 0),
352 },
353 }); err != nil {
354 return 0, fmt.Errorf("send cache_drain: %w", err)
355 }
356
357 for {
358 msg, err := s.decode(ctx)
359 if err != nil {
360 return 0, err
361 }
362 if msg.Id != drainID {
363 continue
364 }
365 p := msg.CacheDrainResult
366 if p == nil {
367 continue
368 }
369 s.l.Info("cache drain complete", "uploaded", p.CacheUploaded, "failed", p.CacheFailed, "queued", p.CacheQueued, "active", p.CacheActive)
370 if p.Error != "" {
371 return 0, fmt.Errorf("cache drain failed: %s", p.Error)
372 }
373 if p.CacheFailed > 0 {
374 return 0, fmt.Errorf("cache drain failed for %d paths", p.CacheFailed)
375 }
376 if p.CacheQueued > 0 || p.CacheActive > 0 {
377 return 0, fmt.Errorf("cache drain incomplete: queued=%d active=%d", p.CacheQueued, p.CacheActive)
378 }
379 return p.CacheUploaded, nil
380 }
381}
382
383func (s *AgentSession) decode(ctx context.Context) (*agentproto.Message, error) {
384 if err := ctx.Err(); err != nil {
385 return nil, err
386 }
387
388 if deadline, ok := ctx.Deadline(); ok {
389 _ = s.conn.SetReadDeadline(deadline)
390 } else {
391 _ = s.conn.SetReadDeadline(time.Time{})
392 }
393
394 // a blocked vsock read wont wake up just from the ctx being cancelled,
395 // only a deadline will wake it up, so if the VM crashes mid-step the read would
396 // hang until workflow timeout. so we will set a deadline in the past to cancel it.
397 //
398 // we set a deadline here instead of closing the connection, this is the long-lived
399 // connection that everything reuses, so we only really want to interrupt it for this
400 // current read. this also lands as a timeout error which the netErr.Timeout() check
401 // below maps to ctx.Err() correctly
402 stop := context.AfterFunc(ctx, func() {
403 _ = s.conn.SetReadDeadline(time.Now())
404 })
405 defer stop()
406
407 msg, err := s.dec.Decode()
408 if err != nil {
409 var netErr net.Error
410 if errors.As(err, &netErr) && netErr.Timeout() && ctx.Err() != nil {
411 return nil, ctx.Err()
412 }
413 return nil, fmt.Errorf("read agent message: %w", err)
414 }
415 return msg, nil
416}
417
418func (s *AgentSession) Close() error {
419 if s == nil || s.conn == nil {
420 return nil
421 }
422 return s.conn.Close()
423}
424
425// this pulls the deadline from the context and converts it to what the
426// agentproto expects
427func timeoutSeconds(ctx context.Context, lead time.Duration) uint32 {
428 deadline, ok := ctx.Deadline()
429 if !ok {
430 return 0
431 }
432 seconds := int64((time.Until(deadline) - lead).Round(time.Second) / time.Second)
433 if seconds < 1 {
434 return 1
435 }
436 if seconds > int64(^uint32(0)) {
437 return ^uint32(0)
438 }
439 return uint32(seconds)
440}