From b75097dd3e2e67e4c82f2701e23a96cdc04677b0 Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Fri, 9 Aug 2024 09:28:32 -0400
Subject: [PATCH] improvement: set context in addition to tenant

---
 lib/ash_authentication/plug/helpers.ex               | 9 ++++++---
 lib/ash_authentication/strategies/magic_link/plug.ex | 4 ++--
 lib/ash_authentication/strategies/oauth2/plug.ex     | 4 ++--
 lib/ash_authentication/strategies/password/plug.ex   | 4 ++--
 4 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/lib/ash_authentication/plug/helpers.ex b/lib/ash_authentication/plug/helpers.ex
index 5b3fa63..c317407 100644
--- a/lib/ash_authentication/plug/helpers.ex
+++ b/lib/ash_authentication/plug/helpers.ex
@@ -86,11 +86,13 @@ defmodule AshAuthentication.Plug.Helpers do
                    "jti" => jti,
                    "purpose" => "user"
                  },
-                 tenant: Ash.PlugHelpers.get_tenant(conn)
+                 tenant: Ash.PlugHelpers.get_tenant(conn),
+                 context: Ash.PlugHelpers.get_context(conn)
                ),
              {:ok, user} <-
                AshAuthentication.subject_to_user(subject, resource,
-                 tenant: Ash.PlugHelpers.get_tenant(conn)
+                 tenant: Ash.PlugHelpers.get_tenant(conn),
+                 context: Ash.PlugHelpers.get_context(conn)
                ) do
           Conn.assign(conn, current_subject_name, user)
         else
@@ -138,7 +140,8 @@ defmodule AshAuthentication.Plug.Helpers do
              validate_token(resource, jti),
            {:ok, user} <-
              AshAuthentication.subject_to_user(subject, resource,
-               tenant: Ash.PlugHelpers.get_tenant(conn)
+               tenant: Ash.PlugHelpers.get_tenant(conn),
+               context: Ash.PlugHelpers.get_context(conn)
              ),
            {:ok, subject_name} <- Info.authentication_subject_name(resource),
            current_subject_name <- current_subject_name(subject_name) do
diff --git a/lib/ash_authentication/strategies/magic_link/plug.ex b/lib/ash_authentication/strategies/magic_link/plug.ex
index 0f41e13..011a83d 100644
--- a/lib/ash_authentication/strategies/magic_link/plug.ex
+++ b/lib/ash_authentication/strategies/magic_link/plug.ex
@@ -7,7 +7,7 @@ defmodule AshAuthentication.Strategy.MagicLink.Plug do
 
   alias AshAuthentication.{Info, Strategy, Strategy.MagicLink}
   alias Plug.Conn
-  import Ash.PlugHelpers, only: [get_actor: 1, get_tenant: 1]
+  import Ash.PlugHelpers, only: [get_actor: 1, get_tenant: 1, get_context: 1]
   import AshAuthentication.Plug.Helpers, only: [store_authentication_result: 2]
 
   @doc """
@@ -55,7 +55,7 @@ defmodule AshAuthentication.Strategy.MagicLink.Plug do
   end
 
   defp opts(conn) do
-    [actor: get_actor(conn), tenant: get_tenant(conn)]
+    [actor: get_actor(conn), tenant: get_tenant(conn), context: get_context(conn)]
     |> Enum.reject(&is_nil(elem(&1, 1)))
   end
 end
diff --git a/lib/ash_authentication/strategies/oauth2/plug.ex b/lib/ash_authentication/strategies/oauth2/plug.ex
index 43f297e..71320a2 100644
--- a/lib/ash_authentication/strategies/oauth2/plug.ex
+++ b/lib/ash_authentication/strategies/oauth2/plug.ex
@@ -7,7 +7,7 @@ defmodule AshAuthentication.Strategy.OAuth2.Plug do
   alias AshAuthentication.{Errors, Info, Strategy, Strategy.OAuth2}
   alias Assent.{Config, HTTPAdapter.Finch}
   alias Plug.Conn
-  import Ash.PlugHelpers, only: [get_actor: 1, get_tenant: 1]
+  import Ash.PlugHelpers, only: [get_actor: 1, get_tenant: 1, get_context: 1]
   import AshAuthentication.Plug.Helpers, only: [store_authentication_result: 2]
   import Plug.Conn
 
@@ -74,7 +74,7 @@ defmodule AshAuthentication.Strategy.OAuth2.Plug do
   end
 
   defp action_opts(conn) do
-    [actor: get_actor(conn), tenant: get_tenant(conn)]
+    [actor: get_actor(conn), tenant: get_tenant(conn), get_context(conn)]
     |> Enum.reject(&is_nil(elem(&1, 1)))
   end
 
diff --git a/lib/ash_authentication/strategies/password/plug.ex b/lib/ash_authentication/strategies/password/plug.ex
index 51828ef..80c4b55 100644
--- a/lib/ash_authentication/strategies/password/plug.ex
+++ b/lib/ash_authentication/strategies/password/plug.ex
@@ -7,7 +7,7 @@ defmodule AshAuthentication.Strategy.Password.Plug do
 
   alias AshAuthentication.{Info, Strategy, Strategy.Password}
   alias Plug.Conn
-  import Ash.PlugHelpers, only: [get_actor: 1, get_tenant: 1]
+  import Ash.PlugHelpers, only: [get_actor: 1, get_tenant: 1, get_context: 1]
   import AshAuthentication.Plug.Helpers, only: [store_authentication_result: 2]
 
   @doc "Handle a registration request"
@@ -65,7 +65,7 @@ defmodule AshAuthentication.Strategy.Password.Plug do
   end
 
   defp opts(conn) do
-    [actor: get_actor(conn), tenant: get_tenant(conn)]
+    [actor: get_actor(conn), tenant: get_tenant(conn), context: get_context(conn)]
     |> Enum.reject(&is_nil(elem(&1, 1)))
   end
 end