From 392bd9ef562a98bb026d49ddd83f108cbe9fdb19 Mon Sep 17 00:00:00 2001
From: csaurus <csaurus@mailbox.org>
Date: Thu, 10 May 2018 22:15:42 -0400
Subject: [PATCH] Stream function to handle direct messages.

---
 lib/pleroma/web/activity_pub/activity_pub.ex | 14 ++++++--------
 lib/pleroma/web/streamer.ex                  | 17 +++++++++++++++--
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/lib/pleroma/web/activity_pub/activity_pub.ex b/lib/pleroma/web/activity_pub/activity_pub.ex
index 38e3a84fb..8c1ba1ea3 100644
--- a/lib/pleroma/web/activity_pub/activity_pub.ex
+++ b/lib/pleroma/web/activity_pub/activity_pub.ex
@@ -42,18 +42,16 @@ def insert(map, local \\ true) when is_map(map) do
   def stream_out(activity) do
     if activity.data["type"] in ["Create", "Announce"] do
       Pleroma.Web.Streamer.stream("user", activity)
-      direct? = activity.data["object"]["visibility"] == "direct"
 
-      cond do
-        direct? ->
-          Pleroma.Web.Streamer.stream("direct", activity)
+      visibility = Pleroma.Web.MastodonAPI.StatusView.get_visibility(activity.data["object"])
 
-        Enum.member?(activity.data["to"], "https://www.w3.org/ns/activitystreams#Public") ->
+      case visibility do
+        "public" ->
           Pleroma.Web.Streamer.stream("public", activity)
+          if activity.local, do: Pleroma.Web.Streamer.stream("public:local", activity)
 
-          if activity.local do
-            Pleroma.Web.Streamer.stream("public:local", activity)
-          end
+        "direct" ->
+          Pleroma.Web.Streamer.stream("direct", activity)
       end
     end
   end
diff --git a/lib/pleroma/web/streamer.ex b/lib/pleroma/web/streamer.ex
index 33041ec12..6aac472dc 100644
--- a/lib/pleroma/web/streamer.ex
+++ b/lib/pleroma/web/streamer.ex
@@ -46,6 +46,19 @@ def handle_cast(%{action: :ping}, topics) do
     {:noreply, topics}
   end
 
+  def handle_cast(%{action: :stream, topic: "direct", item: item}, topics) do
+    recipient_topics =
+      User.get_recipients_from_activity(item)
+      |> Enum.map(fn %{id: id} -> "direct:#{id}" end)
+
+    Enum.each(recipient_topics || [], fn user_topic ->
+      Logger.debug("Trying to push direct message to #{user_topic}\n\n")
+      push_to_socket(topics, user_topic, item)
+    end)
+
+    {:noreply, topics}
+  end
+
   def handle_cast(%{action: :stream, topic: "user", item: %Notification{} = item}, topics) do
     topic = "user:#{item.user_id}"
 
@@ -137,8 +150,8 @@ def push_to_socket(topics, topic, item) do
     end)
   end
 
-  defp internal_topic("user", socket) do
-    "user:#{socket.assigns[:user].id}"
+  defp internal_topic(topic, socket) when topic in ~w[user, direct] do
+    "#{topic}:#{socket.assigns[:user].id}"
   end
 
   defp internal_topic(topic, _), do: topic