#!/bin/bash
# Copyright 2018 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.
#
# A systemd unit generator that outputs .mount units for
# /mnt/stateful_partition and /usr/share/oem.
#
# This generator admits the following kernel command line arguments:
#
# cos.stateful_dev=[path or path prefix]:[mkfs_stateful]
#   This command line argument indicates the block device that should be used
#   for the stateful partition and other options for configuring the stateful
#   partition.
#
#   [path or path prefix]:
#     If empty, COS will use the default stateful partition as the stateful
#     device. Otherwise, the device specified by this value
#     (e.g. /dev/disk/by-id/google-persistent-disk-0) will be used as for the
#     stateful partition. If the last character of this input is a *, this value
#     will be interpreted as a path prefix. If this value is a path prefix, COS
#     will choose the stateful device to be the largest file path according to
#     `sort -V` with the provided prefix. Note that if the provided prefix
#     identifies a unique file path, that path will be chosen as the stateful
#     device.
#
#   [mkfs_stateful]:
#     Can be 0 or 1. Indicates if the stateful device should be reformatted at
#     every boot or not. 1 means that the stateful device will be reformatted at
#     every boot, and 0 means otherwise. If empty, the stateful device will not
#     be reformatted at every boot.
#
# cos.logs_dev=[path]
#   This command line argument indicates the block device that should be used
#   for /var/log.
#
#   [path]:
#     If empty, COS will use the default location on the stateful partition to
#     store logs. Otherwise, the device specified by this value (e.g.
#     /dev/disk/by-id/google-persistent-disk-0) will be formatted if need be and
#     mounted at /var/log.
#
# cos.oem_dev=[path]
#   This command line arguments indicates the block device that should be used
#   for /usr/share/oem.
#
#  [path]:
#    If empty, COS will use the default location as the 8th partition on the
#    disk.
#
# cos.protected_stateful_partition=[protection_level]
#   This command line argument indicates that the stateful partition block
#   device should include additional protection.
#   NOTE: using this flag means that COS cannot resize the stateful partition
#   to use the rest of the disk. COS will instead be limited to the existing
#   partition size on the boot disk.
#
#   [protection_level]:
#     If empty or "n", COS will default to not additionally protecting the
#     stateful partition. If "e", the stateful partition will be
#     confidentiality and integrity protected with an ephemeral key. With this
#     option the stateful partition will not be recoverable on reboot.
#
# When editing this file, keep in mind that we can't use heredocs here because
# using heredocs creates a tempfile. A tmpfs isn't mounted at /tmp this early in
# the boot process, causing /tmp to be read-only.

# Load the image settings.
# shellcheck disable=SC1091
. "/usr/sbin/write_gpt.sh"
load_base_vars

# Flag variables; populated by parse_kernel_args
STATEFUL_DEV=""
MKFS_STATEFUL=""
LOGS_DEV=""
OEM_DEV=""
STATEFUL_PROTECTION_LEVEL="n"

# The stateful partition block device, to be encrypted with integrity.
# This flag is only used when cos.stateful_dev is unused and
# cos.protected_stateful_partition=e.
PROTECTED_BLOCK_DEVICE=""

# Parses arguments from the kernel command line and stores values in flag
# variables.
#
# Args:
# rootdev: The root device of the boot disk (e.g. /dev/sda)
parse_kernel_args() {
  local -r rootdev="$1"
  # shellcheck disable=SC2046
  set -- $(cat /proc/cmdline)
  for arg in "$@"; do
    case "${arg}" in
      cos.stateful_dev=*)
        STATEFUL_DEV="$(echo "${arg#cos.stateful_dev=}" | cut -d : -f 1)"
        MKFS_STATEFUL="$(echo "${arg#cos.stateful_dev=}" | cut -d : -f 2 -s)"
        ;;
      cos.logs_dev=*)
        LOGS_DEV="${arg#cos.logs_dev=}"
        ;;
      cos.oem_dev=*)
        OEM_DEV="${arg#cos.oem_dev=}"
        ;;
      cos.protected_stateful_partition=*)
        STATEFUL_PROTECTION_LEVEL="${arg#cos.protected_stateful_partition=}"
        ;;
      *)
        ;;
    esac
  done
  if [[ -z "${MKFS_STATEFUL}" ]]; then
    MKFS_STATEFUL=0
  fi
  if [[ -z "${STATEFUL_DEV}" ]]; then
    if [[ "${STATEFUL_PROTECTION_LEVEL}" == "e" ]]; then
      STATEFUL_DEV="/dev/mapper/protected_stateful_partition"
      PROTECTED_BLOCK_DEVICE="${rootdev}${PARTITION_NUM_STATE}"
    else
      STATEFUL_DEV="${rootdev}${PARTITION_NUM_STATE}"
    fi
  else
    STATEFUL_DEV="$(readlink -m "${STATEFUL_DEV}")"
  fi
}

