#! /bin/bash
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

set -e
set -u
set -o pipefail

PROG_NAME="$(basename "${0}")"
readonly PROG_NAME
SECURE_BOOT_CMD="/usr/share/cloud/is-secure-boot"
readonly SECURE_BOOT_CMD
ESP_PARTITION="/dev/disk/by-partlabel/EFI-SYSTEM"
readonly ESP_PARTITION
PROC_CMDLINE="/proc/cmdline"
readonly PROC_CMDLINE
CGROUP_HYBRID=1
readonly CGROUP_HYBRID
CGROUP_UNIFIED=2
readonly CGROUP_UNIFIED
CGROUP_UNKNOWN=-1
readonly CGROUP_UNKNOWN

#
# usage <exit_code>
#
# Print usage and exit.
#
usage() {
        local exit_code="${1}"

        cat <<EOF
Usage:
        ${PROG_NAME} {show| set { hybrid | unified} |help}
        show         - Shows the current cgroup type.
        set hybrid   - Changes the cgroup to hybrid
                      (only when secure boot is disabled).
        set unified  - Changes the cgroup to unified
                       (only when secure boot is disabled).
        help         - Prints this message.
EOF
        exit "${exit_code}"
}

#
# get_current_cgroup
#
# Checks /proc/cmdline and existance of /sys/fs/cgroup/cgroup.controllers
# and returns the cgroup accordingly.
#
get_current_cgroup() {
  cgroup=$(sed "s/.*systemd.unified_cgroup_hierarchy=\([a-z]*\).*/\\1/g" \
          "${PROC_CMDLINE}")
  local retval=${CGROUP_UNKNOWN}
  if [[ ( "${cgroup}" == "false" ) &&
        ( ! -f /sys/fs/cgroup/cgroup.controllers ) ]]; then
    retval=${CGROUP_HYBRID}
  fi
  if [[ ( "${cgroup}" != "false" ) &&
        ( -f /sys/fs/cgroup/cgroup.controllers ) ]]; then
    retval=${CGROUP_UNIFIED}
  fi
  echo "${retval}"

}

#
# cgroup_hybrid
#
# Configure cgroup as hybrid by modifying kernel commandline in grub.cfg.
# Takes effect after reboot.
#
cgroup_hybrid() {

  if [[ $(get_current_cgroup) -eq ${CGROUP_HYBRID} ]]; then
    echo "Cgroup is already configured as 'hybrid'."
    exit 0
  fi

  local dir
  dir="$(mktemp -d)"
  mount "${ESP_PARTITION}" "${dir}"
  local search='\(.*\)\(init=\)\(.*\)'
  local replace=" systemd.unified_cgroup_hierarchy=false "
  replace+="systemd.legacy_systemd_cgroup_controller=false"
  sed "s/${search}/\\1\\2\\3${replace}/g" -i "${dir}/efi/boot/grub.cfg"
  sync
  umount "${dir}"
  rm -rf "${dir}"
  echo "Cgroup is changed to 'hybrid' and needs a reboot to take effect."
  check_and_print_secure_msg
}

#
# cgroup_unified
#
# Configure cgroup as unified by modifying kernel commandline in grub.cfg.
# Takes effect after reboot.
#
cgroup_unified() {
  if  [[ $(get_current_cgroup) -eq ${CGROUP_UNIFIED} ]]; then
    echo "Cgroup is already configured as 'unified'."
    exit 0
  fi
  local dir
  dir="$(mktemp -d)"
  mount "${ESP_PARTITION}" "${dir}"
  local search="systemd.unified_cgroup_hierarchy=false "
  search+="systemd.legacy_systemd_cgroup_controller=false"
  local replace=""
  sed "s/${search}/${replace}/g" -i "${dir}/efi/boot/grub.cfg"
  sync
  umount "${dir}"
  rm -rf "${dir}"
  echo "Cgroup is changed to 'unified' and needs a reboot to take effect."
  check_and_print_secure_msg
}

#
# cgroup_show
#
# Show the current cgroup configuration.
#
cgroup_show() {
  cur_cgroup=$(get_current_cgroup)
  if [[ ${cur_cgroup} -eq ${CGROUP_HYBRID} ]]; then
    echo "Configured cgroup is 'hybrid'."
  elif [[ ${cur_cgroup} -eq ${CGROUP_UNIFIED} ]]; then
    echo "Configured cgroup is 'unified'."
  else
    echo "Inconsistency in cgroup configuration."
  fi
}

check_and_print_secure_msg() {
  if ${SECURE_BOOT_CMD}; then
    echo "secure boot is enabled."
    echo "Cgroup changes won't be effective while secure boot is enabled."
    echo -n "Disable secure boot, and reboot for the cgroup changes to be "
    echo "effective."
  fi
}

main() {
  case "${1-help}" in
    show)
    cgroup_show
    ;;
    set)
      case "${2-help}" in
        hybrid)
          cgroup_hybrid
          ;;
        unified)
          cgroup_unified
          ;;
        *)
          usage 1
          ;;
      esac
      ;;
    help)
    usage 0
    ;;
    *)
    usage 1
    ;;
  esac
}

main "$@"