#!/bin/sh
#---------------------------------------------------------------------
# Summary: Grub bootloader logic for Ubuntu Snappy systems.
# Description: This is a heavily modified "10_linux" grub snippet that
#   deals with Snappy dual-rootfs systems.
#
# XXX: Note that this script is called from within a chroot environment
# on snappy systems!
#
#---------------------------------------------------------------------

set -e

prefix="/usr"
exec_prefix="/usr"
datarootdir="/usr/share"

# Utility functions
. "${datarootdir}/grub/grub-mkconfig_lib"

# Globals
machine=`uname -m`

SNAPPY_OS="Ubuntu Core Snappy"
SNAPPY_TYPE=simple
SNAPPY_ARGS="${GRUB_CMDLINE_LINUX} ${GRUB_CMDLINE_LINUX_DEFAULT}"

#---------------------------------------------------------------------

# Display message and exit
die()
{
    msg="$*"
    echo "ERROR: $msg" >&2
    exit 1
}

# Create a grub menu entry by writing a menuentry to stdout.
linux_entry_ext()
{
  local name="$1"
  local os="$2"
  local version="$3"
  local type="$4"
  local args="$5"
  local device="$6"
  local kernel="$7"
  local initrd="$8"

  local boot_device_id=
  local prepare_root_cache=
  local prepare_boot_cache=

  if [ -z "$boot_device_id" ]; then
      boot_device_id="$(grub_get_device_id "${device}")"
  fi

  echo "menuentry '$name' ${CLASS} \$menuentry_id_option 'gnulinux-simple-$boot_device_id' {" | sed "s/^/$submenu_indentation/"

  if [ x$type != xrecovery ] ; then
      save_default_entry | grub_add_tab
  fi

  # Use ELILO's generic "efifb" when it's known to be available.
  # FIXME: We need an interface to select vesafb in case efifb can't be used.
  if [ "x$GRUB_GFXPAYLOAD_LINUX" = x ]; then
      echo "	load_video" | sed "s/^/$submenu_indentation/"
  else
      if [ "x$GRUB_GFXPAYLOAD_LINUX" != xtext ]; then
	  echo "	load_video" | sed "s/^/$submenu_indentation/"
      fi
  fi
  if ([ "$ubuntu_recovery" = 0 ] || [ x$type != xrecovery ]) && \
     ([ "x$GRUB_GFXPAYLOAD_LINUX" != x ] || [ "$gfxpayload_dynamic" = 1 ]); then
      echo "	gfxmode \$linux_gfx_mode" | sed "s/^/$submenu_indentation/"
  fi

  echo "	insmod gzio" | sed "s/^/$submenu_indentation/"
  echo "	if [ x\$grub_platform = xxen ]; then insmod xzio; insmod lzopio; fi" | sed "s/^/$submenu_indentation/"

  # device may be a label (LABEL=name), so convert back to full path
  label_name=$(echo "$device"|sed 's/^LABEL=//g')
  if [ "$device" = "$label_name" ]
  then
    device_path="$device"
  else
    # found a label
    device_path=$(get_partition_from_label "$label_name")
  fi

  if [ x$dirname = x/ ]; then
    if [ -z "${prepare_root_cache}" ]; then

      prepare_root_cache="$(prepare_grub_to_access_device ${device_path} | grub_add_tab)"
    fi
    printf '%s\n' "${prepare_root_cache}" | sed "s/^/$submenu_indentation/"
  else
    if [ -z "${prepare_boot_cache}" ]; then
      prepare_boot_cache="$(prepare_grub_to_access_device ${device_path} | grub_add_tab)"
    fi
    printf '%s\n' "${prepare_boot_cache}" | sed "s/^/$submenu_indentation/"
  fi

  if [ x"$quiet_boot" = x0 ] || [ x"$type" != xsimple ]; then
    message="$(gettext_printf "Loading Linux %s ..." ${version})"
    sed "s/^/$submenu_indentation/" << EOF
	echo	'$(echo "$message" | grub_quote)'
EOF
  fi

    sed "s/^/$submenu_indentation/" << EOF
	linux ${kernel} root=${device} ro ${args}
EOF

  if test -n "${initrd}"; then
    # TRANSLATORS: ramdisk isn't identifier. Should be translated.
    if [ x"$quiet_boot" = x0 ] || [ x"$type" != xsimple ]; then
      message="$(gettext_printf "Loading initial ramdisk ...")"
      sed "s/^/$submenu_indentation/" << EOF
	echo	'$(echo "$message" | grub_quote)'
EOF
    fi
    sed "s/^/$submenu_indentation/" << EOF
	initrd ${initrd}
EOF
  fi

  sed "s/^/$submenu_indentation/" << EOF
}
EOF
}

