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

.gitignore

@@ -177,3 +177,4 @@ $RECYCLE.BIN/
 *.lnk
 
 # End of https://www.gitignore.io/api/c,vim,linux,macos,elixir,windows,visualstudiocode
+received.flv

examples/source.exs

@@ -45,6 +45,8 @@ defmodule Pipeline do
   end
 end
 
+Logger.configure(level: :error)
+
 # Start a pipeline with `Membrane.RTMP.Source` that will spawn an RTMP server waiting for
 # the client connection on given URL
 {:ok, _supervisor, pipeline} = Membrane.Pipeline.start_link(Pipeline)

examples/source_with_standalone_server.exs

@@ -45,25 +45,38 @@ defmodule Pipeline do
   end
 end
 
+Logger.configure(level: :error)
+
 # 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 a message back to parent server.
+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
   )
 
-# 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"
+
+{:ok, client_ref} =
+  receive do
+    {:client_ref, client_ref, ^app, ^stream_key} ->
+      {:ok, client_ref}
+  after
+    5000 -> :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

@@ -81,15 +81,22 @@ 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
       )
 
     state = %{state | app: app, stream_key: stream_key, server: server_pid}
-    {[], state}
+    {[setup: :incomplete], state}
   end
 
   @impl true
@@ -156,19 +163,22 @@ 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}}
+    {[], state}
   end
 
   @impl true

lib/membrane_rtmp_plugin/rtmp_server.ex

@@ -1,6 +1,12 @@
 defmodule Membrane.RTMP.Server do
   @moduledoc """
   A simple RTMP server, which handles each new incoming connection.
+
+  When new client connects to the server, it goes into :client_waiting_queue and its RTMP handshake will remanin unfinished.
+  Only when pipeline tries to pull data from client, its handshake will be finished, and client will be registered.
+
+  Also when new client connects, optional, annonymous function defined by user is triggered.
+  The lambda function is given PID of parent server, app and stream key.
   """
   use GenServer
 
@@ -15,7 +21,9 @@ 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()) | nil
         ]
 
   @type server_identifier :: pid() | atom()
@@ -30,33 +38,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,12 +55,11 @@ defmodule Membrane.RTMP.Server do
 
     {:ok,
      %{
-       subscriptions: %{},
-       client_reference_mapping: %{},
        listener: pid,
        port: nil,
        to_reply: [],
-       use_ssl?: server_options.use_ssl?
+       use_ssl?: server_options.use_ssl?,
+       new_client_callback: server_options.new_client_callback
      }}
   end
 
@@ -92,33 +72,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,9 @@ defmodule Membrane.RTMP.Server.ClientHandler do
        stream_key: nil,
        server: opts.server,
        buffers_demanded: 0,
-       published?: false
+       published?: false,
+       client_register_attempt_made?: false,
+       new_client_callback: opts.new_client_callback
      }}
   end
 
@@ -77,7 +79,7 @@ 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
@@ -96,6 +98,20 @@ 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.client_register_attempt_made? 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)
+        end
+
+        %{state | client_register_attempt_made?: true}
+      else
+        state
+      end
+
     state = Enum.reduce(events, state, &handle_event/2)
 
     request_data(state)
@@ -136,8 +152,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 +173,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,8 @@ 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
       )
 
     case :gen_tcp.controlling_process(client, client_reference) do