fix: make OpenCore NBD injection lifecycle safe

qemu/inject.go

@@ -8,91 +8,265 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"strconv"
+	"strings"
 	"time"
 
 	"github.com/projecteru2/core/log"
 	"howett.net/plist"
 
+	"github.com/cocoonstack/cocoon/lock/flock"
 	"github.com/cocoonstack/cocoon/utils"
 )
 
+const (
+	nbdDeviceCount    = 16
+	nbdLeasePath      = "/run/lock/cocoon-macos-nbd.lock"
+	nbdCleanupTimeout = 15 * time.Second
+	nbdCommandTimeout = 5 * time.Second
+)
+
+type nbdConnection struct {
+	device  string
+	pid     int
+	pidFile string
+	ocPath  string
+}
+
 // InjectConfig mounts the OpenCore qcow2 via qemu-nbd (the only way to edit a FAT partition inside a qcow2; needs root + the nbd module) and patches config.plist with a per-VM identity.
 func InjectConfig(ctx context.Context, ocPath string, sm *SMBIOS) error {
-	_ = exec.CommandContext(ctx, "modprobe", "nbd", "max_part=8").Run()
-	nbd, err := connectFreeNBD(ctx, ocPath)
-	if err != nil {
-		return err
-	}
-	// cleanup stays on plain exec.Command: it must still run after ctx cancellation
-	defer disconnectNBD(ctx, nbd, ocPath)
-	waitForPart(ctx, nbd)
-	mnt, err := os.MkdirTemp("", "oc-efi-")
-	if err != nil {
-		return fmt.Errorf("create mount dir: %w", err)
-	}
-	defer func() { _ = os.RemoveAll(mnt) }()
-	var mountErr error
-	for _, p := range []string{nbd + "p1", nbd + "p2", nbd} {
-		if mountErr = exec.CommandContext(ctx, "mount", p, mnt).Run(); mountErr == nil {
-			break
+	return withNBDLease(ctx, nbdLeasePath, func() (retErr error) {
+		_ = exec.CommandContext(ctx, "modprobe", "nbd", "max_part=8").Run()
+		conn, connectErr := connectFreeNBD(ctx, ocPath)
+		if connectErr != nil {
+			return connectErr
 		}
+		defer func() { retErr = errors.Join(retErr, conn.disconnect()) }()
+		if waitErr := waitForPart(ctx, conn.device); waitErr != nil {
+			return waitErr
+		}
+		mnt, err := os.MkdirTemp("", "oc-efi-")
+		if err != nil {
+			return fmt.Errorf("create mount dir: %w", err)
+		}
+		mounted := false
+		defer func() { retErr = errors.Join(retErr, cleanupNBDMount(mnt, mounted)) }()
+		var mountErr error
+		for _, p := range []string{conn.device + "p1", conn.device + "p2", conn.device} {
+			if mountErr = exec.CommandContext(ctx, "mount", p, mnt).Run(); mountErr == nil {
+				mounted = true
+				break
+			}
+		}
+		if mountErr != nil {
+			return fmt.Errorf("mount OpenCore EFI partition on %s: %w", conn.device, mountErr)
+		}
+		return patchPlist(filepath.Join(mnt, "EFI", "OC", "config.plist"), sm)
+	})
+}
+
+// withNBDLease serializes the complete connect/mount/patch/unmount/disconnect
+// transaction across cocoon-macos processes. A free-device check followed by
+// qemu-nbd --connect is not atomic, so choosing devices concurrently is unsafe.
+func withNBDLease(ctx context.Context, path string, fn func() error) error {
+	l := flock.New(path)
+	if err := l.Lock(ctx); err != nil {
+		return fmt.Errorf("lock qemu-nbd lease: %w", err)
 	}
-	if mountErr != nil {
-		return fmt.Errorf("mount OpenCore EFI partition on %s: %w", nbd, mountErr)
-	}
-	defer func() { _ = exec.Command("umount", mnt).Run() }()
-	return patchPlist(filepath.Join(mnt, "EFI", "OC", "config.plist"), sm)
+	defer func() { _ = l.Unlock(context.Background()) }()
+	return fn()
 }
 
 // waitForPart blocks until the kernel's async partition scan exposes nbdXp1 for the mount.
-func waitForPart(ctx context.Context, nbd string) {
-	_ = utils.WaitFor(ctx, 5*time.Second, 100*time.Millisecond, func() (bool, error) {
+func waitForPart(ctx context.Context, nbd string) error {
+	if err := utils.WaitFor(ctx, 5*time.Second, 100*time.Millisecond, func() (bool, error) {
 		if _, err := os.Stat(nbd + "p1"); err == nil {
 			return true, nil
 		}
-		_ = exec.CommandContext(ctx, "partprobe", nbd).Run()
 		return false, nil
-	})
+	}); err != nil {
+		return fmt.Errorf("wait for partition on %s: %w", nbd, err)
+	}
+	return nil
+}
+
+func cleanupNBDMount(mnt string, mounted bool) error {
+	if !mounted {
+		if err := os.Remove(mnt); err != nil && !os.IsNotExist(err) {
+			return fmt.Errorf("remove mount dir %s: %w", mnt, err)
+		}
+		return nil
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), nbdCleanupTimeout)
+	defer cancel()
+	if out, err := exec.CommandContext(ctx, "umount", mnt).CombinedOutput(); err != nil {
+		return fmt.Errorf("unmount %s (output: %s): %w", mnt, out, err)
+	}
+	if err := os.Remove(mnt); err != nil && !os.IsNotExist(err) {
+		return fmt.Errorf("remove mount dir %s: %w", mnt, err)
+	}
+	return nil
 }
 
-// disconnectNBD waits out qemu-nbd's asynchronous release, or the qemu launch races in and fails with "Failed to get shared write lock".
-func disconnectNBD(ctx context.Context, nbd, ocPath string) {
+// disconnect waits out qemu-nbd's asynchronous release. Cleanup deliberately
+// uses a fresh context because the caller is commonly unwinding after timeout.
+func (c *nbdConnection) disconnect() error {
 	logger := log.WithFunc("qemu.disconnectNBD")
-	_ = exec.Command("qemu-nbd", "--disconnect", nbd).Run()
-	if err := utils.WaitFor(ctx, 10*time.Second, 100*time.Millisecond, func() (bool, error) {
-		return !isFileHeld(ocPath), nil
+	var commandErr error
+	disconnectCtx, disconnectCancel := context.WithTimeout(context.Background(), nbdCommandTimeout)
+	if out, err := exec.CommandContext(disconnectCtx, "qemu-nbd", "--disconnect", c.device).CombinedOutput(); err != nil {
+		logger.Warnf(disconnectCtx, "disconnect %s (output: %s): %v", c.device, out, err)
+		commandErr = fmt.Errorf("disconnect %s (output: %s): %w", c.device, out, err)
+	}
+	disconnectCancel()
+	waitCtx, waitCancel := context.WithTimeout(context.Background(), nbdCleanupTimeout)
+	defer waitCancel()
+	if err := utils.WaitFor(waitCtx, 10*time.Second, 100*time.Millisecond, func() (bool, error) {
+		held, scanErr := isFileHeld(c.ocPath)
+		return !held, scanErr
 	}); err != nil {
-		logger.Warnf(ctx, "qcow2 %s still held after nbd disconnect: %v", ocPath, err)
+		logger.Warnf(waitCtx, "qcow2 %s still held after nbd disconnect: %v", c.ocPath, err)
+		killCtx, killCancel := context.WithTimeout(context.Background(), nbdCleanupTimeout)
+		defer killCancel()
+		var cleanupErrs []error
+		if cleanupErr := utils.TerminateProcess(killCtx, c.pid, "qemu-nbd", c.ocPath, time.Second); cleanupErr != nil {
+			logger.Warnf(killCtx, "terminate qemu-nbd pid %d for %s: %v", c.pid, c.ocPath, cleanupErr)
+			cleanupErrs = append(cleanupErrs, cleanupErr)
+		}
+		if cleanupErr := CleanupNBDForPath(killCtx, c.ocPath); cleanupErr != nil {
+			logger.Warnf(killCtx, "terminate stale qemu-nbd for %s: %v", c.ocPath, cleanupErr)
+			cleanupErrs = append(cleanupErrs, cleanupErr)
+		}
+		commandErr = errors.Join(commandErr, err, errors.Join(cleanupErrs...))
+	}
+	_ = os.Remove(c.pidFile)
+	held, err := isFileHeld(c.ocPath)
+	if err != nil {
+		return errors.Join(commandErr, err)
+	}
+	if !held {
+		return nil
 	}
+	return commandErr
 }
 
 // isFileHeld matches the daemonized qemu-nbd server's cmdline — cheaper than scanning fd tables, and that server is the only holder to wait out.
-func isFileHeld(ocPath string) bool {
+func isFileHeld(ocPath string) (bool, error) {
 	pids, err := utils.FindVMMByCmdline("qemu-nbd", ocPath)
-	return err == nil && len(pids) > 0
+	return len(pids) > 0, err
+}
+
+// CleanupNBDForPath terminates qemu-nbd processes whose command line references
+// path. PID identity is verified before signaling, so an unrelated process is
+// never killed even if a PID was reused.
+func CleanupNBDForPath(ctx context.Context, path string) error {
+	pids, err := utils.FindVMMByCmdline("qemu-nbd", path)
+	if err != nil {
+		return fmt.Errorf("scan qemu-nbd processes for %s: %w", path, err)
+	}
+	var errs []error
+	for _, pid := range pids {
+		if err := utils.TerminateProcess(ctx, pid, "qemu-nbd", path, time.Second); err != nil {
+			errs = append(errs, fmt.Errorf("terminate qemu-nbd pid %d: %w", pid, err))
+		}
+	}
+	return errors.Join(errs...)
 }
 
-// connectFreeNBD claims a device by connecting: the connect itself is the exclusive operation, so a race with another VM create just advances to the next candidate.
-func connectFreeNBD(ctx context.Context, ocPath string) (string, error) {
+// connectFreeNBD is called while holding the host-wide NBD lease. --fork is
+// required: without it qemu-nbd remains in the foreground for the lifetime of
+// the mapping and the invoking CLI never reaches the mount/patch steps.
+func connectFreeNBD(ctx context.Context, ocPath string) (*nbdConnection, error) {
 	var lastErr error
-	for i := range 16 {
+	pidFile := ocPath + ".nbd.pid"
+	_ = os.Remove(pidFile)
+	for i := range nbdDeviceCount {
 		nbd := fmt.Sprintf("/dev/nbd%d", i)
 		if _, err := os.Stat(nbd); err != nil {
 			continue
 		}
 		if _, err := os.Stat(fmt.Sprintf("/sys/block/nbd%d/pid", i)); !os.IsNotExist(err) {
 			continue
 		}
-		out, cerr := exec.CommandContext(ctx, "qemu-nbd", "--connect="+nbd, "-f", "qcow2", ocPath).CombinedOutput()
+		referenced, err := processReferencesBlockDevice("/proc", nbd)
+		if err != nil {
+			return nil, fmt.Errorf("check whether %s is referenced by a process: %w", nbd, err)
+		}
+		if referenced {
+			continue
+		}
+		out, cerr := exec.CommandContext(ctx, "qemu-nbd", qemuNBDConnectArgs(nbd, pidFile, ocPath)...).CombinedOutput()
 		if cerr == nil {
-			return nbd, nil
+			pid, err := utils.ReadPIDFile(pidFile)
+			if err == nil {
+				return &nbdConnection{device: nbd, pid: pid, pidFile: pidFile, ocPath: ocPath}, nil
+			}
+			cleanupCtx, cancel := context.WithTimeout(context.Background(), nbdCleanupTimeout)
+			_ = exec.CommandContext(cleanupCtx, "qemu-nbd", "--disconnect", nbd).Run()
+			_ = CleanupNBDForPath(cleanupCtx, ocPath)
+			cancel()
+			return nil, fmt.Errorf("read qemu-nbd pid file %s: %w", pidFile, err)
 		}
 		lastErr = fmt.Errorf("connect qemu-nbd %s (output: %s): %w", nbd, out, cerr)
+		cleanupCtx, cancel := context.WithTimeout(context.Background(), nbdCleanupTimeout)
+		_ = CleanupNBDForPath(cleanupCtx, ocPath)
+		cancel()
+		if ctx.Err() != nil {
+			return nil, errors.Join(lastErr, ctx.Err())
+		}
 	}
 	if lastErr != nil {
-		return "", lastErr
+		return nil, lastErr
 	}
-	return "", errors.New("no free /dev/nbd device (is the nbd module loaded)")
+	return nil, errors.New("no free /dev/nbd device (is the nbd module loaded)")
+}
+
+// processReferencesBlockDevice catches userspace operations that are still
+// blocked on an NBD device after the kernel has already removed its sysfs pid.
+// Reusing such a device can block the next qemu-nbd attach indefinitely.
+func processReferencesBlockDevice(procRoot, device string) (bool, error) {
+	entries, err := os.ReadDir(procRoot)
+	if err != nil {
+		return false, fmt.Errorf("read %s: %w", procRoot, err)
+	}
+	for _, entry := range entries {
+		if !entry.IsDir() {
+			continue
+		}
+		if _, err := strconv.Atoi(entry.Name()); err != nil {
+			continue
+		}
+		cmdlinePath := filepath.Join(procRoot, entry.Name(), "cmdline")
+		cmdline, err := os.ReadFile(cmdlinePath)
+		if err != nil {
+			if os.IsNotExist(err) {
+				continue
+			}
+			return false, fmt.Errorf("read %s: %w", cmdlinePath, err)
+		}
+		for arg := range strings.SplitSeq(string(cmdline), "\x00") {
+			if nbdArgReferencesDevice(arg, device) {
+				return true, nil
+			}
+		}
+	}
+	return false, nil
+}
+
+func nbdArgReferencesDevice(arg, device string) bool {
+	if arg == device || arg == "--connect="+device {
+		return true
+	}
+	partition := strings.TrimPrefix(arg, device+"p")
+	if partition == arg || partition == "" {
+		return false
+	}
+	_, err := strconv.Atoi(partition)
+	return err == nil
+}
+
+func qemuNBDConnectArgs(nbd, pidFile, ocPath string) []string {
+	return []string{"--fork", "--pid-file=" + pidFile, "--connect=" + nbd, "-f", "qcow2", ocPath}
 }
 
 func patchPlist(path string, sm *SMBIOS) error {