ex_web3

if Code.ensure_loaded?(Plug) do
  defmodule ExWeb3.Guardian.Plug.VerifySignature do
    @moduledoc """
    Looks for and validates a signature found in the `Authorization` header.
    In the case where either:
    1. A signature is already found for `:key`
    2. No signature is found in the `Authorization` header.
    This plug will not do anything.
    This, like all other Guardian plugs, requires a Guardian pipeline to be setup.
    It requires an implementation module, an error handler and a key.
    These can be set either:
    1. Upstream on the connection with `plug Guardian.Pipeline`
    2. Upstream on the connection with `Guardian.Pipeline.{put_module, put_error_handler, put_key}`
    3. Inline with an option of `:module`, `:error_handler`, `:key`
    If a signature is found but is invalid, the error handler will be called with
    `auth_error(conn, {:invalid_token, reason}, opts)`
    Once a signature has been found it will be decoded, the token and claims will
    be put onto the connection.
    They will be available using `Guardian.Plug.current_claims/2` and `Guardian.Plug.current_token/2`
    Options:
    * `resource_by_pub_key` - Function to fetch authorized resource by a given public key in hex encoding.
    * `claims_for_resource` - Function to fetch claims for authorized resource.
    ### Example
    ```elixir
    # setup the upstream pipeline
    plug ExWeb3.Guardian.Plug.VerifySignature,
      resource_by_pub_key: &Filters.fetch_by_pub_key/1,
      claims_for_resource: &(MyApp.ImplementationModule.build_claims(%{}, &1))
    ```
    This will check the authorization header for a token
    `Authorization: Bearer <token>`
    This token will be placed into the connection depending on the key and can be accessed with
    `Guardian.Plug.current_token` and `Guardian.Plug.current_claims`.
    OR
    `MyApp.ImplementationModule.current_token` and `MyApp.ImplementationModule.current_claims`.
    """

    alias Guardian.Plug.Pipeline

    import Plug.Conn
    require Logger

    @behaviour Plug

    @impl Plug
    @spec init(opts :: Keyword.t()) :: Keyword.t()
    def init(opts \\ []) do
      unless Keyword.has_key?(opts, :resource_by_pub_key),
        do: throw "Function :resource_by_pub_key is required"
      unless Keyword.has_key?(opts, :claims_for_resource),
        do: throw "Function :claims_for_resource is required"

      opts
      |> get_scheme()
      |> put_scheme_reg(opts)
    end

    @impl Plug
    @spec call(Plug.Conn.t(), Keyword.t()) :: Plug.Conn.t()
    def call(conn, opts) do
      fn_resource_by_pub_key = Keyword.fetch!(opts, :resource_by_pub_key)
      fn_claims_for_resource = Keyword.fetch!(opts, :claims_for_resource)
      with nil <- Guardian.Plug.current_token(conn, opts),
           {:ok, signature} <- fetch_signature_from_header(conn, opts),
           {:ok, digest} <- fetch_digest_from_body(conn, opts),
           conn <- parse_body(conn),
           key <- storage_key(conn, opts),
           {:ok, bin_pub_key} <- ExWeb3.eth_recover(digest, signature),
           {:ok, resource} <- fn_resource_by_pub_key.(ExWeb3.to_hex(bin_pub_key)),
           {:ok, claims} <- fn_claims_for_resource.(resource) do
        conn
        |> Guardian.Plug.put_current_token(signature, key: key)
        |> Guardian.Plug.put_current_resource(resource, key: key)
        |> Guardian.Plug.put_current_claims(claims, key: key)
      else
        error ->
          parse_body(conn)
          |> handle_error(error, opts)
      end
    end

    defp put_scheme_reg("", opts) do
      opts
    end

    defp put_scheme_reg(:none, opts) do
      opts
    end

    defp put_scheme_reg(scheme, opts) do
      {:ok, reg} = Regex.compile("#{scheme}\:?\s+(.*)$", "i")
      Keyword.put(opts, :scheme_reg, reg)
    end

    defp get_scheme(opts) do
      if Keyword.has_key?(opts, :realm) do
        IO.warn("`:realm` option is deprecated; please rename `:realm` to `:scheme` option instead.")
        Keyword.get(opts, :realm, "Bearer")
      else
        Keyword.get(opts, :scheme, "Bearer")
      end
    end

    defp handle_error(conn, error, opts) do
      if refresh_from_cookie_opts = fetch_refresh_from_cookie_options(opts) do
        Guardian.Plug.VerifyCookie.refresh_from_cookie(conn, refresh_from_cookie_opts)
      else
        apply_error(conn, error, opts)
      end
    end

    defp apply_error(conn, {:error, reason}, opts) do
      conn
      |> Pipeline.fetch_error_handler!(opts)
      |> apply(:auth_error, [conn, {:invalid_token, reason}, opts])
      |> Guardian.Plug.maybe_halt(opts)
    end

    defp apply_error(conn, _, _) do
      conn
    end

    defp fetch_refresh_from_cookie_options(opts) do
      case Keyword.get(opts, :refresh_from_cookie) do
        value when is_list(value) -> value
        true -> []
        _ -> nil
      end
    end

    @spec fetch_digest_from_body(Plug.Conn.t(), Keyword.t()) ::
            :no_digest_found
            | {:ok, binary()}
    defp fetch_digest_from_body(%{private: %{raw_body: message}} = _conn, _opts) do
      try do
        {:ok, ExWeb3.eth_hash_message!(message)}
      catch
        error ->
          Logger.error "[#{__MODULE__}] failed to generate digest to match signature for authentication: #{inspect(error)}"
          {:error, :invalid_token}
      end
    end
    defp fetch_digest_from_body(_conn, _opts) do
      Logger.critical("[#{__MODULE__}] failed to get request body to match signature for authentication. " <>
        "Use ExWeb3.Phoenix.Plug.SaveRawBody to set :raw_body private property for your connection")
      {:error, :invalid_token}
    end

    @spec fetch_signature_from_header(Plug.Conn.t(), Keyword.t()) ::
            :no_token_found
            | {:ok, String.t()}
    defp fetch_signature_from_header(conn, opts) do
      header_name = Keyword.get(opts, :header_name, "authorization")
      headers = get_req_header(conn, header_name)
      fetch_signature_from_header(conn, opts, headers)
    end

    @spec fetch_signature_from_header(Plug.Conn.t(), Keyword.t(), Keyword.t()) ::
            :no_token_found
            | {:ok, String.t()}
    defp fetch_signature_from_header(_, _, []), do: :no_token_found

    defp fetch_signature_from_header(conn, opts, [token | tail]) do
      reg = Keyword.get(opts, :scheme_reg, ~r/^(.*)$/)
      trimmed_token = String.trim(token)

      case Regex.run(reg, trimmed_token) do
        [_, match] ->
          with "0x" <> _ = hex_signature <- String.trim(match) do
            {:ok, hex_signature}
          else
            error ->
              Logger.error "[#{__MODULE__}] failed to extract signature from header: #{inspect(error)}"
              :no_token_found
          end

        _ ->
          fetch_signature_from_header(conn, opts, tail)
      end
    end

    defp parse_body(%{private: %{raw_body: message}, params: params} = conn) do
      case Jason.decode(message) do
        {:ok, data} ->
          %{conn | body_params: data, params: Map.merge(params, data)}

        _ ->
          conn
      end
    end
    defp parse_body(conn), do: conn

    @spec storage_key(Plug.Conn.t(), Keyword.t()) :: String.t()
    defp storage_key(conn, opts), do: Pipeline.fetch_key(conn, opts)
  end
end