#!/bin/bash
# Copyright (c) 2023 by the Flatcar Maintainers.
# Use of this source code is governed by the Apache 2.0 license.

# Helper script for building OS images w/ sysexts included.
# Called by build_image -> prod_image_util.sh.
# This is a separate script mainly so we can trap EXIT and clean up our mounts
#  without interfering with traps set by build_image.

# We're in build_library/, script root is one up
SCRIPT_ROOT="$(cd "$(dirname "$(readlink -f "$0")")/../"; pwd)"
. "${SCRIPT_ROOT}/common.sh" || exit 1

# Script must run inside the chroot
assert_inside_chroot
switch_to_strict_mode

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

# Create a sysext from a package and install it to the OS image.
# Conventions:
# - For each <group>/<package>, <group>_<package>_pkginfo will be built. Can be used in subsequent calls
#   to build dependent sysexts.
# - If ${BUILD_LIBRARY_DIR}/sysext_mangle_<group>_<package> exists it will be used as FS mangle script
#   when building the sysext.
create_prod_sysext() {
  local BOARD="$1"
  local output_dir="$2"
  local workdir="$3"
  local base_sysext="$4"
  local install_root="$5"
  local name="$6"
  local grp_pkgs="$7"
  local pkginfo="${8:-}"

  local -a build_sysext_opts=()

  local -a grp_pkg
  mapfile -t grp_pkg <<<"${grp_pkgs//&/$'\n'}"
  local msg="Installing ${grp_pkg[*]} in sysext ${name}.raw"

  # Include previous sysexts' pkginfo if supplied
  if [[ -n "${pkginfo}" ]] ; then
    if [[ ! -f "${output_dir}/${pkginfo}" ]] ; then
      die "Sysext build '${name}': unable to find package info at '${output_dir}/${pkginfo}'."
    fi
    msg="${msg} w/ package info '${pkginfo}'"
    build_sysext_opts+=( "--base_pkginfo=${output_dir}/${pkginfo}" )
  fi

  # Include FS mangle script if present
  if [[ -x "${BUILD_LIBRARY_DIR}/sysext_mangle_${name}" ]] ; then
    build_sysext_opts+=( "--manglefs_script=${BUILD_LIBRARY_DIR}/sysext_mangle_${name}" )
    msg="${msg}, FS mangle script 'sysext_mangle_${name}'"
  fi

  info "${msg}."

  # Pass the build ID extracted from root FS to build_sysext. This prevents common.sh
  #   in build_sysext to generate a (timestamp based) build ID during a DEV build of a
  #   release tag (which breaks its version check).
  #
  # The --install_root_basename="${name}-base-sysext-rootfs" flag is
  # important - it sets the name of a rootfs directory, which is used
  # to determine the package target in coreos/base/profile.bashrc
  #
  # Built-in sysexts are stored in the compressed /usr partition, so we
  # disable compression to avoid double-compression.
  sudo -E "FLATCAR_BUILD_ID=$FLATCAR_BUILD_ID" "${SCRIPTS_DIR}/build_sysext" \
            --board="${BOARD}" \
            --image_builddir="${workdir}/sysext-build" \
            --squashfs_base="${base_sysext}" \
            --generate_pkginfo \
            --compression=none \
            --install_root_basename="${name}-base-sysext-rootfs" \
            "${build_sysext_opts[@]}" \
            "${name}" "${grp_pkg[@]}"

  sudo mv "${workdir}/sysext-build/${name}.raw" "${workdir}/sysext-build/${name}_pkginfo.raw" \
                              "${workdir}/sysext-build/${name}"_*.txt "${output_dir}"

  sudo mkdir -p "${install_root}"/usr/share/flatcar/sysext
  sudo install -m 0644 -D "${output_dir}/${name}.raw" "${install_root}"/usr/share/flatcar/sysext/

  sudo mkdir -p "${install_root}"/etc/extensions/
  sudo ln -sf "/usr/share/flatcar/sysext/${name}.raw" "${install_root}/etc/extensions/${name}.raw"
}
# --

