src/pkg/provisioner/disk_layout.go

// Copyright 2021 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 provisioner

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"golang.org/x/sys/unix"

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

func switchRoot(deps Deps, runState *state) (err error) {
	if !runState.data.Config.BootDisk.ReclaimSDA3 {
		log.Println("ReclaimSDA3 is not set, not switching root device")
		return nil
	}
	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return err
	}
	sda3DeviceSuffix, err := partutil.PartNumIntToString(bootDev, 3)
	if err != nil {
		return err
	}
	sda3Device := bootDev + sda3DeviceSuffix
	sda5DeviceSuffix, err := partutil.PartNumIntToString(bootDev, 5)
	if err != nil {
		return err
	}
	sda5Device := bootDev + sda5DeviceSuffix
	rootDev, err := exec.Command(deps.RootdevCmd, "-s").Output()
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			return fmt.Errorf("error running rootdev: %v: stderr = %q", exitErr, string(exitErr.Stderr))
		}
		return fmt.Errorf("error running rootdev: %v", err)
	}
	rootDevStr := strings.TrimSpace(string(rootDev))
	log.Printf("Checking root device: current=%q, target=%q", rootDevStr, sda5Device)
	if rootDevStr == sda5Device {
		log.Printf("Current root device is %q, not switching root device", rootDevStr)
		return nil
	}
	log.Println("Need to switch root device")
	log.Println("Copying partition 3 to partition 5...")
	in, err := os.Open(sda3Device)
	if err != nil {
		return err
	}
	defer utils.CheckClose(in, fmt.Sprintf("error closing %q", sda3Device), &err)
	out, err := os.Create(sda5Device)
	if err != nil {
		return err
	}
	defer utils.CheckClose(out, fmt.Sprintf("error closing %q", sda5Device), &err)
	if _, err := io.Copy(out, in); err != nil {
		return fmt.Errorf("error copying from %q to %q: %v", sda3Device, sda5Device, err)
	}
	log.Println("Setting GPT priority...")
	if err := utils.RunCommand([]string{deps.CgptCmd, "prioritize", "-P", "5", "-i", "4", bootDev}, "", nil); err != nil {
		return err
	}
	log.Println("Reboot required to switch root device")
	return ErrRebootRequired
}

func shrinkSDA3(deps Deps, runState *state) error {
	if !runState.data.Config.BootDisk.ReclaimSDA3 {
		log.Println("ReclaimSDA3 is not set, not shrinking partition 3")
		return nil
	}
	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return err
	}
	minimal, err := partutil.IsPartitionMinimal(bootDev, 3)
	if err != nil {
		return fmt.Errorf("error checking partition 3 size: %v", err)
	}
	if minimal {
		log.Println("partition 3 is minimally sized, not shrinking partition 3")
		return nil
	}
	log.Println("ReclaimSDA3 is set, and partition 3 is not minimal; now shrinking partition 3")
	if _, err := partutil.MinimizePartition(bootDev, 3); err != nil {
		return fmt.Errorf("error minimizing partition 3: %v", err)
	}
	log.Println("Reboot required to reload partition table changes")
	return ErrRebootRequired
}

func setupOnShutdownUnit(deps Deps, runState *state) (err error) {
	if err := mountFunc("", filepath.Join(deps.RootDir, "tmp"), "", unix.MS_REMOUNT|unix.MS_NOSUID|unix.MS_NODEV, ""); err != nil {
		return fmt.Errorf("error remounting /tmp as exec: %v", err)
	}
	out := filepath.Join(deps.RootDir, "tmp", "handle_disk_layout.bin")
	if err := utils.CopyFile(deps.HandleDiskLayoutBin, out); err != nil {
		return err
	}
	if err := os.Chmod(out, 755); err != nil {
		return err
	}
	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return err
	}
	ttyPath := "/dev/ttyS2"
	if out, err := exec.Command("uname", "-m").Output(); err == nil && strings.TrimSpace(string(out)) == "aarch64" {
		ttyPath = "/dev/ttyAMA2"
	}
	data := fmt.Sprintf(`[Unit]
Description=Run after everything unmounted
DefaultDependencies=false
Conflicts=shutdown.target
Before=mnt-stateful_partition.mount usr-share-oem.mount
After=tmp.mount

[Service]
Type=oneshot
RemainAfterExit=true
ExecStart=/bin/true
ExecStop=/bin/bash -c '/tmp/handle_disk_layout.bin %s 1 8 "%s" "%t" 2>&1 | sed "s/^/BuildStatus: /"'
TimeoutStopSec=600
StandardOutput=tty
StandardError=tty
TTYPath=%s
`, bootDev, runState.data.Config.BootDisk.OEMSize, runState.data.Config.BootDisk.ReclaimSDA3, ttyPath)
	if err := ioutil.WriteFile(filepath.Join(deps.RootDir, "etc/systemd/system/last-run.service"), []byte(data), 0664); err != nil {
		return err
	}
	systemd := systemdClient{systemctl: deps.SystemctlCmd}
	if err := systemd.start("last-run.service", []string{"--no-block"}); err != nil {
		return err
	}
	// journald needs to be stopped in order for the stateful partition to be
	// unmounted at shutdown. We need the stateful partition to be unmounted so
	// that disk repartitioning can occur.
	if err := systemd.stopJournald(deps.RootDir); err != nil {
		return err
	}
	return nil
}

