Don't accept streams until server is ready to read them

examples/source_with_standalone_server.exs

@@ -47,23 +47,36 @@ end
 
 # The client will connect on `rtmp://localhost:1935/app/stream_key`
 port = 1935
-app = "app"
-stream_key = "stream_key"
+
+# example lambda function that upon launching will send client reference back to parent process.
+parent_process_pid = self()
+
+new_client_callback = fn client_ref, app, stream_key ->
+  send(parent_process_pid, {:client_ref, client_ref, app, stream_key})
+end
 
 # Run the standalone server
 {:ok, server} =
   Membrane.RTMP.Server.start_link(
     handler: %Membrane.RTMP.Source.ClientHandler{controlling_process: self()},
     port: port,
-    use_ssl?: false
+    use_ssl?: false,
+    new_client_callback: new_client_callback,
+    client_timeout: 5_000
   )
 
-# Subscribe to receive client reference that connected to the
-# server with given app id and stream key
-:ok = Membrane.RTMP.Server.subscribe(server, app, stream_key)
+app = "app"
+stream_key = "stream_key"
+
+# Wait max 10s for client to connect on /app/stream_key
+{:ok, client_ref} =
+  receive do
+    {:client_ref, client_ref, ^app, ^stream_key} ->
+      {:ok, client_ref}
+  after
+    10_000 -> :timeout
+  end
 
-# Wait for the client reference
-{:ok, client_ref} = Membrane.RTMP.Server.await_subscription(app, stream_key)
 # Start the pipeline and provide it with the client_ref
 {:ok, _supervisor, pipeline} =
   Membrane.Pipeline.start_link(Pipeline, client_ref: client_ref)

lib/membrane_rtmp_plugin/rtmp/source/client_handler.ex

@@ -26,13 +26,7 @@ defmodule Membrane.RTMP.Source.ClientHandler do
 
   @impl true
   def handle_stream_published(publish_msg, state) do
-    state = %{state | stream_key: publish_msg.stream_key}
-
-    if state.controlling_process do
-      send(state.controlling_process, {:client_connected, state.app, state.stream_key})
-    end
-
-    state
+    %{state | stream_key: publish_msg.stream_key}
   end
 
   @impl true

lib/membrane_rtmp_plugin/rtmp/source/message_handler.ex

@@ -48,6 +48,9 @@ defmodule Membrane.RTMP.MessageHandler do
       header_sent?: false,
       events: [],
       receiver_pid: nil,
+      #
+      publish_msg: nil,
+      publish_header: nil,
       # how many times the Source tries to get control of the socket
       socket_retries: 3,
       # epoch required for performing a handshake with the pipeline
@@ -72,6 +75,17 @@ defmodule Membrane.RTMP.MessageHandler do
     end)
   end
 
+  @spec send_publish_success(map()) :: {map(), list()}
+  def send_publish_success(state) do
+    Responses.publish_success(state.publish_msg.stream_key)
+    |> send_rtmp_payload(state.socket,
+      chunk_stream_id: 3,
+      stream_id: state.publish_header.stream_id
+    )
+
+    {%{state | events: []}, [{:published, state.publish_msg} | state.events]}
+  end
+
   # Expected flow of messages:
   # 1.  [in] c0_c1 handshake -> [out] s0_s1_s2 handshake
   # 2.  [in] c2 handshake -> [out] empty
@@ -141,11 +155,10 @@ defmodule Membrane.RTMP.MessageHandler do
     %Messages.UserControl{event_type: @stream_begin_type, data: <<0, 0, 0, 1>>}
     |> send_rtmp_payload(state.socket, chunk_stream_id: 2)
 
-    Responses.publish_success(publish_msg.stream_key)
-    |> send_rtmp_payload(state.socket, chunk_stream_id: 3, stream_id: header.stream_id)
+    # at this point pause the unfinished handshake until pipeline demands data from this client
+    # (this mechanism prevents accepting streams with no listeners)
 
