Make sure we cleanup properly after update-grub

internal/statemachine/helper.go

@@ -1008,7 +1008,27 @@ func (stateMachine *StateMachine) updateGrub(rootfsVolName string, rootfsPartNum
 
 	// Slice used to store all the commands that need to be run
 	// to properly update grub.cfg in the chroot
+	// updateGrubCmds should be filled as a FIFO list
 	var updateGrubCmds []*exec.Cmd
+	// Slice used to store all the commands that need to be run
+	// to properly cleanup everything after the update of grub.cfg
+	// updateGrubCmds should be filled as a LIFO list (so new entries should added at the start of the slice)
+	var teardownCmds []*exec.Cmd
+
+	defer func() {
+		for _, teardownCmd := range teardownCmds {
+			cmdOutput := helper.SetCommandOutput(teardownCmd, stateMachine.commonFlags.Debug)
+			tmpErr := teardownCmd.Run()
+			if tmpErr != nil {
+				if err != nil {
+					err = fmt.Errorf("Error running command \"%s\". Error is \"%s\". Output is: \n%s",
+						teardownCmd.String(), err.Error(), cmdOutput.String())
+				} else {
+					err = tmpErr
+				}
+			}
+		}
+	}()
 
 	imgPath := filepath.Join(stateMachine.commonFlags.OutputDir, stateMachine.VolumeNames[rootfsVolName])
 
@@ -1017,34 +1037,35 @@ func (stateMachine *StateMachine) updateGrub(rootfsVolName string, rootfsPartNum
 		return err
 	}
 
-	var umounts []*exec.Cmd
+	// detach the loopback device
+	teardownCmds = append(teardownCmds,
+		//nolint:gosec,G204
+		execCommand("losetup", "--detach", loopUsed),
+	)
+
 	updateGrubCmds = append(updateGrubCmds,
 		// mount the rootfs partition in which to run update-grub
 		//nolint:gosec,G204
-		exec.Command("mount",
+		execCommand("mount",
 			fmt.Sprintf("%sp%d", loopUsed, rootfsPartNum),
 			mountDir,
 		),
 	)
 
+	teardownCmds = append([]*exec.Cmd{execCommand("umount", mountDir)}, teardownCmds...)
+
 	// set up the mountpoints
 	mountPoints := []string{"/dev", "/proc", "/sys"}
 	for _, mountPoint := range mountPoints {
 		mountCmds, umountCmds := mountFromHost(mountDir, mountPoint)
 		updateGrubCmds = append(updateGrubCmds, mountCmds...)
-		umounts = append(umounts, umountCmds...)
-		defer func(cmds []*exec.Cmd) {
-			_ = runAll(cmds)
-		}(umountCmds)
-
+		teardownCmds = append(umountCmds, teardownCmds...)
 	}
-	// make sure to unmount the disk too
-	umounts = append(umounts, exec.Command("umount", mountDir))
 
 	// divert GRUB's os-prober as we don't want to scan for other OSes on
 	// the build system
 	updateGrubCmds = append(updateGrubCmds,
-		exec.Command("chroot",
+		execCommand("chroot",
 			mountDir,
 			"dpkg-divert",
 			"--local",
@@ -1055,17 +1076,9 @@ func (stateMachine *StateMachine) updateGrub(rootfsVolName string, rootfsPartNum
 		),
 	)
 
-	// actually run update-grub
-	updateGrubCmds = append(updateGrubCmds,
-		exec.Command("chroot",
-			mountDir,
-			"update-grub",
-		),
-	)
-
 	// undivert GRUB's os-prober
-	updateGrubCmds = append(updateGrubCmds,
-		exec.Command("chroot",
+	teardownCmds = append([]*exec.Cmd{
+		execCommand("chroot",
 			mountDir,
 			"dpkg-divert",
 			"--remove",
@@ -1074,21 +1087,17 @@ func (stateMachine *StateMachine) updateGrub(rootfsVolName string, rootfsPartNum
 			"/etc/grub.d/30_os-prober.dpkg-divert",
 			"--rename",
 			"/etc/grub.d/30_os-prober",
-		),
+		)},
+		teardownCmds...,
 	)
 
-	// unmount /dev /proc and /sys
-	updateGrubCmds = append(updateGrubCmds, umounts...)
-
-	// tear down the loopback
-	teardownCmd := exec.Command("losetup",
-		"--detach",
-		loopUsed,
+	// actually run update-grub
+	updateGrubCmds = append(updateGrubCmds,
+		execCommand("chroot",
+			mountDir,
+			"update-grub",
+		),
 	)
-	defer func() {
-		_ = teardownCmd.Run()
-	}()
-	updateGrubCmds = append(updateGrubCmds, teardownCmd)
 
 	// now run all the commands
 	for _, cmd := range updateGrubCmds {

internal/statemachine/helper_test.go

@@ -11,6 +11,7 @@ import (
 	"os/exec"
 	"path/filepath"
 	"reflect"
+	"regexp"
 	"runtime"
 	"strings"
 	"testing"
@@ -1369,6 +1370,65 @@ func TestFailedGetPreseededSnaps(t *testing.T) {
 	})
 }
 
+// TestStateMachine_updateGrub_checkcmds checks commands to update grub order is ok
+func TestStateMachine_updateGrub_checkcmds(t *testing.T) {
+	asserter := helper.Asserter{T: t}
+	var stateMachine StateMachine
+	stateMachine.commonFlags, stateMachine.stateMachineFlags = helper.InitCommonOpts()
+	stateMachine.commonFlags.Debug = true
+	stateMachine.commonFlags.OutputDir = "/tmp"
+
+	err := stateMachine.makeTemporaryDirectories()
+	asserter.AssertErrNil(err, true)
+
+	t.Cleanup(func() { os.RemoveAll(stateMachine.stateMachineFlags.WorkDir) })
+
+	mockCmder := NewMockExecCommand()
+
+	execCommand = mockCmder.Command
+	t.Cleanup(func() { execCommand = exec.Command })
+
+	stdout, restoreStdout, err := helper.CaptureStd(&os.Stdout)
+	t.Cleanup(func() { restoreStdout() })
+
+	err = stateMachine.updateGrub("", 2)
+	asserter.AssertErrNil(err, true)
+
+	restoreStdout()
+	readStdout, err := io.ReadAll(stdout)
+
+	expectedCmds := []*regexp.Regexp{
+		regexp.MustCompile("mount .*p2 .*/scratch/loopback"),
+		regexp.MustCompile("mount --bind /dev .*/scratch/loopback/dev"),
+		regexp.MustCompile("mount --bind /proc .*/scratch/loopback/proc"),
+		regexp.MustCompile("mount --bind /sys .*/scratch/loopback/sys"),
+		regexp.MustCompile("chroot .*/scratch/loopback dpkg-divert"),
+		regexp.MustCompile("chroot .*/scratch/loopback update-grub"),
+		regexp.MustCompile("chroot .*/scratch/loopback dpkg-divert --remove"),
+		regexp.MustCompile("mount --make-rprivate .*/scratch/loopback/sys"),
+		regexp.MustCompile("umount --recursive .*scratch/loopback/sys"),
+		regexp.MustCompile("mount --make-rprivate .*/scratch/loopback/proc"),
+		regexp.MustCompile("umount --recursive .*scratch/loopback/proc"),
+		regexp.MustCompile("mount --make-rprivate .*/scratch/loopback/dev"),
+		regexp.MustCompile("umount --recursive .*scratch/loopback/dev"),
+		regexp.MustCompile("umount .*scratch/loopback"),
+		regexp.MustCompile("losetup --detach .* /tmp"),
+	}
+
+	gotCmds := strings.Split(strings.TrimSpace(string(readStdout)), "\n")
+	if len(expectedCmds) != len(gotCmds) {
+		t.Fatalf("%v commands to be executed, expected %v", len(gotCmds), len(expectedCmds))
+	}
+
+	for i, gotCmd := range gotCmds {
+		expected := expectedCmds[i]
+
+		if !expected.Match([]byte(gotCmd)) {
+			t.Errorf("Cmd \"%v\" not matching. Expected %v\n", gotCmd, expected.String())
+		}
+	}
+}
+
 // TestFailedUpdateGrub tests failures in the updateGrub function
 func TestFailedUpdateGrub(t *testing.T) {
 	t.Run("test_failed_update_grub", func(t *testing.T) {

internal/statemachine/tests_helper_test.go

@@ -3,6 +3,7 @@ package statemachine
 import (
 	"fmt"
 	"io/fs"
+	"os/exec"
 )
 
 type osMockConf struct {
@@ -63,3 +64,15 @@ func (o *osMock) Truncate(name string, size int64) error {
 func NewOSMock(conf *osMockConf) *osMock {
 	return &osMock{conf: conf}
 }
+
+type mockExecCmd struct{}
+
+func NewMockExecCommand() *mockExecCmd {
+	return &mockExecCmd{}
+}
+
+func (m *mockExecCmd) Command(cmd string, args ...string) *exec.Cmd {
+	// Replace the command with an echo of it
+	//nolint:gosec,G204
+	return exec.Command("echo", append([]string{cmd}, args...)...)
+}