RLS-based persistence

Author: leandrocp

lib/realtime/api/channel.ex

@@ -12,7 +12,7 @@ defmodule Realtime.Api.Message do
   @timestamps_opts [type: :naive_datetime_usec]
   schema "messages" do
     field(:topic, :string)
-    field(:extension, Ecto.Enum, values: [:broadcast, :presence])
+    field(:extension, Ecto.Enum, values: [:broadcast, :presence, :persistence])
     field(:payload, :map)
     field(:event, :string)
     field(:private, :boolean)

lib/realtime/api/message.ex

@@ -121,7 +121,6 @@ defmodule Realtime.Application do
           id: Realtime.LogThrottle
         ),
         Realtime.Tenants.Cache,
-        Realtime.Channels,
         Realtime.FeatureFlags.Cache,
         Realtime.RateCounter.DynamicSupervisor,
         Realtime.Latency,

lib/realtime/application.ex

@@ -4,8 +4,6 @@ defmodule Realtime.Messages do
   """
 
   alias Realtime.Api.Message
-  alias Realtime.Channels
-  alias Realtime.Tenants.Connect
   alias Realtime.Tenants.Repo
 
   import Ecto.Query, only: [from: 2]
@@ -14,50 +12,32 @@ defmodule Realtime.Messages do
   @default_timeout 5_000
 
   @doc """
-  Persists a broadcast sent over WebSocket for `topic` if storage is enabled for such topic.
+  Persists a broadcast sent over WebSocket for `topic` as a `persistence` message.
 
