# This module builds the initial ramdisk, which contains an init
# script that performs the first stage of booting the system: it loads
# the modules necessary to mount the root file system, then calls the
# init in the root file system to start the second boot stage.

{ config, lib, utils, pkgs, ... }:

with lib;

let

  udev = pkgs.eudev;

  kernel-name = config.boot.kernelPackages.kernel.name or "kernel";

  modulesTree = config.system.modulesTree.override { name = kernel-name + "-modules"; };
  firmware = config.hardware.firmware;


  # Determine the set of modules that we need to mount the root FS.
  modulesClosure = pkgs.makeModulesClosure {
    rootModules = config.boot.initrd.availableKernelModules ++ config.boot.initrd.kernelModules;
    kernel = modulesTree;
    firmware = firmware;
    allowMissing = false;
  };


  # The initrd only has to mount `/` or any FS marked as necessary for
  # booting (such as the FS containing `/nix/store`, or an FS needed for
  # mounting `/`, like `/` on a loopback).
  fileSystems = filter utils.fsNeededForBoot config.system.build.fileSystems;

  # A utility for enumerating the shared-library dependencies of a program
  findLibs = pkgs.buildPackages.writeShellScriptBin "find-libs" ''
    set -euo pipefail

    declare -A seen
    left=()

    patchelf="${pkgs.buildPackages.patchelf}/bin/patchelf"

    function add_needed {
      rpath="$($patchelf --print-rpath $1)"
      dir="$(dirname $1)"
      for lib in $($patchelf --print-needed $1); do
        left+=("$lib" "$rpath" "$dir")
      done
    }

    add_needed "$1"

    while [ ''${#left[@]} -ne 0 ]; do
      next=''${left[0]}
      rpath=''${left[1]}
      ORIGIN=''${left[2]}
      left=("''${left[@]:3}")
      if [ -z ''${seen[$next]+x} ]; then
        seen[$next]=1

        # Ignore the dynamic linker which for some reason appears as a DT_NEEDED of glibc but isn't in glibc's RPATH.
        case "$next" in
          ld*.so.?) continue;;
        esac

        IFS=: read -ra paths <<< $rpath
        res=
        for path in "''${paths[@]}"; do
          path=$(eval "echo $path")
          if [ -f "$path/$next" ]; then
              res="$path/$next"
              echo "$res"
              add_needed "$res"
              break
          fi
        done
        if [ -z "$res" ]; then
          echo "Couldn't satisfy dependency $next" >&2
          exit 1
        fi
      fi
    done
  '';

  # Some additional utilities needed in stage 1, like mount, fsck
  # etc.  We don't want to bring in all of those packages, so we just
  # copy what we need.  Instead of using statically linked binaries,
  # we just copy what we need from Glibc and use patchelf to make it
  # work.
  extraUtils = pkgs.runCommandCC "extra-utils"
    { nativeBuildInputs = [pkgs.buildPackages.nukeReferences];
      allowedReferences = [ "out" ]; # prevent accidents like glibc being included in the initrd
    }
    ''
      set +o pipefail

      mkdir -p $out/bin $out/lib
      ln -s $out/bin $out/sbin

      copy_bin_and_libs () {
        [ -f "$out/bin/$(basename $1)" ] && rm "$out/bin/$(basename $1)"
        cp -pdv $1 $out/bin
      }

      # Copy BusyBox.
      for BIN in ${pkgs.busybox}/{s,}bin/*; do
        copy_bin_and_libs $BIN
      done

      # Copy some util-linux stuff.
      copy_bin_and_libs ${pkgs.util-linux}/sbin/blkid

      # Add RAID mdadm tool.
      copy_bin_and_libs ${pkgs.mdadm}/sbin/mdadm
      copy_bin_and_libs ${pkgs.mdadm}/sbin/mdmon

      # Copy udev.
      copy_bin_and_libs ${udev}/bin/udevd
      copy_bin_and_libs ${udev}/bin/udevadm
      for BIN in ${udev}/lib/udev/*_id; do
        copy_bin_and_libs $BIN
      done

      # Copy mdevd.
      for BIN in ${pkgs.mdevd}/bin/*; do
        copy_bin_and_libs $BIN
      done

      # Copy modprobe.
      copy_bin_and_libs ${pkgs.kmod}/bin/kmod
      ln -sf kmod $out/bin/modprobe

      # Dirty hack to make sure the kernel properly loads modules
      # such as ext4 on demand (e.g. on a `mount(2)` syscall). This is necessary
      # because `kmod` isn't linked against `libpthread.so.0` anymore (since
      # it was merged into `libc.so.6` since version `2.34`), but still needs
      # to access it for some reason. This is not an issue in stage-1 itself
      # because of the `LD_LIBRARY_PATH`-variable and anytime later because the rpath of
      # kmod/modprobe points to glibc's `$out/lib` where `libpthread.so.6` exists.
      # However, this is a problem when the kernel calls `modprobe` inside
      # the initial ramdisk because it doesn't know about the
      # `LD_LIBRARY_PATH` and the rpath was nuked.
      #
      # Also, we can't use `makeWrapper` here because `kmod` only does
      # `modprobe` functionality if `argv[0] == "modprobe"`.
      cat >$out/bin/modprobe-kernel <<EOF
      #!$out/bin/ash
      export LD_LIBRARY_PATH=$out/lib
      exec $out/bin/modprobe "\$@"
      EOF
      chmod +x $out/bin/modprobe-kernel

      # Copy resize2fs if any ext* filesystems are to be resized
      ${optionalString (any (fs: fs.autoResize && (lib.hasPrefix "ext" fs.fsType)) fileSystems) ''
        # We need mke2fs in the initrd.
        copy_bin_and_libs ${pkgs.e2fsprogs}/sbin/resize2fs
      ''}

      # Copy multipath.
      ${optionalString config.services.multipath.enable ''
        copy_bin_and_libs ${config.services.multipath.package}/bin/multipath
        copy_bin_and_libs ${config.services.multipath.package}/bin/multipathd
        # Copy lib/multipath manually.
        cp -rpv ${config.services.multipath.package}/lib/multipath $out/lib
      ''}

      # Copy secrets if needed.
      #
      # TODO: move out to a separate script; see #85000.
      ${optionalString (!config.boot.loader.supportsInitrdSecrets)
          (concatStringsSep "\n" (mapAttrsToList (dest: source:
             let source' = if source == null then dest else source; in
               ''
                  mkdir -p $(dirname "$out/secrets/${dest}")
                  # Some programs (e.g. ssh) doesn't like secrets to be
                  # symlinks, so we use `cp -L` here to match the
                  # behaviour when secrets are natively supported.
                  cp -Lr ${source'} "$out/secrets/${dest}"
                ''
          ) config.boot.initrd.secrets))
       }

      ${config.boot.initrd.extraUtilsCommands}

      # Copy ld manually since it isn't detected correctly
      cp -pv ${pkgs.stdenv.cc.libc.out}/lib/ld*.so.? $out/lib

      # Copy all of the needed libraries
      find $out/bin $out/lib -type f | while read BIN; do
        echo "Copying libs for executable $BIN"
        for LIB in $(${findLibs}/bin/find-libs $BIN); do
          TGT="$out/lib/$(basename $LIB)"
          if [ ! -f "$TGT" ]; then
            SRC="$(readlink -e $LIB)"
            cp -pdv "$SRC" "$TGT"
          fi
        done
      done

      # Strip binaries further than normal.
      chmod -R u+w $out
      stripDirs "$STRIP" "lib bin" "-s"

      # Run patchelf to make the programs refer to the copied libraries.
      find $out/bin $out/lib -type f | while read i; do
        if ! test -L $i; then
          nuke-refs -e $out $i
        fi
      done

      find $out/bin -type f | while read i; do
        if ! test -L $i; then
          echo "patching $i..."
          patchelf --set-interpreter $out/lib/ld*.so.? --set-rpath $out/lib $i || true
        fi
      done

      if [ -z "${toString (pkgs.stdenv.hostPlatform != pkgs.stdenv.buildPlatform)}" ]; then
      # Make sure that the patchelf'ed binaries still work.
      echo "testing patched programs..."
      $out/bin/ash -c 'echo hello world' | grep "hello world"
      export LD_LIBRARY_PATH=$out/lib
      $out/bin/mount --help 2>&1 | grep -q "BusyBox"
      $out/bin/blkid -V 2>&1 | grep -q 'libblkid'
      $out/bin/udevadm --version
      $out/bin/mdadm --version
      ${optionalString config.services.multipath.enable ''
        ($out/bin/multipath || true) 2>&1 | grep -q 'need to be root'
        ($out/bin/multipathd || true) 2>&1 | grep -q 'need to be root'
      ''}

      ${config.boot.initrd.extraUtilsCommandsTest}
      fi
    ''; # */

  udevRules = pkgs.runCommand "udev-rules" {
      allowedReferences = [ extraUtils ];
      preferLocalBuild = true;
    } ''
      mkdir -p $out

      echo 'ENV{LD_LIBRARY_PATH}="${extraUtils}/lib"' > $out/00-env.rules

      cp -v ${udev}/var/lib/udev/rules.d/60-cdrom_id.rules $out/
      cp -v ${udev}/var/lib/udev/rules.d/60-persistent-storage.rules $out/
      cp -v ${udev}/var/lib/udev/rules.d/75-net-description.rules $out/
      cp -v ${udev}/var/lib/udev/rules.d/80-drivers.rules $out/
      cp -v ${udev}/var/lib/udev/rules.d/80-net-name-slot.rules $out/
      ${config.boot.initrd.extraUdevRulesCommands}

      for i in $out/*.rules; do
          substituteInPlace $i \
            --replace ata_id ${extraUtils}/bin/ata_id \
            --replace scsi_id ${extraUtils}/bin/scsi_id \
            --replace cdrom_id ${extraUtils}/bin/cdrom_id \
            --replace ${pkgs.coreutils}/bin/basename ${extraUtils}/bin/basename \
            --replace ${pkgs.util-linux}/bin/blkid ${extraUtils}/bin/blkid \
            --replace ${pkgs.mdadm}/sbin ${extraUtils}/sbin \
            --replace ${pkgs.bash}/bin/sh ${extraUtils}/bin/sh \
            --replace ${udev} ${extraUtils}
      done

      # Work around a bug in QEMU, which doesn't implement the "READ
      # DISC INFORMATION" SCSI command:
      #   https://bugzilla.redhat.com/show_bug.cgi?id=609049
      # As a result, `cdrom_id' doesn't print
      # ID_CDROM_MEDIA_TRACK_COUNT_DATA, which in turn prevents the
      # /dev/disk/by-label symlinks from being created.  We need these
      # in the NixOS installation CD, so use ID_CDROM_MEDIA in the
      # corresponding udev rules for now.  This was the behaviour in
      # udev <= 154.  See also
      #   http://www.spinics.net/lists/hotplug/msg03935.html
      substituteInPlace $out/60-persistent-storage.rules \
        --replace ID_CDROM_MEDIA_TRACK_COUNT_DATA ID_CDROM_MEDIA
    ''; # */


  # The init script of boot stage 1 (loading kernel modules for
  # mounting the root FS).
  bootStage1 = pkgs.substituteAll {
    src = ./stage-1-init.sh;

    shell = "${extraUtils}/bin/ash";

    isExecutable = true;

    postInstall = ''
      echo checking syntax
      # check both with bash
      ${pkgs.buildPackages.bash}/bin/sh -n $target
      # and with ash shell, just in case
      ${pkgs.buildPackages.busybox}/bin/ash -n $target
    '';

    inherit udevRules extraUtils modulesClosure;

    inherit (config.boot) resumeDevice;

    inherit (config.system.build) earlyMountScript;

    inherit (config.boot.initrd) checkJournalingFS verbose
      preDeviceCommands postDeviceCommands postMountCommands preFailCommands kernelModules;

    resumeDevices = map (sd: if sd ? device then sd.device else "/dev/disk/by-label/${sd.label}")
                    (filter (sd: hasPrefix "/dev/" sd.device && !sd.randomEncryption.enable
                             # Don't include zram devices
                             && !(hasPrefix "/dev/zram" sd.device)
                            ) config.swapDevices);

    fsInfo =
      let f = fs: [ fs.mountPoint (if fs.device != null then fs.device else "/dev/disk/by-label/${fs.label}") fs.fsType (builtins.concatStringsSep "," fs.options) ];
      in pkgs.writeText "initrd-fsinfo" (concatStringsSep "\n" (concatMap f fileSystems));

    setHostId = optionalString (config.networking.hostId != null) ''
      hi="${config.networking.hostId}"
      ${if pkgs.stdenv.isBigEndian then ''
        echo -ne "\x''${hi:0:2}\x''${hi:2:2}\x''${hi:4:2}\x''${hi:6:2}" > /etc/hostid
      '' else ''
        echo -ne "\x''${hi:6:2}\x''${hi:4:2}\x''${hi:2:2}\x''${hi:0:2}" > /etc/hostid
      ''}
    '';
  };


  # The closure of the init script of boot stage 1 is what we put in
  # the initial RAM disk.
  initialRamdisk = pkgs.makeInitrd {
    name = "initrd-${kernel-name}";
    inherit (config.boot.initrd) compressor compressorArgs prepend;

    contents =
      [ { object = bootStage1;
          symlink = "/init";
        }
        { object = pkgs.writeText "mdadm.conf" config.boot.initrd.services.swraid.mdadmConf;
          symlink = "/etc/mdadm.conf";
        }
        { object = pkgs.runCommand "initrd-kmod-blacklist-ubuntu" {
              src = "${pkgs.kmod-blacklist-ubuntu}/modprobe.conf";
              preferLocalBuild = true;
            } ''
              target=$out
              ${pkgs.buildPackages.perl}/bin/perl -0pe 's/## file: iwlwifi.conf(.+?)##/##/s;' $src > $out
            '';
          symlink = "/etc/modprobe.d/ubuntu.conf";
        }
        { object = config.environment.etc."modprobe.d/nixos.conf".source;
          symlink = "/etc/modprobe.d/nixos.conf";
        }
        { object = pkgs.kmod-debian-aliases;
          symlink = "/etc/modprobe.d/debian.conf";
        }
      ] ++ lib.optionals config.services.multipath.enable [
        { object = pkgs.runCommand "multipath.conf" {
              src = config.environment.etc."multipath.conf".text;
              preferLocalBuild = true;
            } ''
              target=$out
              printf "$src" > $out
              substituteInPlace $out \
                --replace ${config.services.multipath.package}/lib ${extraUtils}/lib
            '';
          symlink = "/etc/multipath.conf";
        }
      ] ++ (lib.mapAttrsToList
        (symlink: options:
          {
            inherit symlink;
            object = options.source;
          }
        )
        config.boot.initrd.extraFiles);
  };

  # Script to add secret files to the initrd at bootloader update time
  initialRamdiskSecretAppender =
    let
      compressorExe = initialRamdisk.compressorExecutableFunction pkgs;
    in pkgs.writeScriptBin "append-initrd-secrets"
      ''
        #!${pkgs.bash}/bin/bash -e
        function usage {
          echo "USAGE: $0 INITRD_FILE" >&2
          echo "Appends this configuration's secrets to INITRD_FILE" >&2
        }

        if [ $# -ne 1 ]; then
          usage
          exit 1
        fi

        if [ "$1"x = "--helpx" ]; then
          usage
          exit 0
        fi

        ${lib.optionalString (config.boot.initrd.secrets == {})
            "exit 0"}

        export PATH=${pkgs.coreutils}/bin:${pkgs.libarchive}/bin:${pkgs.gzip}/bin:${pkgs.findutils}/bin

        function cleanup {
          if [ -n "$tmp" -a -d "$tmp" ]; then
            rm -fR "$tmp"
          fi
        }
        trap cleanup EXIT

        tmp=$(mktemp -d ''${TMPDIR:-/tmp}/initrd-secrets.XXXXXXXXXX)

        ${lib.concatStringsSep "\n" (mapAttrsToList (dest: source:
            let source' = if source == null then dest else toString source; in
              ''
                mkdir -p $(dirname "$tmp/.initrd-secrets/${dest}")
                cp -a ${source'} "$tmp/.initrd-secrets/${dest}"
              ''
          ) config.boot.initrd.secrets)
         }

        (cd "$tmp" && find . -print0 | sort -z | bsdtar --uid 0 --gid 0 -cnf - -T - | bsdtar --null -cf - --format=newc @-) | \
          ${compressorExe} ${lib.escapeShellArgs initialRamdisk.compressorArgs} >> "$1"
      '';
in

{
  config = mkIf config.boot.initrd.enable {

    system.build = lib.mapAttrs (_: lib.mkForce) { inherit bootStage1 initialRamdisk initialRamdiskSecretAppender extraUtils; };

  };
}