From 64ad451a7b8a2ac89079a1bc32680e9cf08ef24e Mon Sep 17 00:00:00 2001
From: Mark Felder <feld@feld.me>
Date: Tue, 13 Feb 2024 19:21:45 -0500
Subject: [PATCH 1/3] Websocket refactor to use Phoenix.Socket.Transport

This will make us compatible with Cowboy and Bandit
---
 config/config.exs                             |   9 +-
 lib/phoenix/transports/web_socket/raw.ex      |  93 -------
 lib/pleroma/web/endpoint.ex                   |   9 +
 .../web/mastodon_api/websocket_handler.ex     | 256 +++++++++---------
 4 files changed, 133 insertions(+), 234 deletions(-)
 delete mode 100644 lib/phoenix/transports/web_socket/raw.ex

diff --git a/config/config.exs b/config/config.exs
index bb17ab145..435387a64 100644
--- a/config/config.exs
+++ b/config/config.exs
@@ -114,14 +114,7 @@
 config :pleroma, Pleroma.Web.Endpoint,
   url: [host: "localhost"],
   http: [
-    ip: {127, 0, 0, 1},
-    dispatch: [
-      {:_,
-       [
-         {"/api/v1/streaming", Pleroma.Web.MastodonAPI.WebsocketHandler, []},
-         {:_, Plug.Cowboy.Handler, {Pleroma.Web.Endpoint, []}}
-       ]}
-    ]
+    ip: {127, 0, 0, 1}
   ],
   protocol: "https",
   secret_key_base: "aK4Abxf29xU9TTDKre9coZPUgevcVCFQJe/5xP/7Lt4BEif6idBIbjupVbOrbKxl",
diff --git a/lib/phoenix/transports/web_socket/raw.ex b/lib/phoenix/transports/web_socket/raw.ex
deleted file mode 100644
index cf4fda79f..000000000
--- a/lib/phoenix/transports/web_socket/raw.ex
+++ /dev/null
@@ -1,93 +0,0 @@
-# Pleroma: A lightweight social networking server
-# Copyright © 2017-2022 Pleroma Authors <https://pleroma.social/>
-# SPDX-License-Identifier: AGPL-3.0-only
-
-defmodule Phoenix.Transports.WebSocket.Raw do
-  import Plug.Conn,
-    only: [
-      fetch_query_params: 1,
-      send_resp: 3
-    ]
-
-  alias Phoenix.Socket.Transport
-
-  def default_config do
-    [
-      timeout: 60_000,
-      transport_log: false,
-      cowboy: Phoenix.Endpoint.CowboyWebSocket
-    ]
-  end
-
-  def init(%Plug.Conn{method: "GET"} = conn, {endpoint, handler, transport}) do
-    {_, opts} = handler.__transport__(transport)
-
-    conn =
-      conn
-      |> fetch_query_params
-      |> Transport.transport_log(opts[:transport_log])
-      |> Transport.check_origin(handler, endpoint, opts)
-
-    case conn do
-      %{halted: false} = conn ->
-        case handler.connect(%{
-               endpoint: endpoint,
-               transport: transport,
-               options: [serializer: nil],
-               params: conn.params
-             }) do
-          {:ok, socket} ->
-            {:ok, conn, {__MODULE__, {socket, opts}}}
-
-          :error ->
-            send_resp(conn, :forbidden, "")
-            {:error, conn}
-        end
-
-      _ ->
-        {:error, conn}
-    end
-  end
-
-  def init(conn, _) do
-    send_resp(conn, :bad_request, "")
-    {:error, conn}
-  end
-
-  def ws_init({socket, config}) do
-    Process.flag(:trap_exit, true)
-    {:ok, %{socket: socket}, config[:timeout]}
-  end
-
-  def ws_handle(op, data, state) do
-    state.socket.handler
-    |> apply(:handle, [op, data, state])
-    |> case do
-      {op, data} ->
-        {:reply, {op, data}, state}
-
-      {op, data, state} ->
-        {:reply, {op, data}, state}
-
-      %{} = state ->
-        {:ok, state}
-
-      _ ->
-        {:ok, state}
-    end
-  end
-
-  def ws_info({_, _} = tuple, state) do
-    {:reply, tuple, state}
-  end
-
-  def ws_info(_tuple, state), do: {:ok, state}
-
-  def ws_close(state) do
-    ws_handle(:closed, :normal, state)
-  end
-
-  def ws_terminate(reason, state) do
-    ws_handle(:closed, reason, state)
-  end
-end
diff --git a/lib/pleroma/web/endpoint.ex b/lib/pleroma/web/endpoint.ex
index 307fa069e..c56362049 100644
--- a/lib/pleroma/web/endpoint.ex
+++ b/lib/pleroma/web/endpoint.ex
@@ -9,6 +9,15 @@ defmodule Pleroma.Web.Endpoint do
 
   alias Pleroma.Config
 
