fix(JWT): Generate tenant claims and validate them.

lib/ash_authentication/jwt/config.ex

@@ -48,6 +48,18 @@ defmodule AshAuthentication.Jwt.Config do
       &Joken.generate_jti/0,
       &validate_jti(&1, &2, &3, opts)
     )
+    |> maybe_add_tenant_claim(opts[:tenant])
+  end
+
+  defp maybe_add_tenant_claim(cfg, nil), do: cfg
+
+  defp maybe_add_tenant_claim(cfg, tenant) do
+    Config.add_claim(
+      cfg,
+      "tenant",
+      fn -> tenant end,
+      &validate_tenant(&1, tenant)
+    )
   end
 
   @doc """
@@ -77,6 +89,12 @@ defmodule AshAuthentication.Jwt.Config do
     "~> #{vsn.major}.#{vsn.minor}"
   end
 
+  @doc """
+  Validate that the "tenant" claim matches the provided tenant option.
+  """
+  @spec validate_tenant(nil | String.t(), nil | String) :: boolean()
+  def validate_tenant(maybe_tenant, tenant), do: maybe_tenant == tenant
+
   @doc """
   The validation function used to validate the "aud" claim.
 

priv/repo/migrations/20250218022011_add_user_with_multitenancy.exs

@@ -0,0 +1,42 @@
+defmodule Example.Repo.Migrations.AddUserWithMultitenancy do
+  @moduledoc """
+  Updates resources based on their most recent snapshots.
+
+  This file was autogenerated with `mix ash_postgres.generate_migrations`
+  """
+
+  use Ecto.Migration
+
+  def up do
+    create table(:user_with_multitenancy, primary_key: false) do
+      add(:id, :uuid, null: false, default: fragment("gen_random_uuid()"), primary_key: true)
+      add(:email, :citext, null: false)
+      add(:hashed_password, :text)
+      add(:tenant, :text, null: false)
+
+      add(:created_at, :utc_datetime_usec,
+        null: false,
+        default: fragment("(now() AT TIME ZONE 'utc')")
+      )
+
+      add(:updated_at, :utc_datetime_usec,
+        null: false,
+        default: fragment("(now() AT TIME ZONE 'utc')")
+      )
+    end
+
+    create unique_index(:user_with_multitenancy, [:tenant, :email],
+             name: "user_with_multitenancy_email_index"
+           )
+  end
+
+  def down do
+    drop_if_exists(
+      unique_index(:user_with_multitenancy, [:tenant, :email],
+        name: "user_with_multitenancy_email_index"
+      )
+    )
+
+    drop(table(:user_with_multitenancy))
+  end
+end

priv/resource_snapshots/repo/user_with_multitenancy/20250218022011.json

@@ -0,0 +1,94 @@
+{
+  "attributes": [
+    {
+      "allow_nil?": false,
+      "default": "fragment(\"gen_random_uuid()\")",
+      "generated?": false,
+      "primary_key?": true,
+      "references": null,
+      "size": null,
+      "source": "id",
+      "type": "uuid"
+    },
+    {
+      "allow_nil?": false,
+      "default": "nil",
+      "generated?": false,
+      "primary_key?": false,
+      "references": null,
+      "size": null,
+      "source": "email",
+      "type": "citext"
+    },
+    {
+      "allow_nil?": true,
+      "default": "nil",
+      "generated?": false,
+      "primary_key?": false,
+      "references": null,
+      "size": null,
+      "source": "hashed_password",
+      "type": "text"
+    },
+    {
+      "allow_nil?": false,
+      "default": "nil",
+      "generated?": false,
+      "primary_key?": false,
+      "references": null,
+      "size": null,
+      "source": "tenant",
+      "type": "text"
+    },
+    {
+      "allow_nil?": false,
+      "default": "fragment(\"(now() AT TIME ZONE 'utc')\")",
+      "generated?": false,
+      "primary_key?": false,
+      "references": null,
+      "size": null,
+      "source": "created_at",
+      "type": "utc_datetime_usec"
+    },
+    {
+      "allow_nil?": false,
+      "default": "fragment(\"(now() AT TIME ZONE 'utc')\")",
+      "generated?": false,
+      "primary_key?": false,
+      "references": null,
+      "size": null,
+      "source": "updated_at",
+      "type": "utc_datetime_usec"
+    }
+  ],
+  "base_filter": null,
+  "check_constraints": [],
+  "custom_indexes": [],
+  "custom_statements": [],
+  "has_create_action": true,
+  "hash": "A78FBE6B466C501B7ADFCA26D65967EBBBC9D5503BE161C21455EA3DFEADC486",
+  "identities": [
+    {
+      "all_tenants?": false,
+      "base_filter": null,
+      "index_name": "user_with_multitenancy_email_index",
+      "keys": [
+        {
+          "type": "atom",
+          "value": "email"
+        }
+      ],
+      "name": "email",
+      "nils_distinct?": true,
+      "where": null
+    }
+  ],
+  "multitenancy": {
+    "attribute": "tenant",
+    "global": false,
+    "strategy": "attribute"
+  },
+  "repo": "Elixir.Example.Repo",
+  "schema": null,
+  "table": "user_with_multitenancy"
+}

test/ash_authentication/jwt/config_test.exs

@@ -49,6 +49,18 @@ defmodule AshAuthentication.Jwt.ConfigTest do
     end
   end
 