-  Bypass RLS because the check has already been done on the message broadcast,
-  and then set `broadcasted_at` to avoid re-broadcasting the message twice.
+  Only called after the sender's `persistence` policy authorized it, so the persisted row is always
+  private. Bypasses RLS because that authorization already happened on the broadcast, and sets
+  `broadcasted_at` to avoid re-broadcasting the message twice.
 
   Automatically uses RPC if the database connection is not on the same node.
   """
-  @spec store(DBConnection.conn(), String.t(), String.t(), String.t(), term(), boolean()) ::
-          {:ok, binary()} | {:error, any()} | {:error, :rpc_error, term} | {:error, :storage_disabled}
-  def store(conn, tenant_id, topic, event, payload, private) when node(conn) == node() do
-    if Channels.storage_enabled?(conn, tenant_id, topic) do
-      insert(conn, topic, event, payload, private)
-    else
-      {:error, :storage_disabled}
-    end
+  @spec store(DBConnection.conn(), String.t(), String.t(), term()) ::
+          {:ok, binary()} | {:error, any()} | {:error, :rpc_error, term}
+  def store(conn, topic, event, payload) when node(conn) == node() do
+    insert(conn, topic, event, payload)
   end
 
-  def store(conn, tenant_id, topic, event, payload, private) do
-    Realtime.GenRpc.call(node(conn), __MODULE__, :store, [conn, tenant_id, topic, event, payload, private], key: topic)
-  end
-
-  @doc """
-  Similar to `store/6` but runs asynchronously for callers that don't already hold a connection.
-  """
-  @spec store_async(String.t(), String.t(), String.t(), term()) :: :ok
-  def store_async(tenant_id, topic, event, payload) do
-    Task.Supervisor.start_child(Realtime.TaskSupervisor, fn ->
-      case Connect.lookup_or_start_connection(tenant_id) do
-        {:ok, conn} -> store(conn, tenant_id, topic, event, payload, false)
-        {:error, _reason} -> :skip
-      end
-    end)
-
-    :ok
+  def store(conn, topic, event, payload) do
+    Realtime.GenRpc.call(node(conn), __MODULE__, :store, [conn, topic, event, payload], key: topic)
   end
 
-  defp insert(conn, topic, event, payload, private) do
+  defp insert(conn, topic, event, payload) do
     changeset =
       Message.changeset(%Message{}, %{
         topic: topic,
-        extension: :broadcast,
+        extension: :persistence,
         event: event,
         payload: payload,
-        private: private,
+        private: true,
         broadcasted_at: NaiveDateTime.utc_now(:microsecond)
       })
 
@@ -110,7 +90,7 @@ defmodule Realtime.Messages do
         where:
           m.topic == ^topic and
             m.private == true and
-            m.extension == :broadcast and
+            m.extension in [:broadcast, :persistence] and
             m.inserted_at >= ^since and
             m.inserted_at < ^now,
         limit: ^limit,

lib/realtime/channels.ex

@@ -320,26 +320,35 @@ defmodule Realtime.Tenants.Authorization do
   end
 
   defp check_write_policies(conn, authorization_context, extensions, policies) do
+    with {:ok, policies} <- check_extension_write_policies(conn, authorization_context, extensions, policies),
+         {:ok, persistence_write?} <- authorize_write(conn, authorization_context, :persistence) do
+      {:ok, Policies.update_policies(policies, :persistence, :write, persistence_write?)}
+    end
+  end
+
+  defp check_extension_write_policies(conn, authorization_context, extensions, policies) do
     Enum.reduce_while(@all_extensions, {:ok, policies}, fn extension, {:ok, acc} ->
       if extension in extensions do
-        changeset = Message.changeset(%Message{}, %{topic: authorization_context.topic, extension: extension})
-
-        case Repo.insert(conn, changeset, Message, mode: :savepoint, returning: false) do
-          {:ok, _} ->
-            {:cont, {:ok, Policies.update_policies(acc, extension, :write, true)}}
-
-          {:error, %Postgrex.Error{postgres: %{code: :insufficient_privilege}}} ->
-            {:cont, {:ok, Policies.update_policies(acc, extension, :write, false)}}
-
-          {:error, reason} ->
-            {:halt, {:error, reason}}
+        case authorize_write(conn, authorization_context, extension) do
+          {:ok, allowed?} -> {:cont, {:ok, Policies.update_policies(acc, extension, :write, allowed?)}}
+          {:error, reason} -> {:halt, {:error, reason}}
         end
       else
         {:cont, {:ok, Policies.update_policies(acc, extension, :write, false)}}
       end
     end)
   end
 
+  defp authorize_write(conn, authorization_context, extension) do
+    changeset = Message.changeset(%Message{}, %{topic: authorization_context.topic, extension: extension})
+
+    case Repo.insert(conn, changeset, Message, mode: :savepoint, returning: false) do
+      {:ok, _} -> {:ok, true}
+      {:error, %Postgrex.Error{postgres: %{code: :insufficient_privilege}}} -> {:ok, false}
+      {:error, reason} -> {:error, reason}
+    end
+  end
+
   defp rate_counter(tenant_id) do
     %Tenant{} = tenant = Realtime.Tenants.Cache.get_tenant_by_external_id(tenant_id)
     rate_counter = Realtime.Tenants.authorization_errors_per_second_rate(tenant)

lib/realtime/messages.ex

@@ -2,20 +2,24 @@ defmodule Realtime.Tenants.Authorization.Policies do
   @moduledoc """
   Policies structure that holds the required authorization information for a given connection.
 
-  Currently there are two types of policies:
+  Currently there are three types of policies:
   * Realtime.Tenants.Authorization.Policies.BroadcastPolicies - Used to store the access to Broadcast feature on a given Topic
   * Realtime.Tenants.Authorization.Policies.PresencePolicies - Used to store the access to Presence feature on a given Topic
+  * Realtime.Tenants.Authorization.Policies.PersistencePolicies - Used to store whether messages may be persisted on a given Topic
   """
 
   alias Realtime.Tenants.Authorization.Policies.BroadcastPolicies
   alias Realtime.Tenants.Authorization.Policies.PresencePolicies
+  alias Realtime.Tenants.Authorization.Policies.PersistencePolicies
 
   defstruct broadcast: %BroadcastPolicies{},
-            presence: %PresencePolicies{}
+            presence: %PresencePolicies{},
+            persistence: %PersistencePolicies{}
 
   @type t :: %__MODULE__{
           broadcast: BroadcastPolicies.t(),
-          presence: PresencePolicies.t()
+          presence: PresencePolicies.t(),
+          persistence: PersistencePolicies.t()
         }
 
   @doc """

