diff options
| -rw-r--r-- | lib/pleroma/web/mastodon_api/websocket_handler.ex | 8 | ||||
| -rw-r--r-- | lib/pleroma/web/o_auth/token/strategy/revoke.ex | 1 | ||||
| -rw-r--r-- | lib/pleroma/web/streamer.ex | 24 | ||||
| -rw-r--r-- | test/pleroma/web/streamer_test.exs | 54 |
4 files changed, 81 insertions, 6 deletions
diff --git a/lib/pleroma/web/mastodon_api/websocket_handler.ex b/lib/pleroma/web/mastodon_api/websocket_handler.ex index 0d1faffbd..ffbc2c4de 100644 --- a/lib/pleroma/web/mastodon_api/websocket_handler.ex +++ b/lib/pleroma/web/mastodon_api/websocket_handler.ex @@ -32,7 +32,7 @@ defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do req end - {:cowboy_websocket, req, %{user: user, topic: topic, 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} -> @@ -54,7 +54,7 @@ defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do }, topic #{state.topic}" ) - Streamer.add_socket(state.topic, state.user) + Streamer.add_socket(state.topic, state.oauth_token) {:ok, %{state | timer: timer()}} end @@ -100,6 +100,10 @@ defmodule Pleroma.Web.MastodonAPI.WebsocketHandler 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 diff --git a/lib/pleroma/web/o_auth/token/strategy/revoke.ex b/lib/pleroma/web/o_auth/token/strategy/revoke.ex index 8d6572704..03a0b91ae 100644 --- a/lib/pleroma/web/o_auth/token/strategy/revoke.ex +++ b/lib/pleroma/web/o_auth/token/strategy/revoke.ex @@ -22,5 +22,6 @@ defmodule Pleroma.Web.OAuth.Token.Strategy.Revoke do @spec revoke(Token.t()) :: {:ok, Token.t()} | {:error, Ecto.Changeset.t()} def revoke(%Token{} = token) do Repo.delete(token) + Pleroma.Web.Streamer.close_streams_by_oauth_token(token) end end diff --git a/lib/pleroma/web/streamer.ex b/lib/pleroma/web/streamer.ex index fc3bbb130..8bf70d99b 100644 --- a/lib/pleroma/web/streamer.ex +++ b/lib/pleroma/web/streamer.ex @@ -37,7 +37,7 @@ defmodule Pleroma.Web.Streamer do {:ok, topic :: String.t()} | {:error, :bad_topic} | {:error, :unauthorized} def get_topic_and_add_socket(stream, user, oauth_token, params \\ %{}) do with {:ok, topic} <- get_topic(stream, user, oauth_token, params) do - add_socket(topic, user) + add_socket(topic, oauth_token) end end @@ -120,10 +120,10 @@ defmodule Pleroma.Web.Streamer do end @doc "Registers the process for streaming. Use `get_topic/3` to get the full authorized topic." - def add_socket(topic, user) do + def add_socket(topic, oauth_token) do if should_env_send?() do - auth? = if user, do: true - Registry.register(@registry, topic, auth?) + oauth_token_id = if oauth_token, do: oauth_token.id, else: false + Registry.register(@registry, topic, oauth_token_id) end {:ok, topic} @@ -320,6 +320,22 @@ defmodule Pleroma.Web.Streamer do end end + def close_streams_by_oauth_token(oauth_token) do + if should_env_send?() do + Registry.select( + @registry, + [ + { + {:"$1", :"$2", :"$3"}, + [{:==, :"$3", oauth_token.id}], + [:"$2"] + } + ] + ) + |> Enum.each(fn pid -> send(pid, :close) end) + end + end + # In test environement, only return true if the registry is started. # In benchmark environment, returns false. # In any other environment, always returns true. diff --git a/test/pleroma/web/streamer_test.exs b/test/pleroma/web/streamer_test.exs index b788a9138..5426467e5 100644 --- a/test/pleroma/web/streamer_test.exs +++ b/test/pleroma/web/streamer_test.exs @@ -813,4 +813,58 @@ defmodule Pleroma.Web.StreamerTest do assert last_status["id"] == to_string(create_activity.id) end end + + describe "stop streaming if token got revoked" do + test "do not revoke other tokens" do + %{user: user, token: token} = oauth_access(["read"]) + %{token: token2} = oauth_access(["read"], user: user) + %{user: user2, token: user2_token} = oauth_access(["read"]) + + post_user = insert(:user) + CommonAPI.follow(user, post_user) + CommonAPI.follow(user2, post_user) + + Streamer.get_topic_and_add_socket("user", user, token) + Streamer.get_topic_and_add_socket("user", user, token2) + Streamer.get_topic_and_add_socket("user", user2, user2_token) + + {:ok, _} = + CommonAPI.post(post_user, %{ + status: "hi" + }) + + assert_receive {:render_with_user, _, "update.json", _} + assert_receive {:render_with_user, _, "update.json", _} + assert_receive {:render_with_user, _, "update.json", _} + + Pleroma.Web.OAuth.Token.Strategy.Revoke.revoke(token) + + assert_receive :close + refute_receive :close + end + + test "revoke all streams for this token" do + %{user: user, token: token} = oauth_access(["read"]) + + post_user = insert(:user) + CommonAPI.follow(user, post_user) + + Streamer.get_topic_and_add_socket("user", user, token) + Streamer.get_topic_and_add_socket("user", user, token) + + {:ok, _} = + CommonAPI.post(post_user, %{ + status: "hi" + }) + + assert_receive {:render_with_user, _, "update.json", _} + assert_receive {:render_with_user, _, "update.json", _} + + Pleroma.Web.OAuth.Token.Strategy.Revoke.revoke(token) + + assert_receive :close + assert_receive :close + refute_receive :close + end + end end |