defmodule Membrane.RTP.H264.Utils do
@moduledoc """
Utility functions for RTP packets containing H264 encoded frames.
"""
@doc """
Checks whether RTP payload contains H264 keyframe.
By default, with option `look_for` set to `:sps`, will in some cases check
whether the payload contains SPS (NALU payload type 7);
if `look_for` is set to `:idr`, will look exclusively for IDR frames
(NALU payload type 5).
"""
# This is rewritten from galene
# https://github.com/jech/galene/blob/6fbdf0eab2c9640e673d9f9ec0331da24cbf2c4c/codecs/codecs.go#L119
# and based on https://datatracker.ietf.org/doc/html/rfc6184#section-5
#
# it is also unclear why we sometimes check against nalu type == 7
# and sometimes against nalu type == 5 but galene does it this way
# and it works
@spec is_keyframe(binary(), :sps | :idr) :: boolean()
def is_keyframe(rtp_payload, look_for \\ :sps)
def is_keyframe(rtp_payload, _look_for) when byte_size(rtp_payload) < 1, do: false
def is_keyframe(rtp_payload, look_for) do
<<_f::1, _nri::2, nalu_type::5, rest::binary>> = rtp_payload
do_is_keyframe(nalu_type, rest, look_for)
end
# reserved
defp do_is_keyframe(0, _payload, _look_for), do: false
# single NALU
defp do_is_keyframe(nalu_type, _nalu, _look_for) when nalu_type in 1..23 do
nalu_type == 5
end
# STAP-A
defp do_is_keyframe(24, aus, look_for) do
# aus - aggregation units
# those might be single-time or multi-time aggreagtion units
check_aggregation_units(24, aus, look_for)
end
# STAP-B, MTAP16 or MTAP24
defp do_is_keyframe(nalu_type, payload, look_for)
when nalu_type in 25..27 and byte_size(payload) >= 2 do
<<_don::16, aus::binary>> = payload
check_aggregation_units(nalu_type, aus, look_for)
end
# FU-A or FU-B
defp do_is_keyframe(nalu_type, payload, look_for)
when nalu_type in 28..29 and byte_size(payload) >= 1 do
# FU indicator has already been cut off
<<fu_header::binary-size(1), _fu_payload::binary>> = payload
<<s::1, _e::1, _r::1, type::5>> = fu_header
case look_for do
:sps ->
s == 1 and type == 7
:idr ->
s == 1 and type == 5
end
end
defp do_is_keyframe(_type, _payload, _look_for), do: false
defp check_aggregation_units(type, aus, look_for) do
case check_first_aggregation_unit(type, aus, look_for) do
{:error, _reason} -> false
{false, <<>>} -> false
{false, remaining_aus} -> check_aggregation_units(type, remaining_aus, look_for)
{true, _remaining_aus} -> true
end
end
defp check_first_aggregation_unit(type, aus, look_for) when byte_size(aus) >= 2 do
<<nalu_size::16, rest::binary>> = aus
if byte_size(rest) < nalu_size do
{:error, :au_too_short}
else
offset =
case type do
# MTAP16 - skip DOND (8 bits) and TS offset (16 bits)
26 -> 3
# MTAP24 - skip DOND (8 bits) and TS offset (24 bits)
27 -> 4
# STAP-A or STAP-B - nothing to skip
_other -> 0
end
<<_x::binary-size(offset), nalu::binary-size(nalu_size), remaining_aus::binary>> = rest
<<0::1, _nal_ref_idc::2, nalu_type::5, _rest::binary>> = nalu
case look_for do
:sps ->
{nalu_type == 7, remaining_aus}
:idr ->
{nalu_type == 5, remaining_aus}
end
end
end
defp check_first_aggregation_unit(_type, _aus, _look_for), do: {:error, :au_too_short}
end