Custom ZFS decryption with Argon2

This might be a bit weird, but what I’m trying to do is to combine ZFS native encryption (using -o keyformat=raw) with generating the decryption key from password using argon2.

To do that I need a custom systemd initrd service that replaces the automatically created NixOS one and handles the key generation and calls zfs load-key.

I found @ElvishJerricco’s script that’s somewhat similar to what I’m trying to achieve, and wanted to ask if I’m approching this correctly?

(I haven’t tried running this yet, it’s meant to go on a remote server and I don’t know of an easy way to test it other than just provisioning it and deploying it.

Also note this is not supposed to be remote decryption over SSH, I’m not interested in having that. I just want to read the password from the tty.)

  boot.initrd = {
    systemd.enable = true;

    systemd.contents."/tmp/secret.key.salt".source = ./secret.key.salt;

    systemd.storePaths = [
      "${pkgs.libargon2}"
      "${pkgs.tinyxxd}"
    ];

    # https://github.com/ElvishJerricco/stage1-tpm-tailscale/blob/2c07f2a531e1557965a0d483ea694fabf9a6d5bb/filesystems.nix
    systemd.services.zfs-import-tank.enable = false; # Disable built-in NixOS ZFS service.

    systemd.services.import-tank-bare = let
      disk = "${utils.escapeSystemdPath "/dev/disk/by-id/..."}.device";
    in {
      requiredBy = ["tank-load-key.service"];
      after = [disk];
      bindsTo = [disk];
      unitConfig.DefaultDependencies = false;
      serviceConfig = {
        Type = "oneshot";
        ExecStart = "${config.boot.zfs.package}/bin/zpool import -f -N -d /dev/disk/by-id tank";
        RemainAfterExit = true;
      };
    };

    systemd.services.tank-load-key = {
      wantedBy = ["sysroot.mount"];
      before = ["sysroot.mount"];
      serviceConfig = {
        Type = "oneshot";
        RemainAfterExit = true;
        StandardInput = "tty";
        StandardOutput = "tty";
      };
      script = ''
        echo "ZFS password:"
        systemd-ask-password | ${pkgs.libargon2}/bin/argon2 ''$(cat /tmp/secret.key.salt) -id -t 4 -m 20 -p 1 -r | ${pkgs.tinyxxd}/bin/tinyxxd -r -p - > /tmp/secret.key
        ${config.boot.zfs.package}/bin/zfs load-key -L file:///tmp/secret.key tank
      '';
      enableStrictShellChecks = true;
    };
  };

Also wanted to ask for a clarification if I can use script in systemd services when using nixos-init. Running nix flake check doesn’t return any error, so I would assume yes…?