Merge branch 'mastodon-websocket-fix' into 'develop'

Fix Mastodon WebSocket authentication

See merge request pleroma/pleroma!4206

2 files changed, 11 insertions, 1 deletions

diff --git a/lib/pleroma/web/endpoint.ex b/lib/pleroma/web/endpoint.ex
index fef907ace..bab3c9fd0 100644
--- a/lib/pleroma/web/endpoint.ex
+++ b/lib/pleroma/web/endpoint.ex
@@ -14,6 +14,7 @@ defmodule Pleroma.Web.Endpoint do
     websocket: [
       path: "/",
       compress: false,
+      connect_info: [:sec_websocket_protocol],
       error_handler: {Pleroma.Web.MastodonAPI.WebsocketHandler, :handle_error, []},
       fullsweep_after: 20
     ]
diff --git a/lib/pleroma/web/mastodon_api/websocket_handler.ex b/lib/pleroma/web/mastodon_api/websocket_handler.ex
index 730295a4c..3ed1cdd6c 100644
--- a/lib/pleroma/web/mastodon_api/websocket_handler.ex
+++ b/lib/pleroma/web/mastodon_api/websocket_handler.ex
@@ -22,7 +22,7 @@ defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do
   # This only prepares the connection and is not in the process yet
   @impl Phoenix.Socket.Transport
   def connect(%{params: params} = transport_info) do
-    with access_token <- Map.get(params, "access_token"),
+    with access_token <- find_access_token(transport_info),
          {:ok, user, oauth_token} <- authenticate_request(access_token),
          {:ok, topic} <-
            Streamer.get_topic(params["stream"], user, oauth_token, params) do
@@ -244,4 +244,13 @@ defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do
   def handle_error(conn, _reason) do
     Plug.Conn.send_resp(conn, 404, "Not Found")
   end
+
+  defp find_access_token(%{
+         connect_info: %{sec_websocket_protocol: [token]}
+       }),
+       do: token
+
+  defp find_access_token(%{params: %{"access_token" => token}}), do: token
+
+  defp find_access_token(_), do: nil
 end