config/profiles/known-hosts.nix

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

let
  inherit (builtins)
    map
    fromTOML
    toString
    readFile
    elem
    attrNames
    attrValues
    ;

  # TODO: test ssh-ng
  # https://discourse.nixos.org/t/wrapper-to-restrict-builder-access-through-ssh-worth-upstreaming/25834
  nix-ssh-wrapper = pkgs.writeShellScript "nix-ssh-wrapper" ''
    case $SSH_ORIGINAL_COMMAND in
      "nix-daemon --stdio")
        exec ${config.nix.package}/bin/nix-daemon --stdio
        ;;
      "nix-store --serve --write")
        exec ${config.nix.package}/bin/nix-store --serve --write
        ;;
      *)
        echo "Access only allowed for using the nix remote builder" 1>&2
        exit
    esac
  '';

  known-hosts = let
    known-hosts' = lib.importTOML ../hosts/known-hosts.toml; # TODO: eww
  in
    lib.pipe known-hosts' [
      (lib.flip lib.removeAttrs ["__default__"])
      (lib.mapAttrs (fqdn: host:
        lib.recursiveUpdate (known-hosts'."__default__" or {}) host
      ))
      (lib.mapAttrsToList (fqdn: host:
        let
          allHostnames = [ fqdn ] ++ host.aliases;
        in
        lib.forEach allHostnames (alias:
          lib.nameValuePair
            alias
            (host // {
              aliases = lib.remove alias allHostnames;
              isAlias = fqdn != alias;
            })
        )
      ))
      lib.flatten
      lib.listToAttrs
    ];
  hostNames = attrNames known-hosts;
  thisHost = known-hosts.${config.networking.fqdn};
  thisHostIsBuilder = thisHost.buildMachine.maxJobs > 0;
  thisHostIsBuildee = thisHost.ssh ? userPublicKey;
  thisHostIsHopHost = elem config.networking.fqdn (lib.forEach (attrValues known-hosts) (host: host.ssh.proxyJump or null));

  mkRemoteConfig = fqdn: let
    thatHost = known-hosts.${fqdn};
    thatJump = known-hosts.${thatHost.ssh.proxyJump};
    buildMachine = thatHost.buildMachine // {
      hostName = fqdn;
      sshUser = thatHost.ssh.listenUser;
    };
    thatHostIsBuilder = thatHost.buildMachine.maxJobs > 0;
    thatHostIsBuildee = thatHost.ssh ? userPublicKey && thisHostIsBuilder;
    thatHostIsThis = elem config.networking.fqdn ([ fqdn ] ++ thatHost.aliases);
  in lib.mkIf (!thatHostIsThis) ( lib.mkMerge [
    # out
    (lib.mkIf (thisHostIsBuildee && thatHostIsBuilder) {

      # TODO: Allow setting speedFactor for local builds, as local is currently fixed to 0
      # https://github.com/NixOS/nix/issues/2457

      nix.distributedBuilds = true;
      # useful when the builder has a faster internet connection than i do
      nix.settings.builders-use-substitutes = true;
      nix.buildMachines = lib.mkIf (!thatHost.isAlias) [ buildMachine ];

    })
    # out or jump
    (lib.mkIf (thisHostIsBuildee && thatHost.ssh ? listenPublicKey) {
      programs.ssh.knownHosts.${fqdn}.publicKey = thatHost.ssh.listenPublicKey;
      # TODO: use nix.buildMachines.*.publicHostKey ?

      # timeouts are great when remote is unresponsive. nix doesn't care, lix is way and tests each remote only once
      programs.ssh.extraConfig = ''
        Host ${fqdn}
          ConnectTimeout ${toString thatHost.ssh.connectTimeout}
          Port ${toString thatHost.ssh.listenPort}
          ${lib.optionalString (thatHost.ssh ? proxyJump) ''
            ProxyJump ${thatJump.ssh.listenUser}@${thatHost.ssh.proxyJump}:${toString thatJump.ssh.listenPort}
          ''}
          ${lib.optionalString (thatHost.ssh ? userPrivateKey) ''
            IdentityFile ${thatHost.ssh.userPrivateKey}
          ''}
      '';

      sops.secrets = lib.mkIf (lib.hasPrefix "/run/secrets/" (thatHost.ssh.userPrivateKey or "")) {
        "${lib.removePrefix "/run/secrets/" thatHost.ssh.userPrivateKey}" = { };
      };

    })
    # in
    (lib.mkIf ((thisHostIsBuilder || thisHostIsHopHost) && thatHostIsBuildee && !thatHost.isAlias) {
      # TODO: ensure the user is "nixbld-remote"?
      users.groups.${thisHost.ssh.listenUser} = { };
      users.users.${thisHost.ssh.listenUser} = {
        isSystemUser = lib.mkDefault (!config.users.users.${thisHost.ssh.listenUser}.isNormalUser);
        useDefaultShell = lib.mkDefault true;
        openssh.authorizedKeys.keys = [
          # https://man.archlinux.org/man/core/openssh/sshd.8.en#AUTHORIZED_KEYS_FILE_FORMAT
          # TODO: lib.getExe
          ''restrict,pty,command="${nix-ssh-wrapper}" ${thatHost.ssh.userPublicKey}''
        ];
        group = lib.mkOverride 1499 "${thisHost.ssh.listenUser}"; # mkOptionDefault - 1,
      };
    })
    (lib.mkIf (thisHostIsBuilder && thatHostIsBuildee && !thatHost.isAlias) {
      nix.settings.allowed-users = [ thisHost.ssh.listenUser ];
      nix.settings.trusted-users = [ thisHost.ssh.listenUser ];
    })
  ]);

in {

  imports = lib.forEach hostNames mkRemoteConfig;

}