mkvm/common.go

package main

import (
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"net/http"
	"strings"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"libvirt.org/go/libvirt"
	"libvirt.org/go/libvirtxml"

	"mkvm/config"
	"mkvm/libvirtx"
	"mkvm/volumes/pools"
)

var (
	argMemoryMB   int
	argCPUs       int
	argDiskSizeGB int

	// cloudinit args
	argSSHKeys    []string
	argSSHKeyURLs []string
	argPackages   []string
)

type domainModifier func(*libvirtxml.Domain) error

func getLibvirtAndPool() (*libvirt.Connect, pools.StoragePool, error) {
	conn, err := libvirtx.New()
	if err != nil {
		return nil, nil, err
	}
	defer conn.Close()

	pool, err := pools.GetPool(conn, config.C.StoragePool)
	if err != nil {
		return nil, nil, err
	}

	return conn, pool, nil
}

func getServerListenAddress(conn *libvirt.Connect) (string, error) {
	serverInterfaceName := ""
	if config.C.Network != "" {
		nicSource.Network = &libvirtxml.DomainInterfaceSourceNetwork{Network: config.C.Network}
		libvirtnet, err := conn.LookupNetworkByName(config.C.Network)
		if err != nil {
			logrus.WithField("network", config.C.Network).Error("error finding libvirt network")
			return "", err
		}

		xmlstr, err := libvirtnet.GetXMLDesc(0)
		if err != nil {
			logrus.WithField("network", config.C.Network).Error("error getting network xml description")
			return "", err
		}

		var net libvirtxml.Network
		if err := net.Unmarshal(xmlstr); err != nil {
			logrus.WithField("network", config.C.Network).Error("error parsing network xml description")
			return "", err
		}

		serverInterfaceName = net.Bridge.Name
	} else if config.C.Bridge != "" {
		nicSource.Bridge = &libvirtxml.DomainInterfaceSourceBridge{Bridge: config.C.Bridge}
		serverInterfaceName = config.C.Bridge
	} else {
		return "", errors.New("no network or bridge configured")
	}

	serverInterface, err := net.InterfaceByName(serverInterfaceName)
	if err != nil {
		logrus.Error("error finding local network interface to run server on")
		return "", err
	}

	serverInterfaceAddrs, err := serverInterface.Addrs()
	if err != nil {
		logrus.Error("error finding local network interface's IP")
		return "", err
	}

	if len(serverInterfaceAddrs) == 0 {
		return "", fmt.Errorf("bridge interface %s does not have an IP on this machine", serverInterfaceName)
	}

	serverBindIP, _, err := net.ParseCIDR(serverInterfaceAddrs[0].String())
	if err != nil {
		logrus.WithField("interface", serverInterfaceName).WithField("address", serverInterfaceAddrs[0].String()).Error("error parsing local address")
		return "", err
	}

	port := rand.Intn(65535-1025) + 1025
	return fmt.Sprintf("%s:%d", serverBindIP, port), nil
}

func createDomain(conn *libvirt.Connect, pool pools.StoragePool, name string, modifiers ...domainModifier) error {
	interfaces := []libvirtxml.DomainInterface{
		{
			Model:  &libvirtxml.DomainInterfaceModel{Type: "virtio"},
			Source: &nicSource,
		},
	}

	// Create the domain
	domainXML := &libvirtxml.Domain{
		Type:   "kvm",
		Name:   name,
		Memory: &libvirtxml.DomainMemory{Value: uint(argMemoryMB), Unit: "MiB"},
		VCPU:   &libvirtxml.DomainVCPU{Value: uint(argCPUs)},
		OS: &libvirtxml.DomainOS{
			Type:        &libvirtxml.DomainOSType{Arch: "x86_64", Type: "hvm"},
			BootDevices: []libvirtxml.DomainBootDevice{{Dev: "hd"}},
		},
		Features: &libvirtxml.DomainFeatureList{
			ACPI:   &libvirtxml.DomainFeature{},
			APIC:   &libvirtxml.DomainFeatureAPIC{},
			VMPort: &libvirtxml.DomainFeatureState{State: "off"},
		},
		CPU: &libvirtxml.DomainCPU{Mode: "host-model"},
		Devices: &libvirtxml.DomainDeviceList{
			Emulator: "/usr/bin/kvm",
			Disks:    []libvirtxml.DomainDisk{pool.GetDomainDiskXML(name)},
			Channels: []libvirtxml.DomainChannel{
				{
					Source: &libvirtxml.DomainChardevSource{
						UNIX: &libvirtxml.DomainChardevSourceUNIX{Path: "/var/lib/libvirt/qemu/f16x86_64.agent", Mode: "bind"},
					},
					Target: &libvirtxml.DomainChannelTarget{
						VirtIO: &libvirtxml.DomainChannelTargetVirtIO{Name: "org.qemu.guest_agent.0"},
					},
				},
			},
			Consoles:   []libvirtxml.DomainConsole{{Target: &libvirtxml.DomainConsoleTarget{}}},
			Serials:    []libvirtxml.DomainSerial{{}},
			Interfaces: interfaces,
		},
	}

	for _, modifier := range modifiers {
		if err := modifier(domainXML); err != nil {
			return err
		}
	}

	domainXMLString, err := domainXML.Marshal()
	if err != nil {
		return err
	}

	logrus.Debug("defining domain from xml")
	domain, err := conn.DomainDefineXML(domainXMLString)
	if err != nil {
		return fmt.Errorf("error defining domain from xml description: %v", err)
	}

	logrus.Debug("booting domain")
	err = domain.Create()
	if err != nil {
		return fmt.Errorf("error creating domain: %v", err)
	}

	return nil
}

func setSMBIOS(smbios map[int]map[string]string) domainModifier {
	return func(d *libvirtxml.Domain) error {
		qemuArgs := []libvirtxml.DomainQEMUCommandlineArg{}
		if d.QEMUCommandline != nil {
			qemuArgs = d.QEMUCommandline.Args
		}

		for smbiosType, values := range smbios {
			arg := libvirtxml.DomainQEMUCommandlineArg{
				Value: fmt.Sprintf("type=%d", smbiosType),
			}
			for key, value := range values {
				arg.Value = fmt.Sprintf("%s,%s=%s", arg.Value, key, value)
			}
			qemuArgs = append(qemuArgs, libvirtxml.DomainQEMUCommandlineArg{Value: "-smbios"}, arg)
		}

		d.QEMUCommandline = &libvirtxml.DomainQEMUCommandline{Args: qemuArgs}

		return nil
	}
}

func downloadSSHKeys(url string) error {
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusOK {
		logrus.WithFields(logrus.Fields{
			"url":    url,
			"status": resp.Status,
			"body":   string(body),
		}).Error("non-200 response from SSH key URL")
		return fmt.Errorf("non-200 response from SSH key URL: %s", resp.Status)
	}

	count := 0
	for _, key := range strings.Split(string(body), "\n") {
		key = strings.TrimSpace(key)
		if key == "" {
			continue
		}
		argSSHKeys = append(argSSHKeys, key)
		count++
	}

	logrus.WithField("url", url).WithField("keys", count).Debug("downloaded SSH authorized keys")

	return nil
}

func registerGlobalFlags(cmd cobra.Command) {
	cmd.Flags().IntVarP(&argMemoryMB, "memory", "m", 1024, "amount of memory (in MB) to assign to the VM")
	cmd.Flags().IntVarP(&argCPUs, "cpu", "c", 2, "the number of vCPU cores to assign to the VM")
	cmd.Flags().IntVarP(&argDiskSizeGB, "disk", "d", 25, "disk size (in GB)")
}