+  describe "validate_tenant/2" do
+    test "when the provided tenant matches the expected tenant it is valid" do
+      assert Config.validate_tenant("banana", "banana")
+    end
+
+    test "when the provided tenant does not match the expected tenant, it is invalid" do
+      refute Config.validate_tenant("apple", "banana")
+      refute Config.validate_tenant(nil, "banana")
+      refute Config.validate_tenant("apple", nil)
+    end
+  end
+
   describe "validate_jti/3" do
     test "is true when the token has not been revoked" do
       TokenResource

test/ash_authentication/jwt_test.exs

@@ -29,7 +29,7 @@ defmodule AshAuthentication.JwtTest do
     end
   end
 
-  describe "token_for_user/1" do
+  describe "token_for_user/3" do
     test "correctly generates and signs tokens" do
       user = build_user()
       assert {:ok, token, claims} = Jwt.token_for_user(user)
@@ -45,6 +45,18 @@ defmodule AshAuthentication.JwtTest do
       assert_in_delta(claims["nbf"], now, 1.5)
       assert claims["sub"] == "user?id=#{user.id}"
     end
+
+    test "it encodes the tenant when passed one" do
+      user = build_user()
+      assert {:ok, _token, claims} = Jwt.token_for_user(user, %{}, tenant: "banana")
+      assert claims["tenant"] == "banana"
+    end
+
+    test "it doesn't encode the tenant otherwise" do
+      user = build_user()
+      assert {:ok, _token, claims} = Jwt.token_for_user(user, %{})
+      refute is_map_key(claims, "tenant")
+    end
   end
 
   describe "verify/2" do

test/ash_authentication_test.exs

@@ -6,7 +6,12 @@ defmodule AshAuthenticationTest do
 
   describe "authenticated_resources/0" do
     test "it correctly locates all authenticatable resources" do
-      assert [Example.User, Example.UserWithTokenRequired, Example.UserWithRegisterMagicLink] =
+      assert [
+               Example.User,
+               Example.UserWithMultitenancy,
+               Example.UserWithRegisterMagicLink,
+               Example.UserWithTokenRequired
+             ] =
                authenticated_resources(:ash_authentication)
     end
   end

test/support/data_case.ex

@@ -116,4 +116,30 @@ defmodule DataCase do
       Ash.Resource.put_metadata(user, field, value)
     end)
   end
+
+  @doc "User with multitenancy enabled factory"
+  @spec build_user_with_multitenancy(keyword) ::
+          Example.UserWithMultitenancy.t() | no_return
+  def build_user_with_multitenancy(attrs \\ []) do
+    password = password()
+
+    attrs =
+      attrs
+      |> Map.new()
+      |> Map.put_new(:email, Faker.Internet.email())
+      |> Map.put_new(:password, password)
+      |> Map.put_new(:password_confirmation, password)
+      |> Map.put_new(:tenant, Faker.Lorem.word())
+
+    user =
+      Example.UserWithMultitenancy
+      |> Ash.Changeset.new()
+      |> Ash.Changeset.for_create(:register_with_password, attrs)
+      |> Ash.create!(tenant: attrs[:tenant])
+
+    attrs
+    |> Enum.reduce(user, fn {field, value}, user ->
+      Ash.Resource.put_metadata(user, field, value)
+    end)
+  end
 end

test/support/example.ex

@@ -3,11 +3,12 @@ defmodule Example do
   use Ash.Domain, otp_app: :ash_authentication, extensions: [AshGraphql.Domain, AshJsonApi.Domain]
 
   resources do
-    resource Example.User
-    resource Example.UserWithTokenRequired
     resource Example.Token
+    resource Example.User
     resource Example.UserIdentity
+    resource Example.UserWithMultitenancy
     resource Example.UserWithRegisterMagicLink
+    resource Example.UserWithTokenRequired
   end
 
   json_api do

test/support/example/user_with_multitenancy.ex

@@ -0,0 +1,74 @@
+defmodule Example.UserWithMultitenancy do
+  @moduledoc false
+  use Ash.Resource,
+    data_layer: AshPostgres.DataLayer,
+    extensions: [AshAuthentication],
+    domain: Example
+
+  require Logger
+
+  attributes do
+    uuid_primary_key :id, writable?: true
+    attribute :email, :ci_string, allow_nil?: false, public?: true
+    attribute :hashed_password, :string, allow_nil?: true, sensitive?: true, public?: false
+    attribute :tenant, :string, allow_nil?: false, public?: true
+    create_timestamp :created_at
+    update_timestamp :updated_at
+  end
+
+  authentication do
+    tokens do
+      enabled? true
+      store_all_tokens? true
+      require_token_presence_for_authentication? true
+      token_resource Example.Token
+      signing_secret &get_config/2
+    end
+
+    strategies do
+      password do
+        identity_field :email
+        register_action_accept [:tenant]
+
+        resettable do
+          sender fn user, token, _opts ->
+            Logger.debug(
+              "Password reset request for user #{user.username}, token #{inspect(token)}"
+            )
+          end
+        end
+      end
+    end
+  end
+
+  actions do
+    defaults [:create, :read, :update, :destroy]
+  end
+
+  identities do
+    identity :email, [:email]
+  end
+
+  multitenancy do
+    strategy :attribute
+    attribute :tenant
+  end
+
+  postgres do
+    table "user_with_multitenancy"
+    repo(Example.Repo)
+
+    manage_tenant do
+      template([:tenant])
+    end
+  end
+
+  def get_config(path, _resource) do
+    value =
+      :ash_authentication
+      |> Application.get_all_env()
+      |> get_in(path)
+
+    {:ok, value}
+  end
+end