+  socket("/api/v1/streaming", Pleroma.Web.MastodonAPI.WebsocketHandler,
+    longpoll: false,
+    websocket: [
+      path: "/",
+      compress: false,
+      error_handler: {Pleroma.Web.MastodonAPI.WebsocketHandler, :handle_error, []}
+    ]
+  )
+
   socket("/socket", Pleroma.Web.UserSocket,
     websocket: [
       path: "/websocket",
diff --git a/lib/pleroma/web/mastodon_api/websocket_handler.ex b/lib/pleroma/web/mastodon_api/websocket_handler.ex
index 07c2b62e3..bb27d806d 100644
--- a/lib/pleroma/web/mastodon_api/websocket_handler.ex
+++ b/lib/pleroma/web/mastodon_api/websocket_handler.ex
@@ -11,28 +11,21 @@ defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do
   alias Pleroma.Web.Streamer
   alias Pleroma.Web.StreamerView
 
-  @behaviour :cowboy_websocket
+  @behaviour Phoenix.Socket.Transport
 
   # Client ping period.
   @tick :timer.seconds(30)
-  # Cowboy timeout period.
-  @timeout :timer.seconds(60)
-  # Hibernate every X messages
-  @hibernate_every 100
 
-  def init(%{qs: qs} = req, state) do
-    with params <- Enum.into(:cow_qs.parse_qs(qs), %{}),
-         sec_websocket <- :cowboy_req.header("sec-websocket-protocol", req, nil),
-         access_token <- Map.get(params, "access_token"),
-         {:ok, user, oauth_token} <- authenticate_request(access_token, sec_websocket),
-         {:ok, topic} <- Streamer.get_topic(params["stream"], user, oauth_token, params) do
-      req =
-        if sec_websocket do
-          :cowboy_req.set_resp_header("sec-websocket-protocol", sec_websocket, req)
-        else
-          req
-        end
+  @impl Phoenix.Socket.Transport
+  def child_spec(_opts), do: :ignore
 
+  # 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"),
+         {:ok, user, oauth_token} <- authenticate_request(access_token),
+         {:ok, topic} <-
+           Streamer.get_topic(params["stream"], user, oauth_token, params) do
       topics =
         if topic do
           [topic]
@@ -40,41 +33,40 @@ def init(%{qs: qs} = req, state) do
           []
         end
 
-      {:cowboy_websocket, req,
-       %{user: user, topics: topics, oauth_token: oauth_token, count: 0, timer: nil},
-       %{idle_timeout: @timeout}}
+      state = %{
+        user: user,
+        topics: topics,
+        oauth_token: oauth_token,
+        count: 0,
+        timer: nil
+      }
+
+      {:ok, state}
     else
       {:error, :bad_topic} ->
-        Logger.debug("#{__MODULE__} bad topic #{inspect(req)}")
-        req = :cowboy_req.reply(404, req)
-        {:ok, req, state}
+        Logger.debug("#{__MODULE__} bad topic #{inspect(transport_info)}")
+
+        {:error, :bad_topic}
 
       {:error, :unauthorized} ->
-        Logger.debug("#{__MODULE__} authentication error: #{inspect(req)}")
-        req = :cowboy_req.reply(401, req)
-        {:ok, req, state}
+        Logger.debug("#{__MODULE__} authentication error: #{inspect(transport_info)}")
+        {:error, :unauthorized}
     end
   end
 
-  def websocket_init(state) do
-    Logger.debug(
-      "#{__MODULE__} accepted websocket connection for user #{(state.user || %{id: "anonymous"}).id}, topics #{state.topics}"
-    )
-
+  # All subscriptions/links and messages cannot be created
+  # until the processed is launched with init/1
+  @impl Phoenix.Socket.Transport
+  def init(state) do
     Enum.each(state.topics, fn topic -> Streamer.add_socket(topic, state.oauth_token) end)
-    {:ok, %{state | timer: timer()}}
+
+    Process.send_after(self(), :ping, @tick)
+
+    {:ok, state}
   end
 
-  # Client's Pong frame.
-  def websocket_handle(:pong, state) do
-    if state.timer, do: Process.cancel_timer(state.timer)
-    {:ok, %{state | timer: timer()}}
-  end
-
-  # We only receive pings for now
-  def websocket_handle(:ping, state), do: {:ok, state}
-
-  def websocket_handle({:text, text}, state) do
+  @impl Phoenix.Socket.Transport
+  def handle_in({text, [opcode: :text]}, state) do
     with {:ok, %{} = event} <- Jason.decode(text) do
       handle_client_event(event, state)
     else
@@ -84,50 +76,47 @@ def websocket_handle({:text, text}, state) do
     end
   end
 
-  def websocket_handle(frame, state) do
+  def handle_in(frame, state) do
     Logger.error("#{__MODULE__} received frame: #{inspect(frame)}")
     {:ok, state}
   end
 
-  def websocket_info({:render_with_user, view, template, item, topic}, state) do
+  @impl Phoenix.Socket.Transport
+  def handle_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, topic)}, %{state | user: user})
+      message = view.render(template, item, user, topic)
+      {:push, {:text, message}, %{state | user: user}}
     else
       {:ok, state}
     end
   end
 
