AV1 payloader

lib/ex_webrtc/rtp/av1/leb128.ex

@@ -0,0 +1,36 @@
+defmodule ExWebRTC.RTP.AV1.LEB128 do
+  @moduledoc false
+  # Utilities for handling unsigned Little Endian Base 128 integers
+
+  import Bitwise
+
+  # see https://chromium.googlesource.com/external/webrtc/+/HEAD/modules/rtp_rtcp/source/rtp_packetizer_av1.cc#61
+  @spec encode(non_neg_integer(), [bitstring()]) :: binary()
+  def encode(value, acc \\ [])
+
+  def encode(value, acc) when value < 0x80 do
+    for group <- Enum.reverse([value | acc]), into: <<>> do
+      <<group>>
+    end
+  end
+
+  def encode(value, acc) do
+    group = 0x80 ||| (value &&& 0x7F)
+    encode(value >>> 7, [group | acc])
+  end
+
+  # see https://source.chromium.org/chromium/chromium/src/+/main:third_party/webrtc/rtc_base/byte_buffer.cc;drc=8e78783dc1f7007bad46d657c9f332614e240fd8;l=107
+  @spec read(binary(), non_neg_integer(), non_neg_integer(), non_neg_integer()) ::
+          {:ok, pos_integer(), non_neg_integer()} | {:error, :invalid_leb128_data}
+  def read(data, read_bits \\ 0, leb128_size \\ 0, value \\ 0)
+
+  def read(<<0::1, group::7, _rest::binary>>, read_bits, leb128_size, value) do
+    {:ok, leb128_size + 1, value ||| group <<< read_bits}
+  end
+
+  def read(<<1::1, group::7, rest::binary>>, read_bits, leb128_size, value) do
+    read(rest, read_bits + 7, leb128_size + 1, value ||| group <<< read_bits)
+  end
+
+  def read(_, _, _, _), do: {:error, :invalid_leb128_data}
+end

lib/ex_webrtc/rtp/av1/obu.ex