lib/realtime/tenants/authorization.ex

@@ -0,0 +1,10 @@
+defmodule Realtime.Tenants.Authorization.Policies.PersistencePolicies do
+  @moduledoc """
+  PersistencePolicies structure that holds whether a connection is authorized to persist messages for a given Topic.
+  """
+  defstruct write: nil
+
+  @type t :: %__MODULE__{
+          write: boolean() | nil
+        }
+end

lib/realtime/tenants/authorization/policies.ex

@@ -106,14 +106,6 @@ defmodule Realtime.Tenants.Repo do
     end
   end
 
-  defp result_to_single_struct(
-         {:error, %Postgrex.Error{postgres: %{code: :unique_violation, constraint: "channels_topic_index"}}},
-         _struct,
-         changeset
-       ) do
-    Ecto.Changeset.add_error(changeset, :topic, "has already been taken")
-  end
-
   defp result_to_single_struct({:error, _} = error, _, _), do: error
 
   defp result_to_single_struct({:ok, %Postgrex.Result{rows: []}}, _, _) do

lib/realtime/tenants/authorization/policies/persistence_policies.ex

@@ -4,67 +4,9 @@ defmodule Realtime.Tenants.Migrations.AddBroadcastStorage do
 
   def up do
     execute("ALTER TABLE realtime.messages ADD COLUMN IF NOT EXISTS broadcasted_at timestamp")
-
-    execute("""
-    CREATE TABLE IF NOT EXISTS realtime.channels (
-      id bigserial PRIMARY KEY,
-      topic text NOT NULL,
-      broadcast_storage_enabled_at timestamp,
-      inserted_at timestamp NOT NULL DEFAULT now(),
-      updated_at timestamp NOT NULL DEFAULT now()
-    )
-    """)
-
-    execute("""
-    DO $$ BEGIN
-      IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'channels_topic_index') THEN
-        ALTER TABLE realtime.channels ADD CONSTRAINT channels_topic_index UNIQUE (topic);
-      END IF;
-    END $$;
-    """)
-
-    execute("ALTER TABLE realtime.channels OWNER TO supabase_realtime_admin")
-    execute("GRANT SELECT, INSERT, UPDATE, DELETE ON realtime.channels TO postgres, anon, authenticated, service_role")
-    execute("GRANT USAGE ON SEQUENCE realtime.channels_id_seq TO postgres, anon, authenticated, service_role")
-
-    execute("""
-    CREATE OR REPLACE FUNCTION realtime.enable_broadcast_storage(topic text)
-    RETURNS boolean
-    AS $$
-    BEGIN
-      INSERT INTO realtime.channels (topic, broadcast_storage_enabled_at)
-      VALUES (topic, now())
-      ON CONFLICT ON CONSTRAINT channels_topic_index DO UPDATE
-        SET broadcast_storage_enabled_at = now(),
-            updated_at = now();
-
-      RETURN true;
-    END;
-    $$ LANGUAGE plpgsql;
-    """)
-
-    execute("""
-    CREATE OR REPLACE FUNCTION realtime.disable_broadcast_storage(topic text)
-    RETURNS boolean
-    AS $$
-    BEGIN
-      UPDATE realtime.channels
-      SET broadcast_storage_enabled_at = NULL, updated_at = now()
-      WHERE channels.topic = disable_broadcast_storage.topic;
-
-      RETURN true;
-    END;
-    $$ LANGUAGE plpgsql;
-    """)
-
-    execute("ALTER FUNCTION realtime.enable_broadcast_storage(text) OWNER TO supabase_realtime_admin")
-    execute("ALTER FUNCTION realtime.disable_broadcast_storage(text) OWNER TO supabase_realtime_admin")
   end
 
   def down do
-    execute("DROP FUNCTION IF EXISTS realtime.disable_broadcast_storage(text)")
-    execute("DROP FUNCTION IF EXISTS realtime.enable_broadcast_storage(text)")
-    execute("DROP TABLE IF EXISTS realtime.channels")
     execute("ALTER TABLE realtime.messages DROP COLUMN IF EXISTS broadcasted_at")
   end
 end