From 6dfbf03f111bf85f6a4c04405e140a3368a8afcc Mon Sep 17 00:00:00 2001
From: James Harton <59449+jimsynz@users.noreply.github.com>
Date: Tue, 13 Dec 2022 16:35:30 +1300
Subject: [PATCH] improvement: remove the need for a strategy in
 changeset/query contexts. (#89)

The action -> strategy mapping is now stored directly in the resource DSL.

Closes #84.
---
 .../add_ons/confirmation/actions.ex           |  2 --
 .../add_ons/confirmation/confirm_change.ex    |  7 ++--
 .../add_ons/confirmation/transformer.ex       |  5 ++-
 lib/ash_authentication/info.ex                | 36 +++++++++++++++++--
 .../strategies/oauth2/actions.ex              |  2 --
 .../strategies/oauth2/identity_change.ex      |  8 +++--
 .../strategies/oauth2/sign_in_preparation.ex  |  4 +--
 .../strategies/oauth2/transformer.ex          | 14 ++++++--
 .../strategies/password/actions.ex            |  4 ---
 .../password/hash_password_change.ex          |  3 +-
 .../password_confirmation_validation.ex       |  7 ++--
 .../request_password_reset_preparation.ex     |  4 +--
 .../password/reset_token_validation.ex        |  4 +--
 .../password/sign_in_preparation.ex           |  4 +--
 .../strategies/password/transformer.ex        | 26 ++++++++++++--
 test/support/data_case.ex                     |  3 --
 16 files changed, 98 insertions(+), 35 deletions(-)

diff --git a/lib/ash_authentication/add_ons/confirmation/actions.ex b/lib/ash_authentication/add_ons/confirmation/actions.ex
index c23853b..c636dde 100644
--- a/lib/ash_authentication/add_ons/confirmation/actions.ex
+++ b/lib/ash_authentication/add_ons/confirmation/actions.ex
@@ -26,7 +26,6 @@ defmodule AshAuthentication.AddOn.Confirmation.Actions do
          {:ok, user} <- AshAuthentication.subject_to_user(subject, strategy.resource) do
       user
       |> Changeset.new()
-      |> Changeset.set_context(%{strategy: strategy})
       |> Changeset.for_update(strategy.confirm_action_name, params)
       |> api.update(options)
     else
@@ -53,7 +52,6 @@ defmodule AshAuthentication.AddOn.Confirmation.Actions do
          {:ok, _token_record} <-
            token_resource
            |> Changeset.new()
-           |> Changeset.set_context(%{strategy: strategy})
            |> Changeset.for_create(store_changes_action, %{
              token: token,
              extra_data: changes,
diff --git a/lib/ash_authentication/add_ons/confirmation/confirm_change.ex b/lib/ash_authentication/add_ons/confirmation/confirm_change.ex
index 611d4bc..1e2a0ce 100644
--- a/lib/ash_authentication/add_ons/confirmation/confirm_change.ex
+++ b/lib/ash_authentication/add_ons/confirmation/confirm_change.ex
@@ -4,7 +4,7 @@ defmodule AshAuthentication.AddOn.Confirmation.ConfirmChange do
   """
 
   use Ash.Resource.Change
-  alias AshAuthentication.{AddOn.Confirmation.Actions, Jwt}
+  alias AshAuthentication.{AddOn.Confirmation.Actions, Info, Jwt}
 
   alias Ash.{
     Changeset,
@@ -17,12 +17,13 @@ defmodule AshAuthentication.AddOn.Confirmation.ConfirmChange do
   @impl true
   @spec change(Changeset.t(), keyword, Change.context()) :: Changeset.t()
   def change(changeset, _opts, _context) do
-    case Map.fetch(changeset.context, :strategy) do
+    case Info.strategy_for_action(changeset.resource, changeset.action.name) do
       {:ok, strategy} ->
         do_change(changeset, strategy)
 
       :error ->
-        raise AssumptionFailed, message: "Strategy is missing from the changeset context."
+        raise AssumptionFailed,
+          message: "Action does not correlate with an authentication strategy"
     end
   end
 
diff --git a/lib/ash_authentication/add_ons/confirmation/transformer.ex b/lib/ash_authentication/add_ons/confirmation/transformer.ex
index 09841bf..84af978 100644
--- a/lib/ash_authentication/add_ons/confirmation/transformer.ex
+++ b/lib/ash_authentication/add_ons/confirmation/transformer.ex
@@ -66,13 +66,16 @@ defmodule AshAuthentication.AddOn.Confirmation.Transformer do
          :ok <- validate_confirmed_at_attribute(dsl_state, strategy),
          {:ok, dsl_state} <- maybe_build_change(dsl_state, Confirmation.ConfirmationHookChange),
          {:ok, resource} <- persisted_option(dsl_state, :module) do
+      strategy = %{strategy | resource: resource}
+
       dsl_state =
         dsl_state
         |> Transformer.replace_entity(
           [:authentication, :add_ons],
-          %{strategy | resource: resource},
+          strategy,
           &(&1.name == strategy.name)
         )
+        |> Transformer.persist({:authentication_action, strategy.confirm_action_name}, strategy)
 
       {:ok, dsl_state}
     else
diff --git a/lib/ash_authentication/info.ex b/lib/ash_authentication/info.ex
index ee9ca6f..6c90379 100644
--- a/lib/ash_authentication/info.ex
+++ b/lib/ash_authentication/info.ex
@@ -7,10 +7,15 @@ defmodule AshAuthentication.Info do
     extension: AshAuthentication,
     sections: [:authentication]
 
+  alias AshAuthentication.Strategy
+  alias Spark.Dsl.Extension
+
+  @type dsl_or_resource :: module | map
+
   @doc """
   Retrieve a named strategy from a resource.
   """
-  @spec strategy(dsl_or_resource :: map | module, atom) :: {:ok, strategy} | :error
+  @spec strategy(dsl_or_resource | module, atom) :: {:ok, strategy} | :error
         when strategy: struct
   def strategy(dsl_or_resource, name) do
     dsl_or_resource
@@ -24,7 +29,7 @@ defmodule AshAuthentication.Info do
   @doc """
   Retrieve a named strategy from a resource (raising version).
   """
-  @spec strategy!(dsl_or_resource :: map | module, atom) :: strategy | no_return
+  @spec strategy!(dsl_or_resource | module, atom) :: strategy | no_return
         when strategy: struct
   def strategy!(dsl_or_resource, name) do
     case strategy(dsl_or_resource, name) do
@@ -35,4 +40,31 @@ defmodule AshAuthentication.Info do
         raise "No strategy named `#{inspect(name)}` found on resource `#{inspect(dsl_or_resource)}`"
     end
   end
+
+  @doc """
+  Given an action name, retrieve the strategy it is for from the DSL
+  configuration.
+  """
+  @spec strategy_for_action(dsl_or_resource, atom) :: {:ok, Strategy.t()} | :error
+  def strategy_for_action(dsl_or_resource, action_name) do
+    case Extension.get_persisted(dsl_or_resource, {:authentication_action, action_name}) do
+      nil -> :error
+      value -> {:ok, value}
+    end
+  end
+
+  @doc """
+  Given an action name, retrieve the strategy it is for from the DSL
+  configuration.
+  """
+  @spec strategy_for_action!(dsl_or_resource, atom) :: Strategy.t() | no_return
+  def strategy_for_action!(dsl_or_resource, action_name) do
+    case strategy_for_action(dsl_or_resource, action_name) do
+      {:ok, value} ->
+        value
+
+      :error ->
+        raise "No strategy action named `#{inspect(action_name)}` found on resource `#{inspect(dsl_or_resource)}`"
+    end
+  end
 end
diff --git a/lib/ash_authentication/strategies/oauth2/actions.ex b/lib/ash_authentication/strategies/oauth2/actions.ex
index c100917..bccdb3b 100644
--- a/lib/ash_authentication/strategies/oauth2/actions.ex
+++ b/lib/ash_authentication/strategies/oauth2/actions.ex
@@ -26,7 +26,6 @@ defmodule AshAuthentication.Strategy.OAuth2.Actions do
 
     strategy.resource
     |> Query.new()
-    |> Query.set_context(%{strategy: strategy})
     |> Query.for_read(strategy.sign_in_action_name, params)
     |> api.read(options)
     |> case do
@@ -45,7 +44,6 @@ defmodule AshAuthentication.Strategy.OAuth2.Actions do
 
     strategy.resource
     |> Changeset.new()
-    |> Changeset.set_context(%{strategy: strategy})
     |> Changeset.for_create(strategy.register_action_name, params,
       upsert?: true,
       upsert_identity: action.upsert_identity
diff --git a/lib/ash_authentication/strategies/oauth2/identity_change.ex b/lib/ash_authentication/strategies/oauth2/identity_change.ex
index e694306..a48d2c3 100644
--- a/lib/ash_authentication/strategies/oauth2/identity_change.ex
+++ b/lib/ash_authentication/strategies/oauth2/identity_change.ex
@@ -4,7 +4,7 @@ defmodule AshAuthentication.Strategy.OAuth2.IdentityChange do
   """
 
   use Ash.Resource.Change
-  alias AshAuthentication.UserIdentity
+  alias AshAuthentication.{Info, UserIdentity}
   alias Ash.{Changeset, Error.Framework.AssumptionFailed, Resource.Change}
   import AshAuthentication.Utils, only: [is_falsy: 1]
 
@@ -12,13 +12,15 @@ defmodule AshAuthentication.Strategy.OAuth2.IdentityChange do
   @impl true
   @spec change(Changeset.t(), keyword, Change.context()) :: Changeset.t()
   def change(changeset, _opts, _context) do
-    case Map.fetch(changeset.context, :strategy) do
+    case Info.strategy_for_action(changeset.resource, changeset.action.name) do
       {:ok, strategy} ->
         do_change(changeset, strategy)
 
       :error ->
         {:error,
-         AssumptionFailed.exception(message: "Strategy is missing from the changeset context.")}
+         AssumptionFailed.exception(
+           message: "Action does not correlate with an authentication strategy"
+         )}
     end
   end
 
diff --git a/lib/ash_authentication/strategies/oauth2/sign_in_preparation.ex b/lib/ash_authentication/strategies/oauth2/sign_in_preparation.ex
index 2f18c8b..e92ebdb 100644
--- a/lib/ash_authentication/strategies/oauth2/sign_in_preparation.ex
+++ b/lib/ash_authentication/strategies/oauth2/sign_in_preparation.ex
@@ -11,7 +11,7 @@ defmodule AshAuthentication.Strategy.OAuth2.SignInPreparation do
   """
   use Ash.Resource.Preparation
   alias Ash.{Error.Framework.AssumptionFailed, Query, Resource.Preparation}
-  alias AshAuthentication.{Errors.AuthenticationFailed, Jwt, UserIdentity}
+  alias AshAuthentication.{Errors.AuthenticationFailed, Info, Jwt, UserIdentity}
   require Ash.Query
   import AshAuthentication.Utils, only: [is_falsy: 1]
 
@@ -19,7 +19,7 @@ defmodule AshAuthentication.Strategy.OAuth2.SignInPreparation do
   @impl true
   @spec prepare(Query.t(), keyword, Preparation.context()) :: Query.t()
   def prepare(query, _opts, _context) do
-    case Map.fetch(query.context, :strategy) do
+    case Info.strategy_for_action(query.resource, query.action.name) do
       :error ->
         {:error,
          AssumptionFailed.exception(message: "Strategy is missing from the changeset context.")}
diff --git a/lib/ash_authentication/strategies/oauth2/transformer.ex b/lib/ash_authentication/strategies/oauth2/transformer.ex
index 355bef4..ee1c756 100644
--- a/lib/ash_authentication/strategies/oauth2/transformer.ex
+++ b/lib/ash_authentication/strategies/oauth2/transformer.ex
@@ -52,13 +52,23 @@ defmodule AshAuthentication.Strategy.OAuth2.Transformer do
          :ok <- maybe_validate_register_action(dsl_state, strategy),
          :ok <- maybe_validate_sign_in_action(dsl_state, strategy),
          {:ok, resource} <- persisted_option(dsl_state, :module) do
+      strategy = %{strategy | resource: resource}
+
       dsl_state =
         dsl_state
         |> Transformer.replace_entity(
-          [:authentication, :strategies],
-          %{strategy | resource: resource},
+          ~w[authentication strategies]a,
+          strategy,
           &(&1.name == strategy.name)
         )
+        |> then(fn dsl_state ->
+          ~w[register_action_name sign_in_action_name]a
+          |> Stream.map(&Map.get(strategy, &1))
+          |> Enum.reduce(
+            dsl_state,
+            &Transformer.persist(&2, {:authentication_action, &1}, strategy)
+          )
+        end)
 
       {:ok, dsl_state}
     else
diff --git a/lib/ash_authentication/strategies/password/actions.ex b/lib/ash_authentication/strategies/password/actions.ex
index c6be995..c68720a 100644
--- a/lib/ash_authentication/strategies/password/actions.ex
+++ b/lib/ash_authentication/strategies/password/actions.ex
@@ -19,7 +19,6 @@ defmodule AshAuthentication.Strategy.Password.Actions do
 
     strategy.resource
     |> Query.new()
-    |> Query.set_context(%{strategy: strategy})
     |> Query.for_read(strategy.sign_in_action_name, params)
     |> api.read(options)
     |> case do
@@ -37,7 +36,6 @@ defmodule AshAuthentication.Strategy.Password.Actions do
 
     strategy.resource
     |> Changeset.new()
-    |> Changeset.set_context(%{strategy: strategy})
     |> Changeset.for_create(strategy.register_action_name, params)
     |> api.create(options)
   end
@@ -55,7 +53,6 @@ defmodule AshAuthentication.Strategy.Password.Actions do
 
     strategy.resource
     |> Query.new()
-    |> Query.set_context(%{strategy: strategy})
     |> Query.for_read(resettable.request_password_reset_action_name, params)
     |> api.read(options)
     |> case do
@@ -85,7 +82,6 @@ defmodule AshAuthentication.Strategy.Password.Actions do
 
       user
       |> Changeset.new()
-      |> Changeset.set_context(%{strategy: strategy})
       |> Changeset.for_update(resettable.password_reset_action_name, params)
       |> api.update(options)
     else
diff --git a/lib/ash_authentication/strategies/password/hash_password_change.ex b/lib/ash_authentication/strategies/password/hash_password_change.ex
index 469b7b8..f41a097 100644
--- a/lib/ash_authentication/strategies/password/hash_password_change.ex
+++ b/lib/ash_authentication/strategies/password/hash_password_change.ex
@@ -8,6 +8,7 @@ defmodule AshAuthentication.Strategy.Password.HashPasswordChange do
 
   use Ash.Resource.Change
   alias Ash.{Changeset, Error.Framework.AssumptionFailed, Resource.Change}
+  alias AshAuthentication.Info
 
   @doc false
   @impl true
@@ -15,7 +16,7 @@ defmodule AshAuthentication.Strategy.Password.HashPasswordChange do
   def change(changeset, _opts, _) do
     changeset
     |> Changeset.before_action(fn changeset ->
-      with {:ok, strategy} <- Map.fetch(changeset.context, :strategy),
+      with {:ok, strategy} <- Info.strategy_for_action(changeset.resource, changeset.action.name),
            value when is_binary(value) <-
              Changeset.get_argument(changeset, strategy.password_field),
            {:ok, hash} <- strategy.hash_provider.hash(value) do
diff --git a/lib/ash_authentication/strategies/password/password_confirmation_validation.ex b/lib/ash_authentication/strategies/password/password_confirmation_validation.ex
index 5412f9c..885f6a1 100644
--- a/lib/ash_authentication/strategies/password/password_confirmation_validation.ex
+++ b/lib/ash_authentication/strategies/password/password_confirmation_validation.ex
@@ -7,6 +7,7 @@ defmodule AshAuthentication.Strategy.Password.PasswordConfirmationValidation do
 
   use Ash.Resource.Validation
   alias Ash.{Changeset, Error.Changes.InvalidArgument, Error.Framework.AssumptionFailed}
+  alias AshAuthentication.Info
 
   @doc """
   Validates that the password and password confirmation fields contain
@@ -15,7 +16,7 @@ defmodule AshAuthentication.Strategy.Password.PasswordConfirmationValidation do
   @impl true
   @spec validate(Changeset.t(), keyword) :: :ok | {:error, String.t() | Exception.t()}
   def validate(changeset, _) do
-    case Map.fetch(changeset.context, :strategy) do
+    case Info.strategy_for_action(changeset.resource, changeset.action.name) do
       {:ok, %{confirmation_required?: true} = strategy} ->
         validate_password_confirmation(changeset, strategy)
 
@@ -24,7 +25,9 @@ defmodule AshAuthentication.Strategy.Password.PasswordConfirmationValidation do
 
       :error ->
         {:error,
-         AssumptionFailed.exception(message: "Strategy is missing from the changeset context.")}
+         AssumptionFailed.exception(
+           message: "Action does not correlate with an authentication strategy"
+         )}
     end
   end
 
diff --git a/lib/ash_authentication/strategies/password/request_password_reset_preparation.ex b/lib/ash_authentication/strategies/password/request_password_reset_preparation.ex
index 7650e3a..03c8798 100644
--- a/lib/ash_authentication/strategies/password/request_password_reset_preparation.ex
+++ b/lib/ash_authentication/strategies/password/request_password_reset_preparation.ex
@@ -10,14 +10,14 @@ defmodule AshAuthentication.Strategy.Password.RequestPasswordResetPreparation do
   """
   use Ash.Resource.Preparation
   alias Ash.{Query, Resource.Preparation}
-  alias AshAuthentication.Strategy.Password
+  alias AshAuthentication.{Info, Strategy.Password}
   require Ash.Query
 
   @doc false
   @impl true
   @spec prepare(Query.t(), keyword, Preparation.context()) :: Query.t()
   def prepare(query, _opts, _context) do
-    strategy = Map.fetch!(query.context, :strategy)
+    strategy = Info.strategy_for_action!(query.resource, query.action.name)
 
     if Enum.any?(strategy.resettable) do
       identity_field = strategy.identity_field
diff --git a/lib/ash_authentication/strategies/password/reset_token_validation.ex b/lib/ash_authentication/strategies/password/reset_token_validation.ex
index e85d2a1..2b8b938 100644
--- a/lib/ash_authentication/strategies/password/reset_token_validation.ex
+++ b/lib/ash_authentication/strategies/password/reset_token_validation.ex
@@ -5,13 +5,13 @@ defmodule AshAuthentication.Strategy.Password.ResetTokenValidation do
 
   use Ash.Resource.Validation
   alias Ash.{Changeset, Error.Changes.InvalidArgument}
-  alias AshAuthentication.Jwt
+  alias AshAuthentication.{Info, Jwt}
 
   @doc false
   @impl true
   @spec validate(Changeset.t(), keyword) :: :ok | {:error, Exception.t()}
   def validate(changeset, _) do
-    with {:ok, strategy} <- Map.fetch(changeset.context, :strategy),
+    with {:ok, strategy} <- Info.strategy_for_action(changeset.resource, changeset.action.name),
          token when is_binary(token) <- Changeset.get_argument(changeset, :reset_token),
          {:ok, %{"act" => token_action}, _} <- Jwt.verify(token, changeset.resource),
          {:ok, [resettable]} <- Map.fetch(strategy, :resettable),
diff --git a/lib/ash_authentication/strategies/password/sign_in_preparation.ex b/lib/ash_authentication/strategies/password/sign_in_preparation.ex
index bde2aa9..db7376e 100644
--- a/lib/ash_authentication/strategies/password/sign_in_preparation.ex
+++ b/lib/ash_authentication/strategies/password/sign_in_preparation.ex
@@ -13,7 +13,7 @@ defmodule AshAuthentication.Strategy.Password.SignInPreparation do
   an authentication failed error.
   """
   use Ash.Resource.Preparation
-  alias AshAuthentication.{Errors.AuthenticationFailed, Jwt}
+  alias AshAuthentication.{Errors.AuthenticationFailed, Info, Jwt}
   alias Ash.{Query, Resource.Preparation}
   require Ash.Query
 
@@ -21,7 +21,7 @@ defmodule AshAuthentication.Strategy.Password.SignInPreparation do
   @impl true
   @spec prepare(Query.t(), keyword, Preparation.context()) :: Query.t()
   def prepare(query, _opts, _context) do
-    strategy = Map.fetch!(query.context, :strategy)
+    strategy = Info.strategy_for_action!(query.resource, query.action.name)
     identity_field = strategy.identity_field
     identity = Query.get_argument(query, identity_field)
 
diff --git a/lib/ash_authentication/strategies/password/transformer.ex b/lib/ash_authentication/strategies/password/transformer.ex
index 1b5cd5d..4d8329b 100644
--- a/lib/ash_authentication/strategies/password/transformer.ex
+++ b/lib/ash_authentication/strategies/password/transformer.ex
@@ -71,13 +71,35 @@ defmodule AshAuthentication.Strategy.Password.Transformer do
          :ok <- validate_sign_in_action(dsl_state, strategy),
          {:ok, dsl_state, strategy} <- maybe_transform_resettable(dsl_state, strategy),
          {:ok, resource} <- persisted_option(dsl_state, :module) do
+      strategy = %{strategy | resource: resource}
+
       dsl_state =
         dsl_state
         |> Transformer.replace_entity(
-          [:authentication, :strategies],
-          %{strategy | resource: resource},
+          ~w[authentication strategies]a,
+          strategy,
           &(&1.name == strategy.name)
         )
+        |> then(fn dsl_state ->
+          ~w[sign_in_action_name register_action_name]a
+          |> Stream.map(&Map.get(strategy, &1))
+          |> Enum.reduce(
+            dsl_state,
+            &Transformer.persist(&2, {:authentication_action, &1}, strategy)
+          )
+        end)
+        |> then(fn dsl_state ->
+          strategy
+          |> Map.get(:resettable, [])
+          |> Stream.flat_map(fn resettable ->
+            ~w[request_password_reset_action_name password_reset_action_name]a
+            |> Stream.map(&Map.get(resettable, &1))
+          end)
+          |> Enum.reduce(
+            dsl_state,
+            &Transformer.persist(&2, {:authentication_action, &1}, strategy)
+          )
+        end)
 
       {:ok, dsl_state}
     end
diff --git a/test/support/data_case.ex b/test/support/data_case.ex
index ffb6163..3f4be01 100644
--- a/test/support/data_case.ex
+++ b/test/support/data_case.ex
@@ -79,12 +79,9 @@ defmodule DataCase do
       |> Map.put_new(:password, password)
       |> Map.put_new(:password_confirmation, password)
 
-    {:ok, strategy} = AshAuthentication.Info.strategy(Example.User, :password)
-
     user =
       Example.User
       |> Ash.Changeset.new()
-      |> Ash.Changeset.set_context(%{strategy: strategy})
       |> Ash.Changeset.for_create(:register_with_password, attrs)
       |> Example.create!()