src/pkg/tools/handle_disk_layout.go - cos/tools - Git at Google

 // Copyright 2020 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 package tools

 import (
 	"fmt"
 	"log"
 	"strconv"

 	"cos.googlesource.com/cos/tools.git/src/pkg/tools/partutil"
 )

 // HandleDiskLayout changes the partitions on a COS disk.
 // If the auto-update is disabled, it will shrink sda3 to reclaim the space.
 // It also moves stateful partition and the OEM partition (if extended) by
 // a distance relative to a start point.
 //
 // If sda3 is shrinked, the start point will be the start of sda3 + sda3Margin (2MB).
 // Otherwise, start point will be the original start of the stateful partition.
 //
 // The stateful partition will be moved to leave enough space for the OEM partition,
 // and the OEM partition will be moved to the start point.
 // Finally OEM partition will be resized to 1 sector before the new stateful partition.
 //
 // OEMSize can be the number of sectors (without unit) or size like "3G", "100M", "10000K" or "99999B".
 // If there's no need to extend the OEM partition, `oemSize` in the input will be "", a valid input.
 func HandleDiskLayout(disk string, statePartNum, oemPartNum int, oemSize string, reclaimSDA3 bool) error {
 	if len(disk) <= 0 || statePartNum <= 0 || oemPartNum <= 0 {
 		return fmt.Errorf("empty or non-positive input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q",
 			disk, statePartNum, oemPartNum, oemSize)
 	}

 	if len(oemSize) == 0 {
 		oemSize = "0"
 	}

 	// print the old partition table.
 	table, err := partutil.ReadPartitionTable(disk)
 	if err != nil {
 		return fmt.Errorf("cannot read old partition table of %q, "+
 			"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
 			"error msg: (%v)", disk, disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
 	}
 	log.Printf("\nOld partition table:\n%s\n", table)

 	// read new size of OEM partition.
 	newOEMSizeBytes, err := partutil.ConvertSizeToBytes(oemSize)
 	if err != nil {
 		return fmt.Errorf("error in reading new OEM size, "+
 			"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
 			"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
 	}

 	// read original size of OEM partition.
 	oldOEMSizeSector, err := partutil.ReadPartitionSize(disk, oemPartNum)
 	if err != nil {
 		return fmt.Errorf("error in reading old OEM size, "+
 			"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
 			"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
 	}

 	oldOEMSizeBytes := oldOEMSizeSector << 9 // change unit to bytes.
 	startPointSector, returnAndReboot, err := checkAndReclaimSDA3(disk, statePartNum, reclaimSDA3)
 	if err != nil {
 		return fmt.Errorf("error in reclaiming sda3, "+
 			"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
 			"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
 	}
 	// need to reboot to reread disk partition table.
 	if returnAndReboot {
 		return nil
 	}

 	oemSmaller, err := checkNewOEMSizeSmaller(disk, statePartNum, reclaimSDA3, oldOEMSizeBytes,
 		newOEMSizeBytes, startPointSector)
 	if err != nil {
 		return fmt.Errorf("error in dealing with smaller OEM size, "+
 			"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
 			"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
 	}
 	// No need to resize the OEM partition, everything else done.
 	if oemSmaller {
 		return nil
 	}

 	// leave enough space before the stateful partition for the OEM partition.
 	// and shrink the OEM partition to make the new start 4k aligned.
 	newStateStartSector := partutil.FindLast4KSector(startPointSector + (newOEMSizeBytes >> 9))

 	// Move the stateful partition.
 	if err := partutil.MovePartition(disk, statePartNum, strconv.FormatUint(newStateStartSector, 10)); err != nil {
 		return fmt.Errorf("error in moving stateful partition, "+
 			"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
 			"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
 	}

 	// move OEM partition to the start point.
 	if err := partutil.MovePartition(disk, oemPartNum, strconv.FormatUint(startPointSector, 10)); err != nil {
 		return fmt.Errorf("error in moving OEM partition, "+
 			"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
 			"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
 	}
 	log.Println("Reclaimed /dev/sda3.")

 	// extend the OEM partition.
 	if err = partutil.ExtendPartition(disk, oemPartNum, newStateStartSector-1); err != nil {
 		return fmt.Errorf("error in extending OEM partition, "+
 			"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
 			"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
 	}

 	// print the new partition table.
 	table, err = partutil.ReadPartitionTable(disk)
 	if err != nil {
 		return fmt.Errorf("cannot read new partition table of %q, "+
 			"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
 			"error msg: (%v)", disk, disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
 	}
 	log.Printf("\nCompleted extending OEM partition\n\n New partition table:\n%s\n", table)
 	return nil
 }

 // checkAndReclaimSDA3 checks whether need to reclaim sda3
 // and whether it is done.
 // If sda3 is shrinked, the start point will be the start of sda3 + sda3Margin (2MB).
 // Otherwise, start point will be the original start of the stateful partition.
 // It will return startPointSector, returnAndReboot, error.
 func checkAndReclaimSDA3(disk string, statePartNum int, reclaimSDA3 bool) (uint64, bool, error) {
 	// In some situations, `sfdisk --move-data` requires 1MB free space
 	// in the moving direction. Therefore, leaving 2MB after the start of
 	// sda3 is a safe choice.
 	// Also, this will make sure the start point of the next partition 4K aligned.
 	const sda3Margin = 4096 // 2MB margin
 	var startPointSector uint64
 	var err error
 	if reclaimSDA3 {
 		// check whether sda3 has already been shrinked.
 		minimal, err := partutil.IsPartitionMinimal("/dev/sda", 3)
 		if err != nil {
 			return 0, false, err
 		}
 		// not shrinked yet.
 		if !minimal {
 			_, err = partutil.MinimizePartition("/dev/sda", 3)
 			if err != nil {
 				return 0, false, fmt.Errorf("error in shrinking sda3, "+
 					"error msg: (%v)", err)
 			}
 			log.Println("Shrinked /dev/sda3.")
 			// need to reboot to reload the partition table.
 			return 0, true, nil
 		}
 		// no need to shrink sda3 again.
 		sda3StartSector, err := partutil.ReadPartitionStart("/dev/sda", 3)
 		if err != nil {
 			return 0, false, fmt.Errorf("error in reading the start of sda3, "+
 				"error msg: (%v)", err)
 		}
 		startPointSector = sda3StartSector + sda3Margin // leave enough margin

 	} else {
 		// start point is the original start sector of the stateful partition.
 		startPointSector, err = partutil.ReadPartitionStart(disk, statePartNum)
 		if err != nil {
 			return 0, false, fmt.Errorf("error in reading the start of old stateful partition, "+
 				"error msg: (%v)", err)
 		}
 	}
 	return startPointSector, false, nil
 }

 // checkNewOEMSizeSmaller checks whether the new OEM size is smaller than the old one.
 // If true, move the stateful partition to reclaim sda3 if needed.
 // Otherwise, do nothing.
 // It returns oemSmaller,error.
 func checkNewOEMSizeSmaller(disk string, statePartNum int, reclaimSDA3 bool, oldOEMSizeBytes,
 	newOEMSizeBytes, startPointSector uint64) (bool, error) {
 	if newOEMSizeBytes <= oldOEMSizeBytes {
 		if newOEMSizeBytes != 0 {
 			log.Printf("\n!!!!!!!WARNING!!!!!!!\n"+
 				"oemSize: %d bytes is not larger than the original OEM partition size: %d bytes, "+
 				"nothing is done for the OEM partition.\n ", newOEMSizeBytes, oldOEMSizeBytes)
 		}
 		if !reclaimSDA3 {
 			return true, nil
 		}
 		// move the stateful partition to the start point.
 		if err := partutil.MovePartition(disk, statePartNum, strconv.FormatUint(startPointSector, 10)); err != nil {
 			return true, fmt.Errorf("error in moving stateful partition, "+
 				"input: disk=%q, statePartNum=%d,reclaimSDA3=%t,oldOEMSizeBytes=%d,newOEMSizeBytes=%d, "+
 				"startPointSector=%d, error msg: (%v)", disk, statePartNum, reclaimSDA3,
 				oldOEMSizeBytes, newOEMSizeBytes, startPointSector, err)
 		}
 		log.Println("Reclaimed /dev/sda3.")
 		// print the new partition table.
 		table, err := partutil.ReadPartitionTable(disk)
 		if err != nil {
 			return true, fmt.Errorf("cannot read new partition table of %q, "+
 				"input: disk=%q, statePartNum=%d,reclaimSDA3=%t,oldOEMSizeBytes=%d,newOEMSizeBytes=%d, "+
 				"startPointSector=%d, error msg: (%v)", disk, disk, statePartNum, reclaimSDA3,
 				oldOEMSizeBytes, newOEMSizeBytes, startPointSector, err)
 		}
 		log.Printf("New partition table:\n%s\n", table)
 		return true, nil
 	}
 	return false, nil
 }
	// Copyright 2020 Google LLC
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	package tools

	import (
	"fmt"
	"log"
	"strconv"

	"cos.googlesource.com/cos/tools.git/src/pkg/tools/partutil"
	)

	// HandleDiskLayout changes the partitions on a COS disk.
	// If the auto-update is disabled, it will shrink sda3 to reclaim the space.
	// It also moves stateful partition and the OEM partition (if extended) by
	// a distance relative to a start point.
	//
	// If sda3 is shrinked, the start point will be the start of sda3 + sda3Margin (2MB).
	// Otherwise, start point will be the original start of the stateful partition.
	//
	// The stateful partition will be moved to leave enough space for the OEM partition,
	// and the OEM partition will be moved to the start point.
	// Finally OEM partition will be resized to 1 sector before the new stateful partition.
	//
	// OEMSize can be the number of sectors (without unit) or size like "3G", "100M", "10000K" or "99999B".
	// If there's no need to extend the OEM partition, `oemSize` in the input will be "", a valid input.
	func HandleDiskLayout(disk string, statePartNum, oemPartNum int, oemSize string, reclaimSDA3 bool) error {
	if len(disk) <= 0 \|\| statePartNum <= 0 \|\| oemPartNum <= 0 {
	return fmt.Errorf("empty or non-positive input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q",
	disk, statePartNum, oemPartNum, oemSize)
	}

	if len(oemSize) == 0 {
	oemSize = "0"
	}

	// print the old partition table.
	table, err := partutil.ReadPartitionTable(disk)
	if err != nil {
	return fmt.Errorf("cannot read old partition table of %q, "+
	"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
	"error msg: (%v)", disk, disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
	}
	log.Printf("\nOld partition table:\n%s\n", table)

	// read new size of OEM partition.
	newOEMSizeBytes, err := partutil.ConvertSizeToBytes(oemSize)
	if err != nil {
	return fmt.Errorf("error in reading new OEM size, "+
	"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
	"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
	}

	// read original size of OEM partition.
	oldOEMSizeSector, err := partutil.ReadPartitionSize(disk, oemPartNum)
	if err != nil {
	return fmt.Errorf("error in reading old OEM size, "+
	"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
	"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
	}

	oldOEMSizeBytes := oldOEMSizeSector << 9 // change unit to bytes.
	startPointSector, returnAndReboot, err := checkAndReclaimSDA3(disk, statePartNum, reclaimSDA3)
	if err != nil {
	return fmt.Errorf("error in reclaiming sda3, "+
	"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
	"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
	}
	// need to reboot to reread disk partition table.
	if returnAndReboot {
	return nil
	}

	oemSmaller, err := checkNewOEMSizeSmaller(disk, statePartNum, reclaimSDA3, oldOEMSizeBytes,
	newOEMSizeBytes, startPointSector)
	if err != nil {
	return fmt.Errorf("error in dealing with smaller OEM size, "+
	"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
	"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
	}
	// No need to resize the OEM partition, everything else done.
	if oemSmaller {
	return nil
	}

	// leave enough space before the stateful partition for the OEM partition.
	// and shrink the OEM partition to make the new start 4k aligned.
	newStateStartSector := partutil.FindLast4KSector(startPointSector + (newOEMSizeBytes >> 9))

	// Move the stateful partition.
	if err := partutil.MovePartition(disk, statePartNum, strconv.FormatUint(newStateStartSector, 10)); err != nil {
	return fmt.Errorf("error in moving stateful partition, "+
	"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
	"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
	}

	// move OEM partition to the start point.
	if err := partutil.MovePartition(disk, oemPartNum, strconv.FormatUint(startPointSector, 10)); err != nil {
	return fmt.Errorf("error in moving OEM partition, "+
	"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
	"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
	}
	log.Println("Reclaimed /dev/sda3.")

	// extend the OEM partition.
	if err = partutil.ExtendPartition(disk, oemPartNum, newStateStartSector-1); err != nil {
	return fmt.Errorf("error in extending OEM partition, "+
	"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
	"error msg: (%v)", disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
	}

	// print the new partition table.
	table, err = partutil.ReadPartitionTable(disk)
	if err != nil {
	return fmt.Errorf("cannot read new partition table of %q, "+
	"input: disk=%q, statePartNum=%d, oemPartNum=%d, oemSize=%q, reclaimSDA3=%t, "+
	"error msg: (%v)", disk, disk, statePartNum, oemPartNum, oemSize, reclaimSDA3, err)
	}
	log.Printf("\nCompleted extending OEM partition\n\n New partition table:\n%s\n", table)
	return nil
	}

	// checkAndReclaimSDA3 checks whether need to reclaim sda3
	// and whether it is done.
	// If sda3 is shrinked, the start point will be the start of sda3 + sda3Margin (2MB).
	// Otherwise, start point will be the original start of the stateful partition.
	// It will return startPointSector, returnAndReboot, error.
	func checkAndReclaimSDA3(disk string, statePartNum int, reclaimSDA3 bool) (uint64, bool, error) {
	// In some situations, `sfdisk --move-data` requires 1MB free space
	// in the moving direction. Therefore, leaving 2MB after the start of
	// sda3 is a safe choice.
	// Also, this will make sure the start point of the next partition 4K aligned.
	const sda3Margin = 4096 // 2MB margin
	var startPointSector uint64
	var err error
	if reclaimSDA3 {
	// check whether sda3 has already been shrinked.
	minimal, err := partutil.IsPartitionMinimal("/dev/sda", 3)
	if err != nil {
	return 0, false, err
	}
	// not shrinked yet.
	if !minimal {
	_, err = partutil.MinimizePartition("/dev/sda", 3)
	if err != nil {
	return 0, false, fmt.Errorf("error in shrinking sda3, "+
	"error msg: (%v)", err)
	}
	log.Println("Shrinked /dev/sda3.")
	// need to reboot to reload the partition table.
	return 0, true, nil
	}
	// no need to shrink sda3 again.
	sda3StartSector, err := partutil.ReadPartitionStart("/dev/sda", 3)
	if err != nil {
	return 0, false, fmt.Errorf("error in reading the start of sda3, "+
	"error msg: (%v)", err)
	}
	startPointSector = sda3StartSector + sda3Margin // leave enough margin

	} else {
	// start point is the original start sector of the stateful partition.
	startPointSector, err = partutil.ReadPartitionStart(disk, statePartNum)
	if err != nil {
	return 0, false, fmt.Errorf("error in reading the start of old stateful partition, "+
	"error msg: (%v)", err)
	}
	}
	return startPointSector, false, nil
	}

	// checkNewOEMSizeSmaller checks whether the new OEM size is smaller than the old one.
	// If true, move the stateful partition to reclaim sda3 if needed.
	// Otherwise, do nothing.
	// It returns oemSmaller,error.
	func checkNewOEMSizeSmaller(disk string, statePartNum int, reclaimSDA3 bool, oldOEMSizeBytes,
	newOEMSizeBytes, startPointSector uint64) (bool, error) {
	if newOEMSizeBytes <= oldOEMSizeBytes {
	if newOEMSizeBytes != 0 {
	log.Printf("\n!!!!!!!WARNING!!!!!!!\n"+
	"oemSize: %d bytes is not larger than the original OEM partition size: %d bytes, "+
	"nothing is done for the OEM partition.\n ", newOEMSizeBytes, oldOEMSizeBytes)
	}
	if !reclaimSDA3 {
	return true, nil
	}
	// move the stateful partition to the start point.
	if err := partutil.MovePartition(disk, statePartNum, strconv.FormatUint(startPointSector, 10)); err != nil {
	return true, fmt.Errorf("error in moving stateful partition, "+
	"input: disk=%q, statePartNum=%d,reclaimSDA3=%t,oldOEMSizeBytes=%d,newOEMSizeBytes=%d, "+
	"startPointSector=%d, error msg: (%v)", disk, statePartNum, reclaimSDA3,
	oldOEMSizeBytes, newOEMSizeBytes, startPointSector, err)
	}
	log.Println("Reclaimed /dev/sda3.")
	// print the new partition table.
	table, err := partutil.ReadPartitionTable(disk)
	if err != nil {
	return true, fmt.Errorf("cannot read new partition table of %q, "+
	"input: disk=%q, statePartNum=%d,reclaimSDA3=%t,oldOEMSizeBytes=%d,newOEMSizeBytes=%d, "+
	"startPointSector=%d, error msg: (%v)", disk, disk, statePartNum, reclaimSDA3,
	oldOEMSizeBytes, newOEMSizeBytes, startPointSector, err)
	}
	log.Printf("New partition table:\n%s\n", table)
	return true, nil
	}
	return false, nil
	}