-  def websocket_info({:text, message}, state) do
-    # If the websocket processed X messages, force an hibernate/GC.
-    # We don't hibernate at every message to balance CPU usage/latency with RAM usage.
-    if state.count > @hibernate_every do
-      {:reply, {:text, message}, %{state | count: 0}, :hibernate}
-    else
-      {:reply, {:text, message}, %{state | count: state.count + 1}}
-    end
+  def handle_info({:text, text}, state) do
+    {:push, {:text, text}, state}
   end
 
-  # Ping tick. We don't re-queue a timer there, it is instead queued when :pong is received.
-  # As we hibernate there, reset the count to 0.
-  # If the client misses :pong, Cowboy will automatically timeout the connection after
-  # `@idle_timeout`.
-  def websocket_info(:tick, state) do
-    {:reply, :ping, %{state | timer: nil, count: 0}, :hibernate}
+  def handle_info(:ping, state) do
+    Process.send_after(self(), :ping, @tick)
+
+    {:push, {:ping, ""}, state}
   end
 
-  def websocket_info(:close, state) do
-    {:stop, state}
+  def handle_info(:close, state) do
+    {:stop, {:closed, 'connection closed by server'}, state}
   end
 
-  # State can be `[]` only in case we terminate before switching to websocket,
-  # we already log errors for these cases in `init/1`, so just do nothing here
-  def terminate(_reason, _req, []), do: :ok
+  def handle_info(msg, state) do
+    Logger.debug("#{__MODULE__} received info: #{inspect(msg)}")
 
-  def terminate(reason, _req, state) do
+    {:ok, state}
+  end
+
+  @impl Phoenix.Socket.Transport
+  def terminate(reason, state) do
     Logger.debug(
-      "#{__MODULE__} terminating websocket connection for user #{(state.user || %{id: "anonymous"}).id}, topics #{state.topics || "?"}: #{inspect(reason)}"
+      "#{__MODULE__} terminating websocket connection for user #{(state.user || %{id: "anonymous"}).id}, topics #{state.topics || "?"}: #{inspect(reason)})"
     )
 
     Enum.each(state.topics, fn topic -> Streamer.remove_socket(topic) end)
@@ -135,16 +124,13 @@ def terminate(reason, _req, state) do
   end
 
   # Public streams without authentication.
-  defp authenticate_request(nil, nil) do
+  defp authenticate_request(nil) do
     {:ok, nil, nil}
   end
 
   # Authenticated streams.
