From 56a88f8969a870dccde432f0bff3036d4d3d1224 Mon Sep 17 00:00:00 2001
From: JUN JIE NAN <nanjunjie@gmail.com>
Date: Fri, 24 May 2024 18:58:40 +0800
Subject: [PATCH] Fixed duplicated umount action

Signed-off-by: JUN JIE NAN <nanjunjie@gmail.com>
---
 distrobuilder/main_incus.go |  9 ++++++---
 distrobuilder/vm.go         | 29 ++++++++++++++++++++++++-----
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/distrobuilder/main_incus.go b/distrobuilder/main_incus.go
index 49a18821..0b413870 100644
--- a/distrobuilder/main_incus.go
+++ b/distrobuilder/main_incus.go
@@ -289,7 +289,7 @@ func (c *cmdIncus) run(cmd *cobra.Command, args []string, overlayDir string) err
 	var mounts []shared.ChrootMount
 	var vmDir string
 	var vm *vm
-
+	cleanup := true
 	if c.flagVM {
 		vmDir = filepath.Join(c.global.flagCacheDir, "vm")
 
@@ -340,7 +340,9 @@ func (c *cmdIncus) run(cmd *cobra.Command, args []string, overlayDir string) err
 		}
 
 		defer func() {
-			_ = shared.RunCommand(vm.ctx, nil, nil, "umount", "-R", vmDir)
+			if cleanup {
+				_ = vm.umountPartition(vmDir)
+			}
 		}()
 
 		err = vm.createUEFIFS()
@@ -438,11 +440,12 @@ func (c *cmdIncus) run(cmd *cobra.Command, args []string, overlayDir string) err
 
 	// Unmount VM directory and loop device before creating the image.
 	if c.flagVM {
-		err := shared.RunCommand(vm.ctx, nil, nil, "umount", "-R", vmDir)
+		err := vm.umountPartition(vmDir)
 		if err != nil {
 			return fmt.Errorf("Failed to unmount %q: %w", vmDir, err)
 		}
 
+		cleanup = false
 		err = vm.umountImage()
 		if err != nil {
 			return fmt.Errorf("Failed to unmount image: %w", err)
diff --git a/distrobuilder/vm.go b/distrobuilder/vm.go
index 30e2a4c1..6d2e8ec5 100644
--- a/distrobuilder/vm.go
+++ b/distrobuilder/vm.go
@@ -21,6 +21,7 @@ type vm struct {
 	loopDevice string
 	rootFS     string
 	rootfsDir  string
+	bootfsDir  string
 	size       uint64
 	ctx        context.Context
 }
@@ -250,7 +251,7 @@ func (v *vm) createRootFS() error {
 		}
 
 		defer func() {
-			_ = shared.RunCommand(v.ctx, nil, nil, "umount", v.rootfsDir)
+			_ = v.umountPartition(v.rootfsDir)
 		}()
 
 		return shared.RunCommand(v.ctx, nil, nil, "btrfs", "subvolume", "create", fmt.Sprintf("%s/@", v.rootfsDir))
@@ -289,12 +290,30 @@ func (v *vm) mountUEFIPartition() error {
 		return errors.New("Disk image not mounted")
 	}
 
-	mountpoint := filepath.Join(v.rootfsDir, "boot", "efi")
+	v.bootfsDir = filepath.Join(v.rootfsDir, "boot", "efi")
 
-	err := os.MkdirAll(mountpoint, 0755)
+	err := os.MkdirAll(v.bootfsDir, 0755)
 	if err != nil {
-		return fmt.Errorf("Failed to create directory %q: %w", mountpoint, err)
+		return fmt.Errorf("Failed to create directory %q: %w", v.bootfsDir, err)
 	}
 
-	return shared.RunCommand(v.ctx, nil, nil, "mount", "-t", "vfat", v.getUEFIDevFile(), mountpoint, "-o", "discard")
+	return shared.RunCommand(v.ctx, nil, nil, "mount", "-t", "vfat", v.getUEFIDevFile(), v.bootfsDir, "-o", "discard")
+}
+
+func (v *vm) umountPartition(mountpoint string) (err error) {
+	err = v.checkMountpoint(mountpoint)
+	if err != nil {
+		return
+	}
+
+	return shared.RunCommand(v.ctx, nil, nil, "umount", "-R", mountpoint)
+}
+
+func (v *vm) checkMountpoint(mountpoint string) (err error) {
+	err = shared.RunCommand(v.ctx, nil, nil, "mountpoint", mountpoint)
+	if err != nil {
+		err = fmt.Errorf("%s not mounted: %w", mountpoint, err)
+	}
+
+	return err
 }