-    state = %{state | events: [{:published, publish_msg} | state.events]}
-    {:cont, state}
+    {:halt, %{state | publish_msg: publish_msg, publish_header: header}}
   end
 
   # A message containing stream metadata

lib/membrane_rtmp_plugin/rtmp/source/source.ex

@@ -5,12 +5,14 @@ defmodule Membrane.RTMP.Source do
 
   The source can be used in the following two scenarios:
   * by providing the URL on which the client is expected to connect - note, that if the client doesn't
-  connect on this URL, the source won't complete its setup
-  * by spawning `Membrane.RTMP.Server`, subscribing for a given app and stream key on which the client
-  will connect, waiting for a client reference and passing the client reference to the `#{inspect(__MODULE__)}`.
+  connect on this URL, the source won't complete its setup. Note that all attepted connections to
+  other `app` or `stream_key` than specified ones will be rejected.
+
+  * by spawning `Membrane.RTMP.Server`, receiving a client reference and passing it to the `#{inspect(__MODULE__)}`.
   """
   use Membrane.Source
   require Membrane.Logger
+  require Logger
   alias __MODULE__.ClientHandler, as: SourceClientHandler
   alias Membrane.RTMP.Server.ClientHandler
 
@@ -81,15 +83,23 @@ defmodule Membrane.RTMP.Source do
   def handle_setup(_ctx, %{mode: :builtin_server} = state) do
     {use_ssl?, port, app, stream_key} = parse_url(state.url)
 
+    parent_pid = self()
+
+    new_client_callback = fn client_ref, app, stream_key ->
+      send(parent_pid, {:client_ref, client_ref, app, stream_key})
+    end
+
     {:ok, server_pid} =
       Membrane.RTMP.Server.start_link(
         handler: %SourceClientHandler{controlling_process: self()},
         port: port,
-        use_ssl?: use_ssl?
+        use_ssl?: use_ssl?,
+        new_client_callback: new_client_callback,
+        client_timeout: 100
       )
 
     state = %{state | app: app, stream_key: stream_key, server: server_pid}
-    {[], state}
+    {[setup: :incomplete], state}
   end
 
   @impl true
@@ -156,19 +166,23 @@ defmodule Membrane.RTMP.Source do
   end
 
   @impl true
-  def handle_info({:client_connected, app, stream_key}, _ctx, %{mode: :builtin_server} = state) do
-    :ok = Membrane.RTMP.Server.subscribe(state.server, state.app, state.stream_key)
-    state = %{state | app: app, stream_key: stream_key}
-    {[], state}
+  def handle_info(
+        {:client_ref, client_ref, app, stream_key},
+        _ctx,
+        %{mode: :builtin_server} = state
+      )
+      when app == state.app and stream_key == state.stream_key do
+    {[setup: :complete], %{state | client_ref: client_ref}}
   end
 
   @impl true
   def handle_info(
-        {:client_ref, client_ref_pid, _app, _stream_key},
+        {:client_ref, _client_ref, app, stream_key},
         _ctx,
         %{mode: :builtin_server} = state
       ) do
-    {[redemand: :output], %{state | client_ref: client_ref_pid}}
+    Logger.warning("Unexpected client connected on /#{app}/#{stream_key}")
+    {[], state}
   end
 
   @impl true

lib/membrane_rtmp_plugin/rtmp/source/source_bin.ex

@@ -8,8 +8,8 @@ defmodule Membrane.RTMP.SourceBin do
   The bin can be used in the following two scenarios:
   * by providing the URL on which the client is expected to connect - note, that if the client doesn't
   connect on this URL, the bin won't complete its setup
-  * by spawning `Membrane.RTMP.Server`, subscribing for a given app and stream key on which the client
-  will connect, waiting for a client reference and passing the client reference to the `#{inspect(__MODULE__)}`.
+  * by spawning `Membrane.RTMP.Server`, receiving client reference after client connects on a given `app` and `stream_key`
+  and passing the client reference to the `#{inspect(__MODULE__)}`.
   """
   use Membrane.Bin
 