# Generates the mount unit for the OEM partition.
# The mount unit will not be generated if there is
# more than one device mapper devs
#
# Args:
# dev: The device to use for the OEM partition.
# fstype: The file system type of the OEM partition.
# output_dir: The directory to output the unit to.
# enable_dir[Optional]: The directory containing the symlink
# to the output unit for enabling the output unit. It should
# be an absolute path.
gen_oem_partition_mount() {
  local dev="$1"
  local -r fstype="$2"
  local -r output_dir="$3"
  local -r enable_dir="$4"
  local -r disk="$(basename ${dev})"
  local exec_opt="noexec"
  if [[ -n "${OEM_DEV}" ]]; then
    dev="${OEM_DEV}"
  else
    # After sealing the OEM partition, it will be available through
    # /dev/dm-[0-9]* instead of /dev/sda8
    dev_holder_count=$(ls /sys/class/block/${disk}/holders | wc -l)
    if [[ ${dev_holder_count} -gt 1 ]]; then
      return
    fi

    local dev_holder="$(basename /sys/class/block/${disk}/holders/*)"
    if [[ ${dev_holder} != "*" ]]; then
      dev="/dev/${dev_holder}"
      exec_opt="exec"
    fi
  fi
  echo "
[Unit]
Before=local-fs.target

[Mount]
What=${dev}
Where=/usr/share/oem
Type=${fstype}
Options=ro,nodev,${exec_opt},nosuid
" > "${output_dir}/usr-share-oem.mount"

  if [[ -n "${enable_dir}" ]]; then
    # Creating symlink as required by systemd-239
    ln -s "${output_dir}/usr-share-oem.mount" \
       "${enable_dir}/usr-share-oem.mount"
  fi
}

# Generates the path unit for the input stateful device(s).
# Activates dev-stateful.service, which creates a symlink to the stateful
# device.
#
# Args:
# dev: Device path or path prefix for the stateful device
# output_dir: The directory to output the unit to.
# enable_dir[Optional]: The directory containing the symlink
# to the output unit for enabling the output unit. It should
# be an absolute path.
gen_stateful_devices_path() {
  local -r dev="$1"
  local -r output_dir="$2"
  local -r enable_dir="$3"

  echo "
[Unit]
DefaultDependencies=false
After=local-fs-pre.target

[Path]
PathExistsGlob=${dev}
" > "${output_dir}/stateful-devices.path"

  if [[ -n "${enable_dir}" ]]; then
    # Creating symlink as required by systemd-239
    ln -s "${output_dir}/stateful-devices.path" \
       "${enable_dir}/stateful-devices.path"
  fi
}

# Generates the service that creates the /dev/stateful symlink.
#
# Args:
# dev: Device path or path prefix for the stateful device
# output_dir: The directory to output the unit to.
gen_dev_stateful_service() {
  local -r dev="$1"
  local -r output_dir="$2"
  echo "
[Unit]
Description=Creates the /dev/stateful symlink to the stateful device.
DefaultDependencies=false
After=systemd-udevd.service mount-etc-overlay.service

[Service]
Type=oneshot
RemainAfterExit=true
ExecStart=/usr/share/cloud/stateful-dev-sym-sorted 'stateful' '${dev}'
" > "${output_dir}/stateful-devices.service"
}

# Generates systemd units for creating the /dev/stateful symlink. This symlink
# points to the desired stateful device to mount as the stateful partition.
#
# Args:
# dev: The device or device prefix to use as the stateful device.
# output_dir: The directory to output units to.
configure_dev_stateful() {
  local -r dev="$1"
  local -r output_dir="$2"
  gen_stateful_devices_path "${dev}" "${output_dir}" \
      "${output_dir}/local-fs.target.wants"
  gen_dev_stateful_service "${dev}" "${output_dir}"
}

# Generates a systemd unit for remaking the file system on the stateful device.
#
# Args:
# fstype: The file system type of the stateful partition.
# output_dir: The directory to output the unit to.
gen_stateful_partition_mkfs() {
  local -r fstype="$1"
  local -r output_dir="$2"
  echo "
[Unit]
Description=Remakes the file system on the stateful device.
DefaultDependencies=false
BindsTo=dev-stateful.device
After=dev-stateful.device
Before=local-fs.target

[Service]
Type=oneshot
RemainAfterExit=true
ExecStart=/sbin/mkfs -F -t ${fstype} -E lazy_journal_init /dev/stateful
" > "${output_dir}/stateful-partition-mkfs.service"
}

