Support other OAuth Token Exchanges

defmodule ReqOAuthTokenExchange do
  @moduledoc """
  `Req` plugin for [OAuth 2.0] bearer-token authentication for single-request
  access token exchanges.

  This plugin is intended for OAuth flows where the client can obtain an
  access token by making one request to the token endpoint, such as client
  credentials, JWT bearer authorization grants from the [OAuth 2.0 Assertion
  Framework], or refresh-token exchanges. It accepts token request
  configuration under the `:oauth` key. The configured grant module supplies
  the OAuth `grant_type` and any grant-specific token endpoint request params.
  The access token will be cached and reused for subsequent requests. If the
  response to the authenticated request is a 401, the plugin will request a new
  access token and retry once.

  This is not a full OAuth flow implementation and does not handle flows that
  require redirects, browser/user-agent interaction, device-code polling,
  callback handling, or authorization-code acquisition.

  [OAuth 2.0]: https://www.rfc-editor.org/rfc/rfc6749
  [OAuth 2.0 Assertion Framework]: https://www.rfc-editor.org/rfc/rfc7521
  """

  @callback grant_type(Keyword.t()) :: String.t()
  @callback token_request_params(Keyword.t()) :: Keyword.t()
  @callback cache_key(Keyword.t()) :: term()

  @optional_callbacks cache_key: 1

  defguardp validated_request?(request)
            when is_tuple(request.private.req_oauth_data)

  @doc """
  Runs the plugin.

  ## Usage

      Req.new(base_url: "https://api.example.com", finch: Ev2Finch)
      |> ReqOAuthTokenExchange.attach()
      |> Req.get!(
        url: "/resource",
        oauth: [
          grant: MyApp.OauthGrant,
          url: "https://auth.example.com/oauth/token",
          grant_config: []
        ]
      )
      #=> %Req.Response{}
  """
  def attach(%Req.Request{} = req, opts \\ []) do
    req
    |> Req.Request.register_options([:oauth])
    |> Req.Request.merge_options(opts)
    |> Req.Request.append_request_steps(oauth: &auth/1)
    |> Req.Request.prepend_response_steps(oauth: &refresh_on_unauthorized/1)
  end

  def bust_cache(%Req.Request{} = request) do
    with {:ok, cache_key} <- request_cache_key(request) do
      :persistent_term.erase(cache_key)
    end
  end

  @doc false
  def write_cache(%Req.Request{} = request, token) do
    with {:ok, cache_key} <- request_cache_key(request) do
      :persistent_term.put(cache_key, token)
    end
  end

  defp auth(request) do
    options = Req.Request.get_option(request, :oauth, [])

    with {:ok, request} <- validate(request, options),
         {:ok, token} <- fetch_token(request) do
      request
      |> Req.Request.put_header("authorization", "Bearer " <> token)
      |> Req.Request.put_private(:req_oauth_refreshed?, false)
    else
      {_request, _response_or_exception} = result -> result
      _other -> request
    end
  end

  defp refresh_on_unauthorized({request, response}) when validated_request?(request) and response.status == 401 do
    if Req.Request.get_private(request, :req_oauth_refreshed?) do
      {request, response}
    else
      bust_cache(request)

      case request_token(request) do
        {:ok, token} ->
          %{request | halted: false}
          |> Req.Request.put_header("authorization", "Bearer " <> token)
          |> Req.Request.put_private(:req_oauth_refreshed?, true)
          |> Req.Request.run_request()

        _error ->
          {request, response}
      end
    end
  end

  defp refresh_on_unauthorized({request, response}), do: {request, response}

  defp validate(request, options) do
    with :ok <- validate_url(options),
         {:ok, {options, params, cache_key}} <- validate_params(options) do
      {:ok, Req.Request.put_private(request, :req_oauth_data, {options, params, cache_key})}
    else
      _ -> :error
    end
  end

  defp validate_url(options) do
    if is_binary(options[:url]) and options[:url] != "" do
      :ok
    else
      :error
    end
  end

  defp validate_params(options) do
    grant = options[:grant]
    grant_config = Keyword.get(options, :grant_config, [])

    if is_atom(grant) and Code.ensure_loaded?(grant) and
         function_exported?(grant, :grant_type, 1) and function_exported?(grant, :token_request_params, 1) and
         is_list(grant_config) do
      params = [{:grant_type, grant.grant_type(grant_config)} | grant.token_request_params(grant_config)]
      {:ok, {Keyword.take(options, [:url]), params, cache_key(options)}}
    else
      :error
    end
  end

  defp fetch_token(request) do
    case fetch_cache(request) do
      {:ok, token} -> {:ok, token}
      :error -> request_token(request)
    end
  end

  defp fetch_cache(request) do
    with {:ok, cache_key} <- request_cache_key(request),
         token when is_binary(token) <- :persistent_term.get(cache_key, :error) do
      {:ok, token}
    else
      _other -> :error
    end
  end

  defp request_token(request) do
    {options, params, _cache_key} = Req.Request.get_private(request, :req_oauth_data)
    options = put_in(options[:form], params)

    auth_req =
      Req.Request.new()
      |> Req.Request.append_request_steps(Keyword.delete(request.request_steps, :oauth))
      |> Req.Request.append_response_steps(Keyword.delete(request.response_steps, :oauth))
      |> Req.Request.append_error_steps(request.error_steps)
      |> Req.Request.register_options(Enum.to_list(request.registered_options))
      |> Req.Request.merge_options(
        request.options
        |> Map.drop([:body, :form, :json, :oauth])
        |> Map.to_list()
      )

    # Safe because auth_req inherits the parent request options, including finch: Ev2Finch.
    # credo:disable-for-next-line Ev2.Credo.RequireReqFinch
    case Req.post(auth_req, options) do
      {:ok, %{body: %{"access_token" => token}}} ->
        write_cache(request, token)
        {:ok, token}

      {:ok, %Req.Response{} = response} ->
        {request, response}

      _error ->
        :error
    end
  end

  defp request_cache_key(request) do
    case Req.Request.get_private(request, :req_oauth_data) do
      {_options, _params, cache_key} ->
        {:ok, cache_key}

      nil ->
        options = Req.Request.get_option(request, :oauth, [])

        with :ok <- validate_url(options),
             {:ok, {_options, _params, cache_key}} <- validate_params(options) do
          {:ok, cache_key}
        else
          _other -> :error
        end
    end
  end

  defp cache_key(options) do
    grant = options[:grant]
    grant_config = Keyword.get(options, :grant_config, [])

    grant_cache_key =
      if Code.ensure_loaded?(grant) and function_exported?(grant, :cache_key, 1) do
        grant.cache_key(grant_config)
      else
        grant_config
      end

    {__MODULE__, options[:url], grant, grant_cache_key}
  end
