banger/internal/firecracker/client.go

package firecracker

import (
	"context"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"strings"
	"sync"

	sdk "github.com/firecracker-microvm/firecracker-go-sdk"
	models "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
	"github.com/sirupsen/logrus"

	"banger/internal/vsockagent"
)

type MachineConfig struct {
	BinaryPath      string
	VMID            string
	SocketPath      string
	LogPath         string
	MetricsPath     string
	KernelImagePath string
	InitrdPath      string
	KernelArgs      string
	Drives          []DriveConfig
	TapDevice       string
	VSockPath       string
	VSockCID        uint32
	VCPUCount       int
	MemoryMiB       int
	Logger          *slog.Logger
}

type DriveConfig struct {
	ID       string
	Path     string
	ReadOnly bool
	IsRoot   bool
}

type Machine struct {
	machine   *sdk.Machine
	logFile   *os.File
	closeOnce sync.Once
}

type Client struct {
	client *sdk.Client
}

func NewMachine(ctx context.Context, cfg MachineConfig) (*Machine, error) {
	logFile, err := openLogFile(cfg.LogPath)
	if err != nil {
		return nil, err
	}

	cmd := buildProcessRunner(cfg, logFile)
	machine, err := sdk.NewMachine(
		ctx,
		buildConfig(cfg),
		sdk.WithProcessRunner(cmd),
		sdk.WithLogger(newLogger(cfg.Logger)),
	)
	if err != nil {
		if logFile != nil {
			_ = logFile.Close()
		}
		return nil, err
	}

	return &Machine{machine: machine, logFile: logFile}, nil
}

func (m *Machine) Start(ctx context.Context) error {
	// The caller's ctx is INTENTIONALLY not forwarded to the SDK.
	// firecracker-go-sdk's startVMM (machine.go) spawns a goroutine
	// that SIGTERMs firecracker the instant this ctx cancels, and
	// retains it for the lifetime of the VMM — not just the boot
	// phase. Plumbing an RPC request ctx through would mean
	// firecracker dies the moment the daemon writes its RPC response
	// (daemon.go:handleConn defers cancel). That silently breaks
	// `vm start` on a stopped VM: start "succeeds", the handler
	// returns, ctx cancels, firecracker is SIGTERMed, and the next
	// `vm ssh` hits `vmAlive = false`. `vm.create` sidesteps the bug
	// because BeginVMCreate detaches to a background ctx before
	// calling startVMLocked.
	//
	// We own firecracker lifecycle explicitly — StopVM / KillVM /
	// cleanupRuntime — so losing ctx-driven cancellation here is
	// deliberate. The SDK still enforces its own boot-phase timeouts
	// (socket wait, HTTP) with internal deadlines.
	if err := m.machine.Start(context.Background()); err != nil {
		m.closeLog()
		return err
	}
	_ = ctx

	go func() {
		_ = m.machine.Wait(context.Background())
		m.closeLog()
	}()

	return nil
}

func (m *Machine) PID() (int, error) {
	return m.machine.PID()
}

func New(apiSock string, logger *slog.Logger) *Client {
	return &Client{client: sdk.NewClient(apiSock, newLogger(logger), false)}
}

func (c *Client) SendCtrlAltDel(ctx context.Context) error {
	action := models.InstanceActionInfoActionTypeSendCtrlAltDel
	_, err := c.client.CreateSyncAction(ctx, &models.InstanceActionInfo{
		ActionType: &action,
	})
	return err
}

func openLogFile(path string) (*os.File, error) {
	if path == "" {
		return nil, nil
	}
	return os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
}

func buildConfig(cfg MachineConfig) sdk.Config {
	rootDrive, extraDrives := splitDrives(cfg.Drives)
	drivesBuilder := sdk.NewDrivesBuilder(rootDrive.Path).
		WithRootDrive(rootDrive.Path, sdk.WithDriveID(defaultDriveID(rootDrive, "rootfs")), sdk.WithReadOnly(rootDrive.ReadOnly))
	for _, drive := range extraDrives {
		if strings.TrimSpace(drive.Path) == "" {
			continue
		}
		drivesBuilder = drivesBuilder.AddDrive(drive.Path, drive.ReadOnly, sdk.WithDriveID(defaultDriveID(drive, "drive")))
	}
	drives := drivesBuilder.Build()

	return sdk.Config{
		SocketPath:      cfg.SocketPath,
		LogPath:         cfg.LogPath,
		MetricsPath:     cfg.MetricsPath,
		KernelImagePath: cfg.KernelImagePath,
		InitrdPath:      cfg.InitrdPath,
		KernelArgs:      cfg.KernelArgs,
		Drives:          drives,
		NetworkInterfaces: sdk.NetworkInterfaces{{
			StaticConfiguration: &sdk.StaticNetworkConfiguration{
				HostDevName: cfg.TapDevice,
			},
		}},
		VsockDevices: buildVsockDevices(cfg),
		MachineCfg: models.MachineConfiguration{
			VcpuCount:  sdk.Int64(int64(cfg.VCPUCount)),
			MemSizeMib: sdk.Int64(int64(cfg.MemoryMiB)),
			Smt:        sdk.Bool(false),
		},
		VMID: cfg.VMID,
	}
}

