From 5f03f21bb5cca15be4cc7fc26e16e38668b79e21 Mon Sep 17 00:00:00 2001 From: Johannes Christ Date: Fri, 10 May 2024 21:00:26 +0200 Subject: [PATCH] Remove Application.get_env call from Nostrum.Api.Base Closes #546 --- lib/nostrum/api/base.ex | 32 ++++++++++++++++---------- lib/nostrum/api/ratelimiter.ex | 41 +++++++++++++++++++++++----------- lib/nostrum/application.ex | 2 +- 3 files changed, 49 insertions(+), 26 deletions(-) diff --git a/lib/nostrum/api/base.ex b/lib/nostrum/api/base.ex index 5b9512eea..069d24a8f 100644 --- a/lib/nostrum/api/base.ex +++ b/lib/nostrum/api/base.ex @@ -6,12 +6,23 @@ defmodule Nostrum.Api.Base do import Nostrum.Constants, only: [base_route: 0] require Logger + @typedoc "An anonymous function that returns the bot token." + @type wrapped_token :: (-> String.t()) + @type methods :: :get | :post | :put | :delete - @spec request(pid, methods(), String.t(), iodata(), [{String.t(), String.t()}], Enum.t()) :: + @spec request( + pid, + methods(), + String.t(), + iodata(), + [{String.t(), String.t()}], + Enum.t(), + wrapped_token() + ) :: :gun.stream_ref() - def request(conn, method, route, body, raw_headers, params) do - headers = process_request_headers(raw_headers) + def request(conn, method, route, body, raw_headers, params, wrapped_token) do + headers = process_request_headers(raw_headers, wrapped_token) # Convert method from atom to string for `:gun` method = method @@ -21,24 +32,21 @@ defmodule Nostrum.Api.Base do query_string = URI.encode_query(params) full_route = "#{base_route()}#{route}?#{query_string}" - headers = process_request_headers(headers, body) + headers = finalize_request_headers(headers, body) :gun.request(conn, method, full_route, headers, process_request_body(body)) end - def process_request_headers(headers, ""), do: :proplists.delete("content-type", headers) - def process_request_headers(headers, _body), do: headers + def finalize_request_headers(headers, ""), do: :proplists.delete("content-type", headers) + def finalize_request_headers(headers, _body), do: headers def process_request_body(""), do: "" def process_request_body({:multipart, content}), do: content def process_request_body(body), do: Jason.encode_to_iodata!(body) - def process_request_headers(headers) do - user_agent = [ + def process_request_headers(headers, wrapped_token) do + [ + {"authorization", "Bot #{wrapped_token.()}"}, {"user-agent", "DiscordBot (https://github.com/kraigie/nostrum, #{@version})"} | headers ] - - token = "Bot #{Application.get_env(:nostrum, :token)}" - - [{"authorization", token} | user_agent] end def process_response_body(body) do diff --git a/lib/nostrum/api/ratelimiter.ex b/lib/nostrum/api/ratelimiter.ex index 86bdf89d9..413bca7c3 100644 --- a/lib/nostrum/api/ratelimiter.ex +++ b/lib/nostrum/api/ratelimiter.ex @@ -279,6 +279,10 @@ defmodule Nostrum.Api.Ratelimiter do - `:remaining_in_window`: How many calls we may still make to the API during this time window. Reset automatically via timeouts. + + - `:wrapped_token`: An anonymous function that is internally used to retrieve + the token. This is wrapped to ensure that it is not accidentally exposed in + stacktraces. """ @typedoc since: "0.9.0" @type state :: %{ @@ -294,23 +298,24 @@ defmodule Nostrum.Api.Ratelimiter do body :: String.t()} }, conn: pid() | nil, - remaining_in_window: non_neg_integer() + remaining_in_window: non_neg_integer(), + wrapped_token: Base.wrapped_token() } @doc """ Starts the ratelimiter. """ - @spec start_link([:gen_statem.start_opt()]) :: :gen_statem.start_ret() - def start_link(opts) do - :gen_statem.start_link({:local, @registered_name}, __MODULE__, [], opts) + @spec start_link({String.t(), [:gen_statem.start_opt()]}) :: :gen_statem.start_ret() + def start_link({token, opts}) do + :gen_statem.start_link({:local, @registered_name}, __MODULE__, token, opts) end - def init([]) do + def init(token) when is_binary(token) do # Uncomment the following to trace everything the ratelimiter is doing: # me = self() # spawn(fn -> :sys.trace(me, true) end) # See more examples in the `sys` docs. - {:ok, :disconnected, empty_state()} + {:ok, :disconnected, empty_state(token)} end def callback_mode, do: :state_functions @@ -469,7 +474,12 @@ defmodule Nostrum.Api.Ratelimiter do def connected( :internal, {:run, request, bucket, from}, - %{conn: conn, outstanding: outstanding, remaining_in_window: remaining_for_user} = data + %{ + conn: conn, + outstanding: outstanding, + remaining_in_window: remaining_for_user, + wrapped_token: wrapped_token + } = data ) when remaining_for_user > 0 and is_map_key(outstanding, bucket) do stream = @@ -479,7 +489,8 @@ defmodule Nostrum.Api.Ratelimiter do request.route, request.body, request.headers, - request.params + request.params, + wrapped_token ) data_with_this_running = put_in(data, [:running, stream], {bucket, request, from}) @@ -865,7 +876,10 @@ defmodule Nostrum.Api.Ratelimiter do {:next_event, :internal, {:requeue, {request, from}}}} end - def connected(:info, {:gun_down, conn, _, reason, killed_streams}, %{running: running}) do + def connected(:info, {:gun_down, conn, _, reason, killed_streams}, %{ + running: running, + wrapped_token: wrapped_token + }) do # Even with `retry: 0`, gun seems to try and reconnect, potentially because # of WebSocket. Force the connection to die. :ok = :gun.close(conn) @@ -886,7 +900,7 @@ defmodule Nostrum.Api.Ratelimiter do {:reply, client, {:error, {:connection_died, reason}}} end) - {:next_state, :disconnected, empty_state(), replies} + {:next_state, :disconnected, empty_state(wrapped_token.()), replies} end def global_limit(:state_timeout, next, data) do @@ -936,14 +950,15 @@ defmodule Nostrum.Api.Ratelimiter do defp parse_response(status, headers, buffer), do: {:ok, {status, headers, buffer}} - @spec empty_state :: state() - defp empty_state, + @spec empty_state(String.t()) :: state() + defp empty_state(token), do: %{ outstanding: %{}, running: %{}, inflight: %{}, conn: nil, - remaining_in_window: @bot_calls_per_window + remaining_in_window: @bot_calls_per_window, + wrapped_token: fn -> token end } # Helper functions diff --git a/lib/nostrum/application.ex b/lib/nostrum/application.ex index d776280b0..96906812d 100644 --- a/lib/nostrum/application.ex +++ b/lib/nostrum/application.ex @@ -28,7 +28,7 @@ defmodule Nostrum.Application do children = [ Nostrum.Store.Supervisor, Nostrum.ConsumerGroup, - Nostrum.Api.Ratelimiter, + {Nostrum.Api.Ratelimiter, {Application.fetch_env!(:nostrum, :token), []}}, Nostrum.Shard.Connector, Nostrum.Cache.CacheSupervisor, Nostrum.Shard.Supervisor,