| // 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 |
| } |