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 // bind the host context explicitly. if the local CID resolves to the
32 // loopback CID (which happens on some systems when vsock_loopback is active),
33 // Listen() would bind loopback and never see guest VMs
34 ln, err := vsock.ListenContextID(vsock.Host, port, nil)
35 if err != nil {
36 return nil, fmt.Errorf("listen for agent on vsock port %d: %w", port, err)
37 }
38 h := &agentHub{
39 l: l,
40 ln: ln,
41 pending: make(map[uint32]chan net.Conn),
42 }
43 go h.acceptLoop()
44 return h, nil
45}
46
47func (h *agentHub) expect(cid uint32) (<-chan net.Conn, func(), error) {
48 h.mu.Lock()
49 defer h.mu.Unlock()
50 if _, exists := h.pending[cid]; exists {
51 return nil, nil, fmt.Errorf("already waiting for agent cid %d", cid)
52 }
53 ch := make(chan net.Conn, 1)
54 h.pending[cid] = ch
55 unregister := func() {
56 h.mu.Lock()
57 delete(h.pending, cid)
58 h.mu.Unlock()
59 close(ch)
60 for conn := range ch {
61 if conn != nil {
62 _ = conn.Close()
63 }
64 }
65 }
66 return ch, unregister, nil
67}
68
69func (h *agentHub) acceptLoop() {
70 for {
71 conn, err := h.ln.Accept()
72 if err != nil {
73 h.l.Error("agent vsock accept failed", "error", err)
74 return
75 }
76
77 addr, ok := conn.RemoteAddr().(*vsock.Addr)
78 if !ok {
79 h.l.Warn("agent connection has unexpected remote address", "remote", conn.RemoteAddr())
80 _ = conn.Close()
81 continue
82 }
83
84 h.mu.Lock()
85 ch, ok := h.pending[addr.ContextID]
86 if ok {
87 delete(h.pending, addr.ContextID)
88 }
89 h.mu.Unlock()
90
91 // todo: if / when we add agent recovery (reconnect) we should add a
92 // boot-initialized session credential to prevent random connections...
93 // checking cid here works to ensure for now since we dont attempt to
94 // reconnect, so we block anything else thats not expected (and agent
95 // runs first in the boot sequence always).
96 if !ok {
97 h.l.Warn("dropping agent connection for unknown cid", "cid", addr.ContextID)
98 _ = conn.Close()
99 continue
100 }
101
102 select {
103 case ch <- conn:
104 default:
105 _ = conn.Close()
106 }
107 }
108}
109
110type AgentExec struct {
111 *agentv1.ExecStart
112 ID string
113 Stdout io.Writer
114 Stderr io.Writer
115}
116
117type AgentSession struct {
118 conn net.Conn
119 enc *agentproto.Encoder
120 dec *agentproto.Decoder
121 l *slog.Logger
122 mu sync.Mutex
123}
124
125func NewAgentSession(conn net.Conn, l *slog.Logger) *AgentSession {
126 return &AgentSession{
127 conn: conn,
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 s.l.Info("agent connected", "protocol", helloPayload.ProtocolVersion, "version", helloPayload.AgentVersion, "boot", helloPayload.BootId, "nix", helloPayload.NixVersion)
147
148 if err := s.enc.Encode(&agentproto.Message{
149 Id: "init",
150 Init: init,
151 }); err != nil {
152 return fmt.Errorf("send agent init: %w", err)
153 }
154 return nil
155}
156
157func (s *AgentSession) Exec(ctx context.Context, exec AgentExec) (int, error) {
158 s.mu.Lock()
159 defer s.mu.Unlock()
160
161 if exec.ID == "" {
162 return 0, fmt.Errorf("empty ID passed to Exec")
163 }
164
165 if exec.ExecStart.TimeoutSeconds == 0 {
166 exec.ExecStart.TimeoutSeconds = timeoutSeconds(ctx, guestTimeoutGrace)
167 }
168
169 if err := s.enc.Encode(&agentproto.Message{
170 Id: exec.ID,
171 ExecStart: exec.ExecStart,
172 }); err != nil {
173 return 0, fmt.Errorf("send exec_start: %w", err)
174 }
175
176 for {
177 msg, err := s.decode(ctx)
178 if err != nil {
179 return 0, err
180 }
181 if msg.BuiltPaths == nil && msg.Id != exec.ID {
182 continue
183 }
184
185 if p := msg.ExecStdout; p != nil {
186 _, _ = io.WriteString(exec.Stdout, p.Data)
187 } else if p := msg.ExecStderr; p != nil {
188 _, _ = io.WriteString(exec.Stderr, p.Data)
189 } else if p := msg.BuiltPaths; p != nil {
190 // s.l.Debug("guest built paths", "reason", p.Reason, "count", len(p.Paths))
191 } else if p := msg.ExecExit; p != nil {
192 var err error
193 if p.Error != "" {
194 s.l.Warn("guest exec error", "id", msg.Id, "error", p.Error)
195 err = fmt.Errorf("guest exec error: %s", p.Error)
196 }
197 if p.TimedOut {
198 return int(p.ExitCode), errGuestTimedOut
199 }
200 return int(p.ExitCode), err
201 }
202 }
203}
204
205func (s *AgentSession) ActivateConfig(ctx context.Context, id string, req *agentv1.ActivateConfig, out io.Writer) (*agentv1.ActivateConfigResult, error) {
206 s.mu.Lock()
207 defer s.mu.Unlock()
208
209 if id == "" {
210 return nil, fmt.Errorf("empty ID passed to ActivateConfig")
211 }
212 if req.TimeoutSeconds == 0 {
213 req.TimeoutSeconds = timeoutSeconds(ctx, guestTimeoutGrace)
214 }
215 if err := s.enc.Encode(&agentproto.Message{
216 Id: id,
217 ActivateConfig: req,
218 }); err != nil {
219 return nil, fmt.Errorf("send activate_config: %w", err)
220 }
221
222 for {
223 msg, err := s.decode(ctx)
224 if err != nil {
225 return nil, err
226 }
227 if msg.BuiltPaths == nil && msg.Id != id {
228 continue
229 }
230
231 if p := msg.BuiltPaths; p != nil {
232 // s.l.Debug("guest built paths", "reason", p.Reason, "count", len(p.Paths))
233 } else if p := msg.ExecStderr; p != nil {
234 if out != nil {
235 _, _ = io.WriteString(out, p.Data)
236 }
237 } else if p := msg.ExecStdout; p != nil {
238 if out != nil {
239 _, _ = io.WriteString(out, p.Data)
240 }
241 } else if p := msg.ActivateConfigResult; p != nil {
242 if p.Error != "" {
243 return nil, errors.New(p.Error)
244 }
245 if p.Toplevel == "" {
246 return nil, fmt.Errorf("activate config returned empty toplevel")
247 }
248 return p, nil
249 }
250 }
251}
252
253func (s *AgentSession) Poweroff(ctx context.Context) error {
254 s.mu.Lock()
255 defer s.mu.Unlock()
256
257 id := "poweroff"
258 if err := s.enc.Encode(&agentproto.Message{
259 Id: id,
260 Poweroff: &agentv1.Poweroff{},
261 }); err != nil {
262 return fmt.Errorf("send poweroff: %w", err)
263 }
264
265 for {
266 msg, err := s.decode(ctx)
267 if err != nil {
268 return err
269 }
270 if msg.Id != id {
271 continue
272 }
273 p := msg.PoweroffResult
274 if p == nil {
275 continue
276 }
277 if p.Error != "" {
278 return fmt.Errorf("guest poweroff failed: %s", p.Error)
279 }
280 return nil
281 }
282}
283
284func (s *AgentSession) Drain(ctx context.Context) (uint32, error) {
285 s.mu.Lock()
286 defer s.mu.Unlock()
287
288 drainID := "cache-drain"
289 if err := s.enc.Encode(&agentproto.Message{
290 Id: drainID,
291 CacheDrain: &agentv1.CacheDrain{
292 TimeoutSeconds: timeoutSeconds(ctx, 0),
293 },
294 }); err != nil {
295 return 0, fmt.Errorf("send cache_drain: %w", err)
296 }
297
298 for {
299 msg, err := s.decode(ctx)
300 if err != nil {
301 return 0, err
302 }
303 if msg.Id != drainID {
304 continue
305 }
306 p := msg.CacheDrainResult
307 if p == nil {
308 continue
309 }
310 s.l.Info("cache drain complete", "uploaded", p.CacheUploaded, "failed", p.CacheFailed, "queued", p.CacheQueued, "active", p.CacheActive)
311 if p.Error != "" {
312 return 0, fmt.Errorf("cache drain failed: %s", p.Error)
313 }
314 if p.CacheFailed > 0 {
315 return 0, fmt.Errorf("cache drain failed for %d paths", p.CacheFailed)
316 }
317 if p.CacheQueued > 0 || p.CacheActive > 0 {
318 return 0, fmt.Errorf("cache drain incomplete: queued=%d active=%d", p.CacheQueued, p.CacheActive)
319 }
320 return p.CacheUploaded, nil
321 }
322}
323
324func (s *AgentSession) decode(ctx context.Context) (*agentproto.Message, error) {
325 if err := ctx.Err(); err != nil {
326 return nil, err
327 }
328
329 if deadline, ok := ctx.Deadline(); ok {
330 _ = s.conn.SetReadDeadline(deadline)
331 } else {
332 _ = s.conn.SetReadDeadline(time.Time{})
333 }
334
335 // a blocked vsock read wont wake up just from the ctx being cancelled,
336 // only a deadline will wake it up, so if the VM crashes mid-step the read would
337 // hang until workflow timeout. so we will set a deadline in the past to cancel it.
338 //
339 // we set a deadline here instead of closing the connection, this is the long-lived
340 // connection that everything reuses, so we only really want to interrupt it for this
341 // current read. this also lands as a timeout error which the netErr.Timeout() check
342 // below maps to ctx.Err() correctly
343 stop := context.AfterFunc(ctx, func() {
344 _ = s.conn.SetReadDeadline(time.Now())
345 })
346 defer stop()
347
348 msg, err := s.dec.Decode()
349 if err != nil {
350 var netErr net.Error
351 if errors.As(err, &netErr) && netErr.Timeout() && ctx.Err() != nil {
352 return nil, ctx.Err()
353 }
354 return nil, fmt.Errorf("read agent message: %w", err)
355 }
356 return msg, nil
357}
358
359func (s *AgentSession) Close() error {
360 if s == nil || s.conn == nil {
361 return nil
362 }
363 return s.conn.Close()
364}
365
366// this pulls the deadline from the context and converts it to what the
367// agentproto expects
368func timeoutSeconds(ctx context.Context, lead time.Duration) uint32 {
369 deadline, ok := ctx.Deadline()
370 if !ok {
371 return 0
372 }
373 seconds := int64((time.Until(deadline) - lead).Round(time.Second) / time.Second)
374 if seconds < 1 {
375 return 1
376 }
377 if seconds > int64(^uint32(0)) {
378 return ^uint32(0)
379 }
380 return uint32(seconds)
381}