Return stream attribute in server-sent events

4 files changed, 57 insertions, 20 deletions

diff --git a/lib/pleroma/constants.ex b/lib/pleroma/constants.ex
index ed60bcc37..77bc4bfac 100644
--- a/lib/pleroma/constants.ex
+++ b/lib/pleroma/constants.ex
@@ -94,4 +94,8 @@ defmodule Pleroma.Constants do
       "application/activity+json"
     ]
   )
+
+  const(public_streams,
+    do: ["public", "public:local", "public:media", "public:local:media"]
+  )
 end
diff --git a/lib/pleroma/web/mastodon_api/websocket_handler.ex b/lib/pleroma/web/mastodon_api/websocket_handler.ex
index 6233c3340..2707673ba 100644
--- a/lib/pleroma/web/mastodon_api/websocket_handler.ex
+++ b/lib/pleroma/web/mastodon_api/websocket_handler.ex
@@ -89,11 +89,11 @@ defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do
     {:ok, state}
   end
 
-  def websocket_info({:render_with_user, view, template, item}, state) do
+  def websocket_info({:render_with_user, view, template, item, topic}, state) do
     user = %User{} = User.get_cached_by_ap_id(state.user.ap_id)
 
     unless Streamer.filtered_by_user?(user, item) do
-      websocket_info({:text, view.render(template, item, user)}, %{state | user: user})
+      websocket_info({:text, view.render(template, item, user, topic)}, %{state | user: user})
     else
       {:ok, state}
     end
diff --git a/lib/pleroma/web/streamer.ex b/lib/pleroma/web/streamer.ex
index 6b5fdbf0c..70e46617a 100644
--- a/lib/pleroma/web/streamer.ex
+++ b/lib/pleroma/web/streamer.ex
@@ -4,6 +4,7 @@
 
 defmodule Pleroma.Web.Streamer do
   require Logger
+  require Pleroma.Constants
 
   alias Pleroma.Activity
   alias Pleroma.Chat.MessageReference
@@ -24,7 +25,7 @@ defmodule Pleroma.Web.Streamer do
 
   def registry, do: @registry
 
-  @public_streams ["public", "public:local", "public:media", "public:local:media"]
+  @public_streams Pleroma.Constants.public_streams()
   @local_streams ["public:local", "public:local:media"]
   @user_streams ["user", "user:notification", "direct", "user:pleroma_chat"]
 
@@ -223,8 +224,8 @@ defmodule Pleroma.Web.Streamer do
   end
 
   defp do_stream("follow_relationship", item) do
-    text = StreamerView.render("follow_relationships_update.json", item)
     user_topic = "user:#{item.follower.id}"
+    text = StreamerView.render("follow_relationships_update.json", item, user_topic)
 
     Logger.debug("Trying to push follow relationship update to #{user_topic}\n\n")
 
@@ -270,9 +271,11 @@ defmodule Pleroma.Web.Streamer do
 
   defp do_stream(topic, %Notification{} = item)
        when topic in ["user", "user:notification"] do