@@ -0,0 +1,144 @@
+defmodule ExWebRTC.RTP.AV1.OBU do
+  @moduledoc false
+  # Defines the Open Bitstream Unit, the base packetization unit of all structures present in the AV1 bitstream.
+  #
+  # Based on [the AV1 spec](https://aomediacodec.github.io/av1-spec/av1-spec.pdf).
+  #
+  #  OBU syntax:
+  #      0 1 2 3 4 5 6 7
+  #     +-+-+-+-+-+-+-+-+
+  #     |0| type  |X|S|-| (REQUIRED)
+  #     +-+-+-+-+-+-+-+-+
+  #  X: | TID |SID|-|-|-| (OPTIONAL)
+  #     +-+-+-+-+-+-+-+-+
+  #     |1|             |
+  #     +-+ OBU payload |
+  #  S: |1|             | (OPTIONAL, variable length leb128 encoded)
+  #     +-+    size     |
+  #     |0|             |
+  #     +-+-+-+-+-+-+-+-+
+  #     |  OBU payload  |
+  #     |     ...       |
+
+  alias ExWebRTC.RTP.AV1.LEB128
+
+  @obu_sequence_header 1
+  @obu_temporal_delimiter 2
+  @obu_padding 15
+
+  @type t :: %__MODULE__{
+          type: 0..15,
+          x: 0 | 1,
+          s: 0 | 1,
+          tid: 0..7 | nil,
+          sid: 0..3 | nil,
+          payload: binary()
+        }
+
+  @enforce_keys [:type, :x, :s, :payload]
+  defstruct @enforce_keys ++ [:tid, :sid]
+
+  @doc """
+  Parses the low overhead bitstream format defined in AV1 spec section 5.2.
+  On success, returns the parsed OBU as well as the remainder of the AV1 bitstream.
+  """
+  @spec parse(binary()) :: {:ok, t(), binary()} | {:error, :invalid_av1_bitstream}
+  def parse(av1_bitstream_binary)
+
+  def parse(<<0::1, type::4, x::1, s::1, 0::1, rest::binary>>) do
+    with {:ok, tid, sid, rest} <- parse_extension_header(x, rest),
+         {:ok, payload, rest} <- parse_payload(s, rest),
+         :ok <- validate_payload(type, payload) do
+      {:ok,
+       %__MODULE__{
+         type: type,
+         x: x,
+         s: s,
+         tid: tid,
+         sid: sid,
+         payload: payload
+       }, rest}
+    else
+      {:error, _} = err -> err
+    end
+  end
+
+  def parse(_), do: {:error, :invalid_av1_bitstream}
+
+  defp parse_extension_header(0, rest), do: {:ok, nil, nil, rest}
+
+  defp parse_extension_header(1, <<tid::3, sid::2, 0::3, rest::binary>>),
+    do: {:ok, tid, sid, rest}
+
+  defp parse_extension_header(_, _), do: {:error, :invalid_av1_bitstream}
+
+  defp parse_payload(0, rest), do: {:ok, rest, <<>>}
+
+  defp parse_payload(1, rest) do
+    with {:ok, leb128_size, payload_size} <- LEB128.read(rest),
+         <<_::binary-size(leb128_size), payload::binary-size(payload_size), rest::binary>> <- rest do
+      {:ok, payload, rest}
+    else
+      _ -> {:error, :invalid_av1_bitstream}
+    end
+  end
+
+  defp validate_payload(@obu_padding, _), do: :ok
+  defp validate_payload(@obu_temporal_delimiter, <<>>), do: :ok
+  defp validate_payload(type, data) when type != @obu_temporal_delimiter and data != <<>>, do: :ok
+  defp validate_payload(_, _), do: {:error, :invalid_av1_bitstream}
+
+  @spec serialize(t()) :: binary()
+  def serialize(%__MODULE__{type: type, x: x, s: s, payload: payload} = obu) do
+    obu_binary =
+      <<0::1, type::4, x::1, s::1, 0::1>>
+      |> add_extension_header(obu)
+      |> add_payload_size(obu)
+
+    <<obu_binary::binary, payload::binary>>
+  end
+
+  defp add_extension_header(obu_binary, %__MODULE__{x: 0, tid: nil, sid: nil}), do: obu_binary
+
+  defp add_extension_header(obu_binary, %__MODULE__{x: 1, tid: tid, sid: sid})
+       when tid != nil and sid != nil do
+    <<obu_binary::binary, tid::3, sid::2, 0::3>>
+  end
+
+  defp add_extension_header(_obu_binary, _invalid_obu),
+    do: raise("AV1 TID and SID must be set if, and only if X bit is set")
+
+  defp add_payload_size(obu_binary, %__MODULE__{s: 0}), do: obu_binary
+
+  defp add_payload_size(obu_binary, %__MODULE__{s: 1, payload: payload}) do
+    payload_size = payload |> byte_size() |> LEB128.encode()
+    <<obu_binary::binary, payload_size::binary>>
+  end
+
+  @doc """
+  Rewrites a specific case of the sequence header OBU to disable OBU dropping in the AV1 decoder
+  in accordance with av1-rtp-spec sec. 5. Leaves other OBUs unchanged.
+  """
+  @spec disable_dropping_in_decoder_if_applicable(t()) :: t()
+  def disable_dropping_in_decoder_if_applicable(obu)
+
+  # We're handling the following case:
+  #   - still_picture = 0
+  #   - reduced_still_picture_header = 0
+  #   - timing_info_present_flag = 0
+  #   - operating_points_cnt_minus_1 = 0
+  #   - seq_level_idx[0] = 0
+  # and setting operating_point_idc[0] = 0xFFF
+  #
+  # For the sequence header OBU syntax, refer to the AV1 spec sec. 5.5.
+  def disable_dropping_in_decoder_if_applicable(
+        %__MODULE__{
+          type: @obu_sequence_header,
+          payload: <<seq_profile::3, 0::3, iddpf::1, 0::5, _op_idc_0::12, 0::5, rest::bitstring>>
+        } = obu
+      ) do
+    %{obu | payload: <<seq_profile::3, 0::3, iddpf::1, 0::5, 0xFFF::12, 0::5, rest::bitstring>>}
+  end
+
+  def disable_dropping_in_decoder_if_applicable(obu), do: obu
+end

