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

lib/ash_authentication/jwt/config.ex

@@ -7,7 +7,7 @@ defmodule AshAuthentication.Jwt.Config do
   revocation.
   """
 
-  alias Ash.Resource
+  alias Ash.{Error.Invalid.TenantRequired, Resource}
   alias AshAuthentication.{Info, TokenResource}
   alias Joken.{Config, Signer}
 
@@ -48,6 +48,11 @@ defmodule AshAuthentication.Jwt.Config do
       &Joken.generate_jti/0,
       &validate_jti(&1, &2, &3, opts)
     )
+    |> Config.add_claim(
+      "tenant",
+      fn -> generate_tenant(resource, opts) end,
+      &validate_tenant(&1, resource, opts)
+    )
   end
 
   @doc """
@@ -77,6 +82,51 @@ defmodule AshAuthentication.Jwt.Config do
     "~> #{vsn.major}.#{vsn.minor}"
   end
 
+  @doc """
+  The generator function used to generate the "tenant" claim.
+
+  If the token resource or the user resource have tenancy enabled then we
+  validate the presence of the `tenant` option.
+  """
+  @spec generate_tenant(Resource.t(), Keyword.t()) :: nil | String.t()
+  def generate_tenant(resource, options) do
+    if multitenancy_required?(resource) do
+      case Keyword.fetch(options, :tenant) do
+        {:ok, tenant} -> tenant
+        :error -> raise TenantRequired, resource: resource
+      end
+    end
+  end
+
+  @doc """
+  Validate that the "tenant" claim matches the provided tenant option.
+
+  If the token resource or the user resource have tenancy enabled then we
+  validate the value of the "token" claim.
+  """
+  @spec validate_tenant(String.t(), Resource.t(), Keyword.t()) :: boolean()
+  def validate_tenant(tenant, resource, options) do
+    if multitenancy_required?(resource) do
+      case Keyword.fetch(options, :tenant) do
+        {:ok, ^tenant} -> true
+        _ -> false
+      end
+    else
+      true
+    end
+  end
+
+  defp multitenancy_required?(user_resource) do
+    if Resource.Info.multitenancy_strategy(user_resource) do
+      true
+    else
+      case Info.authentication_tokens_token_resource(user_resource) do
+        {:ok, token_resource} -> not is_nil(Resource.Info.multitenancy_strategy(token_resource))
+        _ -> false
+      end
+    end
+  end
+
   @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,43 @@ defmodule AshAuthentication.Jwt.ConfigTest do
     end
   end
 
+  describe "generate_tenant/2" do
+    test "when multitenancy is enabled it returns the tenant option" do
+      assert "banana" = Config.generate_tenant(Example.UserWithMultitenancy, tenant: "banana")
+    end
+
+    test "when multitenancy is enabled and the tenant option is missing, it raises" do
+      assert_raise Ash.Error.Invalid.TenantRequired, fn ->
+        Config.generate_tenant(Example.UserWithMultitenancy, [])
+      end
+    end
+
+    test "when multitenancy is not enabled it returns nil" do
+      refute Config.generate_tenant(Example.UserWithTokenRequired, tenant: "banana")
+    end
+  end
+
+  describe "validate_tenant/2" do
+    test "when multitenancy is enabled and the tenant matches it validates" do
+      assert Config.validate_tenant("banana", Example.UserWithMultitenancy, tenant: "banana")
+    end
+
+    test "when multitenancy is enabled and the tenant doesn't match it doesn't validate" do
+      refute Config.validate_tenant("banana", Example.UserWithMultitenancy, tenant: "apple")
+      refute Config.validate_tenant(nil, Example.UserWithMultitenancy, tenant: "apple")
+    end
+
+    test "when multitenancy is enabled and no tenant option is present it doesn't validate" do
+      refute Config.validate_tenant("banana", Example.UserWithMultitenancy, [])
+      refute Config.validate_tenant(nil, Example.UserWithMultitenancy, [])
+    end
+
+    test "when multitenancy is not enabled it is always valid" do
+      assert Config.validate_tenant("banana", Example.UserWithTokenRequired, tenant: "apple")
+      assert Config.validate_tenant(nil, Example.UserWithTokenRequired, [])
+    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 required" do
+      user = build_user_with_multitenancy(tenant: "banana")
+      assert {:ok, _token, claims} = Jwt.token_for_user(user, %{}, tenant: "banana")
+      assert claims["tenant"] == "banana"
+    end
+
+    test "it doesn't encode the tenant when multitenancy is disabled" do
+      user = build_user()
+      assert {:ok, _token, claims} = Jwt.token_for_user(user, %{}, tenant: "banana")
+      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