Add ability to look for IDR in H264 RTP payload. Bump to 0.15.1 (#39)

README.md

@@ -24,7 +24,7 @@ The package can be installed by adding `membrane_rtp_h264_plugin` to your list o
 ```elixir
 def deps do
   [
-    {:membrane_rtp_h264_plugin, "~> 0.15.0"}
+    {:membrane_rtp_h264_plugin, "~> 0.15.1"}
   ]
 end
 ```

lib/rtp_h264/utils.ex

@@ -5,6 +5,11 @@ defmodule Membrane.RTP.H264.Utils do
 
   @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
@@ -13,57 +18,66 @@ defmodule Membrane.RTP.H264.Utils do
   # 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()) :: boolean()
-  def is_keyframe(rtp_payload) when byte_size(rtp_payload) < 1, do: false
+  @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) do
+  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)
+    do_is_keyframe(nalu_type, rest, look_for)
   end
 
   # reserved
-  defp do_is_keyframe(0, _payload), do: false
+  defp do_is_keyframe(0, _payload, _look_for), do: false
 
   # single NALU
-  defp do_is_keyframe(nalu_type, _nalu) when nalu_type in 1..23 do
+  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) do
+  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)
+    check_aggregation_units(24, aus, look_for)
   end
 
   # STAP-B, MTAP16 or MTAP24
-  defp do_is_keyframe(nalu_type, payload)
+  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)
+    check_aggregation_units(nalu_type, aus, look_for)
   end
 
   # FU-A or FU-B
-  defp do_is_keyframe(nalu_type, payload)
+  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
-    s == 1 and type == 7
+
+    case look_for do
+      :sps ->
+        s == 1 and type == 7
+
+      :idr ->
+        s == 1 and type == 5
+    end
   end
 
-  defp do_is_keyframe(_type, _payload), do: false
+  defp do_is_keyframe(_type, _payload, _look_for), do: false
 
-  defp check_aggregation_units(type, aus) do
-    case check_first_aggregation_unit(type, aus) do
+  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)
+      {false, remaining_aus} -> check_aggregation_units(type, remaining_aus, look_for)
       {true, _remaining_aus} -> true
     end
   end
 
-  defp check_first_aggregation_unit(type, aus) when byte_size(aus) >= 2 do
+  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
@@ -81,9 +95,16 @@ defmodule Membrane.RTP.H264.Utils do
 
       <<_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
-      {nalu_type == 7, remaining_aus}
+
+      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), do: {:error, :au_too_short}
+  defp check_first_aggregation_unit(_type, _aus, _look_for), do: {:error, :au_too_short}
 end

mix.exs

@@ -1,7 +1,7 @@
 defmodule Membrane.RTP.H264.MixProject do
   use Mix.Project
 
-  @version "0.15.0"
+  @version "0.15.1"
   @github_url "https://github.com/membraneframework/membrane_rtp_h264_plugin"
 
   def project do
@@ -22,7 +22,7 @@ defmodule Membrane.RTP.H264.MixProject do
       name: "Membrane RTP H264 Plugin",
       source_url: @github_url,
       docs: docs(),
-      homepage_url: "https://membraneframework.org"
+      homepage_url: "https://membrane.stream"
     ]
   end
 
@@ -54,7 +54,7 @@ defmodule Membrane.RTP.H264.MixProject do
       licenses: ["Apache-2.0"],
       links: %{
         "GitHub" => @github_url,
-        "Membrane Framework Homepage" => "https://membraneframework.org"
+        "Membrane Framework Homepage" => "https://membrane.stream"
       }
     ]
   end