# Returns a list of the currently available kernels.
# $1: If set, look for kernel below "$1/boot/".
get_kernels()
{
    local prefix_dir="$1"
    local list

    case "x$machine" in
    xi?86 | xx86_64)
    list=`for i in $prefix_dir/boot/vmlinuz-* \
                         $prefix_dir/vmlinuz-* \
                         $prefix_dir/boot/kernel-* ; do
                  if grub_file_is_not_garbage "$i" ; then echo -n "$i " ; fi
              done` ;;
    *) 
    list=`for i in $prefix_dir/boot/vmlinuz-* \
                         $prefix_dir/boot/vmlinux-* \
                         $prefix_dir/vmlinuz-* \
                         $prefix_dir/vmlinux-* \
                         $prefix_dir/boot/kernel-* ; do
                  if grub_file_is_not_garbage "$i" ; then echo -n "$i " ; fi
         done` ;;
    esac
    echo "$list"
}

# Returns the path to the initrd for the kernel specified by $1.
# $1: kernel version.
# $2: directory to look in.
get_initrd()
{
  local version="$1"
  local dir="$2"

  local alt_version=`echo $version | sed -e "s,\.old$,,g"`
  local initrd=
  local i=

  for i in "initrd.img-${version}" "initrd-${version}.img" "initrd-${version}.gz" \
           "initrd-${version}" "initramfs-${version}.img" \
           "initrd.img-${alt_version}" "initrd-${alt_version}.img" \
           "initrd-${alt_version}" "initramfs-${alt_version}.img" \
           "initramfs-genkernel-${version}" \
           "initramfs-genkernel-${alt_version}" \
           "initramfs-genkernel-${GENKERNEL_ARCH}-${version}" \
           "initramfs-genkernel-${GENKERNEL_ARCH}-${alt_version}"; do
    if test -e "${dir}/${i}" ; then
      initrd="${dir}/${i}"
      break
    fi
  done
  echo "$initrd"
}

# Determine full path to disk partition given a filesystem label.
get_partition_from_label()
{
    local label="$1"
    local part=
    local path=

    [ -n "$label" ] || grub_warn "need FS label"

    part=$(find /dev -name "$label"|tail -1)
    [ -z "$part" ] && return
    path=$(readlink -f "$part")
    [ -n "$path" ] && echo "$path"
}

# Return the partition label for the given partition device.
# $1: full path to partition device.
get_label_from_device()
{
    local root="$1"

    local label=
    local std_label=
    local label_rootfs=

    for std_label in system-a system-b; do
        label_rootfs=$(findfs "PARTLABEL=$std_label" 2>/dev/null || :)
        if [ "$label_rootfs" = "$root" ]; then
            label="$std_label"
            break
        fi
    done

    echo "$label"
}

# Return the full path to the device corresponding to the given
# partition label.
#
# $1: partition label.
get_device_from_label()
{
    local label="$1"
    local device=

    device=$(findfs "PARTLABEL=$label" 2>/dev/null || :)
    echo "$device"
}

# Given a rootfs label, return the rootfs label corresponding to the
# "other" rootfs partition.
get_other_rootfs_label()
{
    local label="$1"

    if [ "$label" = "system-a" ]; then
        echo "system-b"
    else
        echo "system-a"
    fi
}

# Given a mountpoint, return the corresponding device path
# $1: mountpoint.
get_device_from_mountpoint()
{
    local mountpoint="$1"
    local device=

    # XXX: Parse mount output rather than looking in /proc/mounts to
    # avoid seeing the mounts outside the chroot.
    device=$(/bin/mount | grep "^/dev/.* on ${mountpoint}[[:space:]]" 2>/dev/null |\
        awk '{print $1}' || :)

    echo "$device"
}

# Convert a partition label name to a menuentry name
make_name()
{
    local label="$1"

    echo "$SNAPPY_OS $label rootfs" | grub_quote
}