-    Registry.dispatch(@registry, "#{topic}:#{item.user_id}", fn list ->
+    user_topic = "#{topic}:#{item.user_id}"
+
+    Registry.dispatch(@registry, user_topic, fn list ->
       Enum.each(list, fn {pid, _auth} ->
-        send(pid, {:render_with_user, StreamerView, "notification.json", item})
+        send(pid, {:render_with_user, StreamerView, "notification.json", item, user_topic})
       end)
     end)
   end
@@ -281,7 +284,7 @@ defmodule Pleroma.Web.Streamer do
        when topic in ["user", "user:pleroma_chat"] do
     topic = "#{topic}:#{user.id}"
 
-    text = StreamerView.render("chat_update.json", %{chat_message_reference: cm_ref})
+    text = StreamerView.render("chat_update.json", %{chat_message_reference: cm_ref}, topic)
 
     Registry.dispatch(@registry, topic, fn list ->
       Enum.each(list, fn {pid, _auth} ->
@@ -309,7 +312,7 @@ defmodule Pleroma.Web.Streamer do
   end
 
   defp push_to_socket(topic, %Participation{} = participation) do
-    rendered = StreamerView.render("conversation.json", participation)
+    rendered = StreamerView.render("conversation.json", participation, topic)
 
     Registry.dispatch(@registry, topic, fn list ->
       Enum.each(list, fn {pid, _} ->
@@ -337,12 +340,15 @@ defmodule Pleroma.Web.Streamer do
       Pleroma.Activity.get_create_by_object_ap_id(item.object.data["id"])
       |> Map.put(:object, item.object)
 
-    anon_render = StreamerView.render("status_update.json", create_activity)
+    anon_render = StreamerView.render("status_update.json", create_activity, topic)
 
     Registry.dispatch(@registry, topic, fn list ->
       Enum.each(list, fn {pid, auth?} ->
         if auth? do
-          send(pid, {:render_with_user, StreamerView, "status_update.json", create_activity})
+          send(
+            pid,
+            {:render_with_user, StreamerView, "status_update.json", create_activity, topic}
+          )
         else
           send(pid, {:text, anon_render})
         end
@@ -351,12 +357,13 @@ defmodule Pleroma.Web.Streamer do
   end
 
   defp push_to_socket(topic, item) do
-    anon_render = StreamerView.render("update.json", item)
+    Logger.debug("topic=#{topic}")
+    anon_render = StreamerView.render("update.json", item, topic)
 
     Registry.dispatch(@registry, topic, fn list ->
       Enum.each(list, fn {pid, auth?} ->
         if auth? do
-          send(pid, {:render_with_user, StreamerView, "update.json", item})
+          send(pid, {:render_with_user, StreamerView, "update.json", item, topic})
         else
           send(pid, {:text, anon_render})
         end
diff --git a/lib/pleroma/web/views/streamer_view.ex b/lib/pleroma/web/views/streamer_view.ex
index 0cdcb1918..f591da9a6 100644
--- a/lib/pleroma/web/views/streamer_view.ex
+++ b/lib/pleroma/web/views/streamer_view.ex
@@ -11,8 +11,11 @@ defmodule Pleroma.Web.StreamerView do
   alias Pleroma.User
   alias Pleroma.Web.MastodonAPI.NotificationView
 
-  def render("update.json", %Activity{} = activity, %User{} = user) do
+  require Pleroma.Constants
+
+  def render("update.json", %Activity{} = activity, %User{} = user, topic) do
     %{
+      stream: render("stream.json", %{topic: topic}),
       event: "update",
       payload:
         Pleroma.Web.MastodonAPI.StatusView.render(
@@ -25,8 +28,9 @@ defmodule Pleroma.Web.StreamerView do
     |> Jason.encode!()
   end
 
-  def render("status_update.json", %Activity{} = activity, %User{} = user) do
+  def render("status_update.json", %Activity{} = activity, %User{} = user, topic) do
     %{
+      stream: render("stream.json", %{topic: topic}),
       event: "status.update",
       payload:
         Pleroma.Web.MastodonAPI.StatusView.render(
@@ -39,8 +43,9 @@ defmodule Pleroma.Web.StreamerView do
     |> Jason.encode!()
   end
 
-  def render("notification.json", %Notification{} = notify, %User{} = user) do
+  def render("notification.json", %Notification{} = notify, %User{} = user, topic) do
     %{
+      stream: render("stream.json", %{topic: topic}),
       event: "notification",
       payload:
         NotificationView.render(
@@ -52,8 +57,9 @@ defmodule Pleroma.Web.StreamerView do
     |> Jason.encode!()
   end
 
-  def render("update.json", %Activity{} = activity) do
+  def render("update.json", %Activity{} = activity, topic) do
     %{
+      stream: render("stream.json", %{topic: topic}),
       event: "update",
       payload:
         Pleroma.Web.MastodonAPI.StatusView.render(
@@ -65,8 +71,9 @@ defmodule Pleroma.Web.StreamerView do
     |> Jason.encode!()
   end
 
-  def render("status_update.json", %Activity{} = activity) do
+  def render("status_update.json", %Activity{} = activity, topic) do
     %{
+      stream: render("stream.json", %{topic: topic}),
       event: "status.update",
       payload:
         Pleroma.Web.MastodonAPI.StatusView.render(
@@ -78,7 +85,7 @@ defmodule Pleroma.Web.StreamerView do
     |> Jason.encode!()
   end
 
-  def render("chat_update.json", %{chat_message_reference: cm_ref}) do
+  def render("chat_update.json", %{chat_message_reference: cm_ref}, topic) do
     # Explicitly giving the cmr for the object here, so we don't accidentally
     # send a later 'last_message' that was inserted between inserting this and
     # streaming it out
@@ -93,6 +100,7 @@ defmodule Pleroma.Web.StreamerView do
       )
 
     %{
+      stream: render("stream.json", %{topic: topic}),
       event: "pleroma:chat_update",
       payload:
         representation
@@ -101,8 +109,9 @@ defmodule Pleroma.Web.StreamerView do
     |> Jason.encode!()
   end
 
-  def render("follow_relationships_update.json", item) do
+  def render("follow_relationships_update.json", item, topic) do
     %{
+      stream: render("stream.json", %{topic: topic}),
       event: "pleroma:follow_relationships_update",
       payload:
         %{
@@ -123,8 +132,9 @@ defmodule Pleroma.Web.StreamerView do
     |> Jason.encode!()
   end
 
-  def render("conversation.json", %Participation{} = participation) do
+  def render("conversation.json", %Participation{} = participation, topic) do
     %{
+      stream: render("stream.json", %{topic: topic}),
       event: "conversation",
       payload:
         Pleroma.Web.MastodonAPI.ConversationView.render("participation.json", %{
@@ -150,6 +160,22 @@ defmodule Pleroma.Web.StreamerView do
     |> Jason.encode!()
   end
 
+  def render("stream.json", %{topic: "user:pleroma_chat:" <> _}), do: ["user:pleroma_chat"]
+  def render("stream.json", %{topic: "user:notification:" <> _}), do: ["user:notification"]
+  def render("stream.json", %{topic: "user:" <> _}), do: ["user"]
+  def render("stream.json", %{topic: "direct:" <> _}), do: ["direct"]
+  def render("stream.json", %{topic: "list:" <> id}), do: ["list", id]
+  def render("stream.json", %{topic: "hashtag:" <> tag}), do: ["hashtag", tag]
+
+  def render("stream.json", %{topic: "public:remote:media:" <> instance}),
+    do: ["public:remote:media", instance]
+
+  def render("stream.json", %{topic: "public:remote:" <> instance}),
+    do: ["public:remote", instance]
+
+  def render("stream.json", %{topic: stream}) when stream in Pleroma.Constants.public_streams(),
+    do: [stream]
+
   defp maybe_error(%{error: :bad_topic}), do: %{error: "bad_topic"}
   defp maybe_error(%{error: :unauthorized}), do: %{error: "unauthorized"}
   defp maybe_error(%{error: :already_authenticated}), do: %{error: "already_authenticated"}