ParseDiskUtilPlistOutput: relax the requirements w.r.t. last partition (

#146)
cirruslabs · Apr 4, 2024 · ebc6ff2 · ebc6ff2
1 parent 20106e6
commit ebc6ff2
Show file tree

Hide file tree

Showing 2 changed files with 90 additions and 24 deletions.
diff --git a/builder/tart/diskutil_parser.go b/builder/tart/diskutil_parser.go
@@ -5,7 +5,7 @@ import (
 	"howett.net/plist"
 )
 
-const expectedLastPartitionContent = "Apple_APFS"
+const expectedPartitionContent = "Apple_APFS"
 
 type DiskUtilOutput struct {
 	AllDisksAndPartitions []DiskWithPartitions `plist:"AllDisksAndPartitions"`
@@ -22,10 +22,8 @@ type Partition struct {
 }
 
 // ParseDiskUtilPlistOutput parses "diskutil list -plist" output,
-// makes sure there's only one disk on the system and returns
-// its name and the name of the last partition, additionally
-// validating that the last partition is not a recovery one
-// (which we should've deleted for the disk expansion to work).
+// makes sure there's only one disk on the system and returns its
+// name and the name of a single APFS partition or errors otherwise.
 func ParseDiskUtilPlistOutput(input []byte) (string, string, error) {
 	var diskUtilOutput DiskUtilOutput
 
@@ -42,33 +40,27 @@ func ParseDiskUtilPlistOutput(input []byte) (string, string, error) {
 	var candidates []Candidate
 
 	for _, diskWithPartitions := range diskUtilOutput.AllDisksAndPartitions {
-		// Skip disks without partitions
-		if len(diskWithPartitions.Partitions) == 0 {
-			continue
+		for _, partition := range diskWithPartitions.Partitions {
+			if partition.Content != expectedPartitionContent {
+				continue
+			}
+
+			candidates = append(candidates, Candidate{
+				DiskName:      diskWithPartitions.DeviceIdentifier,
+				PartitionName: partition.DeviceIdentifier,
+			})
 		}
-
-		// Add a candidate if this disk's last partition is expectedLastPartitionContent
-		lastPartition := diskWithPartitions.Partitions[len(diskWithPartitions.Partitions)-1]
-
-		if lastPartition.Content != expectedLastPartitionContent {
-			continue
-		}
-
-		candidates = append(candidates, Candidate{
-			DiskName:      diskWithPartitions.DeviceIdentifier,
-			PartitionName: lastPartition.DeviceIdentifier,
-		})
 	}
 
 	if len(candidates) == 0 {
-		return "", "", fmt.Errorf("found no disks on which the last partition's \"Content\" "+
-			"is %q, make sure that the macOS is installed", expectedLastPartitionContent)
+		return "", "", fmt.Errorf("found no disks on which the partition's \"Content\" "+
+			"is %q, make sure that the macOS is installed", expectedPartitionContent)
 	}
 
 	if len(candidates) > 1 {
-		return "", "", fmt.Errorf("found more than one disk on which the last partition's \"Content\" "+
+		return "", "", fmt.Errorf("found more than one disk on which the partition's \"Content\" "+
 			"is %q, please only mount a single disk that contains APFS partitions otherwise it's hard "+
-			"to tell on which disk the macOS is installed", expectedLastPartitionContent)
+			"to tell on which disk the macOS is installed", expectedPartitionContent)
 	}
 
 	return candidates[0].DiskName, candidates[0].PartitionName, nil

diff --git a/builder/tart/diskutil_parser_test.go b/builder/tart/diskutil_parser_test.go
@@ -70,6 +70,80 @@ func TestSingleDisk(t *testing.T) {
 	require.Equal(t, "disk0s2", partitionName)
 }
 
+func TestSingleDiskWithRecoveryPartition(t *testing.T) {
+	plistBytes := `<?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">
+<dict>
+	<key>AllDisks</key>
+	<array>
+		<string>disk0</string>
+		<string>disk0s1</string>
+		<string>disk0s2</string>
+		<string>disk0s3</string>
+	</array>
+	<key>AllDisksAndPartitions</key>
+	<array>
+		<dict>
+			<key>Content</key>
+			<string>GUID_partition_scheme</string>
+			<key>DeviceIdentifier</key>
+			<string>disk0</string>
+			<key>OSInternal</key>
+			<false/>
+			<key>Partitions</key>
+			<array>
+				<dict>
+					<key>Content</key>
+					<string>Apple_APFS_ISC</string>
+					<key>DeviceIdentifier</key>
+					<string>disk0s1</string>
+					<key>DiskUUID</key>
+					<string>22A4F5CC-CB5B-4EF1-B4F5-EBE9F79C0EE6</string>
+					<key>Size</key>
+					<integer>524288000</integer>
+				</dict>
+				<dict>
+					<key>Content</key>
+					<string>Apple_APFS</string>
+					<key>DeviceIdentifier</key>
+					<string>disk0s2</string>
+					<key>DiskUUID</key>
+					<string>E59D951F-E0B0-42B5-BE26-DF3611AD673C</string>
+					<key>Size</key>
+					<integer>44107001856</integer>
+				</dict>
+				<dict>
+					<key>Content</key>
+					<string>Apple_APFS_Recovery</string>
+					<key>DeviceIdentifier</key>
+					<string>disk0s3</string>
+					<key>DiskUUID</key>
+					<string>7684137C-F165-45C4-948B-681012742DAC</string>
+					<key>Size</key>
+					<integer>5368672256</integer>
+				</dict>
+			</array>
+			<key>Size</key>
+			<integer>50000000000</integer>
+		</dict>
+	</array>
+	<key>VolumesFromDisks</key>
+	<array/>
+	<key>WholeDisks</key>
+	<array>
+		<string>disk0</string>
+	</array>
+</dict>
+</plist>
+`
+
+	diskName, partitionName, err := tart.ParseDiskUtilPlistOutput([]byte(plistBytes))
+	require.NoError(t, err)
+	require.Equal(t, "disk0", diskName)
+	require.Equal(t, "disk0s2", partitionName)
+}
+
 func TestMultipleDisks(t *testing.T) {
 	plistBytes := `<?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">