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 Stdout io.Writer
111 Stderr io.Writer
112}
113
114type AgentSession struct {
115 conn net.Conn
116 enc *agentproto.Encoder
117 dec *agentproto.Decoder
118 l *slog.Logger
119 mu sync.Mutex
120}
121
122func NewAgentSession(conn net.Conn, l *slog.Logger) *AgentSession {
123 return &AgentSession{
124 conn: conn,
125 enc: agentproto.NewEncoder(conn),
126 dec: agentproto.NewDecoder(conn),
127 l: l,
128 }
129}
130
131func (s *AgentSession) Init(ctx context.Context, init *agentv1.Init) error {
132 s.mu.Lock()
133 defer s.mu.Unlock()
134
135 hello, err := s.decode(ctx)
136 if err != nil {
137 return fmt.Errorf("read agent hello: %w", err)
138 }
139 helloPayload := hello.Hello
140 if helloPayload == nil {
141 return fmt.Errorf("expected agent hello, got nil")
142 }
143 s.l.Info("agent connected", "protocol", helloPayload.ProtocolVersion, "version", helloPayload.AgentVersion, "boot", helloPayload.BootId, "nix", helloPayload.NixVersion)
144
145 if err := s.enc.Encode(&agentproto.Message{
146 Id: "init",
147 Init: init,
148 }); err != nil {
149 return fmt.Errorf("send agent init: %w", err)
150 }
151 return nil
152}
153
154func (s *AgentSession) Exec(ctx context.Context, exec AgentExec) (int, error) {
155 s.mu.Lock()
156 defer s.mu.Unlock()
157
158 if exec.ID == "" {
159 return 0, fmt.Errorf("empty ID passed to Exec")
160 }
161
162 if exec.ExecStart.TimeoutSeconds == 0 {
163 exec.ExecStart.TimeoutSeconds = timeoutSeconds(ctx, guestTimeoutGrace)
164 }
165
166 if err := s.enc.Encode(&agentproto.Message{
167 Id: exec.ID,
168 ExecStart: exec.ExecStart,
169 }); err != nil {
170 return 0, fmt.Errorf("send exec_start: %w", err)
171 }
172
173 for {
174 msg, err := s.decode(ctx)
175 if err != nil {
176 return 0, err
177 }
178 if msg.BuiltPaths == nil && msg.Id != exec.ID {
179 continue
180 }
181
182 if p := msg.ExecStdout; p != nil {
183 _, _ = io.WriteString(exec.Stdout, p.Data)
184 } else if p := msg.ExecStderr; p != nil {
185 _, _ = io.WriteString(exec.Stderr, p.Data)
186 } else if p := msg.BuiltPaths; p != nil {
187 // s.l.Debug("guest built paths", "reason", p.Reason, "count", len(p.Paths))
188 } else if p := msg.ExecExit; p != nil {
189 var err error
190 if p.Error != "" {
191 s.l.Warn("guest exec error", "id", msg.Id, "error", p.Error)
192 err = fmt.Errorf("guest exec error: %s", p.Error)
193 }
194 if p.TimedOut {
195 return int(p.ExitCode), errGuestTimedOut
196 }
197 return int(p.ExitCode), err
198 }
199 }
200}
201
202func (s *AgentSession) ActivateConfig(ctx context.Context, id string, req *agentv1.ActivateConfig, out io.Writer) (*agentv1.ActivateConfigResult, error) {
203 s.mu.Lock()
204 defer s.mu.Unlock()
205
206 if id == "" {
207 return nil, fmt.Errorf("empty ID passed to ActivateConfig")
208 }
209 if req.TimeoutSeconds == 0 {
210 req.TimeoutSeconds = timeoutSeconds(ctx, guestTimeoutGrace)
211 }
212 if err := s.enc.Encode(&agentproto.Message{
213 Id: id,
214 ActivateConfig: req,
215 }); err != nil {
216 return nil, fmt.Errorf("send activate_config: %w", err)
217 }
218
219 for {
220 msg, err := s.decode(ctx)
221 if err != nil {
222 return nil, err
223 }
224 if msg.BuiltPaths == nil && msg.Id != id {
225 continue
226 }
227
228 if p := msg.BuiltPaths; p != nil {
229 // s.l.Debug("guest built paths", "reason", p.Reason, "count", len(p.Paths))
230 } else if p := msg.ExecStderr; p != nil {
231 if out != nil {
232 _, _ = io.WriteString(out, p.Data)
233 }
234 } else if p := msg.ExecStdout; p != nil {
235 if out != nil {
236 _, _ = io.WriteString(out, p.Data)
237 }
238 } else if p := msg.ActivateConfigResult; p != nil {
239 if p.Error != "" {
240 return nil, errors.New(p.Error)
241 }
242 if p.Toplevel == "" {
243 return nil, fmt.Errorf("activate config returned empty toplevel")
244 }
245 return p, nil
246 }
247 }
248}
249
250func (s *AgentSession) Poweroff(ctx context.Context) error {
251 s.mu.Lock()
252 defer s.mu.Unlock()
253
254 id := "poweroff"
255 if err := s.enc.Encode(&agentproto.Message{
256 Id: id,
257 Poweroff: &agentv1.Poweroff{},
258 }); err != nil {
259 return fmt.Errorf("send poweroff: %w", err)
260 }
261
262 for {
263 msg, err := s.decode(ctx)
264 if err != nil {
265 return err
266 }
267 if msg.Id != id {
268 continue
269 }
270 p := msg.PoweroffResult
271 if p == nil {
272 continue
273 }
274 if p.Error != "" {
275 return fmt.Errorf("guest poweroff failed: %s", p.Error)
276 }
277 return nil
278 }
279}
280
281func (s *AgentSession) Drain(ctx context.Context) (uint32, error) {
282 s.mu.Lock()
283 defer s.mu.Unlock()
284
285 drainID := "cache-drain"
286 if err := s.enc.Encode(&agentproto.Message{
287 Id: drainID,
288 CacheDrain: &agentv1.CacheDrain{
289 TimeoutSeconds: timeoutSeconds(ctx, 0),
290 },
291 }); err != nil {
292 return 0, fmt.Errorf("send cache_drain: %w", err)
293 }
294
295 for {
296 msg, err := s.decode(ctx)
297 if err != nil {
298 return 0, err
299 }
300 if msg.Id != drainID {
301 continue
302 }
303 p := msg.CacheDrainResult
304 if p == nil {
305 continue
306 }
307 s.l.Info("cache drain complete", "uploaded", p.CacheUploaded, "failed", p.CacheFailed, "queued", p.CacheQueued, "active", p.CacheActive)
308 if p.Error != "" {
309 return 0, fmt.Errorf("cache drain failed: %s", p.Error)
310 }
311 if p.CacheFailed > 0 {
312 return 0, fmt.Errorf("cache drain failed for %d paths", p.CacheFailed)
313 }
314 if p.CacheQueued > 0 || p.CacheActive > 0 {
315 return 0, fmt.Errorf("cache drain incomplete: queued=%d active=%d", p.CacheQueued, p.CacheActive)
316 }
317 return p.CacheUploaded, nil
318 }
319}
320
321func (s *AgentSession) decode(ctx context.Context) (*agentproto.Message, error) {
322 if err := ctx.Err(); err != nil {
323 return nil, err
324 }
325
326 if deadline, ok := ctx.Deadline(); ok {
327 _ = s.conn.SetReadDeadline(deadline)
328 } else {
329 _ = s.conn.SetReadDeadline(time.Time{})
330 }
331
332 // a blocked vsock read wont wake up just from the ctx being cancelled,
333 // only a deadline will wake it up, so if the VM crashes mid-step the read would
334 // hang until workflow timeout. so we will set a deadline in the past to cancel it.
335 //
336 // we set a deadline here instead of closing the connection, this is the long-lived
337 // connection that everything reuses, so we only really want to interrupt it for this
338 // current read. this also lands as a timeout error which the netErr.Timeout() check
339 // below maps to ctx.Err() correctly
340 stop := context.AfterFunc(ctx, func() {
341 _ = s.conn.SetReadDeadline(time.Now())
342 })
343 defer stop()
344
345 msg, err := s.dec.Decode()
346 if err != nil {
347 var netErr net.Error
348 if errors.As(err, &netErr) && netErr.Timeout() && ctx.Err() != nil {
349 return nil, ctx.Err()
350 }
351 return nil, fmt.Errorf("read agent message: %w", err)
352 }
353 return msg, nil
354}
355
356func (s *AgentSession) Close() error {
357 if s == nil || s.conn == nil {
358 return nil
359 }
360 return s.conn.Close()
361}
362
363// this pulls the deadline from the context and converts it to what the
364// agentproto expects
365func timeoutSeconds(ctx context.Context, lead time.Duration) uint32 {
366 deadline, ok := ctx.Deadline()
367 if !ok {
368 return 0
369 }
370 seconds := int64((time.Until(deadline) - lead).Round(time.Second) / time.Second)
371 if seconds < 1 {
372 return 1
373 }
374 if seconds > int64(^uint32(0)) {
375 return ^uint32(0)
376 }
377 return uint32(seconds)
378}