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