feat: vsock mmds snapshot

internal/firecracker/api_test.go

@@ -38,6 +38,7 @@ func TestPutSnapshotLoadIncludesNetworkOverrides(t *testing.T) {
 				HostDevName: "fctap7",
 			},
 		},
+		VsockOverride: &VsockOverride{UDSPath: "/run/microagent-personalizer.vsock"},
 	})
 	if err != nil {
 		t.Fatalf("put snapshot load: %v", err)
@@ -47,7 +48,7 @@ func TestPutSnapshotLoadIncludesNetworkOverrides(t *testing.T) {
 		t.Fatalf("request path mismatch: got %q want %q", gotPath, "/snapshot/load")
 	}
 
-	want := "{\"snapshot_path\":\"vmstate.bin\",\"mem_backend\":{\"backend_type\":\"File\",\"backend_path\":\"memory.bin\"},\"resume_vm\":false,\"network_overrides\":[{\"iface_id\":\"net0\",\"host_dev_name\":\"fctap7\"}]}"
+	want := "{\"snapshot_path\":\"vmstate.bin\",\"mem_backend\":{\"backend_type\":\"File\",\"backend_path\":\"memory.bin\"},\"resume_vm\":false,\"network_overrides\":[{\"iface_id\":\"net0\",\"host_dev_name\":\"fctap7\"}],\"vsock_override\":{\"uds_path\":\"/run/microagent-personalizer.vsock\"}}"
 	if gotBody != want {
 		t.Fatalf("request body mismatch:\n got: %s\nwant: %s", gotBody, want)
 	}

internal/firecracker/configure_test.go

@@ -168,6 +168,74 @@ func TestConfigureMachineConfiguresMMDSBeforeStart(t *testing.T) {
 	}
 }
 
+func TestConfigureMachineConfiguresVsockBeforeStart(t *testing.T) {
+	var requests []capturedRequest
+
+	socketPath, shutdown := startUnixSocketServer(t, func(w http.ResponseWriter, r *http.Request) {
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			t.Fatalf("read request body: %v", err)
+		}
+		requests = append(requests, capturedRequest{
+			Method: r.Method,
+			Path:   r.URL.Path,
+			Body:   string(body),
+		})
+		w.WriteHeader(http.StatusNoContent)
+	})
+	defer shutdown()
+
+	client := newAPIClient(socketPath)
+	spec := MachineSpec{
+		ID:              "vm-3",
+		VCPUs:           1,
+		MemoryMiB:       512,
+		KernelImagePath: "/kernel",
+		RootFSPath:      "/rootfs",
+		Vsock: &VsockSpec{
+			ID:   "microagent-personalizer",
+			CID:  42,
+			Path: "/run/microagent-personalizer.vsock",
+		},
+	}
+	paths := machinePaths{JailedSerialLogPath: "/logs/serial.log"}
+	network := NetworkAllocation{
+		InterfaceID: defaultInterfaceID,
+		TapName:     "fctap0",
+		GuestMAC:    "06:00:ac:10:00:02",
+	}
+
+	if err := configureMachine(context.Background(), client, paths, spec, network); err != nil {
+		t.Fatalf("configure machine: %v", err)
+	}
+
+	gotPaths := make([]string, 0, len(requests))
+	for _, request := range requests {
+		gotPaths = append(gotPaths, request.Path)
+	}
+	wantPaths := []string{
+		"/machine-config",
+		"/boot-source",
+		"/drives/root_drive",
+		"/network-interfaces/net0",
+		"/entropy",
+		"/serial",
+		"/vsock",
+		"/actions",
+	}
+	if len(gotPaths) != len(wantPaths) {
+		t.Fatalf("request count mismatch: got %d want %d (%v)", len(gotPaths), len(wantPaths), gotPaths)
+	}
+	for i := range wantPaths {
+		if gotPaths[i] != wantPaths[i] {
+			t.Fatalf("request %d mismatch: got %q want %q", i, gotPaths[i], wantPaths[i])
+		}
+	}
+	if requests[6].Body != "{\"guest_cid\":42,\"uds_path\":\"/run/microagent-personalizer.vsock\"}" {
+		t.Fatalf("vsock body mismatch: got %q", requests[6].Body)
+	}
+}
+
 func startUnixSocketServer(t *testing.T, handler http.HandlerFunc) (string, func()) {
 	t.Helper()
 

internal/firecracker/launch.go

@@ -132,7 +132,7 @@ func stageMachineFiles(spec MachineSpec, paths machinePaths) (MachineSpec, error
 
 	if spec.Vsock != nil {
 		vsock := *spec.Vsock
-		vsock.Path = jailedVSockPath(spec)
+		vsock.Path = jailedVSockDevicePath(*spec.Vsock)
 		staged.Vsock = &vsock
 	}
 
@@ -244,11 +244,8 @@ func waitForPIDFile(ctx context.Context, pidFilePath string) (int, error) {
 	}
 }
 
-func jailedVSockPath(spec MachineSpec) string {
-	if spec.Vsock == nil {
-		return ""
-	}
-	return path.Join(defaultVSockRunDir, filepath.Base(strings.TrimSpace(spec.Vsock.Path)))
+func jailedVSockDevicePath(spec VsockSpec) string {
+	return path.Join(defaultVSockRunDir, filepath.Base(strings.TrimSpace(spec.Path)))
 }
 
 func linkMachineFile(source string, target string) error {

internal/firecracker/runtime.go

@@ -331,6 +331,11 @@ func (r *Runtime) RestoreBoot(ctx context.Context, loadSpec SnapshotLoadSpec, us
 		}
 	}
 
+	var vsockOverride *VsockOverride
+	if loadSpec.Vsock != nil {
+		vsockOverride = &VsockOverride{UDSPath: jailedVSockDevicePath(*loadSpec.Vsock)}
+	}
+
 	// Load snapshot (replaces the full configure+start sequence)
 	if err := client.PutSnapshotLoad(ctx, SnapshotLoadParams{
 		SnapshotPath: chrootStatePath,
@@ -345,6 +350,7 @@ func (r *Runtime) RestoreBoot(ctx context.Context, loadSpec SnapshotLoadSpec, us
 				HostDevName: network.TapName,
 			},
 		},
+		VsockOverride: vsockOverride,
 	}); err != nil {
 		cleanup(network, paths, command, firecrackerPID)
 		return nil, fmt.Errorf("load snapshot: %w", err)
@@ -369,6 +375,11 @@ func (r *Runtime) RestoreBoot(ctx context.Context, loadSpec SnapshotLoadSpec, us
 	return &state, nil
 }
 
+func (r *Runtime) PutMMDS(ctx context.Context, state MachineState, data any) error {
+	client := newAPIClient(state.SocketPath)
+	return client.PutMMDS(ctx, data)
+}
+
 func processExists(pid int) bool {
 	if pid < 1 {
 		return false