MastodonAPI: Stream fixes.

3 files changed, 13 insertions, 3 deletions

diff --git a/lib/pleroma/web/activity_pub/activity_pub.ex b/lib/pleroma/web/activity_pub/activity_pub.ex
index 9af0f553c..5cbf14868 100644
--- a/lib/pleroma/web/activity_pub/activity_pub.ex
+++ b/lib/pleroma/web/activity_pub/activity_pub.ex
@@ -24,6 +24,9 @@ defmodule Pleroma.Web.ActivityPub.ActivityPub do
          :ok <- maybe_federate(activity) do
       if activity.data["type"] == "Create" and Enum.member?(activity.data["to"], "https://www.w3.org/ns/activitystreams#Public") do
         Pleroma.Web.Streamer.stream("public", activity)
+        if local do
+          Pleroma.Web.Streamer.stream("public:local", activity)
+        end
       end
       {:ok, activity}
     end
diff --git a/lib/pleroma/web/mastodon_api/mastodon_socket.ex b/lib/pleroma/web/mastodon_api/mastodon_socket.ex
index f9c8cec32..af76c8701 100644
--- a/lib/pleroma/web/mastodon_api/mastodon_socket.ex
+++ b/lib/pleroma/web/mastodon_api/mastodon_socket.ex
@@ -1,17 +1,24 @@
 defmodule Pleroma.Web.MastodonAPI.MastodonSocket do
   use Phoenix.Socket
 
+  alias Pleroma.Web.OAuth.Token
+  alias Pleroma.{User, Repo}
+
   transport :streaming, Phoenix.Transports.WebSocket.Raw,
     timeout: :infinity # We never receive data.
 
   def connect(params, socket) do
-    if params["stream"] == "public" do
+    with token when not is_nil(token) <- params["access_token"],
+         %Token{user_id: user_id} <- Repo.get_by(Token, token: token),
+         %User{} = user <- Repo.get(User, user_id),
+         stream when stream in ["public", "public:local"] <- params["stream"] do
       socket = socket
       |> assign(:topic, params["stream"])
+      |> assign(:user, user)
       Pleroma.Web.Streamer.add_socket(params["stream"], socket)
       {:ok, socket}
     else
-      :error
+      _e -> :error
     end
   end
 
diff --git a/lib/pleroma/web/streamer.ex b/lib/pleroma/web/streamer.ex
index 3a7b91743..3b2938676 100644
--- a/lib/pleroma/web/streamer.ex
+++ b/lib/pleroma/web/streamer.ex
@@ -43,7 +43,7 @@ defmodule Pleroma.Web.Streamer do
     Enum.each(topics[topic] || [], fn (socket) ->
       json = %{
         event: "update",
-        payload: Pleroma.Web.MastodonAPI.StatusView.render("status.json", activity: item) |> Poison.encode!
+        payload: Pleroma.Web.MastodonAPI.StatusView.render("status.json", activity: item, for: socket.assigns[:user]) |> Poison.encode!
       } |> Poison.encode!
 
       send socket.transport_pid, {:text, json}