# Generates stateful partition mount options.
stateful_mount_options() {
  local -r dirty_expire_centisecs="$(sysctl -n vm.dirty_expire_centisecs)"
  local -r commit_interval="$(( dirty_expire_centisecs / 100 ))"
  echo "nodev,noexec,nosuid,commit=${commit_interval}"
}

# Generates a systemd mount unit for mounting the stateful partition.
#
# Args:
# fstype: The file system type of the stateful partition.
# mkfs_stateful: Indicates if the file system on the stateful partition should
#   be remade before mounting. Can be 0 or 1; 0 is false, 1 is true.
# output_dir: The directory to output the unit to.
# protection_level: The protection level of the stateful partition. If e,
#   requires the stateful block device be modified before mounting the
#   stateful partition.
# enable_dir[Optional]: The directory containing the symlink
# to the output unit for enabling the output unit. It should
# be an absolute path.
gen_stateful_partition_mount() {
  local -r fstype="$1"
  local -r mkfs_stateful="$2"
  local -r output_dir="$3"
  local -r protection_level="$4"
  local -r enable_dir="$5"

  local depend="systemd-fsck@dev-stateful.service"
  if [[ "${mkfs_stateful}" == "1" ]]; then
    depend="stateful-partition-mkfs.service"
  fi
  if [[ "${protection_level}" == "e" ]]; then
    escaped_blk_dev=$(systemd-escape "${PROTECTED_BLOCK_DEVICE}" --path)
    depend="protected-stateful-partition@${escaped_blk_dev}.service ${depend}"
  fi
  echo "
[Unit]
Before=local-fs.target
BindsTo=dev-stateful.device
After=dev-stateful.device ${depend}
Requires=mnt-stateful_partition-make-private.service ${depend}

[Mount]
What=/dev/stateful
Where=/mnt/stateful_partition
Type=${fstype}
Options=$(stateful_mount_options)
" > "${output_dir}/mnt-stateful_partition.mount"

  if [[ -n "${enable_dir}" ]]; then
    # Creating symlink as required by systemd-239
    ln -s "${output_dir}/mnt-stateful_partition.mount" \
       "${enable_dir}/mnt-stateful_partition.mount"
  fi
}

# Generates a systemd mount unit for mounting a logs disk.
#
# Args:
# device: The device to mount.
# output_dir: The directory to output the unit to.
# enable_dir[Optional]: The directory containing the symlink
# to the output unit for enabling the output unit. It should
# be an absolute path.
gen_logs_mount() {
  local -r device="$1"
  local -r output_dir="$2"
  local -r enable_dir="$3"
  echo "
[Unit]
Before=local-fs.target systemd-journal-flush.service
BindsTo=$(systemd-escape "${device#/}").device
After=var.mount $(systemd-escape "${device#/}").device \
prep-logs-dev@$(systemd-escape "${device#/}").service
Requires=prep-logs-dev@$(systemd-escape "${device#/}").service

[Mount]
What=${device}
Where=/var/log
Type=ext4
Options=$(stateful_mount_options)
" > "${output_dir}/var-log.mount"

  if [[ -n "${enable_dir}" ]]; then
    # Creating symlink as required by systemd-239
    ln -s "${output_dir}/var-log.mount" \
       "${enable_dir}/var-log.mount"
  fi
}

main() {
  local -r output_root="$1"
  local -r rootdev="$(rootdev -s | sed 's/[0-9_]*$//')"
  parse_kernel_args "${rootdev}"
  mkdir -p "${output_root}/local-fs.target.wants"
  gen_oem_partition_mount "${rootdev}${PARTITION_NUM_OEM}" "${FS_FORMAT_OEM}" \
    "${output_root}" "${output_root}/local-fs.target.wants"
  configure_dev_stateful "${STATEFUL_DEV}" "${output_root}"
  gen_stateful_partition_mkfs "${FS_FORMAT_STATE}" "${output_root}"
  gen_stateful_partition_mount "${FS_FORMAT_STATE}" "${MKFS_STATEFUL}" \
    "${output_root}" "${STATEFUL_PROTECTION_LEVEL}" \
    "${output_root}/local-fs.target.wants"
  if [[ -n "${LOGS_DEV}" ]]; then
    gen_logs_mount "${LOGS_DEV}" "${output_root}" \
      "${output_root}/local-fs.target.wants"
  fi
}

main "$@" > /dev/ttyS0 2>&1