-  defp authenticate_request(access_token, sec_websocket) do
-    token = access_token || sec_websocket
-
-    with true <- is_bitstring(token),
-         oauth_token = %Token{user_id: user_id} <- Repo.get_by(Token, token: token),
+  defp authenticate_request(access_token) do
+    with oauth_token = %Token{user_id: user_id} <- Repo.get_by(Token, token: access_token),
          user = %User{} <- User.get_cached_by_id(user_id) do
       {:ok, user, oauth_token}
     else
@@ -152,36 +138,32 @@ defp authenticate_request(access_token, sec_websocket) do
     end
   end
 
-  defp timer do
-    Process.send_after(self(), :tick, @tick)
-  end
-
   defp handle_client_event(%{"type" => "subscribe", "stream" => _topic} = params, state) do
     with {_, {:ok, topic}} <-
            {:topic, Streamer.get_topic(params["stream"], state.user, state.oauth_token, params)},
          {_, false} <- {:subscribed, topic in state.topics} do
       Streamer.add_socket(topic, state.oauth_token)
 
-      {[
-         {:text,
-          StreamerView.render("pleroma_respond.json", %{type: "subscribe", result: "success"})}
-       ], %{state | topics: [topic | state.topics]}}
+      message =
+        StreamerView.render("pleroma_respond.json", %{type: "subscribe", result: "success"})
+
+      {:reply, :ok, {:text, message}, %{state | topics: [topic | state.topics]}}
     else
       {:subscribed, true} ->
-        {[
-           {:text,
-            StreamerView.render("pleroma_respond.json", %{type: "subscribe", result: "ignored"})}
-         ], state}
+        message =
+          StreamerView.render("pleroma_respond.json", %{type: "subscribe", result: "ignored"})
+
+        {:reply, :error, {:text, message}, state}
 
       {:topic, {:error, error}} ->
-        {[
-           {:text,
-            StreamerView.render("pleroma_respond.json", %{
-              type: "subscribe",
-              result: "error",
-              error: error
-            })}
-         ], state}
+        message =
+          StreamerView.render("pleroma_respond.json", %{
+            type: "subscribe",
+            result: "error",
+            error: error
+          })
+
+        {:reply, :error, {:text, message}, state}
     end
   end
 
