#!/bin/bash
# Copyright 2017 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

### make a zedboot USB key

## usage: fx mkzedboot [options] <usb device>
##  -f            force writing to a non-usb target
##  -i|--install  include "offline" install

set -e

source "$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"/lib/image_build_vars.sh || exit $?

if [[ "${FUCHSIA_ARCH}" != "x64" ]]; then
  echo >&2 mkzedboot is not supported for ${FUCHSIA_ARCH}
  exit 1
fi

force=false
if [[ "$1" == "-f" ]]; then
  shift
  force=true
fi

include_install=false
if [[ "$1" == "-i" ]] || [[ "$1" == "--install" ]]; then
  shift
  include_install=true
fi

is_usb() {
  if ! ${force}; then
    fx-command-run list-usb-disks | grep "$1"
  fi
}

USB_DEVICE="$1"
if [[ -z "${USB_DEVICE}" ]]; then
  echo >&2 "device argument required"
  echo "USB disks:"
  fx-command-run list-usb-disks
  exit 1
fi
if ! is_usb "${USB_DEVICE}"; then
  echo >&2 "${USB_DEVICE} does not look like a USB device, use -f to force, or pick from below"
  echo "USB disks:"
  fx-command-run list-usb-disks
  exit 1
fi

if ! [[ -w "${USB_DEVICE}" ]] ; then
  echo >&2 "Changing ownership of ${USB_DEVICE} to ${USER}"
  sudo chown "${USER}" "${USB_DEVICE}"
fi

echo >&2 "Opening device..."
# We open the device and hold onto an fd for the duration of our modifications.
# This prevents automounting solutions from observing a final close and
# rescanning the partition table until we're all done making changes -
# particularly important on macOS where users would otherwise receive
# EAGAIN/EBUSY and so on.
open_device() {
  case "$(uname)" in
  Darwin)
    if ! diskutil quiet unmountDisk "${USB_DEVICE}"; then
      echo >&2 "Failed to unmount ${USB_DEVICE}, cannot continue"
      exit 1
    fi
    ;;
  esac
  exec 3>>"${USB_DEVICE}"
}
close_device() {
  echo >&2 "Closing device."
  exec 3>&-
}
open_device

# Destroy any existing GPT/MBR on the device and re-create
echo "Create new GPT partition table... "
"${FUCHSIA_BUILD_DIR}/tools/cgpt" create "${USB_DEVICE}"
"${FUCHSIA_BUILD_DIR}/tools/cgpt" boot -p "${USB_DEVICE}"
echo "done"

echo "Create new partitions... "
# ESP needs to be a FAT compatible size
if [[ ! -z "${IMAGE_ZEDBOOT_ESP}" ]]; then
  esp_size=$(((63*1024*1024)/512))
  esp_offset=2048
  "${FUCHSIA_BUILD_DIR}/tools/cgpt" add -s "${esp_size}" -t efi -b "${esp_offset}" -l esp "${USB_DEVICE}"
else
  reserved_size=2048
  reserved_offset=2048
  vboot_size=$(((64*1024*1024)/512))
  vboot_offset=$(($reserved_offset + $reserved_size))
  "${FUCHSIA_BUILD_DIR}/tools/cgpt" add -s "${reserved_size}" -t reserved -b "${reserved_offset}" -l reserved "${USB_DEVICE}"
  "${FUCHSIA_BUILD_DIR}/tools/cgpt" add -s "${vboot_size}" -t kernel -b "${vboot_offset}" -l zedboot "${USB_DEVICE}"
fi

# NOTE: Ok, so here goes some stuff. I could have written a much smarter "dd"
# (a thing that can operate on block-boundaries for seek and copy, but that
# doesn't do operations ONE BLOCK AT A TIME because it's 2018 yo), or I could
# do what follows. Before this change, adding the install image to a disk via
# DD would take 20 minutes. That's just absurd.
# The stuff:
# Align the install_offset to a 4mb boundary.
# Pad the partition size to a 4mb boundary.
# Set the dd block size to 4mb, even though it really isn't 4mb.
# Seek offset*lba/4mb
# Write with osync

if $include_install; then
  if [[ ! -f "${FUCHSIA_BUILD_DIR}/${IMAGE_INSTALLER_RAW}" ]]; then
    echo >&2 "Install image not found at ${FUCHSIA_BUILD_DIR}/${IMAGE_INSTALLER_RAW} did you build it?"
    exit 1
  fi

  install_image_size=$(wc -c "${FUCHSIA_BUILD_DIR}/${IMAGE_INSTALLER_RAW}" | awk '{print $1}')
  # Add some slack, like the build does, as the file size doesn't represent the
  # volume size and there's no host tool that presently will print the
  # superblock volume size data.
  install_image_size=$((($install_image_size * 14) / 10))
  # It begins. Pad the image size to a 4mb boundary above it's size:
  install_size=$((($install_image_size + 4194303) / 4194304))
  # We need to specify the install size in 512byte lba's:
  install_size=$(($install_size * 8192))

  if [[ ! -z "${IMAGE_ZEDBOOT_ESP}" ]]; then
    install_min_offset=$(($esp_size + $esp_offset))
  else
    install_min_offset=$(($vboot_size + $vboot_offset))
  fi
  # Align the partition offset to a 4mb "block size"
  install_offset=$(( (($install_min_offset * 512) + 4194303) / 4194304))
  # The lba offset of that is:
  install_lba_offset=$(($install_offset * 8192))
  "${FUCHSIA_BUILD_DIR}/tools/cgpt" add -s "${install_size}" -t "48435546-4953-2041-494E-5354414C4C52" -b "${install_lba_offset}" -l install "${USB_DEVICE}"
fi
if [[ ! -z "${IMAGE_ZIRCONR_SIGNEDZBI}" ]]; then
  "${FUCHSIA_BUILD_DIR}/tools/cgpt" add -i 2 -T 1 -S 1 -P 2 "${USB_DEVICE}"
fi
echo "done"

if [[ ! -z "${IMAGE_ZEDBOOT_ESP}" ]]; then
  echo "Writing zedboot for EFI"
  dd if="${FUCHSIA_BUILD_DIR}/${IMAGE_ZEDBOOT_ESP}" of="${USB_DEVICE}" seek=${esp_offset}
else
  echo "Writing zedboot for Cros"
  dd if="${FUCHSIA_BUILD_DIR}/${IMAGE_ZIRCONR_SIGNEDZBI}" of="${USB_DEVICE}" seek=${vboot_offset}
fi
if $include_install; then
  echo "Writing install partition"
  dd if="${FUCHSIA_BUILD_DIR}/${IMAGE_INSTALLER_RAW}" of="${USB_DEVICE}" seek=${install_offset} bs=4194304
fi
echo "done"

close_device

case "$(uname)" in
  Linux)
    eject "${USB_DEVICE}"
    ;;
  Darwin)
    diskutil eject "${USB_DEVICE}"
    ;;
esac