WIP rework grub setup

Signed-off-by: Paul Mars <[email protected]>

internal/partition/partition.go

@@ -60,20 +60,30 @@ func newPartitionTable(volume *gadget.Volume, sectorSize uint64, imgSize uint64)
 
 // GeneratePartitionTable prepares the partition table for structures in a volume and
 // returns it with the partition number of the root partition.
-func GeneratePartitionTable(volume *gadget.Volume, sectorSize uint64, imgSize uint64, isSeeded bool) (partition.Table, int, error) {
-	partitionNumber, rootfsPartitionNumber := 1, -1
+func GeneratePartitionTable(volume *gadget.Volume, sectorSize uint64, imgSize uint64, isSeeded bool) (partition.Table, int, int, error) {
+	partitionNumber, rootfsPartitionNumber, bootPartitionNumber := 1, -1, -1
 	partitionTable := newPartitionTable(volume, sectorSize, imgSize)
 	onDisk := gadget.OnDiskStructsFromGadget(volume)
 
 	for i := range volume.Structure {
 		structure := &volume.Structure[i]
-		if !structure.IsPartition() || helper.ShouldSkipStructure(structure, isSeeded) {
+		if !structure.IsPartition() {
+			continue
+		}
+
+		// Record the actual partition number of the boot partition, as it
+		// might be useful for certain operations (like updating the bootloader)
+		if helper.IsSystemBootStructure(structure) {
+			bootPartitionNumber = partitionNumber
+		}
+
+		if helper.ShouldSkipStructure(structure, isSeeded) {
 			continue
 		}
 
 		// Record the actual partition number of the root partition, as it
 		// might be useful for certain operations (like updating the bootloader)
-		if helper.IsRootfsStructure(structure) { //nolint:gosec,G301
+		if helper.IsRootfsStructure(structure) {
 			rootfsPartitionNumber = partitionNumber
 		}
 
@@ -85,13 +95,13 @@ func GeneratePartitionTable(volume *gadget.Volume, sectorSize uint64, imgSize ui
 		structureType := getStructureType(structure, volume.Schema)
 		err := partitionTable.AddPartition(structurePair, structureType)
 		if err != nil {
-			return nil, rootfsPartitionNumber, err
+			return nil, rootfsPartitionNumber, bootPartitionNumber, err
 		}
 
 		partitionNumber++
 	}
 
-	return partitionTable.GetConcreteTable(), rootfsPartitionNumber, nil
+	return partitionTable.GetConcreteTable(), rootfsPartitionNumber, bootPartitionNumber, nil
 }
 
 // getStructureType extracts the structure type from the structure.Type considering

internal/partition/partition_test.go

@@ -168,6 +168,7 @@ func TestGeneratePartitionTable(t *testing.T) {
 		args                 args
 		wantPartitionTable   partition.Table
 		wantRootfsPartNumber int
+		wantBootPartNumber   int
 		expectedError        string
 	}{
 		{
@@ -178,6 +179,7 @@ func TestGeneratePartitionTable(t *testing.T) {
 				imgSize:    uint64(4 * quantity.SizeKiB),
 			},
 			wantRootfsPartNumber: 2,
+			wantBootPartNumber:   -1,
 			wantPartitionTable: &gpt.Table{
 				LogicalSectorSize:  int(sectorSize512),
 				PhysicalSectorSize: int(sectorSize512),
@@ -206,6 +208,7 @@ func TestGeneratePartitionTable(t *testing.T) {
 				imgSize:    uint64(4 * quantity.SizeKiB),
 			},
 			wantRootfsPartNumber: 2,
+			wantBootPartNumber:   -1,
 			wantPartitionTable: &gpt.Table{
 				LogicalSectorSize:  int(sectorSize4k),
 				PhysicalSectorSize: int(sectorSize4k),
@@ -234,6 +237,7 @@ func TestGeneratePartitionTable(t *testing.T) {
 				imgSize:    uint64(4 * quantity.SizeKiB),
 			},
 			wantRootfsPartNumber: 2,
+			wantBootPartNumber:   -1,
 			wantPartitionTable: &gpt.Table{
 				LogicalSectorSize:  int(sectorSize512),
 				PhysicalSectorSize: int(sectorSize512),
@@ -262,6 +266,7 @@ func TestGeneratePartitionTable(t *testing.T) {
 				imgSize:    uint64(4 * quantity.SizeKiB),
 			},
 			wantRootfsPartNumber: 2,
+			wantBootPartNumber:   -1,
 			expectedError:        `The structure "writable" overlaps GPT header or GPT partition table`,
 		},
 		{
@@ -272,6 +277,7 @@ func TestGeneratePartitionTable(t *testing.T) {
 				imgSize:    uint64(4 * quantity.SizeKiB),
 			},
 			wantRootfsPartNumber: -1,
+			wantBootPartNumber:   -1,
 			wantPartitionTable: &mbr.Table{
 				Partitions: []*mbr.Partition{
 					{
@@ -288,11 +294,12 @@ func TestGeneratePartitionTable(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			asserter := &helper.Asserter{T: t}
-			gotPartitionTable, gotRootfsPartNumber, gotErr := GeneratePartitionTable(tt.args.volume, tt.args.sectorSize, tt.args.imgSize, tt.args.isSeeded)
+			gotPartitionTable, gotRootfsPartNumber, gotBootPartNumber, gotErr := GeneratePartitionTable(tt.args.volume, tt.args.sectorSize, tt.args.imgSize, tt.args.isSeeded)
 
 			if len(tt.expectedError) == 0 {
 				asserter.AssertErrNil(gotErr, true)
 				asserter.AssertEqual(tt.wantRootfsPartNumber, gotRootfsPartNumber)
+				asserter.AssertEqual(tt.wantBootPartNumber, gotBootPartNumber)
 				asserter.AssertEqual(tt.wantPartitionTable, gotPartitionTable)
 			} else {
 				asserter.AssertErrContains(gotErr, tt.expectedError)

internal/statemachine/classic.go

@@ -507,7 +507,7 @@ func (s *StateMachine) addImgStates(states *[]stateFunc) {
 
 	*states = append(*states,
 		makeDiskState,
-		updateBootloaderState,
+		setupBootloaderState,
 	)
 }
 
@@ -522,7 +522,7 @@ func (s *StateMachine) addQcow2States(states *[]stateFunc) {
 	if !found {
 		*states = append(*states,
 			makeDiskState,
-			updateBootloaderState,
+			setupBootloaderState,
 		)
 	}
 	*states = append(*states, makeQcow2ImgState)

internal/statemachine/classic_states.go

@@ -1309,18 +1309,21 @@ func (stateMachine *StateMachine) makeQcow2Img() error {
 	return nil
 }
 
-var updateBootloaderState = stateFunc{"update_bootloader", (*StateMachine).updateBootloader}
+var setupBootloaderState = stateFunc{"setup_bootloader", (*StateMachine).setupBootloader}
 
-// updateBootloader determines the bootloader for each volume
-// and runs the correct helper function to update the bootloader
-func (stateMachine *StateMachine) updateBootloader() error {
+// setupBootloader determines the bootloader for each volume
+// and runs the correct helper function to install/update the bootloader
+func (stateMachine *StateMachine) setupBootloader() error {
 	if stateMachine.RootfsPartNum == -1 || stateMachine.RootfsVolName == "" {
 		return fmt.Errorf("Error: could not determine partition number of the root filesystem")
 	}
+	if stateMachine.BootPartNum == -1 {
+		return fmt.Errorf("Error: could not determine partition number of the boot filesystem")
+	}
 	volume := stateMachine.GadgetInfo.Volumes[stateMachine.RootfsVolName]
 	switch volume.Bootloader {
 	case "grub":
-		err := stateMachine.updateGrub(stateMachine.RootfsVolName, stateMachine.RootfsPartNum)
+		err := stateMachine.setupGrub(stateMachine.RootfsVolName, stateMachine.RootfsPartNum, stateMachine.BootPartNum)
 		if err != nil {
 			return err
 		}

internal/statemachine/classic_test.go

@@ -195,7 +195,7 @@ func TestClassicStateMachine_calculateStates(t *testing.T) {
 				"populate_bootfs_contents",
 				"populate_prepare_partitions",
 				"make_disk",
-				"update_bootloader",
+				"setup_bootloader",
 				"generate_package_manifest",
 			},
 		},
@@ -221,7 +221,7 @@ func TestClassicStateMachine_calculateStates(t *testing.T) {
 				"populate_bootfs_contents",
 				"populate_prepare_partitions",
 				"make_disk",
-				"update_bootloader",
+				"setup_bootloader",
 				"generate_package_manifest",
 			},
 		},
@@ -248,7 +248,7 @@ func TestClassicStateMachine_calculateStates(t *testing.T) {
 				"populate_bootfs_contents",
 				"populate_prepare_partitions",
 				"make_disk",
-				"update_bootloader",
+				"setup_bootloader",
 				"generate_package_manifest",
 			},
 		},
@@ -277,7 +277,7 @@ func TestClassicStateMachine_calculateStates(t *testing.T) {
 				"populate_bootfs_contents",
 				"populate_prepare_partitions",
 				"make_disk",
-				"update_bootloader",
+				"setup_bootloader",
 				"make_qcow2_image",
 				"generate_package_manifest",
 				"generate_filelist",
@@ -302,7 +302,7 @@ func TestClassicStateMachine_calculateStates(t *testing.T) {
 				"populate_bootfs_contents",
 				"populate_prepare_partitions",
 				"make_disk",
-				"update_bootloader",
+				"setup_bootloader",
 				"generate_package_manifest",
 			},
 		},
@@ -323,7 +323,7 @@ func TestClassicStateMachine_calculateStates(t *testing.T) {
 				"populate_bootfs_contents",
 				"populate_prepare_partitions",
 				"make_disk",
-				"update_bootloader",
+				"setup_bootloader",
 				"generate_package_manifest",
 			},
 		},
@@ -349,7 +349,7 @@ func TestClassicStateMachine_calculateStates(t *testing.T) {
 				"populate_bootfs_contents",
 				"populate_prepare_partitions",
 				"make_disk",
-				"update_bootloader",
+				"setup_bootloader",
 				"generate_package_manifest",
 			},
 		},
@@ -371,7 +371,7 @@ func TestClassicStateMachine_calculateStates(t *testing.T) {
 				"populate_bootfs_contents",
 				"populate_prepare_partitions",
 				"make_disk",
-				"update_bootloader",
+				"setup_bootloader",
 				"generate_package_manifest",
 			},
 		},
@@ -400,7 +400,7 @@ func TestClassicStateMachine_calculateStates(t *testing.T) {
 				"populate_bootfs_contents",
 				"populate_prepare_partitions",
 				"make_disk",
-				"update_bootloader",
+				"setup_bootloader",
 				"generate_package_manifest",
 			},
 		},
@@ -428,7 +428,7 @@ func TestClassicStateMachine_calculateStates(t *testing.T) {
 				"populate_bootfs_contents",
 				"populate_prepare_partitions",
 				"make_disk",
-				"update_bootloader",
+				"setup_bootloader",
 				"make_qcow2_image",
 			},
 		},
@@ -564,7 +564,7 @@ func TestDisplayStates(t *testing.T) {
 [17] populate_bootfs_contents
 [18] populate_prepare_partitions
 [19] make_disk
-[20] update_bootloader
+[20] setup_bootloader
 [21] generate_package_manifest
 `
 	if !strings.Contains(string(readStdout), expectedStates) {
@@ -4968,7 +4968,7 @@ func TestFailedUpdateBootloader(t *testing.T) {
 	// has not been found in earlier steps
 	stateMachine.RootfsPartNum = -1
 	stateMachine.RootfsVolName = ""
-	err = stateMachine.updateBootloader()
+	err = stateMachine.setupBootloader()
 	asserter.AssertErrContains(err, "Error: could not determine partition number of the root filesystem")
 
 	// place a test gadget tree in the scratch directory so we don't
@@ -4999,13 +4999,13 @@ func TestFailedUpdateBootloader(t *testing.T) {
 	asserter.AssertErrNil(err, true)
 	err = stateMachine.loadGadgetYaml()
 	asserter.AssertErrNil(err, true)
-	osMkdir = mockMkdir
+	osMkdirAll = mockMkdirAll
 	t.Cleanup(func() {
-		osMkdir = os.Mkdir
+		osMkdirAll = os.MkdirAll
 	})
 
-	err = stateMachine.updateBootloader()
-	asserter.AssertErrContains(err, "Error creating scratch/loopback directory")
+	err = stateMachine.setupBootloader()
+	asserter.AssertErrContains(err, "Error creating scratch/loopback/boot/efi")
 }
 
 // TestUnsupportedBootloader tests that a warning is thrown if the
@@ -5065,7 +5065,7 @@ func TestUnsupportedBootloader(t *testing.T) {
 	defer restoreStdout()
 	asserter.AssertErrNil(err, true)
 
-	err = stateMachine.updateBootloader()
+	err = stateMachine.setupBootloader()
 	asserter.AssertErrNil(err, true)
 
 	// restore stdout and examine what was printed

internal/statemachine/common_states.go

@@ -478,14 +478,15 @@ func (stateMachine *StateMachine) createDiskImage(volumeName string, volume *gad
 
 // partitionDisk generates a partition table and applies it to the disk
 func (stateMachine *StateMachine) partitionDisk(diskImg *diskutils.Disk, volume *gadget.Volume, volumeName string) error {
-	partitionTable, rootfsPartitionNumber, err := partition.GeneratePartitionTable(volume, uint64(stateMachine.SectorSize), uint64(diskImg.Size), stateMachine.IsSeeded)
+	partitionTable, rootfsPartitionNumber, bootPartitionNumber, err := partition.GeneratePartitionTable(volume, uint64(stateMachine.SectorSize), uint64(diskImg.Size), stateMachine.IsSeeded)
 	if err != nil {
 		return err
 	}
 
-	// Save the rootfs partition number, for later use
+	// Save the rootfs/boot partition numbers, for later use
 	// Store in any case, even if value is -1 to make it clear later it was not found
 	stateMachine.RootfsPartNum = rootfsPartitionNumber
+	stateMachine.BootPartNum = bootPartitionNumber
 	if rootfsPartitionNumber != -1 {
 		stateMachine.RootfsVolName = volumeName
 	}