@@ -191,26 +173,26 @@ defp handle_client_event(%{"type" => "unsubscribe", "stream" => _topic} = params
          {_, true} <- {:subscribed, topic in state.topics} do
       Streamer.remove_socket(topic)
 
-      {[
-         {:text,
-          StreamerView.render("pleroma_respond.json", %{type: "unsubscribe", result: "success"})}
-       ], %{state | topics: List.delete(state.topics, topic)}}
+      message =
+        StreamerView.render("pleroma_respond.json", %{type: "unsubscribe", result: "success"})
+
+      {:reply, :ok, {:text, message}, %{state | topics: List.delete(state.topics, topic)}}
     else
       {:subscribed, false} ->
-        {[
-           {:text,
-            StreamerView.render("pleroma_respond.json", %{type: "unsubscribe", result: "ignored"})}
-         ], state}
+        message =
+          StreamerView.render("pleroma_respond.json", %{type: "unsubscribe", result: "ignored"})
+
+        {:reply, :error, {:text, message}, state}
 
       {:topic, {:error, error}} ->
-        {[
-           {:text,
-            StreamerView.render("pleroma_respond.json", %{
-              type: "unsubscribe",
-              result: "error",
-              error: error
-            })}
-         ], state}
+        message =
+          StreamerView.render("pleroma_respond.json", %{
+            type: "unsubscribe",
+            result: "error",
+            error: error
+          })
+
+        {:reply, :error, {:text, message}, state}
     end
   end
 
@@ -219,39 +201,47 @@ defp handle_client_event(
          state
        ) do
     with {:auth, nil, nil} <- {:auth, state.user, state.oauth_token},
-         {:ok, user, oauth_token} <- authenticate_request(access_token, nil) do
-      {[
-         {:text,
-          StreamerView.render("pleroma_respond.json", %{
-            type: "pleroma:authenticate",
-            result: "success"
-          })}
-       ], %{state | user: user, oauth_token: oauth_token}}
+         {:ok, user, oauth_token} <- authenticate_request(access_token) do
+      message =
+        StreamerView.render("pleroma_respond.json", %{
+          type: "pleroma:authenticate",
+          result: "success"
+        })
+
+      {:reply, :ok, {:text, message}, %{state | user: user, oauth_token: oauth_token}}
     else
       {:auth, _, _} ->
-        {[
-           {:text,
-            StreamerView.render("pleroma_respond.json", %{
-              type: "pleroma:authenticate",
-              result: "error",
-              error: :already_authenticated
-            })}
-         ], state}
+        message =
+          StreamerView.render("pleroma_respond.json", %{
+            type: "pleroma:authenticate",
+            result: "error",
+            error: :already_authenticated
+          })
+
+        {:reply, :error, {:text, message}, state}
 
       _ ->
-        {[
-           {:text,
-            StreamerView.render("pleroma_respond.json", %{
-              type: "pleroma:authenticate",
-              result: "error",
-              error: :unauthorized
-            })}
-         ], state}
+        message =
+          StreamerView.render("pleroma_respond.json", %{
+            type: "pleroma:authenticate",
+            result: "error",
+            error: :unauthorized
+          })
+
+        {:reply, :error, {:text, message}, state}
     end
   end
 
   defp handle_client_event(params, state) do
     Logger.error("#{__MODULE__} received unknown event: #{inspect(params)}")
-    {[], state}
+    {:ok, state}
+  end
+
+  def handle_error(conn, :unauthorized) do
+    Plug.Conn.send_resp(conn, 401, "Unauthorized")
+  end
+
+  def handle_error(conn, _reason) do
+    Plug.Conn.send_resp(conn, 404, "Not Found")
   end
 end

From d0f4b2b02fc3aee3f08239d9c188ca5a2e8ad482 Mon Sep 17 00:00:00 2001
From: Mark Felder <feld@feld.me>
Date: Wed, 14 Feb 2024 14:15:24 -0500
Subject: [PATCH 2/3] Remove invalid test

It is not allowed to use the Sec-WebSocket-Protocol header for arbitrary values. This was possible due to the raw websocket handling we were doing with Cowboy, but Phoenix.Socket.Transport does not allow this as the value of this header is compared against a static list of subprotocols.

https://hexdocs.pm/phoenix/Phoenix.Endpoint.html#socket/3-websocket-configuration

Additionally I cannot find anywhere that we depended on this behavior. Setting the Sec-WebSocket-Protocol header does not appear to be a part of PleromaFE.
---
 test/pleroma/integration/mastodon_websocket_test.exs | 11 -----------
 1 file changed, 11 deletions(-)

diff --git a/test/pleroma/integration/mastodon_websocket_test.exs b/test/pleroma/integration/mastodon_websocket_test.exs
index a2c20f0a6..a0ffddf8d 100644
--- a/test/pleroma/integration/mastodon_websocket_test.exs
+++ b/test/pleroma/integration/mastodon_websocket_test.exs
@@ -268,17 +268,6 @@ test "accepts the 'user:notification' stream", %{token: token} = _state do
       end)
     end
 
-    test "accepts valid token on Sec-WebSocket-Protocol header", %{token: token} do
-      assert {:ok, _} = start_socket("?stream=user", [{"Sec-WebSocket-Protocol", token.token}])
-
-      capture_log(fn ->
-        assert {:error, %WebSockex.RequestError{code: 401}} =
-                 start_socket("?stream=user", [{"Sec-WebSocket-Protocol", "I am a friend"}])
-
-        Process.sleep(30)
-      end)
-    end
-
     test "accepts valid token on client-sent event", %{token: token} do
       assert {:ok, pid} = start_socket()
 

From 6be129ead2ff5d6a19edf7230d102aa51a731b03 Mon Sep 17 00:00:00 2001
From: Mark Felder <feld@feld.me>
Date: Wed, 14 Feb 2024 14:19:24 -0500
Subject: [PATCH 3/3] Websocket refactor changelog

---
 changelog.d/websocket-refactor.change | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 changelog.d/websocket-refactor.change

diff --git a/changelog.d/websocket-refactor.change b/changelog.d/websocket-refactor.change
new file mode 100644
index 000000000..3c447832b
--- /dev/null
+++ b/changelog.d/websocket-refactor.change
@@ -0,0 +1 @@
+Refactor the Mastodon /api/v1/streaming websocket handler to use Phoenix.Socket.Transport