#!/bin/bash
#
# Copyright (c) 2010 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.
#
# Script which ensures that a given image has an up-to-date
# kernel partition, rootfs integrity hashes, and legacy bootloader configs.

# --- BEGIN COMMON.SH BOILERPLATE ---
# Load common CrOS utilities.  Inside the chroot this file is installed in
# /usr/lib/crosutils.  Outside the chroot we find it relative to the script's
# location.
find_common_sh() {
  local common_paths=("$(dirname "$(readlink -f "$0")")/.." /usr/lib/crosutils)
  local path

  SCRIPT_ROOT="${common_paths[0]}"
  for path in "${common_paths[@]}"; do
    if [ -r "${path}/common.sh" ]; then
      SCRIPT_ROOT="${path}"
      break
    fi
  done
}

find_common_sh
. "${SCRIPT_ROOT}/common.sh" || exit 1
# --- END COMMON.SH BOILERPLATE ---

# Need to be inside the chroot to load chromeos-common.sh
assert_inside_chroot

# Load functions and constants for chromeos-install
. /usr/share/misc/chromeos-common.sh || exit 1
. "${BUILD_LIBRARY_DIR}/build_image_util.sh" || exit 1
. "${BUILD_LIBRARY_DIR}/disk_layout_util.sh" || exit 1

switch_to_strict_mode