lib/membrane_rtmp_plugin/rtmp_server.ex

@@ -1,6 +1,14 @@
 defmodule Membrane.RTMP.Server do
   @moduledoc """
-  A simple RTMP server, which handles each new incoming connection.
+  A simple RTMP server, which handles each new incoming connection. When a new client connects, the `new_client_callback` is invoked.
+  New connections remain in an incomplete RTMP handshake state, until another process makes demand for their data.
+  If no data is demanded within the client_timeout period, TCP socket is closed.
+
+  Options:
+   - client_timeout: Time (ms) after which an unused client connection is automatically closed.
+   - new_client_callback: An anonymous function called when a new client connects.
+      It receives the client reference, `app` and `stream_key`, allowing custom processing,
+      like sending the reference to another process.
   """
   use GenServer
 
@@ -15,7 +23,12 @@ defmodule Membrane.RTMP.Server do
           handler: ClientHandlerBehaviour.t(),
           port: :inet.port_number(),
           use_ssl?: boolean(),
-          name: atom() | nil
+          name: atom() | nil,
+          new_client_callback: (client_ref :: pid(),
+                                app :: String.t(),
+                                stream_key :: String.t() ->
+                                  any()),
+          client_timeout: non_neg_integer()
         ]
 
   @type server_identifier :: pid() | atom()
@@ -30,33 +43,6 @@ defmodule Membrane.RTMP.Server do
     GenServer.start_link(__MODULE__, server_options, gen_server_opts)
   end
 
-  @doc """
-  Subscribes for the given app and stream key.
-  When a client connects (or has already connected) to the server with given app and stream key,
-  the subscriber will be informed.
-  """
-  @spec subscribe(server_identifier(), String.t(), String.t()) :: :ok
-  def subscribe(server_identifier, app, stream_key) do
-    GenServer.cast(server_identifier, {:subscribe, app, stream_key, self()})
-    :ok
-  end
-
-  @doc """
-  Awaits for the client reference of the connection to which the user has previously subscribed.
-
-  Note: this function call is blocking!
-  Note: first you need to call `#{__MODULE__}.subscribe/3` to subscribe
-  for a given `app` and `stream_key`.
-  """
-  @spec await_subscription(String.t(), String.t(), non_neg_integer()) :: {:ok, pid()} | :error
-  def await_subscription(app, stream_key, timeout \\ 5_000) do
-    receive do
-      {:client_ref, client_ref, ^app, ^stream_key} -> {:ok, client_ref}
-    after
-      timeout -> :error
-    end
-  end
-
   @doc """
   Returns the port on which the server listens for connection.
   """
@@ -74,8 +60,6 @@ defmodule Membrane.RTMP.Server do
 
     {:ok,
      %{
-       subscriptions: %{},
-       client_reference_mapping: %{},
        listener: pid,
        port: nil,
        to_reply: [],
@@ -92,33 +76,9 @@ defmodule Membrane.RTMP.Server do
     end
   end
 
-  @impl true
-  def handle_cast({:subscribe, app, stream_key, subscriber_pid}, state) do
-    state = put_in(state, [:subscriptions, {app, stream_key}], subscriber_pid)
-    maybe_send_subscription(app, stream_key, state)
-    {:noreply, state}
-  end
-
-  @impl true
-  def handle_info({:register_client, app, stream_key, client_reference_pid}, state) do
-    state = put_in(state, [:client_reference_mapping, {app, stream_key}], client_reference_pid)
-    maybe_send_subscription(app, stream_key, state)
-    {:noreply, state}
-  end
-
   @impl true
   def handle_info({:port, port}, state) do
     Enum.each(state.to_reply, &GenServer.reply(&1, port))
     {:noreply, %{state | port: port, to_reply: []}}
   end
-
-  defp maybe_send_subscription(app, stream_key, state) do
-    if state.subscriptions[{app, stream_key}] != nil and
-         state.client_reference_mapping[{app, stream_key}] != nil do
-      send(
-        state.subscriptions[{app, stream_key}],
-        {:client_ref, state.client_reference_mapping[{app, stream_key}], app, stream_key}
-      )
-    end
-  end
 end

lib/membrane_rtmp_plugin/rtmp_server/client_handler.ex

@@ -38,7 +38,10 @@ defmodule Membrane.RTMP.Server.ClientHandler do
        stream_key: nil,
        server: opts.server,
        buffers_demanded: 0,
-       published?: false
+       published?: false,
+       notified_about_client?: false,
+       new_client_callback: opts.new_client_callback,
+       client_timeout: opts.client_timeout
      }}
   end
 