# Arrange for a grub menuentry to be created for the given partition.
#
# $1: full path to rootfs partition device
# $2: partition label associated with $1
# $3: mountpoint of $1.
handle_menu_entry()
{
    local rootfs_device="$1"
    local label="$2"
    local mountpoint="$3"

    local name=
    local device=
    local mount_prefix=
    local list=
    local linux=
    local basename=
    local dirname=
    local rel_dirname=
    local version=
    local initrd=

    # boot by label
    device="LABEL=$label"

    name=$(make_name "$label")

    # avoid double-leading slashes and the subsequent need to call
    # 'readlink -f'.
    if [ "$mountpoint" = "/" ]; then
        mount_prefix=""
    else
        mount_prefix="$mountpoint"
    fi
    list=$(get_kernels "$mount_prefix")

    linux=$(version_find_latest $list)
    basename=$(basename "$linux")
    dirname=$(dirname "$linux")
    rel_dirname=$(make_system_path_relative_to_its_root "$dirname")
    version=$(echo "$basename" | sed -e "s,^[^0-9]*-,,g")

    initrd=$(get_initrd "$version" "$dirname")

    # convert the path to the mounted "other" rootfs back to a
    # a standard one by removing the mountpoint prefix.
    if [ "$mountpoint" != "/" ]; then
        linux=$(echo "$linux" | sed "s!^${mountpoint}!!g")
        initrd=$(echo "$initrd" | sed "s!^${mountpoint}!!g")
    fi

    # Create menu entries for the 2 snappy rootfs's.
    linux_entry_ext "$name" "$SNAPPY_OS" "$version" \
        "$SNAPPY_TYPE" "$SNAPPY_ARGS" "$device" \
        "$linux" "$initrd"
}

#---------------------------------------------------------------------
# main

case "$machine" in
    i?86) GENKERNEL_ARCH="x86" ;;
    mips|mips64) GENKERNEL_ARCH="mips" ;;
    mipsel|mips64el) GENKERNEL_ARCH="mipsel" ;;
    arm*) GENKERNEL_ARCH="arm" ;;
    *) GENKERNEL_ARCH="$machine" ;;
esac

# Determine which partition label is being used for the current root
# directory. This is slightly convoluted but required since this code
# runs within a chroot environment (where lsblk does not work).
#
# XXX: Note that since this code is run from within a chroot (where the
# "other" rootfs is mounted), it might appear that the logic is
# inverted. However, the code below simply 
this_mountpoint="/"
this_root=$(get_device_from_mountpoint "$this_mountpoint")
[ -z "$this_root" ] && {
    die "cannot determine rootfs for $this_mountpoint"
}

this_label=$(get_label_from_device "$this_root")
[ -z "$this_label" ] && {
    die "cannot determine partition label for rootfs $this_root"
}

handle_menu_entry "$this_root" "$this_label" "$this_mountpoint"

other_mountpoint="/writable/cache/system"

# When this script is run on a real snappy system, even if there is only
# a single rootfs provisioned, the other rootfs partition is expected to
# be formatted and mounted.
#
# However this script is also run at provisioning time where
# $other_mountpoint will not be a mountpoint. Therefore in the provisioning
# scenario, only a single menuentry will be generated if only a single
# rootfs is provisioned.
if $(mountpoint -q "$other_mountpoint"); then
    other_label=$(get_other_rootfs_label "$this_label")

    other_root=$(get_device_from_label "$other_label")
    [ -z "$other_root" ] && {
        die "cannot determine rootfs"
    }

    handle_menu_entry "$other_root" "$other_label" "$other_mountpoint"
fi

# Toggle rootfs if previous boot failed.
#
# Since grub sets snappy_trial_boot, if it is _already_ set when grub starts
# and we're in try mode, the previous boot must have failed to unset it,
# so toggle the rootfs.
sed "s/^/$submenu_indentation/" << EOF
    # set defaults
    if [ -z "\$snappy_mode" ]; then
        set snappy_mode=regular
        save_env snappy_mode
    fi
    if [ -z "\$snappy_ab" ]; then
        set snappy_ab=a
        save_env snappy_ab
    fi

    if [ "\$snappy_mode" = "try" ]; then
        if [ "\$snappy_trial_boot" = "1" ]; then
            # Previous boot failed to unset snappy_trial_boot, so toggle
            # rootfs.
            if [ "\$snappy_ab" = "a" ]; then
                set default="$(make_name system-b)"
                set snappy_ab=b
            else
                set snappy_ab=a
                set default="$(make_name system-a)"
            fi
            save_env snappy_ab
        else
            # Trial mode so set the snappy_trial_boot (which snappy is
            # expected to unset).
            #
            # Note: don't use the standard recordfail variable since that forces
            # the menu to be displayed and sets an infinite timeout if set.
            set snappy_trial_boot=1
            save_env snappy_trial_boot

            if [ "\$snappy_ab" = "a" ]; then
                set default="$(make_name system-a)"
            else
                set default="$(make_name system-b)"
            fi
        fi
    else
        if [ "\$snappy_ab" = "a" ]; then
            set default="$(make_name system-a)"
        else
            set default="$(make_name system-b)"
        fi
    fi
EOF
