From 12a5981cc3da65b7f2763d0ec05871b0986234f5 Mon Sep 17 00:00:00 2001
From: Ivan Tashkinov <ivantashkinov@gmail.com>
Date: Wed, 25 Nov 2020 21:47:23 +0300
Subject: [PATCH] Session token setting on token exchange. Auth-related
 refactoring.

---
 lib/pleroma/helpers/auth_helper.ex            | 15 +++++++++++++++
 .../controllers/account_controller.ex         |  3 +--
 .../controllers/auth_controller.ex            |  5 +++--
 lib/pleroma/web/o_auth/mfa_controller.ex      |  3 +--
 lib/pleroma/web/o_auth/o_auth_controller.ex   | 19 +++++++++++++------
 lib/pleroma/web/plugs/o_auth_plug.ex          |  3 ++-
 .../web/plugs/set_user_session_id_plug.ex     |  5 ++---
 lib/pleroma/web/router.ex                     | 14 +++++++-------
 .../web/o_auth/o_auth_controller_test.exs     | 12 +++++++-----
 test/pleroma/web/plugs/o_auth_plug_test.exs   |  3 ++-
 .../plugs/set_user_session_id_plug_test.exs   |  5 +++--
 11 files changed, 56 insertions(+), 31 deletions(-)

diff --git a/lib/pleroma/helpers/auth_helper.ex b/lib/pleroma/helpers/auth_helper.ex
index 878fec346..392fa7d5d 100644
--- a/lib/pleroma/helpers/auth_helper.ex
+++ b/lib/pleroma/helpers/auth_helper.ex
@@ -4,9 +4,12 @@
 
 defmodule Pleroma.Helpers.AuthHelper do
   alias Pleroma.Web.Plugs.OAuthScopesPlug
+  alias Plug.Conn
 
   import Plug.Conn
 
