desktop/tailscale.nix

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

with lib;

let
  cfg = config.services.tailscale-netns;

  instanceOptions = {
    options.extraDaemonFlags = mkOption {
      type = types.listOf types.str;
      default = [ ];
      description = "Extra flags passed to `tailscaled`.";
    };
  };

  mkNetnsSetup =
    index: name: instanceCfg:
    let
      ns = "ts-${name}ns";
      hostAddr = "10.201.${toString index}.1";
      hostIP = "${hostAddr}/24";
      nsIP = "10.201.${toString index}.2/24";
      vethHost = "veth-ts-${name}";
      vethNS = "veth-ts-${name}-ns";
      fallbackLines = concatMapStrings (n: "  echo nameserver ${n}\n") config.networking.nameservers;
    in
    {
      "netns@ts-${name}" = {
        description = "Network namespace for Tailscale instance ${name}";
        before = [ "network.target" ];
        serviceConfig = {
          Type = "oneshot";
          RemainAfterExit = true;
          ExecStart = pkgs.writers.writeBash "ts-${name}-netns-up" ''
            set -euo pipefail
            ${pkgs.coreutils}/bin/mkdir -p /etc/netns/${ns}

            {
              dns_found=0
              if ${pkgs.systemd}/bin/resolvectl dns 2>/dev/null; then
                saved_ifs="$IFS"
                IFS=$'\n'
                for line in $(${pkgs.systemd}/bin/resolvectl dns 2>/dev/null); do
                  case "$line" in
                    *:*)
                      server=$(echo "$line" | ${pkgs.gnused}/bin/sed 's/^[^:]*: //')
                      case "$server" in
                        127.*|::1|"") ;;
                        *) echo "nameserver $server"; dns_found=1 ;;
                      esac
                      ;;
                  esac
                done
                IFS="$saved_ifs"
              fi

              if [ "$dns_found" -eq 0 ]; then
                while IFS= read -r rline; do
                  case "$rline" in
                    nameserver*)
                      ip=$(echo "$rline" | ${pkgs.gnused}/bin/sed 's/^nameserver //')
                      case "$ip" in
                        127.*|::1) ;;
                        *) echo "$rline"; dns_found=1 ;;
                      esac
                      ;;
                  esac
                done < /etc/resolv.conf
              fi
              ${optionalString (fallbackLines != "") ''
                if [ "$dns_found" -eq 0 ]; then
              ${fallbackLines}    fi
              ''}
            } > /etc/netns/${ns}/resolv.conf

            ${pkgs.iproute2}/bin/ip netns add ${ns}
            ${pkgs.iproute2}/bin/ip link add ${vethHost} type veth peer name ${vethNS}
            ${pkgs.iproute2}/bin/ip link set ${vethNS} netns ${ns}

            ${pkgs.iproute2}/bin/ip addr add ${hostIP} dev ${vethHost}
            ${pkgs.iproute2}/bin/ip link set ${vethHost} up

            ${pkgs.iproute2}/bin/ip -n ${ns} addr add ${nsIP} dev ${vethNS}
            ${pkgs.iproute2}/bin/ip -n ${ns} link set ${vethNS} up
            ${pkgs.iproute2}/bin/ip -n ${ns} link set lo up
            ${pkgs.iproute2}/bin/ip -n ${ns} route add default via ${hostAddr}

            ${pkgs.iptables}/bin/iptables -t nat -C POSTROUTING -s ${nsIP} -j MASQUERADE 2>/dev/null \
              || ${pkgs.iptables}/bin/iptables -t nat -A POSTROUTING -s ${nsIP} -j MASQUERADE
          '';
          ExecStop = pkgs.writers.writeBash "ts-${name}-netns-down" ''
            set -euo pipefail
            ${pkgs.iptables}/bin/iptables -t nat -D POSTROUTING -s ${nsIP} -j MASQUERADE 2>/dev/null || true
            ${pkgs.iproute2}/bin/ip link del ${vethHost} 2>/dev/null || true
            ${pkgs.iproute2}/bin/ip netns del ${ns} 2>/dev/null || true
          '';
        };
      };

      "tailscaled-${name}" = {
        description = "Tailscale daemon for instance ${name}";
        bindsTo = [ "netns@ts-${name}.service" ];
        after = [ "netns@ts-${name}.service" ];
        wantedBy = [ "multi-user.target" ];
        serviceConfig = {
          ExecStart = pkgs.writers.writeBash "tailscaled-${name}-start" ''
            exec ${pkgs.iproute2}/bin/ip netns exec ${ns} \
              ${pkgs.tailscale}/bin/tailscaled \
              --statedir=/var/lib/tailscale-${name} \
              --socket=/run/tailscale-${name}/tailscaled.sock \
              --tun=ts-${name} \
              --port=41641 \
              ${escapeShellArgs instanceCfg.extraDaemonFlags}
          '';
          ExecStopPost = pkgs.writers.writeBash "tailscaled-${name}-stop" ''
            ${pkgs.iproute2}/bin/ip netns exec ${ns} \
              ${pkgs.tailscale}/bin/tailscaled --cleanup \
              --statedir=/var/lib/tailscale-${name} \
              --socket=/run/tailscale-${name}/tailscaled.sock \
              2>/dev/null || true
          '';
          Restart = "on-failure";
          RuntimeDirectory = "tailscale-${name}";
          RuntimeDirectoryMode = "0755";
          StateDirectory = "tailscale-${name}";
          StateDirectoryMode = "0700";
          CacheDirectory = "tailscale-${name}";
          CacheDirectoryMode = "0750";

          AmbientCapabilities = "CAP_NET_ADMIN CAP_NET_RAW";
          LimitNOFILE = "infinity";
        };
      };
    };

in
{
  options.services.tailscale-netns = {
    enable = mkEnableOption "Tailscale instances in isolated network namespaces";

    instances = mkOption {
      type = types.attrsOf (types.submodule instanceOptions);
      default = { };
      description = "Tailscale instances configured in separate network namespaces.";
      example = literalExpression ''
        {
          work = {};
        }
      '';
    };
  };

  config = mkIf cfg.enable {
    systemd.services = listToAttrs (
      flatten (
        imap1 (
          index: elem:
          let
            name = elem.name;
            units = mkNetnsSetup index name elem.value;
          in
          [
            { name = "netns@ts-${name}"; value = units."netns@ts-${name}"; }
            { name = "tailscaled-${name}"; value = units."tailscaled-${name}"; }
          ]
        ) (attrsToList cfg.instances)
      )
    );
  };
}