git.squeep.com Git - akkoma/blob - lib/pleroma/web/mastodon_api/websocket_handler.ex

   1 # Pleroma: A lightweight social networking server
   2 # Copyright © 2017-2020 Pleroma Authors <https://pleroma.social/>
   3 # SPDX-License-Identifier: AGPL-3.0-only
   4
   5 defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do
   6   require Logger
   7
   8   alias Pleroma.Repo
   9   alias Pleroma.User
  10   alias Pleroma.Web.OAuth.Token
  11   alias Pleroma.Web.Streamer
  12
  13   @behaviour :cowboy_websocket
  14
  15   # Handled by periodic keepalive in Pleroma.Web.Streamer.Ping.
  16   @timeout :infinity
  17
  18   def init(%{qs: qs} = req, state) do
  19     with params <- Enum.into(:cow_qs.parse_qs(qs), %{}),
  20          sec_websocket <- :cowboy_req.header("sec-websocket-protocol", req, nil),
  21          access_token <- Map.get(params, "access_token"),
  22          {:ok, user} <- authenticate_request(access_token, sec_websocket),
  23          {:ok, topic} <- Streamer.get_topic(Map.get(params, "stream"), user, params) do
  24       req =
  25         if sec_websocket do
  26           :cowboy_req.set_resp_header("sec-websocket-protocol", sec_websocket, req)
  27         else
  28           req
  29         end
  30
  31       {:cowboy_websocket, req, %{user: user, topic: topic}, %{idle_timeout: @timeout}}
  32     else
  33       {:error, :bad_topic} ->
  34         Logger.debug("#{__MODULE__} bad topic #{inspect(req)}")
  35         {:ok, req} = :cowboy_req.reply(404, req)
  36         {:ok, req, state}
  37
  38       {:error, :unauthorized} ->
  39         Logger.debug("#{__MODULE__} authentication error: #{inspect(req)}")
  40         {:ok, req} = :cowboy_req.reply(401, req)
  41         {:ok, req, state}
  42     end
  43   end
  44
  45   def websocket_init(state) do
  46     send(self(), :subscribe)
  47     {:ok, state}
  48   end
  49
  50   # We never receive messages.
  51   def websocket_handle(_frame, state) do
  52     {:ok, state}
  53   end
  54
  55   def websocket_info(:subscribe, state) do
  56     Logger.debug(
  57       "#{__MODULE__} accepted websocket connection for user #{
  58         (state.user || %{id: "anonymous"}).id
  59       }, topic #{state.topic}"
  60     )
  61
  62     Streamer.add_socket(state.topic, streamer_socket(state))
  63     {:ok, state}
  64   end
  65
  66   def websocket_info({:text, message}, state) do
  67     {:reply, {:text, message}, state}
  68   end
  69
  70   def terminate(reason, _req, state) do
  71     Logger.debug(
  72       "#{__MODULE__} terminating websocket connection for user #{
  73         (state.user || %{id: "anonymous"}).id
  74       }, topic #{state.topic || "?"}: #{inspect(reason)}"
  75     )
  76
  77     Streamer.remove_socket(state.topic, streamer_socket(state))
  78     :ok
  79   end
  80
  81   # Public streams without authentication.
  82   defp authenticate_request(nil, nil) do
  83     {:ok, nil}
  84   end
  85
  86   # Authenticated streams.
  87   defp authenticate_request(access_token, sec_websocket) do
  88     token = access_token || sec_websocket
  89
  90     with true <- is_bitstring(token),
  91          %Token{user_id: user_id} <- Repo.get_by(Token, token: token),
  92          user = %User{} <- User.get_cached_by_id(user_id) do
  93       {:ok, user}
  94     else
  95       _ -> {:error, :unauthorized}
  96     end
  97   end
  98
  99   defp streamer_socket(state) do
 100     %{transport_pid: self(), assigns: state}
 101   end
 102 end