# 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

  info "Setting up UEFI booting on ${IMAGE}"
  # Directory for GSetup data.
  sudo mkdir -p "${ESP_DIR}/EFI/Google/GSetup"
  # Boot sequence.
  if [[ "${ARCH}" == "amd64" ]]; then
    echo "\\efi\\boot\\shimx64.efi" | \
        sudo tee "${ESP_DIR}/EFI/Google/GSetup/Boot" > /dev/null
  elif [[ "${ARCH}" == "arm64" ]]; then
    echo "\\efi\\boot\\shimaa64.efi" | \
        sudo tee "${ESP_DIR}/EFI/Google/GSetup/Boot" > /dev/null
  else
    error "Could not set up UEFI booting on unsupported architecture: ${ARCH}"
    return 1
  fi
  info "Successfully set up UEFI booting on ${IMAGE}"

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

  # For x86 and amd64, create_base_image() in
  # src/scripts/build_library/base_image_util.sh runs
  # src/scripts/build_library/create_legacy_bootloader_templates.sh . The script
  # in create_legacy_bootloader_templates.sh generates grub.cfg in the EFI
  # system partition, but that's only for x86 and amd64. Hence, we should
  # generate one here for lakitu for arm64.

  KERNEL_CMDLINE_FILE_ARM64="$(mktemp legacy_config_XXXXXXXXXX.txt)"

  cat <<EOF > "${KERNEL_CMDLINE_FILE_ARM64}"
init=/usr/lib/systemd/systemd
boot=local
rootwait
ro
noresume
noswap
loglevel=7
noinitrd
console=ttyAMA0
EOF

  local -r script_root="$(readlink -f "$(dirname "${BASH_SOURCE[0]}")")"
  . "${script_root}/build_kernel_image.sh" || return 1
  modify_kernel_command_line "${KERNEL_CMDLINE_FILE_ARM64}"
  # Read back the config_file; translate newlines to space
  common_args="$(tr "\n" " " < "${KERNEL_CMDLINE_FILE_ARM64}")"
  install_efi_bootloaders_arm64_cleanup
  trap - EXIT

  # Common verified boot command-line args
  verity_common="dm_verity.error_behavior=${FLAGS_verity_error_behavior}"
  verity_common="${verity_common} dm_verity.max_bios=${FLAGS_verity_max_ios}"
  verity_common="${verity_common} dm_verity.dev_wait=${dev_wait}"

  # Discover last known partition numbers.
  part_num_kern_a="$(get_layout_partition_number \
      "${FLAGS_image_type}" KERN-A)"
  part_num_kern_b="$(get_layout_partition_number \
      "${FLAGS_image_type}" KERN-B)"
  part_num_root_a="$(get_layout_partition_number \
      "${FLAGS_image_type}" ROOT-A)"
  part_num_root_b="$(get_layout_partition_number \
      "${FLAGS_image_type}" ROOT-B)"
  root_a_uuid="PARTUUID=$(part_index_to_uuid "${LOOP_DEV}" ${part_num_root_a})"
  root_b_uuid="PARTUUID=$(part_index_to_uuid "${LOOP_DEV}" ${part_num_root_b})"

  # Make sure we got the kernel partition
  if [ -z "${part_num_kern_a}" ]; then
      error "failed to get a partition number for KERN-A"
      return 1
  fi

  dm_table=
  # get cmdline from any kernel partition
  kernel_cmdline=$(sudo dump_kernel_config "${LOOP_DEV}p${part_num_kern_a}")
  if echo "$kernel_cmdline" | grep -q 'dm="'; then
    dm_table=$(echo "$kernel_cmdline" | sed -s 's/.*dm="\([^"]*\)".*/\1/')
  fi
  old_root='PARTUUID=%U/PARTNROFF=1'
  grub_dm_table_a=${dm_table//${old_root}/${root_a_uuid}}
  grub_dm_table_b=${dm_table//${old_root}/${root_b_uuid}}

  cat <<EOF | sudo dd of="${ESP_DIR}/efi/boot/grub.cfg"
defaultA=0
defaultB=1
gptpriority \$grubdisk ${part_num_kern_a} prioA
gptpriority \$grubdisk ${part_num_kern_b} prioB

if [ \$prioA -lt \$prioB ]; then
  set default=\$defaultB
else
  set default=\$defaultA
fi

set timeout=2

# NOTE: These magic grub variables are a Chrome OS hack. They are not portable.

menuentry "local image A" {
  linux /syslinux/vmlinuz.A ${common_args} cros_efi \
      root=${root_a_uuid}
}

menuentry "local image B" {
  linux /syslinux/vmlinuz.B ${common_args} cros_efi \
      root=${root_b_uuid}
}

menuentry "verified image A" {
  linux /syslinux/vmlinuz.A ${common_args} ${verity_common} \
      cros_efi root=${ROOTDEV} dm="${grub_dm_table_a}"
}

menuentry "verified image B" {
  linux /syslinux/vmlinuz.B ${common_args} ${verity_common} \
      cros_efi root=${ROOTDEV} dm="${grub_dm_table_b}"
}
EOF

  if [[ ${FLAGS_enable_rootfs_verification} -eq ${FLAGS_TRUE} ]]; then
    sudo sed -i \
      -e '/^defaultA=/s:=.*:=2:' \
      -e '/^defaultB=/s:=.*:=3:' \
      "${ESP_DIR}/efi/boot/grub.cfg"
  fi

  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.
  sudo mkdir -p "${ESP_DIR}"/syslinux
  # TODO: Install vmlinuz.B as well, once we reduce the kernel image size.
  sudo cp "${IMAGE_DIR}"/boot_images/vmlinuz-* "${ESP_DIR}"/syslinux/vmlinuz.A
  info "Successfully installed arm64 kernels"
}