prop/lib/microwaveprop/pskr/mqtt.ex

defmodule Microwaveprop.Pskr.Mqtt do
  @moduledoc """
  Minimal MQTT 3.1.1 codec — just the packet types this project
  actually exchanges with `mqtt.pskreporter.info:1883`.

  We're a subscribe-only QoS 0 client, so the wire protocol reduces
  to:

    * outbound: CONNECT, SUBSCRIBE, PINGREQ, DISCONNECT
    * inbound:  CONNACK, PUBLISH (QoS 0), SUBACK, PINGRESP

  No retained messages, no will, no QoS 1/2 inflight tracking, no
  packet ID reuse logic. Anything beyond that is intentionally not
  implemented — adding a dep just to throw away 80% of its surface
  area felt worse than the ~150 lines below.

  ## Frame layout

  Every MQTT packet is `<<type::4, flags::4, remaining_length::vbi,
  body::binary>>` where `vbi` is a variable-byte integer (1-4 bytes,
  high bit = "more bytes follow"). `parse/1` splits a TCP buffer
  into the next decoded packet plus any trailing bytes; on a short
  read it returns `{:incomplete, buffer}` so the caller can
  re-invoke after the next chunk arrives.

  See: https://docs.oasis-open.org/mqtt/mqtt/v3.1.1/os/mqtt-v3.1.1-os.html
  """

  import Bitwise

  @type packet ::
          {:connack, return_code :: 0..5, session_present :: boolean()}
          | {:publish, topic :: String.t(), payload :: binary()}
          | {:suback, packet_id :: non_neg_integer(), qos_results :: [0..2 | 0x80]}
          | :pingresp

  # ---------- Outbound builders ----------

  @doc """
  CONNECT packet for an MQTT 3.1.1 client. Flags are hardcoded to
  `clean_session=1` and no will / no auth — matches the public
  pskreporter broker which doesn't require credentials.
  """
  @spec connect(String.t(), pos_integer()) :: binary()
  def connect(client_id, keepalive_seconds) when is_binary(client_id) do
    body = <<
      # Variable header
      4::16,
      "MQTT",
      # Protocol level: 3.1.1
      0x04,
      # Connect flags: clean_session
      0x02,
      keepalive_seconds::16,
      # Payload: client ID
      byte_size(client_id)::16,
      client_id::binary
    >>

    fixed_header(0x10, body)
  end

  @doc """
  SUBSCRIBE packet. Topics is a list of `{topic_filter, qos}` pairs.
  Returns `{packet_binary, packet_id}` — the caller tracks the
  packet ID for matching the SUBACK on the wire.
  """
  @spec subscribe(non_neg_integer(), [{String.t(), 0..2}]) :: binary()
  def subscribe(packet_id, topics) when is_list(topics) and topics != [] do
    body =
      IO.iodata_to_binary([
        <<packet_id::16>> | Enum.map(topics, fn {filter, qos} -> <<byte_size(filter)::16, filter::binary, qos>> end)
      ])

    # Bits 3-0 of the SUBSCRIBE first byte are reserved as 0010.
    fixed_header(0x82, body)
  end

  @doc "PINGREQ keep-alive packet."
  @spec pingreq() :: binary()
  def pingreq, do: <<0xC0, 0x00>>

  @doc "DISCONNECT packet for clean shutdown."
  @spec disconnect() :: binary()
  def disconnect, do: <<0xE0, 0x00>>

  # ---------- Inbound parser ----------

  @doc """
  Parse the next packet from a buffer of received TCP bytes.

    * `{:ok, packet, rest}` — packet decoded, `rest` is the
      unconsumed tail (may contain more packets).
    * `{:incomplete, buffer}` — short read; caller should append
      the next TCP chunk and call again.
    * `{:error, reason}` — protocol violation; caller should drop
      the connection.
  """
  @spec parse(binary()) :: {:ok, packet(), binary()} | {:incomplete, binary()} | {:error, term()}
  def parse(<<>>), do: {:incomplete, <<>>}

  def parse(<<type_flags::8, rest::binary>> = buffer) do
    case decode_varint(rest) do
      {:ok, len, after_len} when byte_size(after_len) >= len ->
        <<body::binary-size(^len), tail::binary>> = after_len
        decode_packet(type_flags, body, tail)

      {:ok, _len, _} ->
        {:incomplete, buffer}

      :incomplete ->
        {:incomplete, buffer}

      {:error, _} = err ->
        err
    end
  end

  defp decode_packet(0x20, <<flags, return_code, _::binary>>, tail) do
    {:ok, {:connack, return_code, (flags &&& 1) == 1}, tail}
  end

  # PUBLISH: top nibble is 3, low nibble is DUP/QoS/RETAIN. We only
  # subscribe at QoS 0, so the broker won't send us QoS 1/2 — but
  # we still tolerate the QoS 0 form (low nibble all zero or with
  # RETAIN set). Anything else is rejected so a misconfigured
  # subscription surfaces loudly instead of silently dropping bytes.
  defp decode_packet(type_flags, body, tail) when (type_flags &&& 0xF0) == 0x30 do
    qos = type_flags >>> 1 &&& 0b11

    case qos do
      0 ->
        <<topic_len::16, topic::binary-size(topic_len), payload::binary>> = body
        {:ok, {:publish, topic, payload}, tail}

      _ ->
        {:error, {:unsupported_qos, qos}}
    end
  end

  defp decode_packet(0x90, <<packet_id::16, qos_results::binary>>, tail) do
    {:ok, {:suback, packet_id, :erlang.binary_to_list(qos_results)}, tail}
  end

  defp decode_packet(0xD0, <<>>, tail), do: {:ok, :pingresp, tail}

  defp decode_packet(type_flags, _body, _tail) do
    {:error, {:unhandled_packet_type, type_flags}}
  end

  # ---------- Variable-byte integer ----------

  @doc false
  @spec encode_varint(non_neg_integer()) :: binary()
  def encode_varint(n) do
    n
    |> encode_varint_io([])
    |> IO.iodata_to_binary()
  end

  defp encode_varint_io(n, acc) when n <= 127, do: [acc, <<n>>]

  defp encode_varint_io(n, acc) when n <= 268_435_455 do
    encode_varint_io(n >>> 7, [acc, <<(n &&& 0x7F) ||| 0x80>>])
  end

  @doc false
  @spec decode_varint(binary()) :: {:ok, non_neg_integer(), binary()} | :incomplete | {:error, term()}
  def decode_varint(buffer), do: decode_varint(buffer, 0, 0)

  defp decode_varint(_buffer, _shift, count) when count >= 4, do: {:error, :varint_too_long}
  defp decode_varint(<<>>, _shift, _count), do: :incomplete

  defp decode_varint(<<0::1, n::7, rest::binary>>, shift, _count) do
    {:ok, n <<< shift, rest}
  end

  defp decode_varint(<<1::1, n::7, rest::binary>>, shift, count) do
    case decode_varint(rest, shift + 7, count + 1) do
      {:ok, tail, remaining} -> {:ok, n <<< shift ||| tail, remaining}
      other -> other
    end
  end

  defp fixed_header(byte1, body) do
    <<byte1>> <> encode_varint(byte_size(body)) <> body
  end
end