Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove Application.get_env call from Nostrum.Api.Base #593

Merged
merged 1 commit into from
May 25, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 20 additions & 12 deletions lib/nostrum/api/base.ex
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
41 changes: 28 additions & 13 deletions lib/nostrum/api/ratelimiter.ex
Original file line number Diff line number Diff line change
Expand Up @@ -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 :: %{
Expand All @@ -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
Expand Down Expand Up @@ -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 =
Expand All @@ -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})
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion lib/nostrum/application.ex
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down