add errors on broadcast handler

Author: filipecabaco

lib/realtime/tenants.ex

@@ -501,4 +501,19 @@ defmodule Realtime.Tenants do
   @spec region(Tenant.t()) :: String.t() | nil
   def region(%Tenant{extensions: [%{settings: settings}]}), do: Map.get(settings, "region")
   def region(_), do: nil
+
+  @doc """
+  """
+  @spec validate_payload_size(%Tenant{} | binary(), map()) :: :ok | {:error, :payload_size_exceeded}
+  def validate_payload_size(tenant_id, payload) when is_binary(tenant_id) do
+    tenant_id
+    |> Cache.get_tenant_by_external_id()
+    |> validate_payload_size(payload)
+  end
+
+  def validate_payload_size(%Tenant{max_payload_size_in_kb: max_payload_size_in_kb}, payload) do
+    max_payload_size = max_payload_size_in_kb * 1000
+    payload_size = :erlang.external_size(payload)
+    if payload_size > max_payload_size, do: {:error, :payload_size_exceeded}, else: :ok
+  end
 end

lib/realtime/tenants/batch_broadcast.ex

@@ -116,12 +116,11 @@ defmodule Realtime.Tenants.BatchBroadcast do
 
   defp validate_payload_size(changeset, tenant) do
     payload = get_change(changeset, :payload)
-    payload_size = byte_size(Jason.encode!(payload))
-    max_payload_size = tenant.max_payload_size_in_kb * 1000
 
-    if payload_size > max_payload_size,
-      do: add_error(changeset, :payload, "Payload size exceeds tenant limit"),
-      else: changeset
+    case Tenants.validate_payload_size(tenant, payload) do
+      :ok -> changeset
+      _ -> add_error(changeset, :payload, "Payload size exceeds tenant limit")
+    end
   end
 
   @event_type "broadcast"

lib/realtime_web/channels/realtime_channel/broadcast_handler.ex

@@ -6,6 +6,7 @@ defmodule RealtimeWeb.RealtimeChannel.BroadcastHandler do
 
   import Phoenix.Socket, only: [assign: 3]
 
+  alias Realtime.Tenants
   alias RealtimeWeb.RealtimeChannel
   alias RealtimeWeb.TenantBroadcaster
   alias Phoenix.Socket
@@ -38,8 +39,23 @@ defmodule RealtimeWeb.RealtimeChannel.BroadcastHandler do
           |> increment_rate_counter()
 
         %{ack_broadcast: ack_broadcast} = socket.assigns
-        send_message(tenant_id, self_broadcast, tenant_topic, payload)
-        if ack_broadcast, do: {:reply, :ok, socket}, else: {:noreply, socket}
+
+        res =
+          case Tenants.validate_payload_size(tenant_id, payload) do
+            :ok -> send_message(tenant_id, self_broadcast, tenant_topic, payload)
+            {:error, error} -> {:error, error}
+          end
+
+        cond do
+          ack_broadcast && match?({:error, :payload_size_exceeded}, res) ->
+            {:reply, {:error, :payload_size_exceeded}, socket}
+
+          ack_broadcast ->
+            {:reply, :ok, socket}
+
+          true ->
+            {:noreply, socket}
+        end
 
       {:ok, policies} ->
         {:noreply, assign(socket, :policies, policies)}
@@ -65,11 +81,23 @@ defmodule RealtimeWeb.RealtimeChannel.BroadcastHandler do
     } = socket
 
     socket = increment_rate_counter(socket)
-    send_message(tenant_id, self_broadcast, tenant_topic, payload)
 
-    if ack_broadcast,
-      do: {:reply, :ok, socket},
-      else: {:noreply, socket}
+    res =
+      case Tenants.validate_payload_size(tenant_id, payload) do
+        :ok -> send_message(tenant_id, self_broadcast, tenant_topic, payload)
+        {:error, error} -> {:error, error}
+      end
+
+    cond do
+      ack_broadcast && match?({:error, :payload_size_exceeded}, res) ->
+        {:reply, {:error, :payload_size_exceeded}, socket}
+
+      ack_broadcast ->
+        {:reply, :ok, socket}
+
+      true ->
+        {:noreply, socket}
+    end
   end
 
   defp send_message(tenant_id, self_broadcast, tenant_topic, payload) do

test/realtime_web/channels/realtime_channel/broadcast_handler_test.exs

@@ -312,6 +312,82 @@ defmodule RealtimeWeb.RealtimeChannel.BroadcastHandlerTest do
       {:ok, %{avg: avg}} = RateCounter.get(Tenants.events_per_second_rate(tenant))
       assert avg == 0.0
     end
+
+    test "handle payload size excedding limits in private channels", %{topic: topic, tenant: tenant, db_conn: db_conn} do
+      socket =
+        socket_fixture(tenant, topic,
+          policies: %Policies{broadcast: %BroadcastPolicies{write: true}},
+          ack_broadcast: false
+        )
+
+      assert {:noreply, _} =
+               BroadcastHandler.handle(
+                 %{"data" => random_string(tenant.max_payload_size_in_kb * 1000 + 1)},
+                 db_conn,
+                 socket
+               )
+
+      Process.sleep(120)
+
+      refute_received {:socket_push, :text, _}
+    end
+
+    test "handle payload size excedding limits in public channels", %{topic: topic, tenant: tenant, db_conn: db_conn} do
+      socket = socket_fixture(tenant, topic, ack_broadcast: false, private?: false)
+
+      assert {:noreply, _} =
+               BroadcastHandler.handle(
+                 %{"data" => random_string(tenant.max_payload_size_in_kb * 1000 + 1)},
+                 db_conn,
+                 socket
+               )
+
+      Process.sleep(120)
+
+      refute_received {:socket_push, :text, _}
+    end
+
+    test "handle payload size excedding limits in private channel and if ack it will receive error", %{
+      topic: topic,
+      tenant: tenant,
+      db_conn: db_conn
+    } do
+      socket =
+        socket_fixture(tenant, topic,
+          policies: %Policies{broadcast: %BroadcastPolicies{write: true}},
+          ack_broadcast: true
+        )
+
+      assert {:reply, {:error, :payload_size_exceeded}, _} =
+               BroadcastHandler.handle(
+                 %{"data" => random_string(tenant.max_payload_size_in_kb * 1000 + 1)},
+                 db_conn,
+                 socket
+               )
+
+      Process.sleep(120)
+
+      refute_received {:socket_push, :text, _}
+    end
+
+    test "handle payload size excedding limits in public channels and if ack it will receive error", %{
+      topic: topic,
+      tenant: tenant,
+      db_conn: db_conn
+    } do
+      socket = socket_fixture(tenant, topic, ack_broadcast: true, private?: false)
+
+      assert {:reply, {:error, :payload_size_exceeded}, _} =
+               BroadcastHandler.handle(
+                 %{"data" => random_string(tenant.max_payload_size_in_kb * 1000 + 1)},
+                 db_conn,
+                 socket
+               )
+
+      Process.sleep(120)
+
+      refute_received {:socket_push, :text, _}
+    end
   end
 
   defp initiate_tenant(context) do