Provisioning/internal/pxe/supervisor.go

package pxe

import (
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"text/template"

	"provisioning/internal/model"
)

type SupervisorConfig struct {
	Enabled    bool
	Interface  string
	Subnet     string
	RuntimeDir string
	TFTPRoot   string
	DnsmasqBin string
	PublicURL  string
}

type Supervisor struct {
	cfg    SupervisorConfig
	cmd    *exec.Cmd
	cancel context.CancelFunc
}

func NewSupervisor(cfg SupervisorConfig) *Supervisor {
	return &Supervisor{cfg: cfg}
}

func (s *Supervisor) Start(ctx context.Context, hosts []model.Host) error {
	if !s.cfg.Enabled {
		log.Printf("pxe: dnsmasq disabled")
		return nil
	}
	if err := os.MkdirAll(s.cfg.RuntimeDir, 0o755); err != nil {
		return fmt.Errorf("pxe: create runtime dir: %w", err)
	}
	if err := os.MkdirAll(s.cfg.TFTPRoot, 0o755); err != nil {
		return fmt.Errorf("pxe: create tftp root: %w", err)
	}
	s.seedIPXE()
	if err := s.writeConfig(hosts); err != nil {
		return err
	}
	return s.startProcess(ctx)
}

func (s *Supervisor) Reload(hosts []model.Host) error {
	if !s.cfg.Enabled {
		return nil
	}
	if err := s.writeConfig(hosts); err != nil {
		return err
	}
	if s.cmd != nil && s.cmd.Process != nil {
		return signalReload(s.cmd.Process)
	}
	return nil
}

func (s *Supervisor) Shutdown() error {
	if s.cancel != nil {
		s.cancel()
	}
	if s.cmd != nil && s.cmd.Process != nil {
		return s.cmd.Process.Kill()
	}
	return nil
}

func (s *Supervisor) writeConfig(hosts []model.Host) error {
	confPath := filepath.Join(s.cfg.RuntimeDir, "dnsmasq.conf")
	hostsPath := filepath.Join(s.cfg.RuntimeDir, "dhcp-hostsfile")

	var lines []string
	for _, h := range hosts {
		lines = append(lines, h.MAC+",set:known")
	}
	if err := os.WriteFile(hostsPath, []byte(strings.Join(lines, "\n")+"\n"), 0o644); err != nil {
		return fmt.Errorf("pxe: write dhcp-hostsfile: %w", err)
	}

	subnet := s.cfg.Subnet
	if i := strings.Index(subnet, "/"); i != -1 {
		subnet = subnet[:i]
	}

	conf := dnsmasqConf{
		Interface:  s.cfg.Interface,
		Subnet:     subnet,
		TFTPRoot:   s.cfg.TFTPRoot,
		HostsFile:  hostsPath,
		PublicURL:  s.cfg.PublicURL,
		RuntimeDir: s.cfg.RuntimeDir,
	}
	f, err := os.Create(confPath)
	if err != nil {
		return fmt.Errorf("pxe: create dnsmasq.conf: %w", err)
	}
	defer f.Close()
	if err := dnsmasqTmpl.Execute(f, conf); err != nil {
		return fmt.Errorf("pxe: render dnsmasq.conf: %w", err)
	}
	return nil
}

func (s *Supervisor) startProcess(ctx context.Context) error {
	confPath := filepath.Join(s.cfg.RuntimeDir, "dnsmasq.conf")
	procCtx, cancel := context.WithCancel(ctx)
	s.cancel = cancel
	s.cmd = exec.CommandContext(procCtx, s.cfg.DnsmasqBin, "--keep-in-foreground", "--conf-file="+confPath)
	s.cmd.Stdout = os.Stdout
	s.cmd.Stderr = os.Stderr
	if err := s.cmd.Start(); err != nil {
		cancel()
		return fmt.Errorf("pxe: start dnsmasq: %w", err)
	}
	go func() {
		if err := s.cmd.Wait(); err != nil {
			select {
			case <-procCtx.Done():
			default:
				log.Printf("pxe: dnsmasq exited: %v", err)
			}
		}
	}()
	log.Printf("pxe: dnsmasq started (pid %d)", s.cmd.Process.Pid)
	return nil
}

var ipxeFiles = []struct{ src, dst string }{
	{"/usr/lib/ipxe/undionly.kpxe", "undionly.kpxe"},
	{"/usr/lib/ipxe/ipxe.efi", "ipxe.efi"},
}

var ipxeSymlinks = []struct{ target, link string }{
	{"undionly.kpxe", "undionly.0"},
	{"ipxe.efi", "ipxe.0"},
}

func (s *Supervisor) seedIPXE() {
	for _, f := range ipxeFiles {
		dst := filepath.Join(s.cfg.TFTPRoot, f.dst)
		if _, err := os.Stat(dst); err == nil {
			continue
		}
		if err := copyFile(f.src, dst); err != nil {
			log.Printf("pxe: copy %s → %s: %v", f.src, dst, err)
		} else {
			log.Printf("pxe: seeded %s", f.dst)
		}
	}
	for _, l := range ipxeSymlinks {
		link := filepath.Join(s.cfg.TFTPRoot, l.link)
		if _, err := os.Lstat(link); err == nil {
			continue
		}
		if err := os.Symlink(l.target, link); err != nil {
			log.Printf("pxe: symlink %s → %s: %v", l.link, l.target, err)
		}
	}
}

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	return err
}

type dnsmasqConf struct {
	Interface  string
	Subnet     string
	TFTPRoot   string
	HostsFile  string
	PublicURL  string
	RuntimeDir string
}

var dnsmasqTmpl = template.Must(template.New("dnsmasq.conf").Parse(`
port=0
interface={{.Interface}}
bind-interfaces

dhcp-range={{.Subnet}},proxy
dhcp-hostsfile={{.HostsFile}}
dhcp-ignore=tag:!known

enable-tftp
tftp-root={{.TFTPRoot}}

# Detect iPXE (user-class "iPXE") to break chainload loop
dhcp-userclass=set:ipxe,iPXE

# Initial PXE firmware: serve iPXE bootloader (skip if already iPXE)
pxe-prompt="Provisioning",0
pxe-service=tag:!ipxe,x86PC,"Provisioning",undionly
pxe-service=tag:!ipxe,X86-64_EFI,"Provisioning",ipxe

# iPXE: serve HTTP boot script
dhcp-boot=tag:ipxe,{{.PublicURL}}/ipxe/${mac:hexhyp}

log-dhcp
log-facility={{.RuntimeDir}}/dnsmasq.log
`))