Add support for SRTP in the muxer and demuxer (#177)

* SRTP working

Author: Noarkhh

lib/membrane/rtp/demuxer.ex

@@ -105,6 +105,14 @@ defmodule Membrane.RTP.Demuxer do
                 This option determines the action to be taken after a stream has been announced with a 
                 `t:new_rtp_stream_notification/0` notification but the corresponding pad has not been connected within the specified timeout period.
                 """
+              ],
+              srtp: [
+                spec: false | [ExLibSRTP.Policy.t()],
+                default: false,
+                description: """
+                Specifies whether to use SRTP to decrypt the input streams. Requires adding [srtp](https://github.com/membraneframework/elixir_libsrtp) 
+                dependency to work. If true takes a list of SRTP policies to use for decrypting packets. See `t:ExLibSRTP.Policy.t/0` for details.
+                """
               ]
 
   defmodule State do
@@ -138,17 +146,38 @@ defmodule Membrane.RTP.Demuxer do
     @type t :: %__MODULE__{
             payload_type_mapping: RTP.PayloadFormat.payload_type_mapping(),
             not_linked_pad_handling: %{action: :raise | :ignore, timeout: Membrane.Time.t()},
+            srtp: ExLibSRTP.t() | nil,
             stream_states: %{RTP.ssrc() => StreamState.t()},
             pads_waiting_for_stream: %{Pad.ref() => Membrane.RTP.Demuxer.stream_id()}
           }
 
-    @enforce_keys [:not_linked_pad_handling, :payload_type_mapping]
+    @enforce_keys [:not_linked_pad_handling, :payload_type_mapping, :srtp]
     defstruct @enforce_keys ++ [stream_states: %{}, pads_waiting_for_stream: %{}]
   end
 
   @impl true
   def handle_init(_ctx, opts) do
-    {[], struct(State, Map.from_struct(opts))}
+    srtp =
+      case opts.srtp do
+        false ->
+          nil
+
+        policies ->
+          if not Code.ensure_loaded?(ExLibSRTP) do
+            raise "Optional dependency :ex_libsrtp is required for SRTP"
+          end
+
+          srtp = apply(ExLibSRTP, :new, [])
+          Enum.each(policies, &apply(ExLibSRTP, :add_stream, [srtp, &1]))
+          srtp
+      end
+
+    opts =
+      opts
+      |> Map.from_struct()
+      |> Map.put(:srtp, srtp)
+
+    {[], struct(State, opts)}
   end
 
   @impl true
@@ -290,7 +319,35 @@ defmodule Membrane.RTP.Demuxer do
 
   @spec handle_rtp_packet(binary(), CallbackContext.t(), State.t()) ::
           {[Membrane.Element.Action.t()], State.t()}
-  defp handle_rtp_packet(raw_rtp_packet, ctx, state) do
+  defp handle_rtp_packet(packet, ctx, state) do
+    case unprotect_packet(packet, state.srtp) do
+      nil -> {[], state}
+      raw_packet -> handle_raw_rtp_packet(raw_packet, ctx, state)
+    end
+  end
+
+  @spec unprotect_packet(binary(), ExLibSRTP.t() | nil) :: binary() | nil
+  defp unprotect_packet(rtp_packet, nil) do
+    rtp_packet
+  end
+
+  defp unprotect_packet(protected_rtp_packet, srtp) do
+    case apply(ExLibSRTP, :unprotect, [srtp, protected_rtp_packet]) do
+      {:ok, raw_rtp_packet} ->
+        raw_rtp_packet
+
+      {:error, reason} when reason in [:replay_fail, :replay_old] ->
+        Membrane.Logger.warning("Ignoring packet due to `#{reason}`")
+        nil
+
+      {:error, reason} ->
+        raise "Failed to unprotect packet due to `#{reason}`"
+    end
+  end
+
+  @spec handle_raw_rtp_packet(binary(), CallbackContext.t(), State.t()) ::
+          {[Membrane.Element.Action.t()], State.t()}
+  defp handle_raw_rtp_packet(raw_rtp_packet, ctx, state) do
     {:ok, packet} = ExRTP.Packet.decode(raw_rtp_packet)
 
     {new_stream_actions, state} =

lib/membrane/rtp/muxer.ex

@@ -63,6 +63,15 @@ defmodule Membrane.RTP.Muxer do
 
   def_output_pad :output, accepted_format: %RemoteStream{type: :packetized, content_format: RTP}
 
+  def_options srtp: [
+                spec: false | [ExLibSRTP.Policy.t()],
+                default: false,
+                description: """
+                Specifies whether to use SRTP to encrypt the output stream. Requires adding [srtp](https://github.com/membraneframework/elixir_libsrtp) 
+                dependency to work. If true takes a list of SRTP policies to use for encrypting packets. See `t:ExLibSRTP.Policy.t/0` for details.
+                """
+              ]
+
   defmodule State do
     @moduledoc false
     defmodule StreamState do
@@ -84,16 +93,32 @@ defmodule Membrane.RTP.Muxer do
     end
 
     @type t :: %__MODULE__{
-            stream_states: %{Pad.ref() => StreamState.t()}
+            stream_states: %{Pad.ref() => StreamState.t()},
+            srtp: ExLibSRTP.t() | nil
           }
 
-    @enforce_keys []
+    @enforce_keys [:srtp]
     defstruct @enforce_keys ++ [stream_states: %{}]
   end
 
   @impl true
-  def handle_init(_ctx, _opts) do
-    {[], %State{}}
+  def handle_init(_ctx, opts) do
+    srtp =
+      case opts.srtp do
+        false ->
+          nil
+
+        policies ->
+          if not Code.ensure_loaded?(ExLibSRTP) do
+            raise "Optional dependency :ex_libsrtp is required for SRTP"
+          end
+
+          srtp = apply(ExLibSRTP, :new, [])
+          Enum.each(policies, &apply(ExLibSRTP, :add_stream, [srtp, &1]))
+          srtp
+      end
+
+    {[], %State{srtp: srtp}}
   end
 
   @impl true
@@ -105,7 +130,7 @@ defmodule Membrane.RTP.Muxer do
   def handle_stream_format(pad, stream_format, ctx, state) do
     pad_options = ctx.pads[pad].options
 
-    ssrc = get_stream_ssrc(pad_options, state)
+    ssrc = get_stream_ssrc(pad_options.ssrc, state)
 
     encoding_name =
       @payload_format_to_encoding_name[stream_format.payload_format] ||
@@ -141,7 +166,7 @@ defmodule Membrane.RTP.Muxer do
 
   @impl true
   def handle_buffer(Pad.ref(:input, _ref) = pad_ref, buffer, _ctx, state) do
-    {rtp_metadata, metadata} = Map.pop(buffer.metadata, :rtp, %{})
+    rtp_metadata = Map.get(buffer.metadata, :rtp, %{})
     stream_state = state.stream_states[pad_ref]
 
     rtp_offset =
@@ -165,15 +190,25 @@ defmodule Membrane.RTP.Muxer do
         marker: Map.get(rtp_metadata, :marker, false)
       )
 
-    raw_packet = ExRTP.Packet.encode(packet)
+    buffer_action =
+      packet
+      |> ExRTP.Packet.encode()
+      |> protect_packet(state.srtp)
+      |> case do
+        nil ->
+          []
+
+        raw_packet ->
+          buffer = %Membrane.Buffer{
+            buffer
+            | payload: raw_packet,
+              metadata: Map.put(buffer.metadata, :rtp, %{packet | payload: <<>>})
+          }
 
-    buffer = %Membrane.Buffer{
-      buffer
-      | payload: raw_packet,
-        metadata: Map.put(metadata, :rtp, %{packet | payload: <<>>})
-    }
+          [buffer: {:output, buffer}]
+      end
 
-    {[buffer: {:output, buffer}], state}
+    {buffer_action, state}
   end
 
   @impl true
@@ -187,10 +222,30 @@ defmodule Membrane.RTP.Muxer do
     end
   end
 
-  defp get_stream_ssrc(pad_options, state) do
+  @spec protect_packet(binary(), ExLibSRTP.t() | nil) :: binary() | nil
+  defp protect_packet(rtp_packet, nil) do
+    rtp_packet
+  end
+
+  defp protect_packet(rtp_packet, srtp) do
+    case apply(ExLibSRTP, :protect, [srtp, rtp_packet]) do
+      {:ok, protected_rtp_packet} ->
+        protected_rtp_packet
+
+      {:error, reason} when reason in [:replay_fail, :replay_old] ->
+        Membrane.Logger.warning("Ignoring packet due to `#{reason}`")
+        nil
+
+      {:error, reason} ->
+        raise "Failed to unprotect packet due to `#{reason}`"
+    end
+  end
+
+  @spec get_stream_ssrc(RTP.ssrc() | :random, State.t()) :: RTP.ssrc()
+  defp get_stream_ssrc(ssrc, state) do
     assigned_ssrcs = Enum.map(state.stream_states, fn {_pad_ref, %{ssrc: ssrc}} -> ssrc end)
 
-    case pad_options.ssrc do
+    case ssrc do
       :random ->
         Stream.repeatedly(fn -> Enum.random(0..@max_ssrc) end)
         |> Enum.find(&(&1 not in assigned_ssrcs))

test/membrane/rtp/muxer_demuxer_integration_test.exs

@@ -19,8 +19,8 @@ defmodule Membrane.RTP.MuxerDemuxerTest do
         })
         |> child(:realtimer, Membrane.Realtimer)
         |> child(:rtp_h264_payloader, Membrane.RTP.H264.Payloader)
-        |> child(:rtp_muxer, Membrane.RTP.Muxer)
-        |> child(:rtp_demuxer, Membrane.RTP.Demuxer)
+        |> child(:rtp_muxer, %Membrane.RTP.Muxer{srtp: opts.srtp})
+        |> child(:rtp_demuxer, %Membrane.RTP.Demuxer{srtp: opts.srtp})
         |> via_out(:output,
           options: [stream_id: {:encoding_name, :H264}]
         )
@@ -32,14 +32,26 @@ defmodule Membrane.RTP.MuxerDemuxerTest do
     end
   end
 
-  @tag :tmp_dir
-  test "Muxed and demuxed stream is the same as unchanged one", %{tmp_dir: tmp_dir} do
+  describe "Muxed and demuxed stream is the same as unchanged one" do
+    @tag :tmp_dir
+    test "when not using SRTP encryption", %{tmp_dir: tmp_dir} do
+      perform_test(false, tmp_dir)
+    end
+
+    @tag :tmp_dir
+    test "when using SRTP encryption", %{tmp_dir: tmp_dir} do
+      policy = %ExLibSRTP.Policy{ssrc: :any_inbound, key: String.duplicate("a", 30)}
+      perform_test([policy], tmp_dir)
+    end
+  end
+
+  defp perform_test(srtp, tmp_dir) do
     output_path = Path.join(tmp_dir, "output.h264")
 
     pipeline =
       Testing.Pipeline.start_supervised!(
         module: MuxerDemuxerPipeline,
-        custom_args: %{input_path: @input_path, output_path: output_path}
+        custom_args: %{input_path: @input_path, output_path: output_path, srtp: srtp}
       )
 
     assert_start_of_stream(pipeline, :sink)

test/membrane/rtp/muxer_test.exs

@@ -13,7 +13,7 @@ defmodule Membrane.RTP.MuxerTest do
     use Membrane.Pipeline
 
     @impl true
-    def handle_init(_ctx, _opts) do
+    def handle_init(_ctx, opts) do
       spec = [
         child(:hackney_source, %Membrane.Hackney.Source{
           location:
@@ -27,7 +27,7 @@ defmodule Membrane.RTP.MuxerTest do
           output_alignment: :nalu
         })
         |> child(:h264_payloader, Membrane.RTP.H264.Payloader)
-        |> child(:rtp_muxer, Membrane.RTP.Muxer)
+        |> child(:rtp_muxer, %Membrane.RTP.Muxer{srtp: opts.srtp})
         |> child(:sink, Membrane.Testing.Sink),
         get_child(:mp4_demuxer)
         |> via_out(:output, options: [kind: :audio])
@@ -40,8 +40,20 @@ defmodule Membrane.RTP.MuxerTest do
     end
   end
 
-  test "Muxer muxes correct amount of packets" do
-    pipeline = Testing.Pipeline.start_supervised!(module: Pipeline)
+  describe "Muxer muxes correct amount of packets" do
+    test "when encrypting the stream with SRTP" do
+      policy = %ExLibSRTP.Policy{ssrc: :any_inbound, key: String.duplicate("b", 30)}
+      perform_test([policy])
+    end
+
+    test "when not encrypting the stream with SRTP" do
+      perform_test(false)
+    end
+  end
+
+  defp perform_test(srtp) do
+    pipeline =
+      Testing.Pipeline.start_supervised!(module: Pipeline, custom_args: %{srtp: srtp})
 
     %{audio: %{payload_type: audio_payload_type}, video: %{payload_type: video_payload_type}} =
       @rtp_output