+  @oauth_token_session_key :oauth_token
+
   @doc """
   Skips OAuth permissions (scopes) checks, assigns nil `:token`.
   Intended to be used with explicit authentication and only when OAuth token cannot be determined.
@@ -22,4 +25,16 @@ defmodule Pleroma.Helpers.AuthHelper do
     |> assign(:user, nil)
     |> assign(:token, nil)
   end
+
+  def get_session_token(%Conn{} = conn) do
+    get_session(conn, @oauth_token_session_key)
+  end
+
+  def put_session_token(%Conn{} = conn, token) when is_binary(token) do
+    put_session(conn, @oauth_token_session_key, token)
+  end
+
+  def delete_session_token(%Conn{} = conn) do
+    delete_session(conn, @oauth_token_session_key)
+  end
 end
diff --git a/lib/pleroma/web/mastodon_api/controllers/account_controller.ex b/lib/pleroma/web/mastodon_api/controllers/account_controller.ex
index 7011b7eb1..b4375872b 100644
--- a/lib/pleroma/web/mastodon_api/controllers/account_controller.ex
+++ b/lib/pleroma/web/mastodon_api/controllers/account_controller.ex
@@ -25,7 +25,6 @@ defmodule Pleroma.Web.MastodonAPI.AccountController do
   alias Pleroma.Web.MastodonAPI.MastodonAPIController
   alias Pleroma.Web.MastodonAPI.StatusView
   alias Pleroma.Web.OAuth.OAuthController
-  alias Pleroma.Web.OAuth.OAuthView
   alias Pleroma.Web.Plugs.EnsurePublicOrAuthenticatedPlug
   alias Pleroma.Web.Plugs.OAuthScopesPlug
   alias Pleroma.Web.Plugs.RateLimiter
@@ -103,7 +102,7 @@ defmodule Pleroma.Web.MastodonAPI.AccountController do
          {:ok, user} <- TwitterAPI.register_user(params),
          {_, {:ok, token}} <-
            {:login, OAuthController.login(user, app, app.scopes)} do
-      json(conn, OAuthView.render("token.json", %{user: user, token: token}))
+      OAuthController.after_token_exchange(conn, %{user: user, token: token})
     else
       {:login, {:account_status, :confirmation_pending}} ->
         json_response(conn, :ok, %{
diff --git a/lib/pleroma/web/mastodon_api/controllers/auth_controller.ex b/lib/pleroma/web/mastodon_api/controllers/auth_controller.ex
index 9cc3984d0..fa582dcfc 100644
--- a/lib/pleroma/web/mastodon_api/controllers/auth_controller.ex
+++ b/lib/pleroma/web/mastodon_api/controllers/auth_controller.ex
@@ -7,6 +7,7 @@ defmodule Pleroma.Web.MastodonAPI.AuthController do
 
   import Pleroma.Web.ControllerHelper, only: [json_response: 3]
 
+  alias Pleroma.Helpers.AuthHelper
   alias Pleroma.User
   alias Pleroma.Web.OAuth.App
   alias Pleroma.Web.OAuth.Authorization
@@ -30,7 +31,7 @@ defmodule Pleroma.Web.MastodonAPI.AuthController do
          {:ok, auth} <- Authorization.get_by_token(app, auth_token),
          {:ok, token} <- Token.exchange_token(app, auth) do
       conn
-      |> put_session(:oauth_token, token.token)
+      |> AuthHelper.put_session_token(token.token)
       |> redirect(to: local_mastodon_root_path(conn))
     end
   end
@@ -53,7 +54,7 @@ defmodule Pleroma.Web.MastodonAPI.AuthController do
   @doc "DELETE /auth/sign_out"
   def logout(conn, _) do
     conn
-    |> clear_session
+    |> clear_session()
     |> redirect(to: "/")
   end
 
diff --git a/lib/pleroma/web/o_auth/mfa_controller.ex b/lib/pleroma/web/o_auth/mfa_controller.ex
index f102c93e7..5d5ec286a 100644
--- a/lib/pleroma/web/o_auth/mfa_controller.ex
+++ b/lib/pleroma/web/o_auth/mfa_controller.ex
@@ -13,7 +13,6 @@ defmodule Pleroma.Web.OAuth.MFAController do
   alias Pleroma.Web.Auth.TOTPAuthenticator
   alias Pleroma.Web.OAuth.MFAView, as: View
   alias Pleroma.Web.OAuth.OAuthController
-  alias Pleroma.Web.OAuth.OAuthView
   alias Pleroma.Web.OAuth.Token
 
   plug(:fetch_session when action in [:show, :verify])
@@ -75,7 +74,7 @@ defmodule Pleroma.Web.OAuth.MFAController do
          {:ok, %{user: user, authorization: auth}} <- MFA.Token.validate(mfa_token),
          {:ok, _} <- validates_challenge(user, params),
          {:ok, token} <- Token.exchange_token(app, auth) do
-      json(conn, OAuthView.render("token.json", %{user: user, token: token}))
+      OAuthController.after_token_exchange(conn, %{user: user, token: token})
     else
       _error ->
         conn
diff --git a/lib/pleroma/web/o_auth/o_auth_controller.ex b/lib/pleroma/web/o_auth/o_auth_controller.ex
index 83a25907d..8103395b3 100644
--- a/lib/pleroma/web/o_auth/o_auth_controller.ex
+++ b/lib/pleroma/web/o_auth/o_auth_controller.ex
@@ -5,6 +5,7 @@
 defmodule Pleroma.Web.OAuth.OAuthController do
   use Pleroma.Web, :controller
 
+  alias Pleroma.Helpers.AuthHelper
   alias Pleroma.Helpers.UriHelper
   alias Pleroma.Maps
   alias Pleroma.MFA
@@ -248,7 +249,7 @@ defmodule Pleroma.Web.OAuth.OAuthController do
     with {:ok, app} <- Token.Utils.fetch_app(conn),
          {:ok, %{user: user} = token} <- Token.get_by_refresh_token(app, token),
          {:ok, token} <- RefreshToken.grant(token) do
-      json(conn, OAuthView.render("token.json", %{user: user, token: token}))
+      after_token_exchange(conn, %{user: user, token: token})
     else
       _error -> render_invalid_credentials_error(conn)
     end
@@ -260,7 +261,7 @@ defmodule Pleroma.Web.OAuth.OAuthController do
          {:ok, auth} <- Authorization.get_by_token(app, fixed_token),
          %User{} = user <- User.get_cached_by_id(auth.user_id),
          {:ok, token} <- Token.exchange_token(app, auth) do
-      json(conn, OAuthView.render("token.json", %{user: user, token: token}))
+      after_token_exchange(conn, %{user: user, token: token})
     else
       error ->
         handle_token_exchange_error(conn, error)
@@ -275,7 +276,7 @@ defmodule Pleroma.Web.OAuth.OAuthController do
          {:ok, app} <- Token.Utils.fetch_app(conn),
          requested_scopes <- Scopes.fetch_scopes(params, app.scopes),
          {:ok, token} <- login(user, app, requested_scopes) do
-      json(conn, OAuthView.render("token.json", %{user: user, token: token}))
+      after_token_exchange(conn, %{user: user, token: token})
     else
       error ->
         handle_token_exchange_error(conn, error)
@@ -298,7 +299,7 @@ defmodule Pleroma.Web.OAuth.OAuthController do
     with {:ok, app} <- Token.Utils.fetch_app(conn),
          {:ok, auth} <- Authorization.create_authorization(app, %User{}),
          {:ok, token} <- Token.exchange_token(app, auth) do
-      json(conn, OAuthView.render("token.json", %{token: token}))
+      after_token_exchange(conn, %{token: token})
     else
       _error ->
         handle_token_exchange_error(conn, :invalid_credentails)
@@ -308,6 +309,12 @@ defmodule Pleroma.Web.OAuth.OAuthController do
   # Bad request
   def token_exchange(%Plug.Conn{} = conn, params), do: bad_request(conn, params)
 
+  def after_token_exchange(%Plug.Conn{} = conn, %{token: token} = view_params) do
+    conn
+    |> AuthHelper.put_session_token(token.token)
+    |> json(OAuthView.render("token.json", view_params))
+  end
+
   defp handle_token_exchange_error(%Plug.Conn{} = conn, {:mfa_required, user, auth, _}) do
     conn
     |> put_status(:forbidden)
@@ -365,9 +372,9 @@ defmodule Pleroma.Web.OAuth.OAuthController do
     with {:ok, app} <- Token.Utils.fetch_app(conn),
          {:ok, %Token{} = oauth_token} <- RevokeToken.revoke(app, params) do
       conn =
-        with session_token = get_session(conn, :oauth_token),
+        with session_token = AuthHelper.get_session_token(conn),
              %Token{token: ^session_token} <- oauth_token do
-          delete_session(conn, :oauth_token)
+          AuthHelper.delete_session_token(conn)
         else
           _ -> conn
         end
diff --git a/lib/pleroma/web/plugs/o_auth_plug.ex b/lib/pleroma/web/plugs/o_auth_plug.ex
index a3b7d42f7..eb287318b 100644
--- a/lib/pleroma/web/plugs/o_auth_plug.ex
+++ b/lib/pleroma/web/plugs/o_auth_plug.ex
@@ -8,6 +8,7 @@ defmodule Pleroma.Web.Plugs.OAuthPlug do
   import Plug.Conn
   import Ecto.Query
 
+  alias Pleroma.Helpers.AuthHelper
   alias Pleroma.Repo
   alias Pleroma.User
   alias Pleroma.Web.OAuth.App
@@ -98,7 +99,7 @@ defmodule Pleroma.Web.Plugs.OAuthPlug do
 
   @spec fetch_token_from_session(Plug.Conn.t()) :: :no_token_found | {:ok, String.t()}
   defp fetch_token_from_session(conn) do
-    case get_session(conn, :oauth_token) do
+    case AuthHelper.get_session_token(conn) do
       nil -> :no_token_found
       token -> {:ok, token}
     end
diff --git a/lib/pleroma/web/plugs/set_user_session_id_plug.ex b/lib/pleroma/web/plugs/set_user_session_id_plug.ex
index d2338c03f..9f4a6b6ac 100644
--- a/lib/pleroma/web/plugs/set_user_session_id_plug.ex
+++ b/lib/pleroma/web/plugs/set_user_session_id_plug.ex
@@ -3,8 +3,7 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 
 defmodule Pleroma.Web.Plugs.SetUserSessionIdPlug do
-  import Plug.Conn
-
+  alias Pleroma.Helpers.AuthHelper
   alias Pleroma.Web.OAuth.Token
 
   def init(opts) do
@@ -12,7 +11,7 @@ defmodule Pleroma.Web.Plugs.SetUserSessionIdPlug do
   end
 
   def call(%{assigns: %{token: %Token{} = oauth_token}} = conn, _) do
-    put_session(conn, :oauth_token, oauth_token.token)
+    AuthHelper.put_session_token(conn, oauth_token.token)
   end
 
   def call(conn, _), do: conn
diff --git a/lib/pleroma/web/router.ex b/lib/pleroma/web/router.ex
index 3a3e63db6..b3462ba00 100644
--- a/lib/pleroma/web/router.ex
+++ b/lib/pleroma/web/router.ex
@@ -320,6 +320,11 @@ defmodule Pleroma.Web.Router do
   end
 
   scope "/oauth", Pleroma.Web.OAuth do
+    get("/registration_details", OAuthController, :registration_details)
+
+    post("/mfa/verify", MFAController, :verify, as: :mfa_verify)
+    get("/mfa", MFAController, :show)
+
     scope [] do
       pipe_through(:oauth)
 
@@ -327,17 +332,12 @@ defmodule Pleroma.Web.Router do
       post("/authorize", OAuthController, :create_authorization)
     end
 
-    post("/token", OAuthController, :token_exchange)
-    get("/registration_details", OAuthController, :registration_details)
-
-    post("/mfa/challenge", MFAController, :challenge)
-    post("/mfa/verify", MFAController, :verify, as: :mfa_verify)
-    get("/mfa", MFAController, :show)
-
     scope [] do
       pipe_through(:fetch_session)
 
+      post("/token", OAuthController, :token_exchange)
       post("/revoke", OAuthController, :token_revoke)
+      post("/mfa/challenge", MFAController, :challenge)
     end
 
     scope [] do
diff --git a/test/pleroma/web/o_auth/o_auth_controller_test.exs b/test/pleroma/web/o_auth/o_auth_controller_test.exs
index a00df8cc7..22cbddce3 100644
--- a/test/pleroma/web/o_auth/o_auth_controller_test.exs
+++ b/test/pleroma/web/o_auth/o_auth_controller_test.exs
@@ -4,8 +4,10 @@
 
 defmodule Pleroma.Web.OAuth.OAuthControllerTest do
   use Pleroma.Web.ConnCase
+
   import Pleroma.Factory
 
+  alias Pleroma.Helpers.AuthHelper
   alias Pleroma.MFA
   alias Pleroma.MFA.TOTP
   alias Pleroma.Repo
@@ -454,7 +456,7 @@ defmodule Pleroma.Web.OAuth.OAuthControllerTest do
 
       conn =
         conn
-        |> put_session(:oauth_token, token.token)
+        |> AuthHelper.put_session_token(token.token)
         |> get(
           "/oauth/authorize",
           %{
@@ -478,7 +480,7 @@ defmodule Pleroma.Web.OAuth.OAuthControllerTest do
 
       conn =
         conn
-        |> put_session(:oauth_token, token.token)
+        |> AuthHelper.put_session_token(token.token)
         |> get(
           "/oauth/authorize",
           %{
@@ -501,7 +503,7 @@ defmodule Pleroma.Web.OAuth.OAuthControllerTest do
 
       conn =
         conn
-        |> put_session(:oauth_token, token.token)
+        |> AuthHelper.put_session_token(token.token)
         |> get(
           "/oauth/authorize",
           %{
@@ -527,7 +529,7 @@ defmodule Pleroma.Web.OAuth.OAuthControllerTest do
 
       conn =
         conn
-        |> put_session(:oauth_token, token.token)
+        |> AuthHelper.put_session_token(token.token)
         |> get(
           "/oauth/authorize",
           %{
@@ -551,7 +553,7 @@ defmodule Pleroma.Web.OAuth.OAuthControllerTest do
 
       conn =
         conn
-        |> put_session(:oauth_token, token.token)
+        |> AuthHelper.put_session_token(token.token)
         |> get(
           "/oauth/authorize",
           %{
diff --git a/test/pleroma/web/plugs/o_auth_plug_test.exs b/test/pleroma/web/plugs/o_auth_plug_test.exs
index ad2aa5d1b..1186cdb14 100644
--- a/test/pleroma/web/plugs/o_auth_plug_test.exs
+++ b/test/pleroma/web/plugs/o_auth_plug_test.exs
@@ -5,6 +5,7 @@
 defmodule Pleroma.Web.Plugs.OAuthPlugTest do
   use Pleroma.Web.ConnCase, async: true
 
+  alias Pleroma.Helpers.AuthHelper
   alias Pleroma.Web.OAuth.Token
   alias Pleroma.Web.OAuth.Token.Strategy.Revoke
   alias Pleroma.Web.Plugs.OAuthPlug
@@ -84,7 +85,7 @@ defmodule Pleroma.Web.Plugs.OAuthPlugTest do
         conn
         |> Session.call(Session.init(session_opts))
         |> fetch_session()
-        |> put_session(:oauth_token, oauth_token.token)
+        |> AuthHelper.put_session_token(oauth_token.token)
 
       %{conn: conn}
     end
diff --git a/test/pleroma/web/plugs/set_user_session_id_plug_test.exs b/test/pleroma/web/plugs/set_user_session_id_plug_test.exs
index a50e80107..21417d0e7 100644
--- a/test/pleroma/web/plugs/set_user_session_id_plug_test.exs
+++ b/test/pleroma/web/plugs/set_user_session_id_plug_test.exs
@@ -5,6 +5,7 @@
 defmodule Pleroma.Web.Plugs.SetUserSessionIdPlugTest do
   use Pleroma.Web.ConnCase, async: true
 
+  alias Pleroma.Helpers.AuthHelper
   alias Pleroma.Web.Plugs.SetUserSessionIdPlug
 
   setup %{conn: conn} do
@@ -28,7 +29,7 @@ defmodule Pleroma.Web.Plugs.SetUserSessionIdPlugTest do
     assert ret_conn == conn
   end
 
-  test "sets :oauth_token in session to :token assign", %{conn: conn} do
+  test "sets session token basing on :token assign", %{conn: conn} do
     %{user: user, token: oauth_token} = oauth_access(["read"])
 
     ret_conn =
@@ -37,6 +38,6 @@ defmodule Pleroma.Web.Plugs.SetUserSessionIdPlugTest do
       |> assign(:token, oauth_token)
       |> SetUserSessionIdPlug.call(%{})
 
-    assert get_session(ret_conn, :oauth_token) == oauth_token.token
+    assert AuthHelper.get_session_token(ret_conn) == oauth_token.token
   end
 end