No OneTemporary
Actions

Size

11 KB

Referenced Files

None

Subscribers

None

View Options

	diff --git a/lib/pleroma/web/mastodon_api/websocket_handler.ex b/lib/pleroma/web/mastodon_api/websocket_handler.ex
	index ffbc2c4de..930e9eb29 100644
	--- a/lib/pleroma/web/mastodon_api/websocket_handler.ex
	+++ b/lib/pleroma/web/mastodon_api/websocket_handler.ex
	@@ -1,143 +1,144 @@
	# Pleroma: A lightweight social networking server
	# Copyright © 2017-2021 Pleroma Authors <https://pleroma.social/>
	# SPDX-License-Identifier: AGPL-3.0-only

	defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do
	require Logger

	alias Pleroma.Repo
	alias Pleroma.User
	alias Pleroma.Web.OAuth.Token
	alias Pleroma.Web.Streamer

	@behaviour :cowboy_websocket

	# 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

	- {:cowboy_websocket, req, %{user: user, topic: topic, oauth_token: oauth_token, count: 0, timer: nil},
	+ {:cowboy_websocket, req,
	+ %{user: user, topic: topic, oauth_token: oauth_token, count: 0, timer: nil},
	%{idle_timeout: @timeout}}
	else
	{:error, :bad_topic} ->
	Logger.debug("#{__MODULE__} bad topic #{inspect(req)}")
	req = :cowboy_req.reply(404, req)
	{:ok, req, state}

	{:error, :unauthorized} ->
	Logger.debug("#{__MODULE__} authentication error: #{inspect(req)}")
	req = :cowboy_req.reply(401, req)
	{:ok, req, state}
	end
	end

	def websocket_init(state) do
	Logger.debug(
	"#{__MODULE__} accepted websocket connection for user #{
	(state.user \|\| %{id: "anonymous"}).id
	}, topic #{state.topic}"
	)

	Streamer.add_socket(state.topic, state.oauth_token)
	{:ok, %{state \| timer: timer()}}
	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(frame, state) do
	Logger.error("#{__MODULE__} received frame: #{inspect(frame)}")
	{:ok, state}
	end

	def websocket_info({:render_with_user, view, template, item}, 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)}, %{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
	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}
	end

	def websocket_info(:close, state) do
	{:stop, 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 terminate(reason, _req, state) do
	Logger.debug(
	"#{__MODULE__} terminating websocket connection for user #{
	(state.user \|\| %{id: "anonymous"}).id
	}, topic #{state.topic \|\| "?"}: #{inspect(reason)}"
	)

	Streamer.remove_socket(state.topic)
	:ok
	end

	# Public streams without authentication.
	defp authenticate_request(nil, 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),
	user = %User{} <- User.get_cached_by_id(user_id) do
	{:ok, user, oauth_token}
	else
	_ -> {:error, :unauthorized}
	end
	end

	defp timer do
	Process.send_after(self(), :tick, @tick)
	end
	end
	diff --git a/test/pleroma/integration/mastodon_websocket_test.exs b/test/pleroma/integration/mastodon_websocket_test.exs
	index adb2d7004..d44033842 100644
	--- a/test/pleroma/integration/mastodon_websocket_test.exs
	+++ b/test/pleroma/integration/mastodon_websocket_test.exs
	@@ -1,146 +1,150 @@
	# Pleroma: A lightweight social networking server
	# Copyright © 2017-2021 Pleroma Authors <https://pleroma.social/>
	# SPDX-License-Identifier: AGPL-3.0-only

	defmodule Pleroma.Integration.MastodonWebsocketTest do
	# Needs a streamer, needs to stay synchronous
	use Pleroma.DataCase

	import ExUnit.CaptureLog
	import Pleroma.Factory

	alias Pleroma.Integration.WebsocketClient
	alias Pleroma.Web.CommonAPI
	alias Pleroma.Web.OAuth

	@moduletag needs_streamer: true, capture_log: true

	@path Pleroma.Web.Endpoint.url()
	\|> URI.parse()
	\|> Map.put(:scheme, "ws")
	\|> Map.put(:path, "/api/v1/streaming")
	\|> URI.to_string()

	def start_socket(qs \\ nil, headers \\ []) do
	path =
	case qs do
	nil -> @path
	qs -> @path <> qs
	end

	WebsocketClient.start_link(self(), path, headers)
	end

	test "refuses invalid requests" do
	capture_log(fn ->
	assert {:error, %WebSockex.RequestError{code: 404}} = start_socket()
	assert {:error, %WebSockex.RequestError{code: 404}} = start_socket("?stream=ncjdk")
	Process.sleep(30)
	end)
	end

	test "requires authentication and a valid token for protected streams" do
	capture_log(fn ->
	- assert {:error, %WebSockex.RequestError{code: 401}} = start_socket("?stream=user&access_token=aaaaaaaaaaaa")
	+ assert {:error, %WebSockex.RequestError{code: 401}} =
	+ start_socket("?stream=user&access_token=aaaaaaaaaaaa")
	+
	assert {:error, %WebSockex.RequestError{code: 401}} = start_socket("?stream=user")
	Process.sleep(30)
	end)
	end

	test "allows public streams without authentication" do
	assert {:ok, _} = start_socket("?stream=public")
	assert {:ok, _} = start_socket("?stream=public:local")
	assert {:ok, _} = start_socket("?stream=public:remote&instance=lain.com")
	assert {:ok, _} = start_socket("?stream=hashtag&tag=lain")
	end

	test "receives well formatted events" do
	user = insert(:user)
	{:ok, _} = start_socket("?stream=public")
	{:ok, activity} = CommonAPI.post(user, %{status: "nice echo chamber"})

	assert_receive {:text, raw_json}, 1_000
	assert {:ok, json} = Jason.decode(raw_json)

	assert "update" == json["event"]
	assert json["payload"]
	assert {:ok, json} = Jason.decode(json["payload"])

	view_json =
	Pleroma.Web.MastodonAPI.StatusView.render("show.json", activity: activity, for: nil)
	\|> Jason.encode!()
	\|> Jason.decode!()

	assert json == view_json
	end

	describe "with a valid user token" do
	setup do
	{:ok, app} =
	Pleroma.Repo.insert(
	OAuth.App.register_changeset(%OAuth.App{}, %{
	client_name: "client",
	scopes: ["read"],
	redirect_uris: "url"
	})
	)

	user = insert(:user)

	{:ok, auth} = OAuth.Authorization.create_authorization(app, user)

	{:ok, token} = OAuth.Token.exchange_token(app, auth)

	%{app: app, user: user, token: token}
	end

	test "accepts valid tokens", state do
	assert {:ok, _} = start_socket("?stream=user&access_token=#{state.token.token}")
	end

	test "accepts the 'user' stream", %{token: token} = _state do
	assert {:ok, _} = start_socket("?stream=user&access_token=#{token.token}")

	capture_log(fn ->
	assert {:error, %WebSockex.RequestError{code: 401}} = start_socket("?stream=user")
	Process.sleep(30)
	end)
	end

	test "accepts the 'user:notification' stream", %{token: token} = _state do
	assert {:ok, _} = start_socket("?stream=user:notification&access_token=#{token.token}")

	capture_log(fn ->
	- assert {:error, %WebSockex.RequestError{code: 401}} = start_socket("?stream=user:notification")
	+ assert {:error, %WebSockex.RequestError{code: 401}} =
	+ start_socket("?stream=user:notification")
	+
	Process.sleep(30)
	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 "disconnect when token is revoked", %{app: app, user: user, token: token} do
	assert {:ok, _} = start_socket("?stream=user:notification&access_token=#{token.token}")
	assert {:ok, _} = start_socket("?stream=user&access_token=#{token.token}")

	{:ok, auth} = OAuth.Authorization.create_authorization(app, user)

	{:ok, token2} = OAuth.Token.exchange_token(app, auth)
	assert {:ok, _} = start_socket("?stream=user&access_token=#{token2.token}")

	OAuth.Token.Strategy.Revoke.revoke(token)

	assert_receive {:close, _}
	assert_receive {:close, _}
	refute_receive {:close, _}
	end
	end
	end
	diff --git a/test/support/websocket_client.ex b/test/support/websocket_client.ex
	index abe7d5eda..70d331999 100644
	--- a/test/support/websocket_client.ex
	+++ b/test/support/websocket_client.ex
	@@ -1,66 +1,66 @@
	# Pleroma: A lightweight social networking server
	# Copyright © 2017-2021 Pleroma Authors <https://pleroma.social/>
	# SPDX-License-Identifier: AGPL-3.0-only

	defmodule Pleroma.Integration.WebsocketClient do
	# https://github.com/phoenixframework/phoenix/blob/master/test/support/websocket_client.exs

	use WebSockex

	@doc """
	Starts the WebSocket server for given ws URL. Received Socket.Message's
	are forwarded to the sender pid
	"""
	def start_link(sender, url, headers \\ []) do
	WebSockex.start_link(
	url,
	__MODULE__,
	- %{ sender: sender },
	+ %{sender: sender},
	extra_headers: headers
	)
	end

	@doc """
	Closes the socket
	"""
	def close(socket) do
	send(socket, :close)
	end

	@doc """
	Sends a low-level text message to the client.
	"""
	def send_text(server_pid, msg) do
	send(server_pid, {:text, msg})
	end

	@doc false
	@impl true
	def handle_frame(frame, state) do
	send(state.sender, frame)
	{:ok, state}
	end

	@impl true
	def handle_disconnect(conn_status, state) do
	send(state.sender, {:close, conn_status})
	{:ok, state}
	end

	@doc false
	@impl true
	def handle_info({:text, msg}, state) do
	{:reply, {:text, msg}, state}
	end

	@impl true
	def handle_info(:close, _state) do
	{:close, <<>>, "done"}
	end

	@doc false
	@impl true
	def terminate(_reason, _state) do
	:ok
	end
	end

File Metadata

Mime Type: text/x-diff
Expires: Sat, Aug 29, 4:03 PM (1 d, 18 h)
Storage Engine: blob
Storage Format: Raw Data
Storage Handle: 1737165
Default Alt Text: (11 KB)

No OneTemporaryActions

View Options

File Metadata

Event Timeline

No OneTemporary
Actions