# Copyright 2016 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# Notes:
# (1) Usage: This script gets used as a "library" from other parts of the
#     build scripts.  Specifically, the board_specific_setup.sh script sources
#     this script as part of the board_make_image_bootable() function.  The
#     "cros_make_image_bootable" script calls the board_make_image_bootable()
#     function.
# (2) In order to install GRUB2 on GPT-partitioned disks, there must be a
#     designated "BIOS boot partition".
#
#     On MBR-only disks, the disk sectors immediately following the MBR are
#     used for storing "stage 2" of GRUB2.  On GPT disks, the sectors
#     immediately after MBR are used up to hold the actual partition table.
#     In such case, a well-known GUID can be used to mark a partition as the
#     BIOS boot partition, which the grub-install command uses to store the
#     second stage of the bootloader.  For more details on BIOS boot partition,
#     see https://en.wikipedia.org/wiki/BIOS_boot_partition.
#
#     The disk layout of Lakitu is inherited from the ChromeOS disk layout,
#     which has a partition 11 for firmware and a partition 12 for legacy
#     bootloader configurations.  Since Lakitu doesn't use a custom firmware,
#     we can reuse partition 11 as the BIOS boot partition without having to
#     change the disk layout.  We continue to use partition 12 for storing
#     bootloader configurations.

IMAGE=

ESP_DIR=""
LOOP_DEV=

# Only let dm-verity block if rootfs verification is configured.
# Also, set which device mapper correspondes to verity
dev_wait=0
ROOTDEV=/dev/dm-0
if [[ ${FLAGS_enable_rootfs_verification} -eq ${FLAGS_TRUE} ]]; then
  dev_wait=1
  if [[ ${FLAGS_enable_bootcache} -eq ${FLAGS_TRUE} ]]; then
    ROOTDEV=/dev/dm-1
  fi
fi

part_index_to_uuid() {
  local image="$1"
  local index="$2"
  sudo cgpt show -i "$index" -u "$image"
}

cleanup() {
  if [[ -d "${ESP_DIR}" ]]; then
    if mountpoint -q "${ESP_DIR}"; then
      sudo umount "${ESP_DIR}"
    fi
    # ESP_DIR should be empty after the umount, so rmdir should work.
    rmdir "${ESP_DIR}"
  fi

  if [[ -b "${LOOP_DEV}" ]]; then
    sudo losetup --detach "${LOOP_DEV}"
  fi
}

# Installs bootloaders on the given disk image.
#
# Args:
#   $1: absolute path to the disk image.
bootloader_install() {
  trap cleanup EXIT
  IMAGE="$1"

  info "Installing bootloaders on ${IMAGE}"

  LOOP_DEV="$(sudo losetup --find --show --partscan "${IMAGE}")"

  if ! sudo udevadm settle --timeout=10; then
    warn "Error running 'udevadm settle'"
  fi

  local -r esp_node=${LOOP_DEV}p12
  # Wait till udevd finishes the work.
  for i in {1..10}; do
    info "Probing ${esp_node}"
    if [[ -b "${esp_node}" ]]; then
      break
    fi
    sleep 1
    sudo blockdev --rereadpt "${LOOP_DEV}"
  done
  if [[ ! -b "${esp_node}" ]]; then
    error "EFI partition ${esp_node} not available"
    return 1
  fi

  ESP_DIR="$(mktemp --directory)"
  if ! sudo mount -t vfat "${esp_node}" "${ESP_DIR}"; then
    error "Could not mount EFI partition"
    return 1
  fi

  if [[ "${ARCH}" == "amd64" ]]; then
    install_grub2_amd64
  elif [[ "${ARCH}" == "arm64" ]]; then
    install_efi_bootloaders_arm64
  else
    error "Could not install bootloaders on unsupported architecture: ${ARCH}"
    return 1
  fi

  cleanup
  trap - EXIT
  info "Successfully installed bootloaders on ${IMAGE}"
}

