Add connect message validation (#65)

membraneframework · Aug 8, 2023 · e742ebe · e742ebe
1 parent 7426ea3
commit e742ebe
Show file tree

Hide file tree

Showing 7 changed files with 43 additions and 21 deletions.
diff --git a/README.md b/README.md
@@ -28,7 +28,7 @@ The package can be installed by adding `membrane_rtmp_plugin` to your list of de
 ```elixir
 def deps do
   [
-	  {:membrane_rtmp_plugin, "~> 0.14.1"}
+	  {:membrane_rtmp_plugin, "~> 0.15.0"}
   ]
 end
 ```

diff --git a/lib/membrane_rtmp_plugin/rtmp/source/message_handler.ex b/lib/membrane_rtmp_plugin/rtmp/source/message_handler.ex
@@ -76,29 +76,37 @@ defmodule Membrane.RTMP.MessageHandler do
     {:cont, %{state | message_parser: parser}}
   end
 
-  defp do_handle_client_message(%Messages.Connect{}, _header, state) do
-    chunk_size = state.message_parser.chunk_size
+  @validation_stage :connect
+  defp do_handle_client_message(%Messages.Connect{} = msg, _header, state) do
+    case MessageValidator.validate_connect(state.validator, msg) do
+      {:ok, _msg} = result ->
+        chunk_size = state.message_parser.chunk_size
 
-    [
-      %Messages.WindowAcknowledgement{size: @windows_acknowledgment_size},
-      %Messages.SetPeerBandwidth{size: @peer_bandwidth_size},
-      # stream begin type
-      %Messages.UserControl{event_type: 0x00, data: <<0, 0, 0, 0>>},
-      # by default the ffmpeg server uses 128 chunk size
-      %Messages.SetChunkSize{chunk_size: chunk_size}
-    ]
-    |> Enum.each(&send_rtmp_payload(&1, state.socket, chunk_size))
+        [
+          %Messages.WindowAcknowledgement{size: @windows_acknowledgment_size},
+          %Messages.SetPeerBandwidth{size: @peer_bandwidth_size},
+          # stream begin type
+          %Messages.UserControl{event_type: 0x00, data: <<0, 0, 0, 0>>},
+          # by default the ffmpeg server uses 128 chunk size
+          %Messages.SetChunkSize{chunk_size: chunk_size}
+        ]
+        |> Enum.each(&send_rtmp_payload(&1, state.socket, chunk_size))
 
-    {[tx_id], message_parser} = MessageParser.generate_tx_ids(state.message_parser, 1)
+        {[tx_id], message_parser} = MessageParser.generate_tx_ids(state.message_parser, 1)
 
-    tx_id
-    |> Responses.connection_success()
-    |> send_rtmp_payload(state.socket, chunk_size, chunk_stream_id: 3)
+        tx_id
+        |> Responses.connection_success()
+        |> send_rtmp_payload(state.socket, chunk_size, chunk_stream_id: 3)
+
+        Responses.on_bw_done()
+        |> send_rtmp_payload(state.socket, chunk_size, chunk_stream_id: 3)
 
-    Responses.on_bw_done()
-    |> send_rtmp_payload(state.socket, chunk_size, chunk_stream_id: 3)
+        {:cont,
+         validation_action(%{state | message_parser: message_parser}, @validation_stage, result)}
 
-    {:cont, %{state | message_parser: message_parser}}
+      {:error, _reason} = error ->
+        {:halt, {:error, :stream_validation, validation_action(state, @validation_stage, error)}}
+    end
   end
 
   # According to ffmpeg's documentation, this command should make the server release channel for a media stream

diff --git a/lib/membrane_rtmp_plugin/rtmp/source/message_validator.ex b/lib/membrane_rtmp_plugin/rtmp/source/message_validator.ex
@@ -8,6 +8,12 @@ defprotocol Membrane.RTMP.MessageValidator do
 
   @type validation_result_t :: {:ok, term()} | {:error, reason :: any()}
 
+  @doc """
+  Validates the `t:Membrane.RTMP.Messages.Connect.t/0` message.
+  """
+  @spec validate_connect(t(), Messages.Connect.t()) :: validation_result_t()
+  def validate_connect(impl, message)
+
   @doc """
   Validates the `t:Membrane.RTMP.Messages.ReleaseStream.t/0` message.
   """

diff --git a/lib/membrane_rtmp_plugin/rtmp/source/message_validator/default.ex b/lib/membrane_rtmp_plugin/rtmp/source/message_validator/default.ex
@@ -7,6 +7,9 @@ defmodule Membrane.RTMP.MessageValidator.Default do
 end
 
 defimpl Membrane.RTMP.MessageValidator, for: Membrane.RTMP.MessageValidator.Default do
+  @impl true
+  def validate_connect(_impl, _message), do: {:ok, "connect success"}
+
   @impl true
   def validate_release_stream(_impl, _message), do: {:ok, "release stream success"}
 

diff --git a/lib/membrane_rtmp_plugin/rtmp/source/messages/command/connect.ex b/lib/membrane_rtmp_plugin/rtmp/source/messages/command/connect.ex
@@ -1,5 +1,7 @@
 defmodule Membrane.RTMP.Messages.Connect do
-  @moduledoc false
+  @moduledoc """
+  Defines the RTMP `connect` command.
+  """
 
   @behaviour Membrane.RTMP.Message
 

diff --git a/mix.exs b/mix.exs
@@ -1,7 +1,7 @@
 defmodule Membrane.RTMP.Mixfile do
   use Mix.Project
 
-  @version "0.14.1"
+  @version "0.15.0"
   @github_url "https://github.com/membraneframework/membrane_rtmp_plugin"
 
   def project do

diff --git a/test/support/test_validator.ex b/test/support/test_validator.ex
@@ -10,6 +10,9 @@ end
 defimpl Membrane.RTMP.MessageValidator, for: Support.TestValidator do
   alias Membrane.RTMP.Messages
 
+  @impl true
+  def validate_connect(_impl, _message), do: {:ok, "connect success"}
+
   @impl true
   def validate_publish(%Support.TestValidator{stream_key: nil}, _message),
     do: {:ok, "stream allowed"}