lib/ex_webrtc/rtp/av1/payload.ex

@@ -0,0 +1,96 @@
+defmodule ExWebRTC.RTP.AV1.Payload do
+  @moduledoc false
+  # Defines AV1 payload structure stored in RTP packet payload.
+  #
+  # Based on [RTP Payload Format for AV1](https://aomediacodec.github.io/av1-rtp-spec/v1.0.0.html).
+  #
+  #  RTP payload syntax:
+  #      0 1 2 3 4 5 6 7
+  #     +-+-+-+-+-+-+-+-+
+  #     |Z|Y| W |N|-|-|-| (REQUIRED)
+  #     +=+=+=+=+=+=+=+=+ (REPEATED W-1 times, or any times if W = 0)
+  #     |1|             |
+  #     +-+ OBU fragment|
+  #     |1|             | (REQUIRED, leb128 encoded)
+  #     +-+    size     |
+  #     |0|             |
+  #     +-+-+-+-+-+-+-+-+
+  #     |  OBU fragment |
+  #     |     ...       |
+  #     +=+=+=+=+=+=+=+=+
+  #     |     ...       |
+  #     +=+=+=+=+=+=+=+=+ if W > 0, last fragment MUST NOT have size field
+  #     |  OBU fragment |
+  #     |     ...       |
+  #     +=+=+=+=+=+=+=+=+
+
+  @type t :: %__MODULE__{
+          z: 0 | 1,
+          y: 0 | 1,
+          w: 0 | 1 | 2 | 3,
+          n: 0 | 1,
+          payload: binary()
+        }
+
+  @enforce_keys [:z, :y, :w, :n, :payload]
+  defstruct @enforce_keys ++ []
+
+  @doc """
+  Parses RTP payload as AV1 payload.
+  """
+  @spec parse(binary()) :: {:ok, t()} | {:error, :invalid_packet}
+  def parse(rtp_payload)
+
+  def parse(<<z::1, y::1, w::2, n::1, 0::3, payload::binary>>) do
+    if payload == <<>> do
+      {:error, :invalid_packet}
+    else
+      {:ok,
+       %__MODULE__{
+         z: z,
+         y: y,
+         w: w,
+         n: n,
+         payload: payload
+       }}
+    end
+  end
+
+  def parse(_), do: {:error, :invalid_packet}
+
+  @spec serialize(t()) :: binary()
+  def serialize(%__MODULE__{
+        z: z,
+        y: y,
+        w: w,
+        n: n,
+        payload: payload
+      }) do
+    <<z::1, y::1, w::2, n::1, 0::3, payload::binary>>
+  end
+
+  @doc """
+  Payloads chunked fragments of single OBU and sets Z, Y bits.
+  """
+  @spec payload_obu_fragments([binary()], 0 | 1) :: [t()]
+  def payload_obu_fragments(obu_fragments, n_bit \\ 0)
+
+  def payload_obu_fragments([entire_obu], n_bit) do
+    [%__MODULE__{z: 0, y: 0, w: 1, n: n_bit, payload: entire_obu}]
+  end
+
+  def payload_obu_fragments([first_obu_fragment | next_obu_fragments], n_bit) do
+    # First fragment of OBU: set Y bit
+    first_obu_payload = %__MODULE__{z: 0, y: 1, w: 1, n: n_bit, payload: first_obu_fragment}
+
+    next_obu_payloads =
+      next_obu_fragments
+      # Middle fragments of OBU: set Z, Y bits
+      #   av1-rtp-spec sec. 4.4: N is set only for the first packet of the coded video sequence.
+      |> Enum.map(&%__MODULE__{z: 1, y: 1, w: 1, n: 0, payload: &1})
+      # Last fragment of OBU: set Z bit only (unset Y)
+      |> List.update_at(-1, &%{&1 | y: 0})
+
+    [first_obu_payload | next_obu_payloads]
+  end
+end

lib/ex_webrtc/rtp/av1/payloader.ex

