From d133ee84b7ced7d052f701379b78f30768f29aed Mon Sep 17 00:00:00 2001 From: Brian Manifold Date: Wed, 12 Mar 2025 20:21:09 -0700 Subject: [PATCH] feat(portal): Add API rate limiting (#8417) --- elixir/apps/api/lib/api/application.ex | 3 +- elixir/apps/api/lib/api/plugs/rate_limit.ex | 41 +++++++++++ elixir/apps/api/lib/api/rate_limit.ex | 70 ++++++++++++++++++ elixir/apps/api/lib/api/router.ex | 1 + elixir/apps/api/mix.exs | 1 + elixir/apps/api/test/api/rate_limit.exs | 71 +++++++++++++++++++ .../domain/lib/domain/config/definitions.ex | 10 +++ elixir/config/config.exs | 4 ++ elixir/config/runtime.exs | 4 ++ elixir/mix.lock | 1 + 10 files changed, 205 insertions(+), 1 deletion(-) create mode 100644 elixir/apps/api/lib/api/plugs/rate_limit.ex create mode 100644 elixir/apps/api/lib/api/rate_limit.ex create mode 100644 elixir/apps/api/test/api/rate_limit.exs diff --git a/elixir/apps/api/lib/api/application.ex b/elixir/apps/api/lib/api/application.ex index e83082974..fd3785d75 100644 --- a/elixir/apps/api/lib/api/application.ex +++ b/elixir/apps/api/lib/api/application.ex @@ -7,7 +7,8 @@ defmodule API.Application do _ = OpentelemetryPhoenix.setup(adapter: :cowboy2) children = [ - API.Endpoint + API.Endpoint, + API.RateLimit ] opts = [strategy: :one_for_one, name: API.Supervisor] diff --git a/elixir/apps/api/lib/api/plugs/rate_limit.ex b/elixir/apps/api/lib/api/plugs/rate_limit.ex new file mode 100644 index 000000000..0e0824cae --- /dev/null +++ b/elixir/apps/api/lib/api/plugs/rate_limit.ex @@ -0,0 +1,41 @@ +defmodule API.Plugs.RateLimit do + import Plug.Conn + + @refill_rate_default Domain.Config.fetch_env!(:api, API.RateLimit)[:refill_rate] + @capacity_default Domain.Config.fetch_env!(:api, API.RateLimit)[:capacity] + @cost_default API.RateLimit.default_cost() + + def init(opts), do: Keyword.get(opts, :context_type, :api_client) + + def call(conn, _context_type) do + rate_limit_api(conn, []) + end + + defp rate_limit_api(conn, _opts) do + account = conn.assigns.subject.account + key = "api:#{account.id}" + refill_rate = refill_rate(account) + capacity = capacity(account) + + case API.RateLimit.hit(key, refill_rate, capacity, @cost_default) do + {:allow, _count} -> + conn + + {:deny, _refill_time} -> + conn + |> put_resp_header("retry-after", Integer.to_string(div(@cost_default, refill_rate))) + |> put_status(429) + |> Phoenix.Controller.put_view(json: API.ErrorJSON) + |> Phoenix.Controller.render(:"429") + |> halt() + end + end + + defp refill_rate(account) do + Map.get(account.limits, :api_refill_rate, @refill_rate_default) + end + + defp capacity(account) do + Map.get(account.limits, :api_capacity, @capacity_default) + end +end diff --git a/elixir/apps/api/lib/api/rate_limit.ex b/elixir/apps/api/lib/api/rate_limit.ex new file mode 100644 index 000000000..dc8b8b902 --- /dev/null +++ b/elixir/apps/api/lib/api/rate_limit.ex @@ -0,0 +1,70 @@ +defmodule API.RateLimit do + @moduledoc """ + Distributed, eventually consistent rate limiter using `Domain.PubSub` and `Hammer`. + + This module provides a rate-limiting mechanism for requests using a distributed, + eventually consistent approach. It combines local in-memory counting with a + broadcasting mechanism to keep counters in sync across nodes in a cluster. + """ + + @default_cost 10 + + def default_cost do + @default_cost + end + + def hit(key, refill_rate, capacity, cost \\ @default_cost) do + :ok = broadcast({:hit, key, refill_rate, capacity, cost, Node.self()}) + API.RateLimit.Local.hit(key, refill_rate, capacity, cost) + end + + defmodule Local do + @moduledoc false + use Hammer, backend: :ets, algorithm: :token_bucket + end + + defmodule Listener do + @moduledoc false + use GenServer + + @doc false + def start_link(opts) do + topic = Keyword.fetch!(opts, :topic) + GenServer.start_link(__MODULE__, topic) + end + + @impl true + def init(topic) do + :ok = Domain.PubSub.subscribe(topic) + {:ok, []} + end + + @impl true + def handle_info({:hit, key, refill_rate, capacity, cost, node}, state) do + if node != Node.self() do + {_allow_deny, _int} = Local.hit(key, refill_rate, capacity, cost) + end + + {:noreply, state} + end + end + + @topic "__ratelimit" + + defp broadcast(message) do + Domain.PubSub.broadcast(@topic, message) + end + + def child_spec(opts) do + %{ + id: __MODULE__, + start: {__MODULE__, :start_link, [opts]}, + type: :supervisor + } + end + + def start_link(opts) do + children = [{Local, opts}, {Listener, topic: @topic}] + Supervisor.start_link(children, strategy: :one_for_one) + end +end diff --git a/elixir/apps/api/lib/api/router.ex b/elixir/apps/api/lib/api/router.ex index 35d02310a..94a654474 100644 --- a/elixir/apps/api/lib/api/router.ex +++ b/elixir/apps/api/lib/api/router.ex @@ -9,6 +9,7 @@ defmodule API.Router do plug :accepts, ["json"] plug API.Plugs.Auth + plug API.Plugs.RateLimit end pipeline :public do diff --git a/elixir/apps/api/mix.exs b/elixir/apps/api/mix.exs index d495cfce9..d406d7784 100644 --- a/elixir/apps/api/mix.exs +++ b/elixir/apps/api/mix.exs @@ -61,6 +61,7 @@ defmodule API.MixProject do {:remote_ip, "~> 1.1"}, {:open_api_spex, "~> 3.21.2"}, {:ymlr, "~> 5.0"}, + {:hammer, "~> 7.0.0"}, # Test deps {:credo, "~> 1.5", only: [:dev, :test], runtime: false}, diff --git a/elixir/apps/api/test/api/rate_limit.exs b/elixir/apps/api/test/api/rate_limit.exs new file mode 100644 index 000000000..7144029a2 --- /dev/null +++ b/elixir/apps/api/test/api/rate_limit.exs @@ -0,0 +1,71 @@ +defmodule API.RateLimitTest do + use API.ConnCase, async: true + + setup do + account = Fixtures.Accounts.create_account() + actor = Fixtures.Actors.create_actor(type: :api_client, account: account) + identity = Fixtures.Auth.create_identity(account: account, actor: actor) + subject = Fixtures.Auth.create_subject(identity: identity) + + %{ + account: account, + actor: actor, + identity: identity, + subject: subject + } + end + + describe "verify rate limit" do + test "allows requests under rate limit", %{conn: conn, account: account, actor: actor} do + api_request_cost = API.RateLimit.default_cost() + capacity = div(Domain.Config.fetch_env!(:api, API.RateLimit)[:capacity], api_request_cost) + _resources = for _ <- 1..3, do: Fixtures.Resources.create_resource(%{account: account}) + + for _ <- 0..(capacity - 1) do + resp_conn = call_api(conn, actor) + assert %{"data" => _data, "metadata" => _metadata} = json_response(resp_conn, 200) + end + end + + test "returns 429 after hitting rate limit", %{conn: conn, account: account, actor: actor} do + api_request_cost = API.RateLimit.default_cost() + capacity = div(Domain.Config.fetch_env!(:api, API.RateLimit)[:capacity], api_request_cost) + _resources = for _ <- 1..3, do: Fixtures.Resources.create_resource(%{account: account}) + + for _ <- 0..(capacity - 1) do + resp_conn = call_api(conn, actor) + assert %{"data" => _data, "metadata" => _metadata} = json_response(resp_conn, 200) + end + + resp_conn = call_api(conn, actor) + assert %{"error" => %{"reason" => "Too Many Requests"}} = json_response(resp_conn, 429) + end + + test "allows requests after time window", %{conn: conn, account: account, actor: actor} do + api_request_cost = API.RateLimit.default_cost() + capacity = div(Domain.Config.fetch_env!(:api, API.RateLimit)[:capacity], api_request_cost) + + _resources = for _ <- 1..3, do: Fixtures.Resources.create_resource(%{account: account}) + + for _ <- 0..(capacity - 1) do + resp_conn = call_api(conn, actor) + assert %{"data" => _data, "metadata" => _metadata} = json_response(resp_conn, 200) + end + + resp_conn = call_api(conn, actor) + assert %{"error" => %{"reason" => "Too Many Requests"}} = json_response(resp_conn, 429) + + :timer.sleep(:timer.seconds(1)) + + resp_conn = call_api(conn, actor) + assert %{"data" => _data, "metadata" => _metadata} = json_response(resp_conn, 200) + end + end + + def call_api(conn, actor) do + conn + |> authorize_conn(actor) + |> put_req_header("content-type", "application/json") + |> get("/resources") + end +end diff --git a/elixir/apps/domain/lib/domain/config/definitions.ex b/elixir/apps/domain/lib/domain/config/definitions.ex index 60210e802..309965f0f 100644 --- a/elixir/apps/domain/lib/domain/config/definitions.ex +++ b/elixir/apps/domain/lib/domain/config/definitions.ex @@ -179,6 +179,16 @@ defmodule Domain.Config.Definitions do end ) + @doc """ + The API rate limiter uses a token bucket algorithm. This field sets the rate the bucket is refilled. + """ + defconfig(:api_refill_rate, :integer, default: 10) + + @doc """ + The API rate limiter uses a token bucket algorithm. This field sets the capacity of the bucket. + """ + defconfig(:api_capacity, :integer, default: 200) + @doc """ Enable or disable requiring secure cookies. Required for HTTPS. """ diff --git a/elixir/config/config.exs b/elixir/config/config.exs index ccd62477f..4ddc8b50f 100644 --- a/elixir/config/config.exs +++ b/elixir/config/config.exs @@ -206,6 +206,10 @@ config :api, private_clients: [%{__struct__: Postgrex.INET, address: {172, 28, 0, 0}, netmask: 16}], relays_presence_debounce_timeout_ms: 3_000 +config :api, API.RateLimit, + refill_rate: 10, + capacity: 200 + ############################### ##### Third-party configs ##### ############################### diff --git a/elixir/config/runtime.exs b/elixir/config/runtime.exs index 10baee89a..036f4a633 100644 --- a/elixir/config/runtime.exs +++ b/elixir/config/runtime.exs @@ -197,6 +197,10 @@ if config_env() == :prod do external_trusted_proxies: compile_config!(:phoenix_external_trusted_proxies), private_clients: compile_config!(:phoenix_private_clients) + config :api, API.RateLimit, + refill_rate: compile_config!(:api_refill_rate), + capacity: compile_config!(:api_capacity) + config :web, api_external_url: api_external_url end diff --git a/elixir/mix.lock b/elixir/mix.lock index c3052f7b4..9ccc1e711 100644 --- a/elixir/mix.lock +++ b/elixir/mix.lock @@ -40,6 +40,7 @@ "gproc": {:hex, :gproc, "0.9.1", "f1df0364423539cf0b80e8201c8b1839e229e5f9b3ccb944c5834626998f5b8c", [:rebar3], [], "hexpm", "905088e32e72127ed9466f0bac0d8e65704ca5e73ee5a62cb073c3117916d507"}, "grpcbox": {:hex, :grpcbox, "0.17.1", "6e040ab3ef16fe699ffb513b0ef8e2e896da7b18931a1ef817143037c454bcce", [:rebar3], [{:acceptor_pool, "~> 1.0.0", [hex: :acceptor_pool, repo: "hexpm", optional: false]}, {:chatterbox, "~> 0.15.1", [hex: :ts_chatterbox, repo: "hexpm", optional: false]}, {:ctx, "~> 0.6.0", [hex: :ctx, repo: "hexpm", optional: false]}, {:gproc, "~> 0.9.1", [hex: :gproc, repo: "hexpm", optional: false]}], "hexpm", "4a3b5d7111daabc569dc9cbd9b202a3237d81c80bf97212fbc676832cb0ceb17"}, "hackney": {:hex, :hackney, "1.23.0", "55cc09077112bcb4a69e54be46ed9bc55537763a96cd4a80a221663a7eafd767", [:rebar3], [{:certifi, "~> 2.14.0", [hex: :certifi, repo: "hexpm", optional: false]}, {:idna, "~> 6.1.0", [hex: :idna, repo: "hexpm", optional: false]}, {:metrics, "~> 1.0.0", [hex: :metrics, repo: "hexpm", optional: false]}, {:mimerl, "~> 1.1", [hex: :mimerl, repo: "hexpm", optional: false]}, {:parse_trans, "3.4.1", [hex: :parse_trans, repo: "hexpm", optional: false]}, {:ssl_verify_fun, "~> 1.1.0", [hex: :ssl_verify_fun, repo: "hexpm", optional: false]}, {:unicode_util_compat, "~> 0.7.0", [hex: :unicode_util_compat, repo: "hexpm", optional: false]}], "hexpm", "6cd1c04cd15c81e5a493f167b226a15f0938a84fc8f0736ebe4ddcab65c0b44e"}, + "hammer": {:hex, :hammer, "7.0.1", "136edcd81af44becbe6b73a958c109e2364ab0dc026d7b19892037dc2632078c", [:mix], [], "hexpm", "796edf14ab2aa80df72080210fcf944ee5e8868d8ece7a7511264d802f58cc2d"}, "hpack": {:hex, :hpack_erl, "0.3.0", "2461899cc4ab6a0ef8e970c1661c5fc6a52d3c25580bc6dd204f84ce94669926", [:rebar3], [], "hexpm", "d6137d7079169d8c485c6962dfe261af5b9ef60fbc557344511c1e65e3d95fb0"}, "hpax": {:hex, :hpax, "1.0.2", "762df951b0c399ff67cc57c3995ec3cf46d696e41f0bba17da0518d94acd4aac", [:mix], [], "hexpm", "2f09b4c1074e0abd846747329eaa26d535be0eb3d189fa69d812bfb8bfefd32f"}, "httpoison": {:hex, :httpoison, "2.2.1", "87b7ed6d95db0389f7df02779644171d7319d319178f6680438167d7b69b1f3d", [:mix], [{:hackney, "~> 1.17", [hex: :hackney, repo: "hexpm", optional: false]}], "hexpm", "51364e6d2f429d80e14fe4b5f8e39719cacd03eb3f9a9286e61e216feac2d2df"},