Adapt plugin for general Membrane use cases

lib/membrane_rtsp_plugin/source.ex

@@ -21,10 +21,15 @@ defmodule Membrane.RTSP.Source do
 
   require Membrane.Logger
 
+  alias __MODULE__
   alias __MODULE__.{ConnectionManager, ReadyNotifier}
-  alias Membrane.RTSP.Decapsulator
+  alias Membrane.RTSP.TCP.Decapsulator
   alias Membrane.Time
 
+  @type transport ::
+          {:udp, port_range_start :: non_neg_integer(), port_range_end :: non_neg_integer()}
+          | :tcp
+
   def_options stream_uri: [
                 spec: binary(),
                 description: "The RTSP URI of the resource to stream."
@@ -37,10 +42,7 @@ defmodule Membrane.RTSP.Source do
                 """
               ],
               transport: [
-                spec:
-                  {:udp, port_range_start :: non_neg_integer(),
-                   port_range_end :: non_neg_integer()}
-                  | :tcp,
+                spec: transport(),
                 default: :tcp,
                 description: """
                 Transport protocol that will be used in the established RTSP stream. In case of
@@ -67,8 +69,32 @@ defmodule Membrane.RTSP.Source do
     accepted_format: _any,
     availability: :on_request
 
+  defmodule State do
+    @moduledoc false
+    @type t :: %__MODULE__{
+            stream_uri: binary(),
+            allowed_media_types: ConnectionManager.media_types(),
+            transport: Source.transport(),
+            timeout: Time.t(),
+            keep_alive_interval: Time.t(),
+            connection_manager: ConnectionManager.t(),
+            tracks: [ConnectionManager.track()],
+            ssrc_to_track: %{non_neg_integer() => ConnectionManager.track()}
+          }
+
+    @enforce_keys [:stream_uri, :allowed_media_types, :transport, :timeout, :keep_alive_interval]
+    defstruct @enforce_keys ++
+                [
+                  connection_manager: nil,
+                  tracks: [],
+                  ssrc_to_track: %{}
+                ]
+  end
+
   defmodule ReadyNotifier do
     @moduledoc false
+    # This element's only purpose is to send a notification to it's parent when it's entered playing
+    # state, meaning all it's siblings also did.
     use Membrane.Source
 
     @impl true
@@ -79,10 +105,7 @@ defmodule Membrane.RTSP.Source do
 
   @impl true
   def handle_init(_ctx, options) do
-    state =
-      options
-      |> Map.from_struct()
-      |> Map.merge(%{connection_manager: nil, tracks: [], ssrc_to_track: %{}})
+    state = struct(State, Map.from_struct(options))
 
     {[], state}
   end
@@ -139,7 +162,7 @@ defmodule Membrane.RTSP.Source do
   end
 
   @impl true
-  def handle_info(%{tracks: tracks, transport_info: transport_info}, _ctx, state) do
+  def handle_info(%{tracks: tracks}, _ctx, state) do
     fmt_mapping =
       Enum.map(tracks, fn %{rtpmap: rtpmap} ->
         {rtpmap.payload_type, {String.to_atom(rtpmap.encoding), rtpmap.clock_rate}}
@@ -151,14 +174,16 @@ defmodule Membrane.RTSP.Source do
       child(:ready_notifier, ReadyNotifier)
     ]
 
-    case transport_info do
-      {:tcp, tcp_socket} ->
+    case state.transport do
+      :tcp ->
+        {:tcp, socket} = List.first(tracks).transport
+
         spec =
           common_spec ++
             [
               child(:tcp_source, %Membrane.TCP.Source{
                 connection_side: :client,
-                local_socket: tcp_socket
+                local_socket: socket
               })
               |> child(:tcp_depayloader, Decapsulator)
               |> via_in(Pad.ref(:rtp_input, make_ref()))
@@ -167,13 +192,20 @@ defmodule Membrane.RTSP.Source do
 
         {[spec: spec], %{state | tracks: tracks}}
 
-      {:udp, min_port, max_port} ->
+      {:udp, _port_range_start, _port_range_end} ->
         spec =
           common_spec ++
-            Enum.map(min_port..max_port, fn port ->
-              child({:udp_source, make_ref()}, %Membrane.UDP.Source{local_port_no: port})
-              |> via_in(Pad.ref(:rtp_input, make_ref()))
-              |> get_child(:rtp_session)
+            Enum.flat_map(tracks, fn track ->
+              {:udp, rtp_port, rtcp_port} = track.transport
+
+              [
+                child({:udp, make_ref()}, %Membrane.UDP.Source{local_port_no: rtp_port})
+                |> via_in(Pad.ref(:rtp_input, make_ref()))
+                |> get_child(:rtp_session),
+                child({:udp, make_ref()}, %Membrane.UDP.Source{local_port_no: rtcp_port})
+                |> via_in(Pad.ref(:rtp_input, make_ref()))
+                |> get_child(:rtp_session)
+              ]
             end)
 
         {[spec: spec], %{state | tracks: tracks}}
@@ -204,10 +236,12 @@ defmodule Membrane.RTSP.Source do
     {[terminate: :normal], state}
   end
 
+  @spec get_rtp_depayloader(ConnectionManager.track()) :: module() | nil
   defp get_rtp_depayloader(%{rtpmap: %{encoding: "H264"}}), do: Membrane.RTP.H264.Depayloader
   defp get_rtp_depayloader(%{rtpmap: %{encoding: "H265"}}), do: Membrane.RTP.H265.Depayloader
   defp get_rtp_depayloader(_track), do: nil
 
+  @spec parser(ChildrenSpec.builder(), ConnectionManager.track()) :: ChildrenSpec.builder()
   defp parser(link_builder, %{rtpmap: %{encoding: "H264"}} = track) do
     sps = track.fmtp.sprop_parameter_sets && track.fmtp.sprop_parameter_sets.sps
     pps = track.fmtp.sprop_parameter_sets && track.fmtp.sprop_parameter_sets.pps
@@ -232,6 +266,7 @@ defmodule Membrane.RTSP.Source do
 
   # a strange issue with one of Milesight camera where the parameter sets has
   # <<0, 0, 0, 1>> at the end
+  @spec clean_parameter_set(binary()) :: binary()
   defp clean_parameter_set(ps) do
     case :binary.part(ps, byte_size(ps), -4) do
       <<0, 0, 0, 1>> -> :binary.part(ps, 0, byte_size(ps) - 4)