From 3e518d50aec8ea435345674e064602053f08301d Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Fri, 10 Apr 2026 08:02:10 -0700
Subject: [PATCH] tstest/build-macos-base-vm: fix waitFor, add --rebuild, fix
 MkdirAll

- waitFor now runs the try func at least once before checking the
  deadline, instead of potentially returning nil on a short timeout
- Add --rebuild flag to delete and recreate an existing VM
- Exit early with a message (not an error) when the VM already exists
- Create intermediate directories before writing .AppleSetupDone
  (freshly installed VMs may not have /private/var/db yet)
---
 tstest/build-macos-base-vm/main.go | 60 +++++++++++++++++++++---------
 1 file changed, 43 insertions(+), 17 deletions(-)

diff --git a/tstest/build-macos-base-vm/main.go b/tstest/build-macos-base-vm/main.go
index 465e04cfa..347487662 100644
--- a/tstest/build-macos-base-vm/main.go
+++ b/tstest/build-macos-base-vm/main.go
@@ -31,7 +31,8 @@ import (
 )
 
 var (
-	vmName = flag.String("name", "llmacstation", "VM name (directory under ~/VM.bundle/)")
+	vmName  = flag.String("name", "llmacstation", "VM name (directory under ~/VM.bundle/)")
+	rebuild = flag.Bool("rebuild", false, "delete existing VM and recreate it")
 )
 
 func main() {
@@ -49,7 +50,14 @@ func main() {
 	ipswPath := filepath.Join(bundleDir, "RestoreImage.ipsw")
 
 	if _, err := os.Stat(filepath.Join(vmDir, "Disk.img")); err == nil {
-		log.Fatalf("VM %q already exists at %s. Delete it first or choose a different --name.", *vmName, vmDir)
+		if !*rebuild {
+			log.Printf("VM %q already exists at %s; nothing to do. Use --rebuild to recreate.", *vmName, vmDir)
+			return
+		}
+		log.Printf("Removing existing VM %q...", *vmName)
+		if err := os.RemoveAll(vmDir); err != nil {
+			log.Fatalf("Removing %s: %v", vmDir, err)
+		}
 	}
 
 	os.MkdirAll(bundleDir, 0755)
@@ -185,23 +193,23 @@ func applyPostInstallFixups(vmDir string) error {
 		exec.Command("hdiutil", "detach", diskDev, "-force").Run()
 	}()
 
-	// Wait for APFS volumes to synthesize.
-	time.Sleep(2 * time.Second)
-
-	// Find the APFS Data volume. It's on a synthesized disk derived from
-	// the physical APFS container.
+	// Wait for the APFS Data volume to appear. After hdiutil attach,
+	// the kernel synthesizes APFS volumes asynchronously.
 	var dataVolDev string
-	allDisks, _ := exec.Command("diskutil", "list").CombinedOutput()
-	for _, line := range strings.Split(string(allDisks), "\n") {
-		if strings.Contains(line, "APFS Volume") && strings.Contains(line, "Data") {
-			fields := strings.Fields(line)
-			if len(fields) > 0 {
-				dataVolDev = fields[len(fields)-1]
+	if err := waitFor(10*time.Second, func() error {
+		out, _ := exec.Command("diskutil", "list").CombinedOutput()
+		for _, line := range strings.Split(string(out), "\n") {
+			if strings.Contains(line, "APFS Volume") && strings.Contains(line, "Data") {
+				fields := strings.Fields(line)
+				if len(fields) > 0 {
+					dataVolDev = fields[len(fields)-1]
+					return nil
+				}
 			}
 		}
-	}
-	if dataVolDev == "" {
-		return fmt.Errorf("no APFS Data volume found:\n%s", allDisks)
+		return fmt.Errorf("APFS Data volume not yet available")
+	}); err != nil {
+		return fmt.Errorf("waiting for APFS Data volume: %w", err)
 	}
 
 	// Mount the Data volume via diskutil (handles APFS permissions correctly).
@@ -218,7 +226,11 @@ func applyPostInstallFixups(vmDir string) error {
 	defer exec.Command("diskutil", "unmount", mountPoint).Run()
 
 	// Create .AppleSetupDone to skip the Setup Assistant.
-	setupDone := filepath.Join(mountPoint, "private", "var", "db", ".AppleSetupDone")
+	dbDir := filepath.Join(mountPoint, "private", "var", "db")
+	if err := os.MkdirAll(dbDir, 0755); err != nil {
+		return fmt.Errorf("creating var/db: %v", err)
+	}
+	setupDone := filepath.Join(dbDir, ".AppleSetupDone")
 	if err := os.WriteFile(setupDone, nil, 0644); err != nil {
 		return fmt.Errorf("creating .AppleSetupDone: %v", err)
 	}
@@ -227,6 +239,20 @@ func applyPostInstallFixups(vmDir string) error {
 	return nil
 }
 
+func waitFor(timeout time.Duration, try func() error) error {
+	deadline := time.Now().Add(timeout)
+	for {
+		err := try()
+		if err == nil {
+			return nil
+		}
+		if time.Now().After(deadline) {
+			return err
+		}
+		time.Sleep(200 * time.Millisecond)
+	}
+}
+
 const entitlementsPlist = `<?xml version="1.0" encoding="UTF-8"?>
 <!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
 <plist version="1.0">