147 lines
4.1 KiB
Elixir
147 lines
4.1 KiB
Elixir
defmodule WhoNeedHelp.Trust.RateLimiter do
|
|
@moduledoc """
|
|
Shared PostgreSQL-backed action limits.
|
|
|
|
Policies are supplied as a map whose values contain positive `limit` and
|
|
`window_seconds` integers. The application ships a pilot product policy,
|
|
while runtime configuration can replace the complete map or explicitly use
|
|
an empty map for an isolated test environment.
|
|
"""
|
|
|
|
import Ecto.Query
|
|
|
|
alias WhoNeedHelp.Repo
|
|
alias WhoNeedHelp.Trust.RateLimitBucket
|
|
|
|
def check(action, scope) do
|
|
case check_many([{action, scope}]) do
|
|
{:ok, [limit]} -> {:ok, limit}
|
|
{:ok, []} -> {:ok, :not_configured}
|
|
{:error, :rate_limited} = error -> error
|
|
end
|
|
end
|
|
|
|
@doc """
|
|
Atomically increments every configured action/scope pair in one PostgreSQL
|
|
statement.
|
|
|
|
Unconfigured actions are ignored. The call fails when any configured bucket
|
|
exceeds its policy, while still counting the whole attempt consistently.
|
|
"""
|
|
def check_many(action_scopes) when is_list(action_scopes) do
|
|
now = DateTime.utc_now(:second)
|
|
|
|
prepared =
|
|
action_scopes
|
|
|> Enum.map(fn {action, scope} -> {normalize_action(action), scope} end)
|
|
|> Enum.uniq()
|
|
|> Enum.flat_map(fn {action, scope} ->
|
|
case policy(action) do
|
|
nil -> []
|
|
policy -> [prepare_bucket(action, scope, policy, now)]
|
|
end
|
|
end)
|
|
|
|
case prepared do
|
|
[] -> {:ok, []}
|
|
buckets -> increment_many(buckets)
|
|
end
|
|
end
|
|
|
|
def prune_expired do
|
|
now = DateTime.utc_now(:second)
|
|
Repo.delete_all(from bucket in RateLimitBucket, where: bucket.expires_at <= ^now)
|
|
end
|
|
|
|
defp policy(action) do
|
|
policies = Application.get_env(:who_need_help, :rate_limit_policies, %{})
|
|
|
|
raw =
|
|
if is_map(policies) do
|
|
Map.get(policies, action) || Map.get(policies, String.to_existing_atom(action))
|
|
end
|
|
|
|
case raw do
|
|
%{limit: limit, window_seconds: window}
|
|
when is_integer(limit) and limit > 0 and is_integer(window) and window > 0 ->
|
|
%{limit: limit, window_seconds: window}
|
|
|
|
%{"limit" => limit, "window_seconds" => window}
|
|
when is_integer(limit) and limit > 0 and is_integer(window) and window > 0 ->
|
|
%{limit: limit, window_seconds: window}
|
|
|
|
_ ->
|
|
nil
|
|
end
|
|
rescue
|
|
ArgumentError -> nil
|
|
end
|
|
|
|
defp prepare_bucket(action, scope, %{limit: limit, window_seconds: window_seconds}, now) do
|
|
unix = DateTime.to_unix(now)
|
|
window_start = (div(unix, window_seconds) * window_seconds) |> DateTime.from_unix!()
|
|
expires_at = DateTime.add(window_start, window_seconds, :second)
|
|
scope_hash = :crypto.hash(:sha256, to_string(scope))
|
|
|
|
%{
|
|
action: action,
|
|
scope_hash: scope_hash,
|
|
window_started_at: window_start,
|
|
expires_at: expires_at,
|
|
limit: limit,
|
|
now: now
|
|
}
|
|
end
|
|
|
|
defp increment_many(buckets) do
|
|
now = buckets |> hd() |> Map.fetch!(:now)
|
|
|
|
rows =
|
|
Enum.map(buckets, fn bucket ->
|
|
%{
|
|
id: Ecto.UUID.generate(),
|
|
action: bucket.action,
|
|
scope_hash: bucket.scope_hash,
|
|
window_started_at: bucket.window_started_at,
|
|
count: 1,
|
|
expires_at: bucket.expires_at,
|
|
inserted_at: now,
|
|
updated_at: now
|
|
}
|
|
end)
|
|
|
|
{_count, returned} =
|
|
Repo.insert_all(
|
|
RateLimitBucket,
|
|
rows,
|
|
on_conflict: [inc: [count: 1], set: [updated_at: now]],
|
|
conflict_target: [:action, :scope_hash, :window_started_at],
|
|
returning: [:action, :scope_hash, :count, :expires_at]
|
|
)
|
|
|
|
limits =
|
|
Map.new(buckets, fn bucket ->
|
|
{{bucket.action, bucket.scope_hash}, bucket.limit}
|
|
end)
|
|
|
|
results =
|
|
Enum.map(returned, fn bucket ->
|
|
%{
|
|
action: bucket.action,
|
|
count: bucket.count,
|
|
limit: Map.fetch!(limits, {bucket.action, bucket.scope_hash}),
|
|
resets_at: bucket.expires_at
|
|
}
|
|
end)
|
|
|
|
if Enum.all?(results, &(&1.count <= &1.limit)) do
|
|
{:ok, results}
|
|
else
|
|
{:error, :rate_limited}
|
|
end
|
|
end
|
|
|
|
defp normalize_action(action) when is_atom(action), do: Atom.to_string(action)
|
|
defp normalize_action(action) when is_binary(action), do: action
|
|
end
|