func calcSDA3End(device string) (uint64, error) {
	sda3Start, err := partutil.ReadPartitionStart(device, 3)
	if err != nil {
		return 0, err
	}
	sda3Size, err := partutil.ReadPartitionSize(device, 3)
	if err != nil {
		return 0, err
	}
	sda3End := sda3Start + sda3Size - 1
	return sda3End, nil
}

func waitForDiskResize(deps Deps, runState *state) error {
	if !runState.data.Config.BootDisk.WaitForDiskResize {
		log.Println("WaitForDiskResize is not set, not waiting for a boot disk resize")
		return nil
	}
	if runState.data.DiskResizeComplete {
		log.Println("Already finished waiting for disk resize, not waiting again")
		return nil
	}
	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return err
	}
	startSize, err := ioutil.ReadFile(filepath.Join(deps.RootDir, fmt.Sprintf("sys/class/block/%s/size", filepath.Base(bootDev))))
	if err != nil {
		return err
	}
	log.Println("WaitForDiskResize is set; waiting for the boot disk size to change. Timeout is 3 minutes")
	start := time.Now()
	end := start.Add(3 * time.Minute)
	for time.Now().Before(end) {
		curSize, err := ioutil.ReadFile(filepath.Join(deps.RootDir, fmt.Sprintf("sys/class/block/%s/size", filepath.Base(bootDev))))
		if err != nil {
			return err
		}
		if string(curSize) != string(startSize) {
			log.Printf("Boot disk size has changed: start %q, end %q", strings.TrimSpace(string(startSize)), strings.TrimSpace(string(curSize)))
			runState.data.DiskResizeComplete = true
			return runState.write()
		}
		time.Sleep(time.Second)
	}
	return errors.New("timed out waiting for disk resize")
}

func relocatePartitions(deps Deps, runState *state) error {
	if !runState.data.Config.BootDisk.ReclaimSDA3 && runState.data.Config.BootDisk.OEMSize == "" {
		log.Println("ReclaimSDA3 is not set, OEM resize not requested, not relocating partitions")
		return nil
	}
	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return err
	}
	if runState.data.Config.BootDisk.OEMSize != "" {
		// Check if OEM partition is after sda3; if so, then we're done
		oemStart, err := partutil.ReadPartitionStart(bootDev, 8)
		if err != nil {
			return err
		}
		sda3End, err := calcSDA3End(bootDev)
		if err != nil {
			return err
		}
		if oemStart > sda3End {
			log.Println("OEM resize requested, OEM appears to be relocated after sda3. Partition relocation is complete")
			return nil
		}
	} else {
		// Check two things:
		// 1. sda3 is minimal
		// 2. Stateful partition is located immediately after sda3
		//
		// If both are true, we are done.
		minimal, err := partutil.IsPartitionMinimal(bootDev, 3)
		if err != nil {
			return err
		}
		statefulStart, err := partutil.ReadPartitionStart(bootDev, 1)
		if err != nil {
			return err
		}
		sda3Start, err := partutil.ReadPartitionStart(bootDev, 3)
		// The stateful partition is relocated 4096 sectors after the start of sda3.
		// See src/pkg/tools/handle_disk_layout.go for details.
		if minimal && statefulStart == sda3Start+4096 {
			log.Println("ReclaimSDA3 is set, sda3 appears to have been reclaimed. Partition relocation is complete")
			return nil
		}
	}
	// Partition relocation must be done. Prepare for disk relocation to happen on
	// the next reboot
	log.Println("Partition relocation is required. Preparing for partition relocation to occur on the next reboot")
	if err := setupOnShutdownUnit(deps, runState); err != nil {
		return err
	}
	log.Println("Reboot required to relocate partitions")
	return ErrRebootRequired
}

