GPT header rework - Refacto 2

internal/statemachine/common_states.go

@@ -133,66 +133,128 @@ func (stateMachine *StateMachine) generateDiskInfo() error {
 	return nil
 }
 
-var calculateRootfsSizeState = stateFunc{"calculate_rootfs_size", (*StateMachine).calculateRootfsSize}
+// Semi-arbitrary value, probably larger than needed but big enough to not have issues
+// and low enough to stay in the same magnitude
+const ext4FudgeFactor = 1.5
 
-// calculateRootfsSize calculates the size of the root filesystem.
 // On a 100MiB filesystem, ext4 takes a little over 7MiB for the
-// metadata, so use 8MB as a minimum padding.
+// metadata, so use 8MB as a minimum padding
+const ext4Padding = 8 * quantity.SizeMiB
+
+var calculateRootfsSizeState = stateFunc{"calculate_rootfs_size", (*StateMachine).calculateRootfsSize}
+
+// calculateRootfsSize calculates the size needed by the root filesystem.
+// If an image size was specified, make sure it is big enough to contain the
+// rootfs and try to allocate it to the rootfs
 func (stateMachine *StateMachine) calculateRootfsSize() error {
-	rootfsSize, err := helper.Du(stateMachine.tempDirs.rootfs)
+	rootfsMinSize, err := stateMachine.getRootfsMinSize()
 	if err != nil {
-		return fmt.Errorf("Error getting rootfs size: %s", err.Error())
+		return err
 	}
 
-	// fudge factor for incidentals
-	rootfsPadding := 8 * quantity.SizeMiB
-	rootfsSize = quantity.Size(math.Ceil(float64(rootfsSize) * 1.5))
-	rootfsSize += rootfsPadding
-
-	stateMachine.RootfsSize = stateMachine.alignToSectorSize(rootfsSize)
+	stateMachine.RootfsSize = rootfsMinSize
 
-	if stateMachine.commonFlags.Size != "" {
-		rootfsVolume, rootfsVolumeName := stateMachine.findRootfsVolume()
-		desiredSize := stateMachine.ImageSizes[rootfsVolumeName]
+	rootfsVolume, rootfsVolumeName, rootfsStructure := stateMachine.findRootfsVolumeStructure()
 
-		// subtract the size and offsets of the existing volumes
-		if rootfsVolume != nil {
-			for _, structure := range rootfsVolume.Structure {
-				desiredSize = helper.SafeQuantitySubtraction(desiredSize, structure.Size)
-				if structure.Offset != nil {
-					desiredSize = helper.SafeQuantitySubtraction(desiredSize,
-						quantity.Size(*structure.Offset))
-				}
-			}
+	gadgetRootfsMinSize := stateMachine.getGadgetRootfsMinSize(rootfsStructure)
+	if gadgetRootfsMinSize > stateMachine.RootfsSize {
+		stateMachine.RootfsSize = gadgetRootfsMinSize
+	}
 
-			desiredSize = stateMachine.alignToSectorSize(desiredSize)
-
-			if desiredSize < stateMachine.RootfsSize {
-				return fmt.Errorf("Error: calculated rootfs partition size %d is smaller "+
-					"than actual rootfs contents (%d). Try using a larger value of "+
-					"--image-size",
-					desiredSize, stateMachine.RootfsSize,
-				)
-			}
+	desiredSize, foundDesiredSize := stateMachine.getRootfsDesiredSize(rootfsVolume, rootfsVolumeName)
 
+	if foundDesiredSize {
+		if stateMachine.RootfsSize > desiredSize {
+			fmt.Printf("WARNING: rootfs content %d is bigger "+
+				"than requested image size (%d). Try using a larger value of "+
+				"--image-size",
+				stateMachine.RootfsSize, desiredSize,
+			)
+		} else {
 			stateMachine.RootfsSize = desiredSize
 		}
 	}
 
-	stateMachine.syncGadgetStructureRootfsSize()
+	if rootfsStructure != nil {
+		rootfsStructure.Size = stateMachine.RootfsSize
+	}
+
 	return nil
 }
 
-// findRootfsVolume finds the volume associated to the rootfs
-func (stateMachine *StateMachine) findRootfsVolume() (*gadget.Volume, string) {
+// getRootfsMinSize gets the minimum size needed to hold the built rootfs directory
+func (stateMachine *StateMachine) getRootfsMinSize() (quantity.Size, error) {
+	rootfsMinSize, err := helper.Du(stateMachine.tempDirs.rootfs)
+	if err != nil {
+		return quantity.Size(0), fmt.Errorf("Error getting rootfs size: %s", err.Error())
+	}
+
+	// Take into account ext4 filesystems metadata size
+	rootfsPadding := ext4Padding
+	rootfsMinSize = quantity.Size(math.Ceil(float64(rootfsMinSize) * ext4FudgeFactor))
+	rootfsMinSize += rootfsPadding
+
+	return stateMachine.alignToSectorSize(rootfsMinSize), nil
+}
+
+// getGadgetRootfsMinSize gets the minimum size of the rootfs requested in the
+// gadget YAML
+func (stateMachine *StateMachine) getGadgetRootfsMinSize(rootfsStructure *gadget.VolumeStructure) quantity.Size {
+	if rootfsStructure == nil {
+		return quantity.Size(0)
+	}
+
+	return stateMachine.alignToSectorSize(rootfsStructure.MinSize)
+}
+
+// getRootfsDesiredSize subtracts the size and offsets of the existing structures
+// from the requested image size to determine how much room is left for the rootfs
+func (stateMachine *StateMachine) getRootfsDesiredSize(rootfsVolume *gadget.Volume, rootfsVolumeName string) (quantity.Size, bool) {
+	if rootfsVolume == nil {
+		return quantity.Size(0), false
+	}
+
+	desiredSize, found := stateMachine.ImageSizes[rootfsVolumeName]
+	if !found {
+		// So far we do not know of a desired size for the rootfs
+		return quantity.Size(0), found
+	}
+
+	reservedSize := calculateNoRootfsSize(rootfsVolume)
+	desiredSize = helper.SafeQuantitySubtraction(desiredSize, reservedSize)
+	desiredSize = stateMachine.alignToSectorSize(desiredSize)
+
+	return desiredSize, found
+}
+
+// calculateNoRootfsSize determines the needed space for existing structures
+// except for the rootfs
+func calculateNoRootfsSize(v *gadget.Volume) quantity.Size {
+	var size quantity.Size
+	for _, s := range v.Structure {
+		if isRootfsStructure(&s) { //nolint:gosec,G301
+			continue
+		}
+		if s.Offset != nil && quantity.Size(*s.Offset) > size {
+			size = quantity.Size(*s.Offset)
+		}
+		size += s.MinSize
+	}
+
+	return size
+}
+
+// findRootfsVolumeStructure finds the volume and the structure associated to the rootfs
+func (stateMachine *StateMachine) findRootfsVolumeStructure() (*gadget.Volume, string, *gadget.VolumeStructure) {
 	for volumeName, volume := range stateMachine.GadgetInfo.Volumes {
-		for _, structure := range volume.Structure {
-			if structure.Size == 0 {
-				return volume, volumeName
+		for i := range volume.Structure {
+			s := &volume.Structure[i]
+			if isRootfsStructure(s) { //nolint:gosec,G301
+				return volume, volumeName, s
 			}
 		}
 	}
-	return nil, ""
+	return nil, "", nil
 }
 
 // alignToSectorSize align the given size to the SectorSize of the stateMachine
@@ -201,20 +263,6 @@ func (stateMachine *StateMachine) alignToSectorSize(size quantity.Size) quantity
 		quantity.Size(stateMachine.SectorSize)
 }
 
-// syncGadgetStructureRootfsSize synchronizes size of the gadget.Structure that
-// represents the rootfs with the RootfsSize value of the statemachine
-// This functions assumes stateMachine.RootfsSize was previously correctly updated.
-func (stateMachine *StateMachine) syncGadgetStructureRootfsSize() {
-	for _, volume := range stateMachine.GadgetInfo.Volumes {
-		for structIndex, structure := range volume.Structure {
-			if structure.Size == 0 {
-				structure.Size = stateMachine.RootfsSize
-			}
-			volume.Structure[structIndex] = structure
-		}
-	}
-}
-
 var populateBootfsContentsState = stateFunc{"populate_bootfs_contents", (*StateMachine).populateBootfsContents}
 
 // Populate the Bootfs Contents by using snapd's MountedFilesystemWriter
@@ -299,9 +347,10 @@ func (stateMachine *StateMachine) populateBootfsLayoutStructure(laidOutStructure
 
 var populatePreparePartitionsState = stateFunc{"populate_prepare_partitions", (*StateMachine).populatePreparePartitions}
 
-// Populate and prepare the partitions. For partitions without "filesystem:" specified in
-// gadget.yaml, this involves using dd to copy the content blobs into a .img file. For
-// partitions that do have "filesystem:" specified, we use the Mkfs functions from snapd.
+// populatePreparePartitions populates and prepares the partitions. For partitions without
+// "filesystem:" specified in gadget.yaml, this involves using dd to copy the content blobs
+// into a .img file. For partitions that do have "filesystem:" specified, we use the Mkfs
+// functions from snapd.
 // Throughout this process, the offset is tracked to ensure partitions are not overlapping.
 func (stateMachine *StateMachine) populatePreparePartitions() error {
 	for _, volumeName := range stateMachine.VolumeOrder {
@@ -311,7 +360,7 @@ func (stateMachine *StateMachine) populatePreparePartitions() error {
 		}
 		for structIndex, structure := range volume.Structure {
 			var contentRoot string
-			if structure.Role == gadget.SystemData || structure.Role == gadget.SystemSeed {
+			if isRootfsStructure(&structure) || structure.Role == gadget.SystemSeed { //nolint:gosec,G301
 				contentRoot = stateMachine.tempDirs.rootfs
 			} else {
 				contentRoot = filepath.Join(stateMachine.tempDirs.volumes, volumeName,
@@ -338,7 +387,7 @@ func (stateMachine *StateMachine) populatePreparePartitions() error {
 
 var makeDiskState = stateFunc{"make_disk", (*StateMachine).makeDisk}
 
-// Make the disk
+// makeDisk makes the disk image
 func (stateMachine *StateMachine) makeDisk() error {
 	for volumeName, volume := range stateMachine.GadgetInfo.Volumes {
 		_, found := stateMachine.VolumeNames[volumeName]