@@ -77,11 +80,21 @@ defmodule Membrane.RTMP.Server.ClientHandler do
 
   @impl true
   def handle_info({:demand_data, how_many_buffers_demanded}, state) do
-    state = %{state | buffers_demanded: how_many_buffers_demanded}
+    state = finish_handshake(state) |> Map.replace!(:buffers_demanded, how_many_buffers_demanded)
     request_data(state)
     {:noreply, state}
   end
 
+  @impl true
+  def handle_info({:client_timeout, app, stream_key}, state) do
+    if not state.published? do
+      Logger.warning("No demand made for client /#{app}/#{stream_key}, terminating connection.")
+      :gen_tcp.close(state.socket)
+    end
+
+    {:noreply, state}
+  end
+
   @impl true
   def handle_info(other_msg, state) do
     handler_state = state.handler.handle_info(other_msg, state.handler_state)
@@ -96,6 +109,24 @@ defmodule Membrane.RTMP.Server.ClientHandler do
     {message_handler_state, events} =
       MessageHandler.handle_client_messages(messages, state.message_handler_state)
 
+    state =
+      if message_handler_state.publish_msg != nil and not state.notified_about_client? do
+        %{publish_msg: %Membrane.RTMP.Messages.Publish{stream_key: stream_key}} =
+          message_handler_state
+
+        if is_function(state.new_client_callback) do
+          state.new_client_callback.(self(), state.app, stream_key)
+        else
+          raise "new_client_callback is not a function"
+        end
+
+        Process.send_after(self(), {:client_timeout, state.app, stream_key}, state.client_timeout)
+
+        %{state | notified_about_client?: true}
+      else
+        state
+      end
+
     state = Enum.reduce(events, state, &handle_event/2)
 
     request_data(state)
@@ -136,8 +167,6 @@ defmodule Membrane.RTMP.Server.ClientHandler do
         %{state | handler_state: new_handler_state, app: connected_msg.app}
 
       {:published, publish_msg} ->
-        send(state.server, {:register_client, state.app, publish_msg.stream_key, self()})
-
         new_handler_state =
           state.handler.handle_stream_published(publish_msg, state.handler_state)
 
@@ -159,4 +188,14 @@ defmodule Membrane.RTMP.Server.ClientHandler do
       end
     end
   end
+
+  defp finish_handshake(state) when not state.published? do
+    {message_handler_state, events} =
+      MessageHandler.send_publish_success(state.message_handler_state)
+
+    state = Enum.reduce(events, state, &handle_event/2)
+    %{state | message_handler_state: message_handler_state}
+  end
+
+  defp finish_handshake(state), do: state
 end

lib/membrane_rtmp_plugin/rtmp_server/listener.ex

@@ -52,7 +52,9 @@ defmodule Membrane.RTMP.Server.Listener do
         socket: client,
         use_ssl?: options.use_ssl?,
         handler: options.handler,
-        server: options.server
+        server: options.server,
+        new_client_callback: options.new_client_callback,
+        client_timeout: options.client_timeout
       )
 
     case :gen_tcp.controlling_process(client, client_reference) do