#
# Functions for mktirfs
#

run_hooks() {
  hook_type="$1"
  hooks=$({
    if [ -d /usr/share/tiny-initramfs/hooks."${hook_type}" ] ; then run-parts --list /usr/share/tiny-initramfs/hooks."${hook_type}" ; fi
    if [ -d /etc/tiny-initramfs/hooks."${hook_type}" ]       ; then run-parts --list /etc/tiny-initramfs/hooks."${hook_type}"       ; fi
  } | sed 's%.*/%%' | sort -u)
  for hook in $hooks ; do
    if [ -f /etc/tiny-initramfs/hooks."${hook_type}"/"${hook}" ] ; then
      . /etc/tiny-initramfs/hooks."${hook_type}"/"${hook}"
    elif [ -f /usr/share/tiny-initramfs/hooks."${hook_type}"/"${hook}" ] ; then
      . /usr/share/tiny-initramfs/hooks."${hook_type}"/"${hook}"
    fi
  done
}

add_modules() {
  add_modules_for /
  add_modules_for /usr
}

add_modules_for() {
  look_for="$1"
  blockdev_id=
  blockdev_name=
  p_blockdev_name=
  fstype_id=
  mount_source=

  while read id parent_id majmin devpath mountpoint popts rest ; do
    set -- $rest
    while [ x"$1" != x"-" ] && [ -n "$1" ]; do
      shift
    done
    if [ x"$1" != x"-" ] ; then
      continue
    fi
    shift
    if [ x"$mountpoint" = x"$look_for" ] ; then
      blockdev_id="$majmin"
      fstype_id="$1"
      # parse filesystem specific information
      if [ x"$fstype_id" = x"btrfs" ] ; then
        mount_source="$2"
      fi
      break
    fi
  done < /proc/self/mountinfo
  
  # not a mount point
  if [ -z "$fstype_id" ] ; then
    if [ x"$look_for" = x"/" ] ; then
      echo "$0: WARNING: unable to find mount point information while adding modules for root file system" >&2
      echo "YOUR SYSTEM MIGHT NOT BOOT WITH THIS INITRAMFS." >&2
    fi
    return
  fi

  # special casing ubifs
  if [ x"$fstype_id" = x"ubifs" ] ; then
    printf "ubifs\n" >> "$modules_list"
    return
  fi

  # the major:minor value of st_dev for a btrfs filesystem does not map to a
  # real device, so let's try to figure out the actual device ID (note that
  # this will work for a btrfs backed by a single device, the multi-device
  # case is not tested!)
  if [ x"$fstype_id" = x"btrfs" ] ; then
     uuid="$(blkid --output export $mount_source | grep ^UUID= | cut -f 2 -d =)"
     if [ -d "/sys/fs/btrfs/$uuid" ]; then
       for device in /sys/fs/btrfs/$uuid/devices/* ; do
         if [ -f "$device/dev" ] ; then
	   blockdev_id=$(cat "$device/dev")
	   break
	 fi
       done
     fi
  fi

  # find kernel block device name
  for blk in /sys/block/* ; do
    # handle partitions
    for blk2 in "$blk"/"${blk##*/}"* ; do
      if [ -f "$blk2/dev" ] && [ x"$(cat $blk2/dev)" = x"$blockdev_id" ] ; then
        # use device containing the partitions
        blockdev_name="${blk##*/}"
        p_blockdev_name="${blk2##*/}"
        break
      fi
    done
    if [ -f "$blk/dev" ] && [ x"$(cat $blk/dev)" = x"$blockdev_id" ] ; then
      blockdev_name="${blk##*/}"
    fi
    if [ -n "$blockdev_name" ] ; then
      break
    fi
  done
  if [ -z "$p_blockdev_name" ] ; then
    p_blockdev_name="$blockdev_name"
  fi

  if [ -z "$blockdev_name" ] && [ x"$fstype_id" != x"nfs4" ] ; then
    echo "$0: WARNING: unable to determine kernel block device name for $look_for file system" >&2
    echo "YOUR SYSTEM MIGHT NOT BOOT WITH THIS INITRAMFS." >&2
  else
    # find file system module responsible (for example,
    # ext4 often handles ext2 and ext3 file systems)
    real_fstype_id=
    for fstype in /sys/fs/* ; do
      if [ -d "$fstype/$p_blockdev_name" ] ; then
        real_fstype_id=${fstype##*/}
        break
      fi
    done
  fi

  if [ -z "$real_fstype_id" ] ; then
    real_fstype_id="$fstype_id"
  fi

  # special case module name variance
  # (note that this is not tested!)
  case "$real_fstype_id" in
    nfs4)       real_fstype_id="nfsv4" ;;
  esac

  printf "%s\n" "$real_fstype_id" >> "$modules_list"

  # determine block device
  case "$blockdev_name" in
    dm-*)
      echo "$0: WARNING: $look_for file system on device mapper, not supported by tiny-initramfs" >&2
      echo "YOUR SYSTEM WILL NOT BOOT WITH THIS INITRAMFS." >&2
      blockdev_name=
      ;;
    md/*|md-*)
      echo "$0: WARNING: $look_for file system on md array, not directly supported by tiny-initramfs" >&2
      echo "(please use explicit kernel option for assembling the array and add any non-builtin modules manually)" >&2
      echo "YOUR SYSTEM MIGHT NOT BOOT WITH THIS INITRAMFS." >&2
      blockdev_name=
      ;;
    loop*)
      echo "$0: WARNING: $look_for file system on loop device, not supported by tiny-initramfs" >&2
      echo "YOUR SYSTEM WILL NOT BOOT WITH THIS INITRAMFS." >&2
      blockdev_name=
      ;;
  esac

  if [ -n "$blockdev_name" ] ; then
    # determine modules
    sysfs_add_modules $(readlink -f /sys/block/"${blockdev_name}"/device)
  fi

  # load additional drivers if necessary
  # (Logic taken from initramfs-tools.)
  if [ -e /sys/bus/scsi/devices ] ; then
    printf "sd_mod\n" >> "$modules_list"
  fi
  if [ -e /sys/bus/mmc/devices ] ; then
    printf "mmc_block\n" >> "$modules_list"
  fi
  if [ -e /sys/bus/virtio ] ; then
    printf "virtio_pci\n" >> "$modules_list"
    printf "virtio_mmio\n" >> "$modules_list"
  fi
}

sysfs_add_modules() {
  path="$1"
  while [ -n "$path" ] && [ x"$path" != x"/" ]  && [ x"$path" != x"." ] && [ x"$path" != x"/sys" ] ; do
    modalias="$(cat "$path"/modalias 2>/dev/null || :)"
    if [ -n "$modalias" ] ; then
      printf "%s\n" "$modalias" >> "$modules_list"
    fi
    driver="$(readlink -f "$path"/driver/module || :)"
    if [ -e "$driver" ] ; then
      printf "%s\n" "$(basename "$driver")" >> "$modules_list"
    fi
    path="$(dirname "$path")"
  done
}

process_modules() {
  > "${initramfs_dir}/modules"
  process_modules_helper $(cat "$modules_list")
  if grep -q ^/libcrc32c\.ko "${initramfs_dir}/modules" || grep -q ^/btrfs\.ko "${initramfs_dir}/modules" ; then
    process_modules_helper crc32c
  fi
  if grep -q ^/ubifs\.ko "${initramfs_dir}/modules" ; then
    process_modules_helper deflate lzip lzo
  fi
  if grep -q ^/virtio_blk\.ko "${initramfs_dir}/modules" ; then
    process_modules_helper virtio_pci virtio_mmio
  fi
  # remove file if it's empty
  if ! [ -s "${initramfs_dir}/modules" ] ; then
    rm -f "${initramfs_dir}/modules"
  fi
}

process_modules_helper() {
  if [ $# -eq 0 ] ; then
    return
  fi
  /sbin/modprobe --all --ignore-install --set-version="${VERSION}" --quiet --show-depends "$@" | \
    awk '$1 == "insmod" { print; }' | while read dummy_type mod_file mod_options ; do
    mod_name=$(basename "$mod_file" | sed -E 's/(.*\.ko)(\..*)?/\1/')
    mod_compression=$(basename "$mod_file" | sed -E 's/(.*\.ko)(\..*)?/\2/')
    if ! grep -q ^/"${mod_name}" "${initramfs_dir}/modules" ; then
      if [ "$mod_compression" = .xz ] ; then
        xzcat "${mod_file}" > "${initramfs_dir}/${mod_name}"
      elif [ -z "$mod_compression" ] ; then
        cp "${mod_file}" "${initramfs_dir}/${mod_name}"
      else
        echo "$0: WARNING: unable to determine compression for modules while adding modules" >&2
        echo "YOUR SYSTEM MIGHT NOT BOOT WITH THIS INITRAMFS." >&2
        cp "${mod_file}" "${initramfs_dir}/${mod_name}"
      fi
      printf "%s\n" "/${mod_name}${mod_options:+ $mod_options}" >> "${initramfs_dir}/modules"
    fi
  done
}

# This function is from dracut (dracut-functions.sh, GPL2+)
# get_cpu_vendor
# Only two values are returned: AMD or Intel
get_cpu_vendor ()
{
    if grep -qE AMD /proc/cpuinfo; then
        printf "AMD"
    fi
    if grep -qE Intel /proc/cpuinfo; then
        printf "Intel"
    fi
}

# This function is from dracut (dracut-functions.sh, GPL2+)
# get_host_ucode
# Get the hosts' ucode file based on the /proc/cpuinfo
get_ucode_file ()
{
    local family=`grep -E "cpu family" /proc/cpuinfo | head -1 | sed s/.*:\ //`
    local model=`grep -E "model" /proc/cpuinfo |grep -v name | head -1 | sed s/.*:\ //`
    local stepping=`grep -E "stepping" /proc/cpuinfo | head -1 | sed s/.*:\ //`

    if [ x"$(get_cpu_vendor)" = x"AMD" ]; then
        # If family greater or equal than 0x16, 0x15
        if [ $family -ge 22 ]; then
            printf "microcode_amd_fam16h.bin"
        elif [ $family -ge 21 ]; then
            printf "microcode_amd_fam15h.bin"
        else
            printf "microcode_amd.bin"
        fi
    fi
    if [ x"$(get_cpu_vendor)" = x"Intel" ]; then
        # The /proc/cpuinfo are in decimal.
        printf "%02x-%02x-%02x" ${family} ${model} ${stepping}
    fi
}

# See comment in add_microcode
add_intel_microcode()
{
  local IUCODE_DIRS=""
  for firmware_dir in /lib/firmware/updates /lib/firmware /lib/firmware/${VERSION} ; do
    if [ -d "$firmware_dir/intel-ucode" ] ; then
      IUCODE_DIRS="${IUCODE_DIRS:+$IUCODE_DIRS }${firmware_dir}/intel-ucode"
    fi
  done
  local IUCODE_TOOL_OPTIONS=""
  if [ x"$1" != x"generic" ] ; then
    IUCODE_TOOL_OPTIONS="-S"
  fi
  iucode_tool -q $IUCODE_TOOL_OPTIONS --write-earlyfw="$extra_early_dir/intel_microcode.img" --overwrite $IUCODE_DIRS
  if [ -f "$extra_early_dir/intel_microcode.img" ] ; then
    echo "intel_microcode.img" >> "$extra_early_list"
  fi
}

# This is somewhat based on dracut's microcode logic (dracut.sh, GPL2+)
add_microcode() {
  # only x86 microcodes supported so far
  case "$(uname -m)" in
    i?86|x86_64) ;;
    *)           return ;;
  esac

  destdir="$early_dir/kernel/x86/microcode"

  if [ x"$1" = x"generic" ]; then
    WANTED_VENDORS="Intel AMD"
    specific_file="no"
    firmware_file=""
  else
    WANTED_VENDORS="$(get_cpu_vendor)"
    specific_file="yes"
    firmware_file="$(get_ucode_file)"
    if [ -z "$firmware_file" ] ; then
      echo "$0: WARNING: couldn't determine microcode file name for current CPU, adding all" $WANTED_VENDORS "microcodes." >&2
      specific_file="no"
      if [ -z "$WANTED_VENDORS" ]; then
        WANTED_VENDORS="Intel AMD"
      fi
    fi
  fi
  
  for vendor in $WANTED_VENDORS ; do
    if [ x"$vendor" = x"Intel" ] && /usr/bin/which iucode_tool >/dev/null 2>&1 ; then
      # In case of Intel add the microcode differently if iucode_tool
      # is installed. (Will be the case if the intel-microcode package
      # is installed, because it Depends: iucode-tool, but may not be
      # the case if the user has installed the Intel microcode
      # manually, so keep the old code as a fallback.)
      add_intel_microcode "$1"
      continue
    fi
    case $vendor in
      AMD)   dir="amd-ucode"   ; bin="AuthenticAMD.bin" ;;
      Intel) dir="intel-ucode" ; bin="GenuineIntel.bin" ;;
    esac
    for firmware_dir in /lib/firmware/updates /lib/firmware /lib/firmware/${VERSION} ; do
      if [ -d "$firmware_dir/$dir" ] ; then
        if [ "$specific_file" = "yes" ] ; then
          if [ -r "$firmware_dir/$dir/$firmware_file" ] ; then
            mkdir -p "$destdir"
            cat "$firmware_dir/$dir/$firmware_file" >> "$destdir/$bin"
          elif [ -r "$firmware_dir/$dir/$firmware_file.initramfs" ] ; then
            mkdir -p "$destdir"
            cat "$firmware_dir/$dir/$firmware_file.initramfs" >> "$destdir/$bin"
          fi
        else
          for fn in "$firmware_dir/$dir"/* ; do
            if ! [ -r "$fn" ] ; then
              continue
            fi
            # ignore gpg signatures
            if [ x"${fn%%.asc}" != x"${fn}" ] ; then
              continue
            fi
            mkdir -p "$destdir"
            cat "$fn" >> "$destdir/$bin"
          done
        fi
      fi
    done
  done
}