BOARD="$1"
BUILD_DIR="$2"
root_fs_dir="$3"

merged_rootfs_dir="$4"
sysext_output_dir="$5"

sysexts_list="$6"

grp_pkg=""
prev_pkginfo=""
sysext_workdir="${BUILD_DIR}/prod-sysext-work"
sysext_mountdir="${BUILD_DIR}/prod-sysext-work/mounts"
sysext_base="${sysext_workdir}/base-os.squashfs"

function cleanup() {
  IFS=':' read -r -a mounted_sysexts <<< "$sysext_lowerdirs"
  # skip the rootfs
  mounted_sysexts=("${mounted_sysexts[@]:1}")

  for sysext in "${mounted_sysexts[@]}"; do
    sudo systemd-dissect --umount --rmdir "$sysext"
  done

  sudo umount "${sysext_mountdir}"/* || true
  rm -rf "${sysext_workdir}" || true
}
# --

trap cleanup EXIT

rm -rf "${sysext_workdir}" "${sysext_output_dir}"
mkdir "${sysext_workdir}" "${sysext_output_dir}"

info "creating temporary base OS squashfs"
sudo mksquashfs "${root_fs_dir}" "${sysext_base}" -noappend -xattrs-exclude '^btrfs.'

# Build sysexts on top of root fs and mount sysexts' squashfs + pkginfo squashfs
# for combined overlay later.
prev_pkginfo=""
sysext_lowerdirs="${sysext_mountdir}/rootfs-lower"
mkdir -p "${sysext_mountdir}"
for sysext in ${sysexts_list//,/ }; do
  # format is "<name>:<group>/<package>"
  name="${sysext%|*}"
  grp_pkg="${sysext#*|}"
  create_prod_sysext "${BOARD}" \
                     "${sysext_output_dir}" \
                     "${sysext_workdir}" \
                     "${sysext_base}" \
                     "${root_fs_dir}"\
                     "${name}" \
                     "${grp_pkg}" \
                     "${prev_pkginfo}"

  sudo systemd-dissect \
    --read-only \
    --mount \
    --mkdir \
    --image-policy='root=encrypted+unprotected+absent:usr=encrypted+unprotected+absent' \
    "${sysext_output_dir}/${name}.raw" \
    "${sysext_mountdir}/${name}"

  sudo systemd-dissect \
    --read-only \
    --mount \
    --mkdir \
    --image-policy='root=encrypted+unprotected+absent:usr=encrypted+unprotected+absent' \
    "${sysext_output_dir}/${name}_pkginfo.raw" \
    "${sysext_mountdir}/${name}_pkginfo"

  sysext_lowerdirs="${sysext_lowerdirs}:${sysext_mountdir}/${name}"
  sysext_lowerdirs="${sysext_lowerdirs}:${sysext_mountdir}/${name}_pkginfo"

  prev_pkginfo="${name}_pkginfo.raw"
done

# Mount the combined overlay (base OS, sysexts, and syset pkginfos) and copy a snapshot
# into the designated output dir for upper layers to process.
mkdir -p "${sysext_mountdir}/rootfs-lower" 
sudo mount -rt squashfs -o loop,nodev "${sysext_base}" "${sysext_mountdir}/rootfs-lower"

# Mount overlay for report generation
mkdir -p "${sysext_workdir}/.work"
mkdir -p "${sysext_mountdir}/rootfs-upper" 
sudo mount -t overlay overlay \
         -o lowerdir="${sysext_lowerdirs}",upperdir="${sysext_mountdir}/rootfs-upper",workdir="${sysext_workdir}/.work" \
         "${sysext_mountdir}/rootfs-upper"


sudo rm -rf "${merged_rootfs_dir}"
sudo cp -a "${sysext_mountdir}/rootfs-upper" "${merged_rootfs_dir}"


cleanup
trap -- EXIT