end

defmodule ReqOAuthTokenExchange.DocusignJWTBearerGrant do
  @moduledoc """
  DocuSign JWT bearer grant for `ReqOAuthTokenExchange`.

  Implements DocuSign's JWT grant token request as described in
  [DocuSign's JWT access token guide](https://developers.docusign.com/platform/auth/jwt-get-token/).

  DocuSign expects the requested scope in the signed JWT assertion. This is
  different from the generic OAuth assertion profile, where `scope` is an
  OAuth token endpoint request parameter.
  """
  @behaviour ReqOAuthTokenExchange

  alias ReqOAuthTokenExchange

  @grant_type "urn:ietf:params:oauth:grant-type:jwt-bearer"

  @config_schema [
    auth_host: [type: :string, required: true],
    integration_key: [type: :string, required: true],
    user_id: [type: :string, required: true],
    private_key: [type: :string, required: true],
    scope: [type: :string, required: true],
    expires_in_seconds: [type: :pos_integer, required: true]
  ]

  @impl ReqOAuthTokenExchange
  def grant_type(_config), do: @grant_type

  @impl ReqOAuthTokenExchange
  def token_request_params(config), do: [assertion: jwt(config)]

  @impl ReqOAuthTokenExchange
  def cache_key(config) do
    config = config(config)

    {
      config[:auth_host],
      config[:integration_key],
      config[:user_id],
      config[:scope],
      config[:expires_in_seconds],
      :crypto.hash(:sha256, config[:private_key])
    }
  end

  defp jwt(config) do
    config = config(config)
    now = Joken.current_time()

    claims = %{
      "iss" => config[:integration_key],
      "sub" => config[:user_id],
      "aud" => config[:auth_host],
      "iat" => now,
      "exp" => now + config[:expires_in_seconds],
      "scope" => config[:scope]
    }

    signer = Joken.Signer.create("RS256", %{"pem" => config[:private_key]})
    {:ok, jwt} = Joken.Signer.sign(claims, signer)
    jwt
  end

  defp config(config), do: NimbleOptions.validate!(config, @config_schema)
end

defmodule ReqOAuthTokenExchange.ClientCredentialsGrant do
  @moduledoc """
  OAuth 2.0 client credentials grant for `ReqOAuthTokenExchange`.

  Implements the token request params for the client credentials grant. `scope`
  is optional per OAuth 2.0. `audience` and `extra_params` are included for
  providers that require additional token endpoint params.
  """
  @behaviour ReqOAuthTokenExchange

  alias ReqOAuthTokenExchange

  @grant_type "client_credentials"

  @config_schema [
    client_id: [type: :string, required: true],
    client_secret: [type: :string, required: true],
    scope: [type: :string],
    audience: [type: :string],
    extra_params: [type: :keyword_list, default: []]
  ]

  @impl ReqOAuthTokenExchange
  def grant_type(_config), do: @grant_type

  @impl ReqOAuthTokenExchange
  def token_request_params(config) do
    config = config(config)

    config[:extra_params]
    |> Keyword.put_new(:audience, config[:audience])
    |> Keyword.put_new(:scope, config[:scope])
    |> Keyword.put(:client_secret, config[:client_secret])
    |> Keyword.put(:client_id, config[:client_id])
    |> Enum.reject(fn {_key, value} -> is_nil(value) end)
  end

  defp config(config), do: NimbleOptions.validate!(config, @config_schema)
end