ash_authentication

# SPDX-FileCopyrightText: 2022 Alembic Pty Ltd
#
# SPDX-License-Identifier: MIT

defmodule AshAuthentication.TokenResource.Actions do
  @moduledoc """
  The code interface for interacting with the token resource.
  """

  alias Ash.{Changeset, Query, Resource}
  alias AshAuthentication.{Errors.InvalidToken, Jwt, TokenResource, TokenResource.Info}

  import AshAuthentication.Utils

  require Logger

  @doc false
  @spec read_expired(Resource.t(), keyword) :: {:ok, [Resource.record()]} | {:error, any}
  def read_expired(resource, opts \\ []) do
    with :ok <- assert_resource_has_extension(resource, TokenResource),
         {:ok, domain} <- Info.token_domain(resource),
         {:ok, read_expired_action_name} <- Info.token_read_expired_action_name(resource) do
      resource
      |> Query.new()
      |> Query.set_context(%{
        private: %{
          ash_authentication?: true
        }
      })
      |> Query.for_read(read_expired_action_name, opts)
      |> Ash.read(domain: domain)
    end
  end

  @doc """
  Remove all expired records.
  """
  @spec expunge_expired(Resource.t(), keyword) :: :ok | {:error, any}
  def expunge_expired(resource, opts \\ []) do
    expunge_expired_action_name = Info.token_expunge_expired_action_name!(resource)

    with :ok <- assert_resource_has_extension(resource, TokenResource),
         {:ok, read_expired_action_name} <- Info.token_read_expired_action_name(resource) do
      opts =
        opts
        |> Keyword.put_new_lazy(:domain, fn -> Info.token_domain!(resource) end)

      authorize_with =
        if Ash.DataLayer.data_layer_can?(resource, :expr_error) do
          :error
        else
          :filter
        end

      resource
      |> Query.new()
      |> Query.set_context(%{private: %{ash_authentication?: true}})
      |> Query.for_read(read_expired_action_name, %{}, opts)
      |> Ash.bulk_destroy(
        expunge_expired_action_name,
        %{},
        opts
        |> Keyword.update(
          :context,
          %{private: %{ash_authentication?: true}},
          &Ash.Helpers.deep_merge_maps(&1, %{private: %{ash_authentication?: true}})
        )
        |> Keyword.merge(
          strategy: [:atomic, :atomic_batches, :stream],
          return_errors?: true,
          notify?: true,
          return_records?: false,
          return_notifications?: false,
          authorize_with: authorize_with
        )
      )
      |> case do
        %{status: :success} -> :ok
        %{errors: errors} -> {:error, Ash.Error.to_class(errors)}
      end
    end
  end

  @doc """
  Has the token been revoked?

  Similar to `jti_revoked?/2..3` except that it extracts the JTI from the token,
  rather than relying on it to be passed in.
  """
  @spec token_revoked?(Resource.t(), String.t(), keyword) :: boolean
  def token_revoked?(resource, token, opts \\ []) do
    with :ok <- assert_resource_has_extension(resource, TokenResource),
         {:ok, domain} <- Info.token_domain(resource),
         {:ok, is_revoked_action_name} <- Info.token_revocation_is_revoked_action_name(resource) do
      resource
      |> Ash.ActionInput.for_action(
        is_revoked_action_name,
        %{"token" => token},
        Keyword.put(opts, :domain, domain)
      )
      |> Ash.ActionInput.set_context(%{
        private: %{
          ash_authentication?: true
        }
      })
      |> Ash.run_action()
      |> case do
        {:ok, value} ->
          value

        {:error, error} ->
          Logger.error("""
          Error while checking if token is revoked.
          We must assume that it is revoked for security purposes.

          #{Exception.format(:error, error)}
          """)

          true
      end
    end
  end

  @doc """
  Has the token been revoked?

  Similar to `token-revoked?/2..3` except that rather than extracting the JTI
  from the token, assumes that it's being passed in directly.
  """
  @spec jti_revoked?(Resource.t(), String.t(), keyword) :: boolean
  def jti_revoked?(resource, jti, opts \\ []) do
    with :ok <- assert_resource_has_extension(resource, TokenResource),
         {:ok, domain} <- Info.token_domain(resource),
         {:ok, is_revoked_action_name} <- Info.token_revocation_is_revoked_action_name(resource) do
      resource
      |> Ash.ActionInput.for_action(
        is_revoked_action_name,
        %{"jti" => jti},
        Keyword.take(Keyword.put(opts, :domain, domain), [
          :actor,
          :authorize?,
          :context,
          :tenant,
          :tracer,
          :domain
        ])
      )
      |> Ash.ActionInput.set_context(%{
        private: %{
          ash_authentication?: true
        }
      })
      |> Ash.run_action()
      |> case do
        {:ok, value} ->
          value

        {:error, error} ->
          Logger.error("""
          Error while checking if token is revoked.
          We must assume that it is revoked for security purposes.

          #{Exception.format(:error, error)}
          """)

          true
      end
    end
  end

  @doc false
  @spec valid_jti?(Resource.t(), String.t(), keyword) :: boolean
  def valid_jti?(resource, jti, opts \\ []), do: !jti_revoked?(resource, jti, opts)

  @doc """
  Revoke a token.

  Extracts the JTI from the provided token and uses it to generate a revocation
  record.

  ## Options

    * `:store_all_tokens?` — describes whether the token was generated by an
      authentication resource with `store_all_tokens?` enabled. When provided,
      the revocation is performed atomically so that concurrent revocations of
      the same token cannot both succeed:

        * `true` — the token row is expected to exist; it is locked with
          `SELECT … FOR UPDATE` and its purpose flipped to `"revocation"`.
          Returns `{:error, _}` if the row is already a revocation record.
        * `false` — the token row is expected not to exist; a plain insert is
          performed. A concurrent duplicate results in a primary key conflict
          which is surfaced as `{:error, _}`.

      If omitted, the legacy upsert behaviour is used. This path is NOT safe
      against concurrent duplicate revocations and is retained only for
      backwards compatibility.
  """
  @spec revoke(Resource.t(), String.t(), keyword) :: :ok | {:error, any}
  def revoke(resource, token, opts \\ []) do
    {store_all_tokens?, opts} = Keyword.pop(opts, :store_all_tokens?)

    case store_all_tokens? do
      true -> revoke_locked(resource, token, opts)
      false -> revoke_atomic_insert(resource, token, opts)
      nil -> revoke_legacy_upsert(resource, token, opts)
    end
  end

  defp revoke_atomic_insert(resource, token, opts) do
    with :ok <- assert_resource_has_extension(resource, TokenResource),
         {:ok, domain} <- Info.token_domain(resource),
         {:ok, insert_action_name} <-
           Info.token_revocation_revoke_token_insert_action_name(resource) do
      resource
      |> Changeset.new()
      |> Changeset.set_context(%{private: %{ash_authentication?: true}})
      |> Changeset.for_create(insert_action_name, %{"token" => token}, opts)
      |> Ash.create(domain: domain)
      |> case do
        {:ok, _} -> :ok
        {:error, reason} -> {:error, translate_jti_conflict(reason)}
      end
    end
  end

  defp revoke_locked(resource, token, opts) do
    with :ok <- assert_resource_has_extension(resource, TokenResource),
         {:ok, %{"jti" => jti}} <- Jwt.peek(token),
         {:ok, domain} <- Info.token_domain(resource),
         lookup_opts =
           opts
           |> Keyword.take([:actor, :authorize?, :context, :tenant, :tracer])
           |> Keyword.merge(
             domain: domain,
             authorize?: false,
             error?: false
           )
           |> maybe_lock_opt(resource, :for_update),
         {:ok, existing} <- Ash.get(resource, [jti: jti], lookup_opts) do
      case existing && existing.purpose do
        "revocation" ->
          {:error,
           InvalidToken.exception(type: :revocation, reason: "token has already been revoked")}

        nil ->
          # Row missing despite the caller expecting it to exist (e.g. stored
          # token was deleted, or config changed). Fall back to atomic insert
          # so the race-safety guarantee is preserved.
          revoke_atomic_insert(resource, token, opts)

        _ ->
          revoke_legacy_upsert(resource, token, opts)
      end
    end
  end

  defp revoke_legacy_upsert(resource, token, opts) do
    with :ok <- assert_resource_has_extension(resource, TokenResource),
         {:ok, domain} <- Info.token_domain(resource),
         {:ok, revoke_token_action_name} <-
           Info.token_revocation_revoke_token_action_name(resource) do
      resource
      |> Changeset.new()
      |> Changeset.set_context(%{private: %{ash_authentication?: true}})
      |> Changeset.for_create(
        revoke_token_action_name,
        %{"token" => token},
        Keyword.merge(opts, upsert?: true)
      )
      |> Ash.create(domain: domain)
      |> case do
        {:ok, _} -> :ok
        {:error, reason} -> {:error, reason}
      end
    end
  end

  @doc """
  Revoke a token by JTI.

  If you have the token, you should use `revoke/2` instead.

  Accepts the same `:store_all_tokens?` option as `revoke/3`.
  """
  @spec revoke_jti(Resource.t(), String.t(), String.t(), keyword) ::
          :ok | {:error, any}
  def revoke_jti(resource, jti, subject, opts \\ []) do
    {store_all_tokens?, opts} = Keyword.pop(opts, :store_all_tokens?)

    case store_all_tokens? do
      true -> revoke_jti_locked(resource, jti, subject, opts)
      false -> revoke_jti_atomic_insert(resource, jti, subject, opts)
      nil -> revoke_jti_legacy_upsert(resource, jti, subject, opts)
    end
  end

  defp revoke_jti_atomic_insert(resource, jti, subject, opts) do
    with :ok <- assert_resource_has_extension(resource, TokenResource),
         {:ok, domain} <- Info.token_domain(resource),
         {:ok, insert_action_name} <-
           Info.token_revocation_revoke_jti_insert_action_name(resource) do
      resource
      |> Changeset.new()
      |> Changeset.set_context(%{private: %{ash_authentication?: true}})
      |> Changeset.for_create(insert_action_name, %{"jti" => jti, "subject" => subject}, opts)
      |> Ash.create(domain: domain)
      |> case do
        {:ok, _} -> :ok
        {:error, reason} -> {:error, translate_jti_conflict(reason)}
      end
    end
  end

  defp revoke_jti_locked(resource, jti, subject, opts) do
    with :ok <- assert_resource_has_extension(resource, TokenResource),
         {:ok, domain} <- Info.token_domain(resource),
         lookup_opts =
           opts
           |> Keyword.take([:actor, :authorize?, :context, :tenant, :tracer])
           |> Keyword.merge(
             domain: domain,
             authorize?: false,
             error?: false
           )
           |> maybe_lock_opt(resource, :for_update),
         {:ok, existing} <- Ash.get(resource, [jti: jti], lookup_opts) do
      case existing && existing.purpose do
        "revocation" ->
          {:error,
           InvalidToken.exception(type: :revocation, reason: "token has already been revoked")}

        nil ->
          revoke_jti_atomic_insert(resource, jti, subject, opts)

        _ ->
          revoke_jti_legacy_upsert(resource, jti, subject, opts)
      end
    end
  end

  defp revoke_jti_legacy_upsert(resource, jti, subject, opts) do
    with :ok <- assert_resource_has_extension(resource, TokenResource),
         {:ok, domain} <- Info.token_domain(resource),
         {:ok, revoke_jti_action_name} <-
           Info.token_revocation_revoke_jti_action_name(resource) do
      resource
      |> Changeset.new()
      |> Changeset.set_context(%{private: %{ash_authentication?: true}})
      |> Changeset.for_create(
        revoke_jti_action_name,
        %{"jti" => jti, "subject" => subject},
        Keyword.merge(opts, upsert?: true)
      )
      |> Ash.create(domain: domain)
      |> case do
        {:ok, _} -> :ok
        {:error, reason} -> {:error, reason}
      end
    end
  end

  defp translate_jti_conflict(%Ash.Error.Invalid{errors: errors} = error) do
    if Enum.any?(errors, &jti_unique_conflict?/1) do
      InvalidToken.exception(type: :revocation, reason: "token has already been revoked")
    else
      error
    end
  end

  defp translate_jti_conflict(error), do: error

  defp jti_unique_conflict?(%Ash.Error.Changes.InvalidAttribute{field: :jti, private_vars: vars}),
    do: Keyword.get(vars, :constraint_type) == :unique

  defp jti_unique_conflict?(_), do: false

  @doc """
  Store a token.

  Stores a token for any purpose.
  """
  @spec store_token(Resource.t(), map, keyword) :: :ok | {:error, any}
  def store_token(resource, params, opts \\ []) do
    with :ok <- assert_resource_has_extension(resource, TokenResource),
         {:ok, domain} <- Info.token_domain(resource),
         {:ok, store_token_action_name} <- Info.token_store_token_action_name(resource) do
      resource
      |> Changeset.new()
      |> Changeset.set_context(%{
        private: %{
          ash_authentication?: true
        }
      })
      |> Changeset.for_create(
        store_token_action_name,
        params,
        Keyword.merge(opts, upsert?: true)
      )
      |> Ash.create(domain: domain)
      |> case do
        {:ok, _} -> :ok
        {:error, reason} -> {:error, reason}
      end
    end
  end

  @doc """
  Retrieve a token by token or JTI optionally filtering by purpose.
  """
  @spec get_token(Resource.t(), map, keyword) :: {:ok, [Resource.record()]} | {:error, any}
  def get_token(resource, params, opts \\ []) do
    with :ok <- assert_resource_has_extension(resource, TokenResource),
         {:ok, domain} <- Info.token_domain(resource),
         {:ok, get_token_action_name} <- Info.token_get_token_action_name(resource) do
      resource
      |> Query.new()
      |> Query.set_context(%{private: %{ash_authentication?: true}})
      |> Query.for_read(
        get_token_action_name,
        params,
        Keyword.take(Keyword.put(opts, :domain, domain), [
          :actor,
          :authorize?,
          :tenant,
          :tracer,
          :domain
        ])
      )
      |> Ash.read()
    end
  end
end