auth_cache.ex

defmodule ExAws.Config.AuthCache do
  @moduledoc false

  use GenServer

  # http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html

  @refresh_lead_time Application.compile_env(:ex_aws, :auth_refresh_lead_time, 300_000)
  @instance_auth_key :aws_instance_auth
  @pod_identity_auth_key :aws_pod_identity_auth

  defmodule AuthConfigAdapter do
    @moduledoc false

    @doc "Compute the awscli auth information."
    @callback adapt_auth_config(auth :: map, profile :: String.t(), expiration :: integer) :: any
  end

  def start_link(opts \\ []) do
    GenServer.start_link(__MODULE__, :ok, Keyword.put(opts, :name, __MODULE__))
  end

  def get(config) do
    :ets.lookup(__MODULE__, @instance_auth_key)
    |> refresh_auth_if_required(config)
  end

  def get(:pod_identity, config) do
    :ets.lookup(__MODULE__, @pod_identity_auth_key)
    |> refresh_pod_identity_if_required(config)
  end

  def get(profile, expiration) do
    case :ets.lookup(__MODULE__, {:awscli, profile}) do
      [{{:awscli, ^profile}, auth_config}] ->
        auth_config

      [] ->
        GenServer.call(__MODULE__, {:refresh_awscli_config, profile, expiration}, 30_000)
    end
  end

  ## Callbacks

  def init(:ok) do
    ets = :ets.new(__MODULE__, [:named_table, read_concurrency: true])
    {:ok, ets}
  end

  def handle_call({:refresh_auth, config}, _from, ets) do
    auth = refresh_auth(config, ets)
    {:reply, auth, ets}
  end

  def handle_call({:refresh_pod_identity_auth, config}, _from, ets) do
    auth = refresh_pod_identity_auth(config, ets)
    {:reply, auth, ets}
  end

  def handle_call({:refresh_awscli_config, profile, expiration}, _from, ets) do
    case :ets.lookup(ets, {:awscli, profile}) do
      [{{:awscli, ^profile}, auth}] ->
        {:reply, auth, ets}

      [] ->
        auth = refresh_awscli_config(profile, expiration, ets)
        {:reply, auth, ets}
    end
  end

  def handle_info({:refresh_auth, config}, ets) do
    refresh_auth(config, ets)
    {:noreply, ets}
  end

  def handle_info({:refresh_pod_identity_auth, config}, ets) do
    refresh_pod_identity_auth(config, ets)
    {:noreply, ets}
  end

  def handle_info({:refresh_awscli_config, profile, expiration}, ets) do
    refresh_awscli_config(profile, expiration, ets)
    {:noreply, ets}
  end

  def refresh_awscli_config(profile, expiration, ets) do
    auth =
      try do
        ExAws.Config.awscli_auth_credentials(profile)
      rescue
        e ->
          require Logger

          Logger.error(
            "Failed to refresh AWS CLI config for profile #{profile}: #{Exception.format(:error, e, __STACKTRACE__)}"
          )

          nil
      end

    auth =
      if auth do
        case ExAws.Config.awscli_auth_adapter() do
          nil ->
            auth

          adapter ->
            attempt_credentials_refresh(adapter, auth, profile, expiration)
        end
      else
        nil
      end

    refresh_interval =
      if auth && Map.get(auth, :expiration) do
        interval = next_refresh_in(auth)

        if interval > 0 do
          max(interval, expiration)
        else
          expiration
        end
      else
        expiration
      end

    Process.send_after(self(), {:refresh_awscli_config, profile, expiration}, refresh_interval)

    if auth do
      :ets.insert(ets, {{:awscli, profile}, auth})
    end

    auth
  end

  defp attempt_credentials_refresh(adapter, auth, profile, expiration, retries \\ 6) do
    case adapter.adapt_auth_config(auth, profile, expiration) do
      {:error, error} when retries == 1 ->
        Process.send_after(self(), {:refresh_awscli_config, profile, expiration}, expiration)

        raise "Couldn't get credentials from auth adapter after 6 retries, last error was #{inspect(error)}"

      {:error, _error} ->
        Process.sleep(:rand.uniform(5_000))
        attempt_credentials_refresh(adapter, auth, profile, expiration, retries - 1)

      # Always store a map on AuthCache
      auth when is_map(auth) ->
        auth
    end
  end

  defp refresh_auth_if_required([], config) do
    GenServer.call(__MODULE__, {:refresh_auth, config}, 30_000)
  end

  defp refresh_auth_if_required([{_key, cached_auth}], config) do
    if next_refresh_in(cached_auth) > 0 do
      cached_auth
    else
      GenServer.call(__MODULE__, {:refresh_auth, config}, 30_000)
    end
  end

  defp refresh_auth(config, ets) do
    :ets.lookup(__MODULE__, @instance_auth_key)
    |> refresh_auth_if_stale(config, ets)
  end

  defp refresh_auth_if_stale([], config, ets) do
    refresh_auth_now(config, ets)
  end

  defp refresh_auth_if_stale([{_key, cached_auth}], config, ets) do
    if next_refresh_in(cached_auth) > @refresh_lead_time do
      # we still have a valid auth token, so simply return that
      cached_auth
    else
      refresh_auth_now(config, ets)
    end
  end

  defp refresh_auth_if_stale(_, config, ets), do: refresh_auth_now(config, ets)

  defp refresh_auth_now(config, ets) do
    auth = ExAws.InstanceMeta.security_credentials(config)
    :ets.insert(ets, {@instance_auth_key, auth})
    Process.send_after(__MODULE__, {:refresh_auth, config}, next_refresh_in(auth))
    auth
  end

  defp refresh_pod_identity_if_required([], config) do
    GenServer.call(__MODULE__, {:refresh_pod_identity_auth, config}, 30_000)
  end

  defp refresh_pod_identity_if_required([{_key, cached_auth}], config) do
    if next_refresh_in(cached_auth) > 0 do
      cached_auth
    else
      GenServer.call(__MODULE__, {:refresh_pod_identity_auth, config}, 30_000)
    end
  end

  defp refresh_pod_identity_auth(config, ets) do
    :ets.lookup(__MODULE__, @pod_identity_auth_key)
    |> refresh_pod_identity_if_stale(config, ets)
  end

  defp refresh_pod_identity_if_stale([], config, ets) do
    refresh_pod_identity_now(config, ets)
  end

  defp refresh_pod_identity_if_stale([{_key, cached_auth}], config, ets) do
    if next_refresh_in(cached_auth) > @refresh_lead_time do
      cached_auth
    else
      refresh_pod_identity_now(config, ets)
    end
  end

  defp refresh_pod_identity_if_stale(_, config, ets), do: refresh_pod_identity_now(config, ets)

  defp refresh_pod_identity_now(config, ets) do
    # Only attempt pod identity if the environment is properly configured
    if ExAws.PodIdentity.available?() do
      auth = ExAws.PodIdentity.security_credentials(config)
      :ets.insert(ets, {@pod_identity_auth_key, auth})
      Process.send_after(__MODULE__, {:refresh_pod_identity_auth, config}, next_refresh_in(auth))
      auth
    else
      # Return empty map if pod identity is not available
      %{}
    end
  end

  defp next_refresh_in(%{expiration: expiration}) when is_integer(expiration) do
    try do
      expires_in_ms = (expiration - System.os_time(:second)) * 1000

      # refresh lead_time before auth expires, unless the time has passed
      # otherwise refresh needed now
      max(0, expires_in_ms - @refresh_lead_time)
    rescue
      _e -> 0
    end
  end

  defp next_refresh_in(%{expiration: expiration}) when is_binary(expiration) do
    try do
      expires_in_ms =
        expiration
        |> NaiveDateTime.from_iso8601!()
        |> NaiveDateTime.diff(NaiveDateTime.utc_now(), :millisecond)

      # refresh lead_time before auth expires, unless the time has passed
      # otherwise refresh needed now
      max(0, expires_in_ms - @refresh_lead_time)
    rescue
      _e -> 0
    end
  end

  defp next_refresh_in(_), do: 0
end