func resizeOEMFileSystem(deps Deps, runState *state) error {
	if !runState.data.Config.BootDisk.ReclaimSDA3 && runState.data.Config.BootDisk.OEMSize == "" {
		log.Println("ReclaimSDA3 is not set, OEM resize not requested, partition relocation did not occur, FS resize unnecessary")
		return nil
	}
	// Check if OEM partition is after sda3; if so, then relocation occurred and
	// we need to resize the file system.
	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return err
	}
	sda3End, err := calcSDA3End(bootDev)
	if err != nil {
		return err
	}
	oemStart, err := partutil.ReadPartitionStart(bootDev, 8)
	if err != nil {
		return err
	}
	if oemStart < sda3End {
		log.Println("OEM partition is before sda3; relocation did not occur, FS resize unnecessary")
		return nil
	}
	log.Println("Partition relocation appears to have occurred, resizing the OEM file system")
	systemd := systemdClient{systemctl: deps.SystemctlCmd}
	if err := systemd.stop("usr-share-oem.mount"); err != nil {
		return err
	}
	sda8DeviceSuffix, err := partutil.PartNumIntToString(bootDev, 8)
	if err != nil {
		return err
	}
	sda8Device := bootDev + sda8DeviceSuffix
	if err := utils.RunCommand([]string{deps.E2fsckCmd, "-fp", sda8Device}, "", nil); err != nil {
		return err
	}
	resizeArgs := []string{deps.Resize2fsCmd, sda8Device}
	if runState.data.Config.BootDisk.OEMFSSize4K != 0 {
		resizeArgs = append(resizeArgs, strconv.FormatUint(runState.data.Config.BootDisk.OEMFSSize4K, 10))
	}
	if err := utils.RunCommand(resizeArgs, "", nil); err != nil {
		return err
	}
	if err := systemd.start("usr-share-oem.mount", nil); err != nil {
		return err
	}
	log.Println("OEM file system resized to account for available space")
	return nil
}

func patchGrubForARM64(deps Deps) error {
	out, err := exec.Command("uname", "-m").Output()
	if err != nil {
		return fmt.Errorf("failed to run uname: %v", err)
	}
	arch := strings.TrimSpace(string(out))
	if arch != "aarch64" {
		return nil
	}
	log.Println("ARM64 architecture detected; patching grub.cfg to load vmlinuz.A and target ROOT-B for slot B")
	efiPath, err := partutil.MountEFIPartition()
	if err != nil {
		return fmt.Errorf("error mounting EFI partition: %v", err)
	}
	defer func() {
		if err := partutil.UnmountEFIPartition(); err != nil {
			log.Printf("Warning: failed to unmount EFI partition: %v", err)
		}
	}()

	grubPath := filepath.Join(efiPath, partutil.GRUBCfgEFIPath)
	grubContentBytes, err := ioutil.ReadFile(grubPath)
	if err != nil {
		return fmt.Errorf("error reading grub.cfg: %v", err)
	}
	grubContent := string(grubContentBytes)

	bootDev, err := partutil.BootDev(deps.RootdevCmd)
	if err != nil {
		return fmt.Errorf("error getting boot dev: %v", err)
	}
	partuuidA, err := partutil.ReadPartitionUUID(bootDev, 3)
	if err != nil {
		return fmt.Errorf("failed to get ROOT-A partition UUID: %v", err)
	}
	partuuidB, err := partutil.ReadPartitionUUID(bootDev, 5)
	if err != nil {
		return fmt.Errorf("failed to get ROOT-B partition UUID: %v", err)
	}

	newContent := grubContent
	if strings.Contains(newContent, "/syslinux/vmlinuz.B") {
		log.Println("Replacing /syslinux/vmlinuz.B with /syslinux/vmlinuz.A in grub.cfg")
		newContent = strings.ReplaceAll(newContent, "/syslinux/vmlinuz.B", "/syslinux/vmlinuz.A")
	}

	log.Printf("Replacing PARTUUID=%s with PARTUUID=%s in slot B menu entry", partuuidA, partuuidB)
	newContent = cosboot.ReplaceInSlotEntry(newContent, true /* isSlotB */, "PARTUUID="+partuuidA, "PARTUUID="+partuuidB)

	if err := ioutil.WriteFile(grubPath, []byte(newContent), 0755); err != nil {
		return fmt.Errorf("error writing patched grub.cfg: %v", err)
	}
	return nil
}

// repartitionBootDisk executes all behaviors related to repartitioning the boot
// disk. Most of these behaviors require a reboot. To keep reboots simple (e.g.
// we don't want to initiate a reboot when deferred statements are unresolved),
// we handle reboots by returning ErrRebootRequired and asking the caller to
// initiate the reboot.
func repartitionBootDisk(deps Deps, runState *state) error {
	if err := patchGrubForARM64(deps); err != nil {
		log.Printf("Warning: failed to patch GRUB config: %v", err)
	}
	if err := switchRoot(deps, runState); err != nil {
		return err
	}
	if err := shrinkSDA3(deps, runState); err != nil {
		return err
	}
	if err := waitForDiskResize(deps, runState); err != nil {
		return err
	}
	if err := relocatePartitions(deps, runState); err != nil {
		return err
	}
	if err := resizeOEMFileSystem(deps, runState); err != nil {
		return err
	}
	return nil
}