if [ $# -lt 2 ]; then
  echo "Usage: ${0} /PATH/TO/IMAGE IMAGE.BIN [shflags overrides]"
  exit 1
fi

IMAGE_DIR="$(readlink -f "${1}")"
BOOT_DESC_FILE="${IMAGE_DIR}/boot.desc"
IMAGE="${IMAGE_DIR}/${2}"
shift
shift
FLAG_OVERRIDES="${@}"

if [ ! -r "${BOOT_DESC_FILE}" ]; then
  warn "${BOOT_DESC_FILE} cannot be read!"
  warn "Falling back to command line parsing"
  BOOT_DESC="${@}"
else
  BOOT_DESC="$(cat ${BOOT_DESC_FILE} | tr -s '\n' ' ')"
  info "Boot-time configuration for $(dirname "${IMAGE}"): "
  cat ${BOOT_DESC_FILE} | while read line; do
    info "  ${line}"
  done
fi

if [ ! -r "${IMAGE}" ]; then
  die "${IMAGE} cannot be read!"
fi


locate_gpt
set +e

# Now parse the build settings from ${OUTPUT_DIR}/boot.desc
DEFINE_string adjust_part "" \
  "Adjustments to apply to the partition table"
DEFINE_string board "${DEFAULT_BOARD}" \
  "Board we're building for."
DEFINE_string image_type "base" \
  "Type of image we're building for (base/factory_install)."
DEFINE_string output_dir "/tmp" \
  "Directory to place output in."
DEFINE_string image "chromiumos_base.img" \
  "Full path to the chromiumos image to make bootable."
DEFINE_string arch "x86" \
  "Architecture to make bootable for: arm, mips, x86, or amd64"
DEFINE_boolean cleanup_dirs ${FLAGS_TRUE} \
  "Whether the mount dirs should be removed on completion."

DEFINE_string boot_args "noinitrd" \
  "Additional boot arguments to pass to the commandline"

DEFINE_integer rootfs_size 720 \
  "rootfs filesystem size in MBs."
# ceil(0.1 * rootfs_size) is a good minimum.
DEFINE_integer rootfs_hash_pad 8 \
  "MBs reserved at the end of the rootfs image."

DEFINE_string rootfs_hash "/tmp/rootfs.hash" \
  "Path where the rootfs hash should be stored."
# TODO(taysom): when we turn on boot cache, both verification and
# bootcache should have their default be FLAGS_TRUE.
DEFINE_boolean enable_rootfs_verification ${FLAGS_FALSE} \
  "Default all bootloaders to NOT use kernel-based root fs integrity checking."
DEFINE_boolean enable_bootcache ${FLAGS_FALSE} \
  "Default all bootloaders to NOT use bootcache."
DEFINE_integer verity_error_behavior 3 \
  "Kernel verified boot error behavior (0: I/O errors, 1: reboot, 2: nothing)"
DEFINE_integer verity_max_ios -1 \
  "Number of outstanding I/O operations dm-verity caps at."
DEFINE_string verity_algorithm "sha1" \
  "Cryptographic hash algorithm used for kernel vboot."
DEFINE_string verity_salt "" \
  "Salt for rootfs hash tree."

DEFINE_string keys_dir "/usr/share/vboot/devkeys" \
  "Directory containing the signing keys."

DEFINE_string rootfs_mountpoint "/tmp/rootfs" \
  "Path where the rootfs can be safely mounted"
DEFINE_string statefulfs_mountpoint "/tmp/statefulfs" \
  "Path where the statefulfs can be safely mounted"
DEFINE_string espfs_mountpoint "/tmp/espfs" \
  "Path where the espfs can be safely mounted"

DEFINE_boolean use_dev_keys ${FLAGS_FALSE} \
  "Use developer keys for signing. (Default: false)"

DEFINE_boolean fsck_rootfs ${FLAGS_FALSE} \
  "Check integrity of the rootfs on the modified image."

# TODO(pkumar): Remove once known that no images are using this flag
DEFINE_boolean crosbug12352_arm_kernel_signing ${FLAGS_FALSE} \
  "This flag is deprecated but the bots still need parse old images."

# TODO(sosa):  Remove once known images no longer use this in their config.
DEFINE_string arm_extra_bootargs "" "DEPRECATED FLAG.  Do not use."

DEFINE_boolean force_developer_mode ${FLAGS_FALSE} \
  "Add cros_debug to boot args."

DEFINE_boolean enable_squashfs ${FLAGS_FALSE} \
  "Make the rootfs of the image squashfs."
DEFINE_string squash_sort_file "" \
  "Specify the priority of files when squashing the rootfs."

DEFINE_string enable_serial "" \
  "Enable serial port for printks. Example values: ttyS0"

# Parse the boot.desc and any overrides
eval set -- "${BOOT_DESC} ${FLAG_OVERRIDES}"
FLAGS "${@}" || exit 1

[ -z "${FLAGS_verity_salt}" ] && FLAGS_verity_salt=$(make_salt)

. "${BUILD_LIBRARY_DIR}/board_options.sh" || exit 1

# Only now can we die on error.  shflags functions leak non-zero error codes,
# so will die prematurely if 'switch_to_strict_mode' is specified before now.
switch_to_strict_mode -u

# $1 - Directory where developer rootfs is mounted.
# $2 - Directory where developer stateful_partition is mounted.
# $3 - Directory where the ESP partition is mounted.
mount_gpt_cleanup() {
  local rootfs="${1-$FLAGS_rootfs_mountpoint}"
  local statefs="${2-$FLAGS_statefulfs_mountpoint}"
  local espfs="${3-$FLAGS_espfs_mountpoint}"
  "${SCRIPTS_DIR}/mount_gpt_image.sh" \
    -u -r "${rootfs}" -s "${statefs}" -e "${espfs}"
}

append_boot_flag() {
  local file="$1"
  local pattern="$2"
  local base_pattern="$3"

  [ -f "${file}" ] || return ${FLAGS_TRUE}
  grep -wq "${pattern}" "${file}" && return ${FLAGS_TRUE}
  sudo sed -i "s/\b${base_pattern}\b/& ${pattern}/g" "${file}"
}

check_kernel_size() {
  local kernel_image_size="$1"
  local kernel_part="$2"
  local kernel_slot="$3"
  local kernel_partition_size=$(get_partition_size ${FLAGS_image_type} \
                                                   ${kernel_part})
  local kernel_partition_size_90=$(( kernel_partition_size * 90 / 100 ))
  info "Kernel partition ${kernel_slot} size is ${kernel_partition_size} bytes."
  if [[ ${kernel_image_size} -gt ${kernel_partition_size} ]]; then
    die "Kernel image won't fit in partition ${kernel_slot}!"
  elif [[ ${kernel_image_size} -gt ${kernel_partition_size_90} ]]; then
    warn "Kernel partition ${kernel_slot} is more than 90% full!"
  fi
}

build_img() {
  local image_name="$1"
  local root_dev="$2"
  local keyblock="$3"
  local private="$4"
  local public="$5"
  local vblock=${6:-""}

  # Default to non-verified
  local enable_rootfs_verification_flag=--noenable_rootfs_verification
  if [[ ${FLAGS_enable_rootfs_verification} -eq ${FLAGS_TRUE} ]]; then
    enable_rootfs_verification_flag=--enable_rootfs_verification
  fi
  local enable_bootcache_flag=--noenable_bootcache
  if [[ ${FLAGS_enable_bootcache} -eq ${FLAGS_TRUE} ]]; then
    enable_bootcache_flag=--enable_bootcache
  fi

  if [[ -n "${vblock}" ]]; then
    vblock=--hd_vblock="${FLAGS_output_dir}/${vblock}"
  fi

  ${SCRIPTS_DIR}/build_kernel_image.sh \
    --board="${FLAGS_board}" \
    --arch="${FLAGS_arch}" \
    --to="${FLAGS_output_dir}/${image_name}" \
    --vmlinuz="${FLAGS_rootfs_mountpoint}/boot/vmlinuz" \
    --working_dir="${FLAGS_output_dir}" \
    --boot_args="${FLAGS_boot_args}" \
    --keep_work \
    --rootfs_image=${root_dev} \
    --rootfs_hash=${FLAGS_rootfs_hash} \
    --verity_hash_alg=${FLAGS_verity_algorithm} \
    --verity_max_ios=${FLAGS_verity_max_ios} \
    --verity_error_behavior=${FLAGS_verity_error_behavior} \
    --verity_salt=${FLAGS_verity_salt} \
    --keys_dir="${FLAGS_keys_dir}" \
    --keyblock="${keyblock}" \
    --private="${private}" \
    --public="${public}" \
    --enable_serial="${FLAGS_enable_serial}" \
    ${vblock} \
    ${enable_rootfs_verification_flag} \
    ${enable_bootcache_flag}
}

make_image_bootable() {
  local image="$1"

  trap "mount_gpt_cleanup" EXIT
  "${SCRIPTS_DIR}/mount_gpt_image.sh" --from "$(dirname "${image}")" \
    --image "$(basename ${image})" -r "${FLAGS_rootfs_mountpoint}" \
    -s "${FLAGS_statefulfs_mountpoint}"

  # Update legacy boot config templates (in rootfs) before rootfs is locked.
  # This is required because postinst will copy new legacy boot configurations
  # from rootfs partition instead of modifying existing entries in partition 12
  # (EFI partition).
  if [ ${FLAGS_force_developer_mode} -eq ${FLAGS_TRUE} ]; then
    append_boot_flag "${FLAGS_rootfs_mountpoint}/boot/syslinux/root.A.cfg" \
      "cros_debug" "cros_legacy"
    append_boot_flag "${FLAGS_rootfs_mountpoint}/boot/syslinux/root.B.cfg" \
      "cros_debug" "cros_legacy"
    append_boot_flag "${FLAGS_rootfs_mountpoint}/boot/syslinux/usb.A.cfg" \
      "cros_debug" "cros_legacy"
    append_boot_flag "${FLAGS_rootfs_mountpoint}/boot/efi/boot/grub.cfg" \
      "cros_debug" "cros_efi"
  fi

  # The rootfs should never be mounted rw again after this point without
  # re-calling make_image_bootable.
  sudo mount -o remount,ro "${FLAGS_rootfs_mountpoint}"
  # Newer `mount` will decode the filename backing the loop device,
  # so we need to dig deeper and find the answer ourselves.
  root_dev=$(awk -v mnt="${FLAGS_rootfs_mountpoint}" \
             '$2 == mnt { print $1 }' /proc/mounts)

  # Make the filesystem un-mountable as read-write.
  # mount_gpt_image.sh will undo this as needed.
  # TODO(wad) make sure there is parity in the signing scripts.
  if [ ${FLAGS_enable_rootfs_verification} -eq ${FLAGS_TRUE} ]; then
    # TODO(wad) this would be a good place to reset any other ext2 metadata.
    warn "Disabling r/w mount of the root filesystem"
    disable_rw_mount "$root_dev"
  fi

  # We sign the image with the recovery_key, because this is what goes onto the
  # USB key. We can only boot from the USB drive in recovery mode.
  # For dev install shim, we need to use the installer keyblock instead of
  # the recovery keyblock because of the difference in flags.
  local keyblock
  if [ ${FLAGS_use_dev_keys} -eq ${FLAGS_TRUE} ]; then
    keyblock=installer_kernel.keyblock
    info "DEBUG: use dev install keyblock"
  else
    keyblock=recovery_kernel.keyblock
    info "DEBUG: use recovery keyblock"
  fi

  if [ ${FLAGS_force_developer_mode} -eq ${FLAGS_TRUE} ]; then
    FLAGS_boot_args="${FLAGS_boot_args} cros_debug"
  fi

  local squash_sort_flag=
  if [ -n "${FLAGS_squash_sort_file}" ]; then
    squash_sort_flag="-sort ${FLAGS_squash_sort_file}"
  fi

  if [ $FLAGS_enable_squashfs -eq $FLAGS_TRUE ]; then
    local squashfs_img="${FLAGS_output_dir}/squashfs.image"
    sudo mksquashfs "${FLAGS_rootfs_mountpoint}" ${squashfs_img} -comp lzo \
      -noI -noF -ef ${SCRIPTS_DIR}/exclude-list -wildcards ${squash_sort_flag}
    root_dev=$squashfs_img
  fi
  # Builds the kernel partition image.
  build_img "vmlinuz.image" "${root_dev}" "${keyblock}" \
            "recovery_kernel_data_key.vbprivk" "recovery_key.vbpubk"
  build_img "hd_vmlinuz.image" "${root_dev}" "kernel.keyblock" \
            "kernel_data_key.vbprivk" "kernel_subkey.vbpubk" \
            "vmlinuz_hd.vblock"

  # Check the size of kernel image and issue warning when image size is
  # near the limit.
  local kernel_image_size_A=$(stat -c '%s' ${FLAGS_output_dir}/vmlinuz.image)
  info "Kernel image A   size is ${kernel_image_size_A} bytes."
  local kernel_image_size_B=$(stat -c '%s' ${FLAGS_output_dir}/hd_vmlinuz.image)
  info "Kernel image B   size is ${kernel_image_size_B} bytes."
  check_kernel_size ${kernel_image_size_A} 2 A
  check_kernel_size ${kernel_image_size_B} 4 B

  local rootfs_hash_size=$(stat -c '%s' ${FLAGS_rootfs_hash})
  local rootfs_fs_size=$(get_filesystem_size ${FLAGS_image_type} 3)
  local rootfs_partition_size=$(get_partition_size ${FLAGS_image_type} 3)
  local rootfs_hash_pad=$(( rootfs_partition_size - rootfs_fs_size ))
  info "Appending rootfs.hash (${rootfs_hash_size} bytes) to the root fs"
  if [[ ${rootfs_hash_size} -gt ${rootfs_hash_pad} ]]
  then
    die "rootfs_partition_size - rootfs_fs_size is less than the needed " \
        "rootfs_hash_size (${rootfs_hash_size}), update your disk layout " \
        "configuration"
  fi
  # Unfortunately, mount_gpt_image uses mount and not losetup to create the
  # loop devices.  This means that they are not the correct size.  We have to
  # write directly to the image to append the hash tree data.
  local hash_offset="$(partoffset ${image} 3)"
  if [ $FLAGS_enable_squashfs -eq $FLAGS_TRUE ]; then
    rootfs_file_size=$(stat -c '%s' ${root_dev})
    hash_offset=$((hash_offset + (${rootfs_file_size} / 512)))
  else
    hash_offset=$((hash_offset + (${rootfs_fs_size} / 512)))
  fi
  sudo dd bs=512 \
          seek=${hash_offset} \
          if="${FLAGS_rootfs_hash}" \
          of="${image}" \
          conv=notrunc \
          status=none

  # Move the verification block needed for the hard disk install to the
  # stateful partition. Mount stateful fs, copy file, and umount fs.
  # In original CL: http://codereview.chromium.org/2868044, this was done in
  # create_base_image(). However, it could break the build if it is a clean
  # build because vmlinuz_hd.vblock hasn't been created by build_kernel_image.sh
  # In some builds that don't use vboot to verify the kernel, this file might
  # not get created as part of the build, so only copy them if they were.
  if [ -f "${FLAGS_output_dir}/vmlinuz_hd.vblock" ]; then
    sudo cp "${FLAGS_output_dir}/vmlinuz_hd.vblock" \
            "${FLAGS_statefulfs_mountpoint}"
  fi

  # Install the kernel to both slots A and B so there will always be a regular
  # kernel in slot B on recovery and non-recovery images.
  local koffset="$(partoffset ${image} 2)"
  sudo dd if="${FLAGS_output_dir}/vmlinuz.image" of="${image}" \
    conv=notrunc bs=512 seek=${koffset} status=none
  koffset="$(partoffset ${image} 4)"
  sudo dd if="${FLAGS_output_dir}/hd_vmlinuz.image" of="${image}" \
    conv=notrunc bs=512 seek=${koffset} status=none

  # Update the bootloaders.  The EFI system partition will be updated.
  local kernel_part=

  # We should update the esp in place in the image.
  local bootloader_to="${image}"
  local esp_offset="$(partoffset ${image} 12)"
  esp_offset=$((esp_offset * 512))  # sectors to bytes
  local esp_size="$(partsize ${image} 12)"
  esp_size=$((esp_size * 512))  # sectors to bytes
  local bootloader_to_flags="--to_offset=${esp_offset} --to_size=${esp_size}"

  if [[ "${FLAGS_arch}" = "x86" || "${FLAGS_arch}" = "amd64" ]]; then
    # Use the kernel partition to acquire configuration flags.
    kernel_part="--kernel_partition='${FLAGS_output_dir}/vmlinuz.image'"
    # Install syslinux on the EFI System Partition.
    kernel_part="${kernel_part} --install_syslinux"
  elif [[ "${FLAGS_arch}" = "arm" || "${FLAGS_arch}" = "mips" ]]; then
    # These flags are not used for ARM / MIPS update_bootloaders.sh
    kernel_part=""
  fi

  if [[ ${esp_size} -gt 0 ]]; then
    # Update partition 12
    ${SCRIPTS_DIR}/update_bootloaders.sh \
      --arch="${FLAGS_arch}" \
      --to="${bootloader_to}" \
      --from="${FLAGS_rootfs_mountpoint}"/boot \
      --vmlinuz="${FLAGS_rootfs_mountpoint}"/boot/vmlinuz \
      ${bootloader_to_flags} \
      $kernel_part
  fi

  # We don't need to keep these files around anymore.
  sudo rm -f "${FLAGS_rootfs_hash}" "${FLAGS_output_dir}/vmlinuz.image" \
             "${FLAGS_output_dir}/hd_vmlinuz.image" \
             "${FLAGS_output_dir}/vmlinuz_hd.vblock"

  trap - EXIT
  ${SCRIPTS_DIR}/mount_gpt_image.sh -u -r "${FLAGS_rootfs_mountpoint}" \
    -s "${FLAGS_statefulfs_mountpoint}"

  # I can only copy the squashfs image to the image only when it is umounted.
  if [ $FLAGS_enable_squashfs -eq $FLAGS_TRUE ]; then
    # copy the squashfs image to the partition
    info "copy the squashfs to the partition"
    local part_offset="$(partoffset ${image} 3)"
    sudo dd bs=512 if="${squashfs_img}" of="${image}" \
            seek=${part_offset} conv=notrunc status=none
    sudo rm "${squashfs_img}"
  fi
}

verify_image_rootfs() {
  local image=$1
  local rootfs_offset="$(partoffset ${image} 3)"

  local rootfs_tmp_file=$(mktemp)
  trap "rm ${rootfs_tmp_file}" EXIT
  sudo dd if="${image}" of="${rootfs_tmp_file}" bs=512 skip="${rootfs_offset}" \
          status=none

  # This flips the read-only compatibility flag, so that
  # e2fsck does not complain about unknown file system capabilities.
  enable_rw_mount "${rootfs_tmp_file}"
  info "Running e2fsck to check root file system for errors"
  sudo e2fsck -fn "${rootfs_tmp_file}" ||
    die "Root file system has errors, please ensure boot.desc and/or \
command line parameters are correct"
}

# Store output and temporary files next to image.
FLAGS_output_dir="${IMAGE_DIR}"
FLAGS_rootfs_hash="${IMAGE_DIR}/rootfs.hash"
FLAGS_rootfs_mountpoint="${IMAGE_DIR}/rootfs_dir"
FLAGS_statefulfs_mountpoint="${IMAGE_DIR}/stateful_dir"
FLAGS_espfs_mountpoint="${IMAGE_DIR}/esp"

# Create the directories if they don't exist.
mkdir -p ${FLAGS_rootfs_mountpoint}
mkdir -p ${FLAGS_statefulfs_mountpoint}
mkdir -p ${FLAGS_espfs_mountpoint}

make_image_bootable "${IMAGE}"
# We can't verify the image if squashfs is enabled because the kernel
# on the host does not support squashfs with LZO
if [ ${FLAGS_fsck_rootfs} -eq ${FLAGS_TRUE} \
    -a ${FLAGS_enable_squashfs} -eq ${FLAGS_FALSE} ]; then
  verify_image_rootfs "${IMAGE}"
fi

if [ ${FLAGS_cleanup_dirs} -eq ${FLAGS_TRUE} ]; then
  rmdir ${FLAGS_rootfs_mountpoint}
  rmdir ${FLAGS_statefulfs_mountpoint}
  rmdir ${FLAGS_espfs_mountpoint}
fi