@@ -0,0 +1,78 @@
+defmodule ExWebRTC.RTP.Payloader.AV1 do
+  @moduledoc false
+  # Encapsulates AV1 video temporal units into RTP packets.
+  #
+  # Resources:
+  # * [RTP Payload Format for AV1 (av1-rtp-spec)](https://aomediacodec.github.io/av1-rtp-spec/v1.0.0.html)
+  # * [AV1 spec](https://aomediacodec.github.io/av1-spec/av1-spec.pdf).
+  # * https://norkin.org/research/av1_decoder_model/index.html
+  # * https://chromium.googlesource.com/external/webrtc/+/HEAD/modules/rtp_rtcp/source/video_rtp_depacketizer_av1.cc
+
+  @behaviour ExWebRTC.RTP.Payloader.Behaviour
+
+  alias ExWebRTC.RTP.AV1.{OBU, Payload}
+  alias ExWebRTC.Utils
+
+  @obu_sequence_header 1
+  @obu_temporal_delimiter 2
+
+  @aggregation_header_size_bytes 1
+
+  @type t :: %__MODULE__{
+          max_payload_size: non_neg_integer()
+        }
+
+  @enforce_keys [:max_payload_size]
+  defstruct @enforce_keys
+
+  @impl true
+  def new(max_payload_size) when max_payload_size > 100 do
+    %__MODULE__{max_payload_size: max_payload_size}
+  end
+
+  @impl true
+  def payload(payloader, temporal_unit) when temporal_unit != <<>> do
+    # In AV1, a temporal unit consists of all OBUs associated with a specific time instant.
+    # Temporal units always start with a temporal delimiter OBU. They may contain multiple AV1 frames.
+    #   av1-rtp-spec sec. 5: The temporal delimiter OBU should be removed when transmitting.
+    obus =
+      case parse_obus(temporal_unit) do
+        [%OBU{type: @obu_temporal_delimiter} | next_obus] ->
+          next_obus
+
+        _ ->
+          raise "Invalid AV1 temporal unit: does not start with temporal delimiter OBU"
+      end
+
+    # With the current implementation, each RTP packet will contain one OBU element.
+    # This element can be an entire OBU, or a fragment of an OBU bigger than max_payload_size.
+    rtp_packets =
+      Stream.flat_map(obus, fn obu ->
+        n_bit = Utils.to_int(obu.type == @obu_sequence_header)
+
+        obu
+        |> OBU.disable_dropping_in_decoder_if_applicable()
+        |> OBU.serialize()
+        |> Utils.chunk(payloader.max_payload_size - @aggregation_header_size_bytes)
+        |> Payload.payload_obu_fragments(n_bit)
+      end)
+      |> Stream.map(&Payload.serialize/1)
+      |> Enum.map(&ExRTP.Packet.new/1)
+      |> List.update_at(-1, &%{&1 | marker: true})
+
+    {rtp_packets, payloader}
+  end
+
+  defp parse_obus(data, obus \\ [])
+  defp parse_obus(<<>>, obus), do: Enum.reverse(obus)
+
+  defp parse_obus(data, obus) do
+    case OBU.parse(data) do
+      {:ok, obu, rest} ->
+        parse_obus(rest, [obu | obus])
+
+      {:error, :invalid_av1_bitstream} ->
+        raise "Invalid AV1 bitstream: unable to parse OBU"
+    end
+  end
+end

lib/ex_webrtc/rtp/payloader.ex

@@ -39,6 +39,7 @@ defmodule ExWebRTC.RTP.Payloader do
   defp to_payloader_module(mime_type) do
     case String.downcase(mime_type) do
       "video/vp8" -> {:ok, ExWebRTC.RTP.Payloader.VP8}
+      "video/av1" -> {:ok, ExWebRTC.RTP.Payloader.AV1}
       "audio/opus" -> {:ok, ExWebRTC.RTP.Payloader.Opus}
       "audio/pcma" -> {:ok, ExWebRTC.RTP.Payloader.G711}
       "audio/pcmu" -> {:ok, ExWebRTC.RTP.Payloader.G711}