install_grub2_amd64() {
  local -r grub_dir="boot/grub"
  local -r grub_target="i386-pc"
  local -r grub_modules="part_gpt gptpriority test fat ext2 normal boot chain
                         configfile linux search search_fs_uuid search_label
                         terminal memdisk echo biosdisk serial"

  # Use partition #11 (originally meant for firmware RW) as BIOS Boot Partition.
  local -r bios_part_num="11"
  # See https://www.gnu.org/software/grub/manual/html_node/BIOS-installation.html
  # for the UUID.
  local -r bios_part_uuid="21686148-6449-6E6F-744E-656564454649"

  info "Installing amd64 GRUB2 ${grub_target}"

  if ! sudo mkdir -p "${ESP_DIR}/${grub_dir}"; then
    error "Could not create dir for grub config"
    return 1
  fi

  if [[ ! -f ${ESP_DIR}/efi/boot/grub.cfg ]]; then
    error "Could not find grub.cfg from EFI installation."
    return 1
  fi

  if ! sudo sed -i -e 's|/sbin/init|/usr/lib/systemd/systemd|' \
                -e '/^set timeout=/s|=.*|=0|' \
                "${ESP_DIR}/efi/boot/grub.cfg"; then
    error "Failed to update grub configuration file."
    return 1
  fi

  # Create a minimal grub.cfg that sets up environment variables and loads the
  # configuration from EFI installation.
  cat <<EOF | sudo dd of="${ESP_DIR}/${grub_dir}/grub.cfg" 2>/dev/null
# The configuration relies on |grubdisk| environment variable to find the GPT
# disk.  When booting with EFI, the efidisk module exports the env var.  In case
# of BIOS, we need to hard-code it and export.
set grubdisk=hd0
export grubdisk
configfile /efi/boot/grub.cfg
EOF

  # Mark the firmware partition as the BIOS boot partition.
  # Note that this must be done prior to running grub-install, or else the
  # grub-install command will fail.
  info "Setting BIOS boot partition ..."
  if ! sudo cgpt add -i "${bios_part_num}" \
                     -t "${bios_part_uuid}" "${LOOP_DEV}"; then
    error "Error running cgpt"
    return 1
  fi

  info "Running grub-install ..."
  if ! sudo grub-install --target="${grub_target}" \
          --boot-directory="${ESP_DIR}/boot" \
          --modules="${grub_modules}" \
          "${LOOP_DEV}"; then
    error "Error running grub-install"
    return 1
  fi

  info "Successfully installed amd64 GRUB2 ${grub_target}"
}

ROOT_FS_DIR_ARM64=""
STATEFUL_FS_DIR_ARM64=""
KERNEL_CMDLINE_FILE_ARM64=""

install_efi_bootloaders_arm64_cleanup() {
  if [[ -d "${ROOT_FS_DIR_ARM64}" ]]; then
    if mountpoint -q "${ROOT_FS_DIR_ARM64}"; then
      unmount_image
    fi
    rmdir "${ROOT_FS_DIR_ARM64}"
  fi
  if [[ -d "${STATEFUL_FS_DIR_ARM64}" ]]; then
    if mountpoint -q "${STATEFUL_FS_DIR_ARM64}"; then
      unmount_image
    fi
    rmdir "${STATEFUL_FS_DIR_ARM64}"
  fi
  rm -f "${KERNEL_CMDLINE_FILE_ARM64}"
}

install_efi_bootloaders_arm64() {
  info "Installing arm64 EFI Shim and GRUB2"

  trap install_efi_bootloaders_arm64_cleanup EXIT

  # The script in src/scripts/update_bootloaders.sh copies the EFI bootloader
  # files to the EFI system partition, only for x86 and amd64. Hence, we should
  # copy them here for lakitu-arm64.
  ROOT_FS_DIR_ARM64="$(mktemp --directory)"
  STATEFUL_FS_DIR_ARM64="$(mktemp --directory)"
  mount_image "${IMAGE}" "${ROOT_FS_DIR_ARM64}" "${STATEFUL_FS_DIR_ARM64}" "" \
              "--read_only"
  sudo mkdir -p "${ESP_DIR}"/efi/boot
  sudo cp "${ROOT_FS_DIR_ARM64}"/boot/efi/boot/*.efi "${ESP_DIR}"/efi/boot/

  if [[ ! -f ${ESP_DIR}/efi/boot/grub.cfg ]]; then
    error "Could not find grub.cfg from EFI installation."
    return 1
  fi

  if ! sudo sed -i -e 's|/sbin/init|/usr/lib/systemd/systemd|' \
                -e '/^set timeout=/s|=.*|=0|' \
                "${ESP_DIR}/efi/boot/grub.cfg"; then
    error "Failed to update grub configuration file."
    return 1
  fi

  unmount_image
  install_efi_bootloaders_arm64_cleanup
  trap - EXIT

  info "Successfully installed arm64 EFI Shim and GRUB2"

  info "Installing arm64 kernels"
  # The function create_base_image() in
  # src/scripts/build_library/base_image_util.sh places the Linux kernel images
  # in ${IMAGE_DIR}/boot_images/ . Then create_base_image() runs
  # src/scripts/bin/cros_make_image_bootable, which runs
  # src/update_bootloaders.sh . The script in update_bootloaders.sh copies the
  # Linux kernel files into the EFI system partition for some architectures, but
  # not for arm64. Hence, we should copy them here for lakitu for arm64.
  # TODO: Install vmlinuz.B as well, once we reduce the kernel image size.
  sudo mkdir -p "${ESP_DIR}"/syslinux
  sudo cp "${IMAGE_DIR}"/boot_images/vmlinuz-* "${ESP_DIR}"/syslinux/vmlinuz.A

  # sign_official_image.sh uses wildcard *.cfg to do some editing
  # and if there is no cfg file it fails. Add dummy one so wildcard
  # is properly expanded
  sudo touch "${ESP_DIR}"/syslinux/dummy.cfg
  info "Successfully installed arm64 kernels"
}