mdns_client.ex

defmodule MDNSClient do
  defmodule Device do
    defstruct ip: nil,
              services: [],
              domain: nil,
              payload: %{}
  end

  use GenServer
  @mdns_group {224, 0, 0, 251}
  @port Application.get_env(:mdns, :port, 5353)

  @query_packet %DNS.Record{
    header: %DNS.Header{},
    qdlist: []
  }

  defmodule State do
    defstruct devices: %{},
              udp: nil,
              handlers: [],
              queries: [],
              from: nil
  end

  @impl true
  def query(namespace) do
    GenServer.call(__MODULE__, {:query, namespace})
  end

  def start_link(args) do
    GenServer.start_link(__MODULE__, args, name: __MODULE__)
  end

  @impl GenServer
  def init(_args) do
    send(self(), :connect)
    {:ok, %State{}}
  end

  @impl GenServer
  def handle_call({:query, namespace}, from, state) do
    packet = %DNS.Record{
      @query_packet
      | :qdlist => [
          %DNS.Query{domain: to_charlist(namespace), type: :ptr, class: :in}
        ]
    }

    p = DNS.Record.encode(packet)
    :ok = :gen_udp.send(state.udp, @mdns_group, @port, p)

    {:noreply,
     %State{state | queries: Enum.uniq([namespace | state.queries]), from: {from, namespace}}}
  end

  @impl GenServer
  def handle_info(:connect, state) do
    udp_options = [
      :binary,
      broadcast: true,
      active: true,
      ip: {0, 0, 0, 0},
      ifaddr: {0, 0, 0, 0},
      add_membership: {@mdns_group, {0, 0, 0, 0}},
      multicast_if: {0, 0, 0, 0},
      multicast_loop: true,
      multicast_ttl: 32,
      reuseaddr: true
    ]

    case :gen_udp.open(@port, udp_options) do
      {:ok, udp} ->
        {:noreply, %State{state | udp: udp}}

      error ->
        {:stop, error, state}
    end
  end

  def handle_info({:udp, _socket, ip, _port, packet}, state) do
    {:noreply, handle_packet(ip, packet, state)}
  end

  defp handle_packet(ip, packet, state) do
    record = DNS.Record.decode(packet)

    case record.header.qr do
      true ->
        handle_response(ip, record, state)

      _ ->
        state
    end
  end

  defp handle_response(ip, record, state) do
    device = get_device(ip, record, state)

    devices =
      Enum.reduce(state.queries, %{:other => []}, fn query, acc ->
        cond do
          Enum.any?(device.services, fn service -> String.ends_with?(service, query) end) ->
            {_namespace, devices} = create_namespace_devices(query, device, acc, state)
            devices

          true ->
            Map.merge(acc, state.devices)
        end
      end)

    state = %State{state | :devices => devices}

    case state.from do
      {from, namespace_query} ->
        GenServer.reply(from, {:ok, devices[String.to_atom(namespace_query)]})
        %{state | from: nil}

      _ ->
        state
    end
  end

  defp handle_device(%DNS.Resource{:type => :ptr} = record, device) do
    %Device{
      device
      | :services =>
          Enum.uniq([to_string(record.data), to_string(record.domain)] ++ device.services)
    }
  end

  defp handle_device(%DNS.Resource{:type => :a} = record, device) do
    %Device{device | :domain => to_string(record.domain)}
  end

  defp handle_device({:dns_rr, _d, :txt, _id, _, _, data, _, _, _}, device) do
    %Device{
      device
      | :payload =>
          Enum.reduce(data, %{}, fn kv, acc ->
            case String.split(to_string(kv), "=", parts: 2) do
              [k, v] -> Map.put(acc, String.downcase(k), String.trim(v))
              _ -> nil
            end
          end)
    }
  end

  defp handle_device(%DNS.Resource{}, device) do
    device
  end

  defp handle_device({:dns_rr, _, _, _, _, _, _, _, _, _}, device) do
    device
  end

  defp handle_device({:dns_rr_opt, _, _, _, _, _, _, _}, device) do
    device
  end

  defp get_device(ip, record, state) do
    orig_device =
      Enum.concat(Map.values(state.devices))
      |> Enum.find(%Device{:ip => ip}, fn device ->
        device.ip == ip
      end)

    Enum.reduce(record.anlist ++ record.arlist, orig_device, fn r, acc ->
      handle_device(r, acc)
    end)
  end

  defp create_namespace_devices(query, device, devices, state) do
    namespace = String.to_atom(query)

    {namespace,
     cond do
       Enum.any?(Map.get(state.devices, namespace, []), fn dev -> dev.ip == device.ip end) ->
         Map.merge(devices, %{namespace => merge_device(device, namespace, state)})

       true ->
         Map.merge(devices, %{namespace => [device | Map.get(state.devices, namespace, [])]})
     end}
  end

  defp merge_device(device, namespace, state) do
    Enum.map(Map.get(state.devices, namespace, []), fn d ->
      cond do
        device.ip == d.ip -> Map.merge(d, device)
        true -> d
      end
    end)
  end
end