From 17c3dc95ca3673182f5dff0f0c8874b9ac0fa7f0 Mon Sep 17 00:00:00 2001
From: Chris Gunn <chrisgun@microsoft.com>
Date: Mon, 21 Oct 2024 14:47:33 -0700
Subject: [PATCH] Toolkit: Add missing `flock` calls.

When making changes to partitions or filesystems, it is recommended to
take a file lock over the disk block device as this informs the host OS
that you are making changes and that it should avoid scanning or
changing the device until you are done. While most of the relevant
operations are covered, there a few places that are missing the lock.
For example, when calling `mkfs` or `resize2fs`.
---
 toolkit/tools/imagegen/diskutils/diskutils.go     | 15 ++++++++-------
 .../pkg/imagecustomizerlib/shrinkfilesystems.go   |  8 +++++---
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/toolkit/tools/imagegen/diskutils/diskutils.go b/toolkit/tools/imagegen/diskutils/diskutils.go
index 5e441206b4c..5bc93baa33a 100644
--- a/toolkit/tools/imagegen/diskutils/diskutils.go
+++ b/toolkit/tools/imagegen/diskutils/diskutils.go
@@ -492,7 +492,7 @@ func CreatePartitions(diskDevPath string, disk configuration.Disk, rootEncryptio
 			return partDevPathMap, partIDToFsTypeMap, encryptedRoot, readOnlyRoot, err
 		}
 
-		partFsType, err := FormatSinglePartition(partDevPath, partition)
+		partFsType, err := formatSinglePartition(diskDevPath, partDevPath, partition)
 		if err != nil {
 			err = fmt.Errorf("failed to format partition:\n%w", err)
 			return partDevPathMap, partIDToFsTypeMap, encryptedRoot, readOnlyRoot, err
@@ -782,12 +782,13 @@ func setGptPartitionType(partition configuration.Partition, timeoutInSeconds, di
 	return
 }
 
-// FormatSinglePartition formats the given partition to the type specified in the partition configuration
-func FormatSinglePartition(partDevPath string, partition configuration.Partition,
+// formatSinglePartition formats the given partition to the type specified in the partition configuration
+func formatSinglePartition(diskDevPath string, partDevPath string, partition configuration.Partition,
 ) (fsType string, err error) {
 	const (
-		totalAttempts = 5
-		retryDuration = time.Second
+		totalAttempts    = 5
+		retryDuration    = time.Second
+		timeoutInSeconds = "5"
 	)
 
 	fsType = partition.FsType
@@ -803,12 +804,12 @@ func FormatSinglePartition(partDevPath string, partition configuration.Partition
 			fsType = "vfat"
 		}
 
-		mkfsArgs := []string{"-t", fsType}
+		mkfsArgs := []string{"--timeout", timeoutInSeconds, diskDevPath, "mkfs", "-t", fsType}
 		mkfsArgs = append(mkfsArgs, mkfsOptions...)
 		mkfsArgs = append(mkfsArgs, partDevPath)
 
 		err = retry.Run(func() error {
-			_, stderr, err := shell.Execute("mkfs", mkfsArgs...)
+			_, stderr, err := shell.Execute("flock", mkfsArgs...)
 			if err != nil {
 				logger.Log.Warnf("Failed to format partition using mkfs: %v", stderr)
 				return err
diff --git a/toolkit/tools/pkg/imagecustomizerlib/shrinkfilesystems.go b/toolkit/tools/pkg/imagecustomizerlib/shrinkfilesystems.go
index aa058dcf422..87fcc0b228d 100644
--- a/toolkit/tools/pkg/imagecustomizerlib/shrinkfilesystems.go
+++ b/toolkit/tools/pkg/imagecustomizerlib/shrinkfilesystems.go
@@ -85,7 +85,8 @@ func shrinkFilesystems(imageLoopDevice string, verityHashPartition *imagecustomi
 		}
 
 		// Shrink the file system with resize2fs -M
-		stdout, stderr, err := shell.Execute("resize2fs", "-M", partitionLoopDevice)
+		stdout, stderr, err := shell.Execute("flock", "--timeout", "5", imageLoopDevice,
+			"resize2fs", "-M", partitionLoopDevice)
 		if err != nil {
 			return fmt.Errorf("failed to resize %s with resize2fs:\n%v", partitionLoopDevice, stderr)
 		}
@@ -103,8 +104,9 @@ func shrinkFilesystems(imageLoopDevice string, verityHashPartition *imagecustomi
 		}
 
 		// Resize the partition with parted resizepart
-		_, stderr, err = shell.ExecuteWithStdin("yes" /*stdin*/, "parted", "---pretend-input-tty",
-			imageLoopDevice, "resizepart", strconv.Itoa(partitionNumber), end)
+		_, stderr, err = shell.ExecuteWithStdin("yes" /*stdin*/, "flock", "--timeout", "5", imageLoopDevice,
+			"parted", "---pretend-input-tty", imageLoopDevice, "resizepart",
+			strconv.Itoa(partitionNumber), end)
 		if err != nil {
 			return fmt.Errorf("failed to resizepart %s with parted:\n%v", partitionLoopDevice, stderr)
 		}