func buildVsockDevices(cfg MachineConfig) []sdk.VsockDevice {
	if strings.TrimSpace(cfg.VSockPath) == "" || cfg.VSockCID == 0 {
		return nil
	}
	return []sdk.VsockDevice{{
		ID:   "vsock",
		Path: cfg.VSockPath,
		CID:  cfg.VSockCID,
	}}
}

func splitDrives(drives []DriveConfig) (DriveConfig, []DriveConfig) {
	root := DriveConfig{ID: "rootfs"}
	var extras []DriveConfig
	for _, drive := range drives {
		if strings.TrimSpace(drive.Path) == "" {
			continue
		}
		if drive.IsRoot {
			root = drive
			if root.ID == "" {
				root.ID = "rootfs"
			}
			continue
		}
		extras = append(extras, drive)
	}
	return root, extras
}

func defaultDriveID(drive DriveConfig, fallback string) string {
	if strings.TrimSpace(drive.ID) != "" {
		return drive.ID
	}
	return fallback
}

// buildProcessRunner constructs the *exec.Cmd the SDK will start. Args are
// passed directly — no shell, no string interpolation — so any future change
// to MachineConfig fields can't smuggle shell metacharacters into the launch.
//
// The daemon and root-helper processes set umask 077 at startup, so the
// API/vsock sockets firecracker creates inherit 0600 mode without needing a
// shell-level `umask` wrapper.
//
// When firecracker has to be launched under sudo (non-root daemon), the
// resulting sockets are root-owned. The caller (LaunchFirecracker) kicks off
// fcproc.EnsureSocketAccessForAsync immediately *before* Machine.Start so the
// chown wins the race against the SDK's HTTP probe over the API socket. That
// replaces the previous in-shell chown_watcher.
func buildProcessRunner(cfg MachineConfig, logFile *os.File) *exec.Cmd {
	args := []string{"--api-sock", cfg.SocketPath, "--id", cfg.VMID}
	var cmd *exec.Cmd
	if os.Geteuid() == 0 {
		cmd = exec.Command(cfg.BinaryPath, args...)
	} else {
		cmd = exec.Command("sudo", append([]string{"-n", "-E", cfg.BinaryPath}, args...)...)
	}
	cmd.Stdin = nil
	if logFile != nil {
		cmd.Stdout = logFile
		cmd.Stderr = logFile
	}
	return cmd
}

func newLogger(base *slog.Logger) *logrus.Entry {
	logger := logrus.New()
	logger.SetOutput(io.Discard)
	logger.SetLevel(logrus.DebugLevel)
	logger.AddHook(slogHook{logger: base})
	return logrus.NewEntry(logger)
}

func HealthVSock(ctx context.Context, logger *slog.Logger, socketPath string) error {
	return vsockagent.Health(ctx, logger, socketPath)
}

func PingVSock(ctx context.Context, logger *slog.Logger, socketPath string) error {
	return HealthVSock(ctx, logger, socketPath)
}

type slogHook struct {
	logger *slog.Logger
}

func (h slogHook) Levels() []logrus.Level {
	return logrus.AllLevels
}

func (h slogHook) Fire(entry *logrus.Entry) error {
	if h.logger == nil {
		return nil
	}
	level := slog.LevelDebug
	switch entry.Level {
	case logrus.PanicLevel, logrus.FatalLevel, logrus.ErrorLevel:
		level = slog.LevelError
	case logrus.WarnLevel:
		level = slog.LevelWarn
	default:
		level = slog.LevelDebug
	}
	attrs := make([]any, 0, len(entry.Data)*2+2)
	attrs = append(attrs, "component", "firecracker_sdk")
	for key, value := range entry.Data {
		attrs = append(attrs, key, value)
	}
	h.logger.Log(context.Background(), level, entry.Message, attrs...)
	return nil
}

func (m *Machine) closeLog() {
	m.closeOnce.Do(func() {
		if m.logFile != nil {
			_ = m.logFile.Close()
		}
	})
}