improvement: assign the token record to the connection assigns (#426)

lib/ash_authentication/plug/helpers.ex

@@ -120,6 +120,9 @@ defmodule AshAuthentication.Plug.Helpers do
   your request.  If a valid bearer token is present then the subject is loaded
   into the assigns under their subject name (with the prefix `current_`).
 
+  If the authentication token is required to be present in the database, it is
+  loaded into the assigns using `current_\#{subject_name}_token_record`
+
   If there is no user present for a resource then the assign is set to `nil`.
   """
   @spec retrieve_from_bearer(Conn.t(), module) :: Conn.t()
@@ -131,7 +134,8 @@ defmodule AshAuthentication.Plug.Helpers do
     |> Enum.reduce(conn, fn token, conn ->
       with {:ok, %{"sub" => subject, "jti" => jti} = claims, resource}
            when not is_map_key(claims, "act") <- Jwt.verify(token, otp_app),
-           :ok <- validate_token(resource, jti),
+           {:ok, token_record} <-
+             validate_token(resource, jti),
            {:ok, user} <-
              AshAuthentication.subject_to_user(subject, resource,
                tenant: Ash.PlugHelpers.get_tenant(conn)
@@ -140,24 +144,40 @@ defmodule AshAuthentication.Plug.Helpers do
            current_subject_name <- current_subject_name(subject_name) do
         conn
         |> Conn.assign(current_subject_name, user)
+        |> maybe_assign_token_record(token_record, subject_name)
       else
         _ -> conn
       end
     end)
   end
 
+  defp maybe_assign_token_record(conn, _token_record, subject_name) when is_nil(subject_name) do
+    conn
+  end
+
+  defp maybe_assign_token_record(conn, token_record, subject_name) do
+    conn
+    |> Conn.assign(
+      current_subject_token_record_name(subject_name),
+      token_record
+    )
+  end
+
   defp validate_token(resource, jti) do
     if Info.authentication_tokens_require_token_presence_for_authentication?(resource) do
-      with {:ok, token_resource} <- Info.authentication_tokens_token_resource(resource),
-           {:ok, [_]} <-
+      with {:ok, token_resource} <-
+             Info.authentication_tokens_token_resource(resource),
+           {:ok, [token_record]} <-
              TokenResource.Actions.get_token(token_resource, %{
                "jti" => jti,
                "purpose" => "user"
              }) do
-        :ok
+        {:ok, token_record}
+      else
+        _ -> :error
       end
     else
-      :ok
+      {:ok, nil}
     end
   end
 
@@ -258,4 +278,8 @@ defmodule AshAuthentication.Plug.Helpers do
   # sobelow_skip ["DOS.StringToAtom"]
   defp current_subject_name(subject_name) when is_atom(subject_name),
     do: String.to_atom("current_#{subject_name}")
+
+  # sobelow_skip ["DOS.StringToAtom"]
+  defp current_subject_token_record_name(subject_name) when is_atom(subject_name),
+    do: String.to_atom("current_#{subject_name}_token_record")
 end

test/ash_authentication/plug/helpers_test.exs

@@ -145,6 +145,14 @@ defmodule AshAuthentication.Plug.HelpersTest do
         |> Helpers.retrieve_from_bearer(:ash_authentication)
 
       assert conn.assigns.current_user_with_token_required.id == user.id
+
+      assert {:ok, user_for_token} =
+               AshAuthentication.subject_to_user(
+                 conn.assigns.current_user_with_token_required_token_record.subject,
+                 Example.UserWithTokenRequired
+               )
+
+      assert user_for_token.id == user.id
     end
 
     test "when token presense is required and the token is not present in the token resource it doesn't load the subjects from the authorization header(s)",
@@ -162,6 +170,7 @@ defmodule AshAuthentication.Plug.HelpersTest do
         |> Helpers.retrieve_from_bearer(:ash_authentication)
 
       refute is_map_key(conn.assigns, :current_user_with_token_required)
+      refute is_map_key(conn.assigns, :current_user_with_token_required_token_record)
     end
 
     test "when token presense is required and the token has been revoked it doesn't lkoad the subjects from the authorization header(s)",