src/pkg/tools/partutil/handle_partition_table.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 partutil

 import (
 	"bytes"
 	"fmt"
 	"log"
 	"os"
 	"os/exec"
 	"strconv"
 	"strings"
 )

 // minPartitionSizeSectors is the minimum partition size on COS in sectors. This
 // size was chosen to maintain 4K alignment of partition start sectors (8
 // sectors = 4K bytes).
 const minPartitionSizeSectors = uint64(8)

 // When we read disk information by dumping the partition table, we get output like the following:
 // sudo sfdisk --dump /dev/sdb
 // label: gpt
 // label-id: 8071096F-DA33-154D-A687-AE097B8252C5
 // device: /dev/sdb
 // unit: sectors
 // first-lba: 2048
 // last-lba: 20971486

 // /dev/sdb1 : start=     4401152, size=     2097152, type=0FC63DAF-8483-4772-8E79-3D69D8477DE4, uuid=3B41256B-E064-544A-9101-D2647C0B3A38
 // /dev/sdb2 : start=      206848, size=     4194304, type=0FC63DAF-8483-4772-8E79-3D69D8477DE4, uuid=60E55EA1-4EEA-9F44-A066-4720F0129089
 // /dev/sdb3 : start=     6498304, size=      204800, type=0FC63DAF-8483-4772-8E79-3D69D8477DE4, uuid=9479C34A-49A6-9442-A56F-956396DFAC20

 // PartContent contains the info of a partition
 type PartContent struct {
 	Start uint64
 	Size  uint64
 }

 // HandlePartitionTable takes a partition table and get the start and size of the target partition.
 // If change==true, it will rebuild the partition table with data passed in by p *PartContent
 // and return the new table
 func HandlePartitionTable(table, partName string, change bool, f func(p *PartContent)) (string, error) {
 	foundPartition := false
 	lines := strings.Split(table, "\n")
 	for idx, line := range lines {
 		// a white space is needed to prevent cases like /dev/sda14 matches /dev/sda1
 		if !strings.HasPrefix(line, partName+" ") {
 			continue
 		}
 		foundPartition = true
 		content := strings.Split(line, ":")
 		partInfo := strings.Split(content[1], ",")
 		startSec := strings.Split(partInfo[0], "=")
 		sizeSec := strings.Split(partInfo[1], "=")
 		var p PartContent
 		var err error
 		p.Start, err = strconv.ParseUint(strings.TrimSpace(startSec[1]), 10, 64)
 		if err != nil {
 			return "", fmt.Errorf("cannot convert %q to int, "+
 				"partition info: %q, error msg: (%v)", strings.TrimSpace(startSec[1]), line, err)
 		}
 		p.Size, err = strconv.ParseUint(strings.TrimSpace(sizeSec[1]), 10, 64)
 		if err != nil {
 			return "", fmt.Errorf("cannot convert %q to int, "+
 				"partition info: %q, error msg: (%v)", strings.TrimSpace(startSec[1]), line, err)
 		}
 		// run reading or changing function on the PartContent struct.
 		f(&p)

 		// need to rebuild the partition table.
 		if change {
 			startSec[1] = strconv.FormatUint(p.Start, 10)
 			partInfo[0] = strings.Join(startSec, "=")
 			sizeSec[1] = strconv.FormatUint(p.Size, 10)
 			partInfo[1] = strings.Join(sizeSec, "=")
 			content[1] = strings.Join(partInfo, ",")
 			lines[idx] = strings.Join(content, ":")
 		}
 		break
 	}
 	if !foundPartition {
 		return table, fmt.Errorf("cannot find the target partition %q, "+
 			"partition table: %s", partName, table)
 	}

 	if change {
 		table = strings.Join(lines, "\n")
 	}
 	return table, nil
 }

 // ReadPartitionTable reads the partition table of a disk.
 func ReadPartitionTable(disk string) (string, error) {
 	table, err := exec.Command("sudo", "sfdisk", "--dump", disk).Output()
 	if err != nil {
 		return "", fmt.Errorf("cannot dump partition table of %q, "+
 			"error msg: (%v)", disk, err)
 	}
 	return string(table), nil
 }

 // ReadPartitionSize reads the size of a partition (unit:sectors of 512 Bytes).
 func ReadPartitionSize(disk string, partNumInt int) (uint64, error) {
 	if len(disk) <= 0 || partNumInt <= 0 {
 		return 0, fmt.Errorf("invalid input: disk=%q, partNumInt=%d", disk, partNumInt)
 	}

 	// get partition number string
 	partNum, err := PartNumIntToString(disk, partNumInt)
 	if err != nil {
 		return 0, fmt.Errorf("error in converting partition number, "+
 			"input: disk=%q, partNumInt=%d, "+
 			"error msg: (%v)", disk, partNumInt, err)
 	}
 	partName := disk + partNum

 	table, err := ReadPartitionTable(disk)
 	if err != nil {
 		return 0, fmt.Errorf("cannot read partition table of %q, "+
 			"input: disk=%q, partNumInt=%d, "+
 			"error msg: (%v)", disk, disk, partNumInt, err)
 	}

 	var size uint64 = 0
 	if _, err = HandlePartitionTable(table, partName, false, func(p *PartContent) { size = p.Size }); err != nil {
 		return 0, fmt.Errorf("error parsing partition table of %q, "+
 			"input: disk=%q, partNumInt=%d, "+
 			"error msg: (%v)", disk, disk, partNumInt, err)
 	}
 	return size, nil
 }

 // ReadPartitionStart reads the start sector of a partition.
 func ReadPartitionStart(disk string, partNumInt int) (uint64, error) {
 	if len(disk) <= 0 || partNumInt <= 0 {
 		return 0, fmt.Errorf("invalid input: disk=%q, partNumInt=%d", disk, partNumInt)
 	}

 	// get partition number string
 	partNum, err := PartNumIntToString(disk, partNumInt)
 	if err != nil {
 		return 0, fmt.Errorf("error in converting partition number, "+
 			"input: disk=%q, partNumInt=%d, "+
 			"error msg: (%v)", disk, partNumInt, err)
 	}
 	partName := disk + partNum

 	table, err := ReadPartitionTable(disk)
 	if err != nil {
 		return 0, fmt.Errorf("cannot read partition table of %q, "+
 			"input: disk=%q, partNumInt=%d, "+
 			"error msg: (%v)", disk, disk, partNumInt, err)
 	}

 	var start uint64 = 0
 	if _, err = HandlePartitionTable(table, partName, false, func(p *PartContent) { start = p.Start }); err != nil {
 		return 0, fmt.Errorf("error parsing partition table of %q, "+
 			"input: disk=%q, partNumInt=%d, "+
 			"error msg: (%v)", disk, disk, partNumInt, err)
 	}
 	return start, nil
 }

 // IsPartitionMinimal determines if a partition is the smallest size it can be.
 // If this function returns true, MinimizePartition can make the given partition
 // smaller.
 func IsPartitionMinimal(disk string, partNumInt int) (bool, error) {
 	numSectors, err := ReadPartitionSize(disk, partNumInt)
 	if err != nil {
 		return false, err
 	}
 	if numSectors > minPartitionSizeSectors {
 		return false, nil
 	}
 	return true, nil
 }

 // MinimizePartition minimizes the input partition and
 // returns the next sector of the end sector.
 // The smallest partition from fdisk is 1 sector partition.
 func MinimizePartition(disk string, partNumInt int) (uint64, error) {
 	minSize := minPartitionSizeSectors
 	if len(disk) == 0 || partNumInt <= 0 {
 		return 0, fmt.Errorf("empty disk name or nonpositive part number, "+
 			"input: disk=%q, partNumInt=%d", disk, partNumInt)
 	}

 	// get partition number string
 	partNum, err := PartNumIntToString(disk, partNumInt)
 	if err != nil {
 		return 0, fmt.Errorf("error in converting partition number, "+
 			"input: disk=%q, partNumInt=%d, "+
 			"error msg: (%v)", disk, partNumInt, err)
 	}

 	partName := disk + partNum

 	var tableBuffer bytes.Buffer
 	var oldSize uint64

 	// dump partition table.
 	table, err := ReadPartitionTable(disk)
 	if err != nil {
 		return 0, fmt.Errorf("cannot read partition table of %q, "+
 			"input: disk=%q, partNumInt=%d, "+
 			"error msg: (%v)", disk, disk, partNumInt, err)
 	}

 	var startSector uint64

 	// edit partition table.
 	table, err = HandlePartitionTable(table, partName, true, func(p *PartContent) {
 		startSector = p.Start
 		oldSize = p.Size
 		p.Size = minSize
 	})
 	if err != nil {
 		return 0, fmt.Errorf("error when editing partition table of %q, "+
 			"input: disk=%q, partNumInt=%d, "+
 			"error msg: (%v)", disk, disk, partNumInt, err)
 	}

 	if oldSize <= minSize {
 		log.Printf("warning: old partition size=%d is smaller than minSize=%d, "+
 			"nothing is done, "+
 			"return value is start sector + minSize (%d)", oldSize, minSize, minSize)
 		return startSector + minSize, nil
 	}

 	tableBuffer.WriteString(table)

 	// write partition table back.
 	writeTableCmd := exec.Command("sudo", "sfdisk", "--no-reread", disk)
 	writeTableCmd.Stdin = &tableBuffer
 	writeTableCmd.Stdout = os.Stdout
 	writeTableCmd.Stderr = os.Stderr
 	if err := writeTableCmd.Run(); err != nil {
 		return 0, fmt.Errorf("error in writing partition table back to %q, "+
 			"input: disk=%q, partNumInt=%d, "+
 			"error msg: (%v)", disk, disk, partNumInt, err)
 	}

 	log.Printf("\nCompleted minimizing %q\n\n", partName)
 	return startSector + minSize, 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 partutil

	import (
	"bytes"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	)

	// minPartitionSizeSectors is the minimum partition size on COS in sectors. This
	// size was chosen to maintain 4K alignment of partition start sectors (8
	// sectors = 4K bytes).
	const minPartitionSizeSectors = uint64(8)

	// When we read disk information by dumping the partition table, we get output like the following:
	// sudo sfdisk --dump /dev/sdb
	// label: gpt
	// label-id: 8071096F-DA33-154D-A687-AE097B8252C5
	// device: /dev/sdb
	// unit: sectors
	// first-lba: 2048
	// last-lba: 20971486

	// /dev/sdb1 : start= 4401152, size= 2097152, type=0FC63DAF-8483-4772-8E79-3D69D8477DE4, uuid=3B41256B-E064-544A-9101-D2647C0B3A38
	// /dev/sdb2 : start= 206848, size= 4194304, type=0FC63DAF-8483-4772-8E79-3D69D8477DE4, uuid=60E55EA1-4EEA-9F44-A066-4720F0129089
	// /dev/sdb3 : start= 6498304, size= 204800, type=0FC63DAF-8483-4772-8E79-3D69D8477DE4, uuid=9479C34A-49A6-9442-A56F-956396DFAC20

	// PartContent contains the info of a partition
	type PartContent struct {
	Start uint64
	Size uint64
	}

	// HandlePartitionTable takes a partition table and get the start and size of the target partition.
	// If change==true, it will rebuild the partition table with data passed in by p *PartContent
	// and return the new table
	func HandlePartitionTable(table, partName string, change bool, f func(p *PartContent)) (string, error) {
	foundPartition := false
	lines := strings.Split(table, "\n")
	for idx, line := range lines {
	// a white space is needed to prevent cases like /dev/sda14 matches /dev/sda1
	if !strings.HasPrefix(line, partName+" ") {
	continue
	}
	foundPartition = true
	content := strings.Split(line, ":")
	partInfo := strings.Split(content[1], ",")
	startSec := strings.Split(partInfo[0], "=")
	sizeSec := strings.Split(partInfo[1], "=")
	var p PartContent
	var err error
	p.Start, err = strconv.ParseUint(strings.TrimSpace(startSec[1]), 10, 64)
	if err != nil {
	return "", fmt.Errorf("cannot convert %q to int, "+
	"partition info: %q, error msg: (%v)", strings.TrimSpace(startSec[1]), line, err)
	}
	p.Size, err = strconv.ParseUint(strings.TrimSpace(sizeSec[1]), 10, 64)
	if err != nil {
	return "", fmt.Errorf("cannot convert %q to int, "+
	"partition info: %q, error msg: (%v)", strings.TrimSpace(startSec[1]), line, err)
	}
	// run reading or changing function on the PartContent struct.
	f(&p)

	// need to rebuild the partition table.
	if change {
	startSec[1] = strconv.FormatUint(p.Start, 10)
	partInfo[0] = strings.Join(startSec, "=")
	sizeSec[1] = strconv.FormatUint(p.Size, 10)
	partInfo[1] = strings.Join(sizeSec, "=")
	content[1] = strings.Join(partInfo, ",")
	lines[idx] = strings.Join(content, ":")
	}
	break
	}
	if !foundPartition {
	return table, fmt.Errorf("cannot find the target partition %q, "+
	"partition table: %s", partName, table)
	}

	if change {
	table = strings.Join(lines, "\n")
	}
	return table, nil
	}

	// ReadPartitionTable reads the partition table of a disk.
	func ReadPartitionTable(disk string) (string, error) {
	table, err := exec.Command("sudo", "sfdisk", "--dump", disk).Output()
	if err != nil {
	return "", fmt.Errorf("cannot dump partition table of %q, "+
	"error msg: (%v)", disk, err)
	}
	return string(table), nil
	}

	// ReadPartitionSize reads the size of a partition (unit:sectors of 512 Bytes).
	func ReadPartitionSize(disk string, partNumInt int) (uint64, error) {
	if len(disk) <= 0 \|\| partNumInt <= 0 {
	return 0, fmt.Errorf("invalid input: disk=%q, partNumInt=%d", disk, partNumInt)
	}

	// get partition number string
	partNum, err := PartNumIntToString(disk, partNumInt)
	if err != nil {
	return 0, fmt.Errorf("error in converting partition number, "+
	"input: disk=%q, partNumInt=%d, "+
	"error msg: (%v)", disk, partNumInt, err)
	}
	partName := disk + partNum

	table, err := ReadPartitionTable(disk)
	if err != nil {
	return 0, fmt.Errorf("cannot read partition table of %q, "+
	"input: disk=%q, partNumInt=%d, "+
	"error msg: (%v)", disk, disk, partNumInt, err)
	}

	var size uint64 = 0
	if _, err = HandlePartitionTable(table, partName, false, func(p *PartContent) { size = p.Size }); err != nil {
	return 0, fmt.Errorf("error parsing partition table of %q, "+
	"input: disk=%q, partNumInt=%d, "+
	"error msg: (%v)", disk, disk, partNumInt, err)
	}
	return size, nil
	}

	// ReadPartitionStart reads the start sector of a partition.
	func ReadPartitionStart(disk string, partNumInt int) (uint64, error) {
	if len(disk) <= 0 \|\| partNumInt <= 0 {
	return 0, fmt.Errorf("invalid input: disk=%q, partNumInt=%d", disk, partNumInt)
	}

	// get partition number string
	partNum, err := PartNumIntToString(disk, partNumInt)
	if err != nil {
	return 0, fmt.Errorf("error in converting partition number, "+
	"input: disk=%q, partNumInt=%d, "+
	"error msg: (%v)", disk, partNumInt, err)
	}
	partName := disk + partNum

	table, err := ReadPartitionTable(disk)
	if err != nil {
	return 0, fmt.Errorf("cannot read partition table of %q, "+
	"input: disk=%q, partNumInt=%d, "+
	"error msg: (%v)", disk, disk, partNumInt, err)
	}

	var start uint64 = 0
	if _, err = HandlePartitionTable(table, partName, false, func(p *PartContent) { start = p.Start }); err != nil {
	return 0, fmt.Errorf("error parsing partition table of %q, "+
	"input: disk=%q, partNumInt=%d, "+
	"error msg: (%v)", disk, disk, partNumInt, err)
	}
	return start, nil
	}

	// IsPartitionMinimal determines if a partition is the smallest size it can be.
	// If this function returns true, MinimizePartition can make the given partition
	// smaller.
	func IsPartitionMinimal(disk string, partNumInt int) (bool, error) {
	numSectors, err := ReadPartitionSize(disk, partNumInt)
	if err != nil {
	return false, err
	}
	if numSectors > minPartitionSizeSectors {
	return false, nil
	}
	return true, nil
	}

	// MinimizePartition minimizes the input partition and
	// returns the next sector of the end sector.
	// The smallest partition from fdisk is 1 sector partition.
	func MinimizePartition(disk string, partNumInt int) (uint64, error) {
	minSize := minPartitionSizeSectors
	if len(disk) == 0 \|\| partNumInt <= 0 {
	return 0, fmt.Errorf("empty disk name or nonpositive part number, "+
	"input: disk=%q, partNumInt=%d", disk, partNumInt)
	}

	// get partition number string
	partNum, err := PartNumIntToString(disk, partNumInt)
	if err != nil {
	return 0, fmt.Errorf("error in converting partition number, "+
	"input: disk=%q, partNumInt=%d, "+
	"error msg: (%v)", disk, partNumInt, err)
	}

	partName := disk + partNum

	var tableBuffer bytes.Buffer
	var oldSize uint64

	// dump partition table.
	table, err := ReadPartitionTable(disk)
	if err != nil {
	return 0, fmt.Errorf("cannot read partition table of %q, "+
	"input: disk=%q, partNumInt=%d, "+
	"error msg: (%v)", disk, disk, partNumInt, err)
	}

	var startSector uint64

	// edit partition table.
	table, err = HandlePartitionTable(table, partName, true, func(p *PartContent) {
	startSector = p.Start
	oldSize = p.Size
	p.Size = minSize
	})
	if err != nil {
	return 0, fmt.Errorf("error when editing partition table of %q, "+
	"input: disk=%q, partNumInt=%d, "+
	"error msg: (%v)", disk, disk, partNumInt, err)
	}

	if oldSize <= minSize {
	log.Printf("warning: old partition size=%d is smaller than minSize=%d, "+
	"nothing is done, "+
	"return value is start sector + minSize (%d)", oldSize, minSize, minSize)
	return startSector + minSize, nil
	}

	tableBuffer.WriteString(table)

	// write partition table back.
	writeTableCmd := exec.Command("sudo", "sfdisk", "--no-reread", disk)
	writeTableCmd.Stdin = &tableBuffer
	writeTableCmd.Stdout = os.Stdout
	writeTableCmd.Stderr = os.Stderr
	if err := writeTableCmd.Run(); err != nil {
	return 0, fmt.Errorf("error in writing partition table back to %q, "+
	"input: disk=%q, partNumInt=%d, "+
	"error msg: (%v)", disk, disk, partNumInt, err)
	}

	log.Printf("\nCompleted minimizing %q\n\n", partName)
	return startSector + minSize, nil
	}