This repository has no description
10 kB
357 lines
1//go:build linux
2
3package main
4
5import (
6 "context"
7 "database/sql"
8 "errors"
9 "fmt"
10 "log/slog"
11 "net"
12 "os"
13 "time"
14
15 "github.com/mdlayher/vsock"
16 "github.com/urfave/cli/v3"
17 agentv1 "tangled.org/core/spindle/agentproto/gen"
18 "tangled.org/core/spindle/db"
19 "tangled.org/core/spindle/engines/microvm"
20)
21
22func SpindleMicroVMRunCommand() *cli.Command {
23 return &cli.Command{
24 Name: "spindle-microvm-run",
25 Usage: "launch the Spindle base microVM and run one command over vsock",
26 Flags: []cli.Flag{
27 &cli.StringFlag{
28 Name: "image-spec",
29 Sources: cli.EnvVars("SPINDLE_MICROVM_IMAGE_SPEC"),
30 Usage: "path to microVM image spec JSON",
31 },
32 &cli.StringFlag{
33 Name: "mkfs-ext4",
34 Usage: "override mkfs.ext4 binary",
35 },
36 &cli.StringFlag{
37 Name: "work-dir",
38 Usage: "directory for per-run disks and sockets",
39 },
40 &cli.UintFlag{
41 Name: "cid",
42 Usage: "guest vsock CID; defaults to a random high CID",
43 },
44 &cli.UintFlag{
45 Name: "port",
46 Value: 10240,
47 Usage: "host vsock port to listen on",
48 },
49 &cli.UintFlag{
50 Name: "memory-mib",
51 Usage: "override the guest memory size in MiB (defaults to the image spec)",
52 },
53 &cli.BoolFlag{
54 Name: "disable-kvm",
55 Usage: "run without -enable-kvm even if /dev/kvm is available",
56 },
57 &cli.BoolFlag{
58 Name: "dev",
59 Usage: "enable dev mode (allows host network access, disables SSL verification)",
60 },
61 &cli.DurationFlag{
62 Name: "qmp-timeout",
63 Value: 10 * time.Second,
64 Usage: "how long to wait for qmp to become ready",
65 },
66 &cli.DurationFlag{
67 Name: "accept-timeout",
68 Value: 15 * time.Second,
69 Usage: "how long to wait for the guest agent after qemu starts",
70 },
71 &cli.DurationFlag{
72 Name: "exec-timeout",
73 Value: 30 * time.Second,
74 Usage: "timeout for the guest command",
75 },
76 &cli.DurationFlag{
77 Name: "cache-drain-timeout",
78 Value: 5 * time.Minute,
79 Usage: "how long to wait for queued cache uploads after the guest command exits",
80 },
81 &cli.DurationFlag{
82 Name: "shutdown-timeout",
83 Value: 10 * time.Second,
84 Usage: "how long to wait for qemu to exit after guest powerdown",
85 },
86 &cli.StringFlag{
87 Name: "cwd",
88 Usage: "guest working directory",
89 },
90 &cli.StringSliceFlag{
91 Name: "cache-read-url",
92 Sources: cli.EnvVars("SPINDLE_NIX_CACHE_READ_URLS"),
93 Usage: "Nix binary cache URL to pass to the guest; repeatable",
94 },
95 &cli.StringSliceFlag{
96 Name: "cache-trusted-public-key",
97 Sources: cli.EnvVars("SPINDLE_NIX_CACHE_TRUSTED_PUBLIC_KEYS"),
98 Usage: "Nix binary cache public key to trust in the guest; repeatable",
99 },
100 &cli.StringFlag{
101 Name: "cache-upload-url",
102 Sources: cli.EnvVars("SPINDLE_NIX_CACHE_UPLOAD_URL"),
103 Usage: "optional cache upload URL for guest-built store paths",
104 },
105 &cli.StringFlag{
106 Name: "activate-config",
107 Usage: "JSON user config to activate before exec (e.g. '{\"services\":{\"openssh\":{\"enable\":true}}}')",
108 },
109 &cli.StringFlag{
110 Name: "db",
111 Usage: "path to sqlite database for config cache",
112 },
113 },
114 Action: runMicroVMRunDev,
115 }
116}
117
118func runMicroVMRunDev(ctx context.Context, cmd *cli.Command) error {
119 imageSpecPath := cmd.String("image-spec")
120 if imageSpecPath == "" {
121 return fmt.Errorf("--image-spec or SPINDLE_MICROVM_IMAGE_SPEC is required")
122 }
123
124 imageSpec, err := microvm.LoadImageSpec(imageSpecPath)
125 if err != nil {
126 return err
127 }
128
129 port := uint32(cmd.Uint("port"))
130 // tell the guest which host vsock port to dial back on. shuttle reads the
131 // cmdline params this is so we can run multiple of this process
132 // concurrently, because otherwise it listens on a specific vsock port, and
133 // we cant bind to the same port twice...
134 imageSpec.BootArgs = fmt.Sprintf("%s shuttle.vsock_port=%d", imageSpec.BootArgs, port)
135 if mib := cmd.Uint("memory-mib"); mib > 0 {
136 imageSpec.MemoryMiB = int(mib)
137 }
138 ln, err := vsock.Listen(port, nil)
139 if err != nil {
140 return fmt.Errorf("listen on vsock port %d: %w", port, err)
141 }
142 defer ln.Close()
143
144 vm, err := microvm.StartVM(ctx, microvm.VMConfig{
145 Image: imageSpec,
146 BootTimeout: cmd.Duration("qmp-timeout"),
147 CID: uint32(cmd.Uint("cid")),
148 EnableKVM: !cmd.Bool("disable-kvm"),
149 MkfsExt4: cmd.String("mkfs-ext4"),
150 WorkDir: cmd.String("work-dir"),
151 Dev: cmd.Bool("dev"),
152 }, slog.Default())
153 if err != nil {
154 return err
155 }
156 defer vm.Close()
157
158 logs := vm.Logs()
159 fmt.Fprintf(os.Stderr, "microvm started: cid=%d work-dir=%s serial-log=%s qemu-log=%s\n",
160 vm.CID(),
161 vm.WorkDir(),
162 logs.Serial,
163 logs.Extra["qemu"],
164 )
165
166 logger := slog.Default()
167
168 if cmd.Duration("accept-timeout") > 0 {
169 if err := ln.SetDeadline(time.Now().Add(cmd.Duration("accept-timeout"))); err != nil {
170 return fmt.Errorf("set accept deadline: %w", err)
171 }
172 }
173
174 argv := cmd.Args().Slice()
175 if len(argv) == 0 {
176 argv = []string{"/run/current-system/sw/bin/echo", "hello-from-spindle"}
177 }
178 jobID := "spindle-microvm-run"
179 execID := "dev-1"
180
181 fmt.Fprintf(os.Stderr, "listening for agent on %s\n", ln.Addr())
182 conn, err := acceptExpectedVsockConn(ln, vm.CID(), logger)
183 if err != nil {
184 return fmt.Errorf("accept agent connection: %w", err)
185 }
186 defer conn.Close()
187
188 upstreams, err := microvm.BuildCacheUpstreams(cmd.StringSlice("cache-read-url"), nil)
189 if err != nil {
190 return fmt.Errorf("build cache upstreams: %w", err)
191 }
192
193 var readCache *microvm.ReadCacheProxy
194 if len(cmd.StringSlice("cache-read-url")) > 0 {
195 var err error
196 readCache, err = microvm.StartReadCacheProxy(ctx, vm.CID(), upstreams, logger)
197 if err != nil {
198 return fmt.Errorf("start read cache proxy: %w", err)
199 }
200 defer readCache.Close()
201 }
202
203 var uploadCache *microvm.UploadCacheProxy
204 if cmd.String("cache-upload-url") != "" {
205 var err error
206 uploadCache, err = microvm.StartUploadCacheProxy(ctx, vm.CID(), cmd.String("cache-upload-url"), upstreams, logger)
207 if err != nil {
208 return fmt.Errorf("start upload cache proxy: %w", err)
209 }
210 defer uploadCache.Close()
211 }
212 dnsProxy, err := microvm.StartDNSProxy(ctx, vm.CID(), logger)
213 if err != nil {
214 return fmt.Errorf("start dns proxy: %w", err)
215 }
216 defer dnsProxy.Close()
217
218 session := microvm.NewAgentSession(conn, logger)
219
220 initCtx, cancelInit := context.WithTimeout(ctx, 30*time.Second)
221 defer cancelInit()
222 if err := session.Init(initCtx, &agentv1.Init{
223 JobId: jobID,
224 CacheTrustedPublicKeys: cmd.StringSlice("cache-trusted-public-key"),
225 CacheReadProxyPort: readCache.Port(),
226 CacheUploadProxyPort: uploadCache.Port(),
227 DnsProxyPort: dnsProxy.Port(),
228 }); err != nil {
229 return fmt.Errorf("init agent: %w", err)
230 }
231
232 execCtx := ctx
233 if cmd.Duration("exec-timeout") > 0 {
234 var cancel context.CancelFunc
235 execCtx, cancel = context.WithTimeout(ctx, cmd.Duration("exec-timeout"))
236 defer cancel()
237 }
238
239 if cmd.String("activate-config") != "" {
240 actCtx := execCtx
241 baseHash, err := microvm.BaseConfigHash(imageSpec)
242 if err != nil {
243 return fmt.Errorf("calculate base config hash: %w", err)
244 }
245
246 var d *db.DB
247 var configKey string
248 var cachedToplevel string
249 if cmd.String("db") != "" {
250 d, err = db.Make(ctx, cmd.String("db"))
251 if err != nil {
252 return fmt.Errorf("failed to open database: %w", err)
253 }
254 defer d.Close()
255
256 configKey, err = microvm.BuildConfigKey(imageSpec, cmd.String("activate-config"))
257 if err != nil {
258 return fmt.Errorf("calculate config key: %w", err)
259 }
260
261 record, err := d.GetNixOSToplevelCacheRecord(configKey)
262 if err != nil {
263 if !errors.Is(err, sql.ErrNoRows) {
264 return fmt.Errorf("lookup config cache: %w", err)
265 }
266 } else {
267 cachedToplevel = record.Toplevel
268 fmt.Printf("realizing cached NixOS config %s\n", cachedToplevel)
269 }
270 }
271
272 result, err := session.ActivateConfig(actCtx, "dev-activate", &agentv1.ActivateConfig{
273 ConfigKey: configKey,
274 BaseConfigHash: baseHash,
275 UserConfig: cmd.String("activate-config"),
276 Toplevel: cachedToplevel,
277 })
278 if err != nil {
279 return fmt.Errorf("activate config: %w", err)
280 }
281 fmt.Fprintf(os.Stderr, "activated config toplevel: %s\n", result.Toplevel)
282
283 if d != nil && cachedToplevel == "" && result.Toplevel != "" && configKey != "" {
284 err = d.SaveNixOSToplevelCacheRecord(configKey, result.Toplevel)
285 if err != nil {
286 return fmt.Errorf("save config cache: %w", err)
287 }
288 }
289 }
290
291 exitCode, err := session.Exec(execCtx, microvm.AgentExec{
292 ID: execID,
293 ExecStart: &agentv1.ExecStart{
294 Argv: argv,
295 Cwd: cmd.String("cwd"),
296 },
297 Stdout: os.Stdout,
298 Stderr: os.Stderr,
299 })
300 if err != nil {
301 return err
302 }
303
304 if uploadCache != nil {
305 drainCtx := ctx
306 if cmd.Duration("cache-drain-timeout") > 0 {
307 var cancel context.CancelFunc
308 drainCtx, cancel = context.WithTimeout(ctx, cmd.Duration("cache-drain-timeout"))
309 defer cancel()
310 }
311 uploaded, err := session.Drain(drainCtx)
312 if err != nil {
313 return err
314 }
315 fmt.Printf("cache uploaded: %d\n", uploaded)
316 }
317
318 // mirror the engine shutdown order: ask the agent to power off first,
319 // then fall back to qemu powerdown / kill
320 shutdownCtx, cancel := context.WithTimeout(context.Background(), cmd.Duration("shutdown-timeout"))
321 defer cancel()
322 poweredOff := false
323 if err := session.Poweroff(shutdownCtx); err != nil {
324 fmt.Fprintf(os.Stderr, "agent poweroff: %s\n", err)
325 } else if err := vm.WaitContext(shutdownCtx); err == nil {
326 poweredOff = true
327 }
328 if !poweredOff {
329 if err := vm.Shutdown(shutdownCtx); err != nil {
330 fmt.Fprintf(os.Stderr, "microvm shutdown fallback: %s\n", err)
331 }
332 }
333
334 if exitCode != 0 {
335 return fmt.Errorf("guest command exited with code %d", exitCode)
336 }
337 return nil
338}
339
340func acceptExpectedVsockConn(ln *vsock.Listener, allowedCID uint32, logger *slog.Logger) (net.Conn, error) {
341 for {
342 conn, err := ln.Accept()
343 if err != nil {
344 return nil, err
345 }
346 if allowedCID == 0 {
347 return conn, nil
348 }
349 addr, ok := conn.RemoteAddr().(*vsock.Addr)
350 if ok && addr.ContextID == allowedCID {
351 return conn, nil
352 }
353 remote := conn.RemoteAddr()
354 _ = conn.Close()
355 logger.Warn("dropped agent connection from unexpected cid", "remote", remote, "expected", allowedCID)
356 }
357}