-
Notifications
You must be signed in to change notification settings - Fork 13
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #22 from rschef/rate-limiter
Rate limiter
- Loading branch information
Showing
8 changed files
with
271 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,71 @@ | ||
| defmodule Rajska.RateLimiter do | ||
| @moduledoc """ | ||
| Rate limiter absinthe middleware. Uses [Hammer](https://github.com/ExHammer/hammer). | ||
| ## Usage | ||
| First configure Hammer, following its documentation. For example: | ||
| config :hammer, | ||
| backend: {Hammer.Backend.ETS, [expiry_ms: 60_000 * 60 * 4, | ||
| cleanup_interval_ms: 60_000 * 10]} | ||
| Add your middleware to the query that should be limited: | ||
| field :default_config, :string do | ||
| middleware Rajska.RateLimiter | ||
| resolve fn _, _ -> {:ok, "ok"} end | ||
| end | ||
| You can also configure it and use multiple rules for limiting in one query: | ||
| field :login_user, :session do | ||
| arg :email, non_null(:string) | ||
| arg :password, non_null(:string) | ||
| middleware Rajska.RateLimiter, limit: 10 # Using the default identifier (user IP) | ||
| middleware Rajska.RateLimiter, keys: :email, limit: 5 # Using the value provided in the email arg | ||
| resolve &AccountsResolver.login_user/2 | ||
| end | ||
| The allowed configuration are: | ||
| * `scale_ms`: The timespan for the maximum number of actions. Defaults to 60_000. | ||
| * `limit`: The maximum number of actions in the specified timespan. Defaults to 10. | ||
| * `id`: An atom or string to be used as the bucket identifier. Note that this will always be the same, so by using this the limit will be global instead of by user. | ||
| * `keys`: An atom or a list of atoms to get a query argument as identifier. Use a list when the argument is nested. | ||
| * `error_msg`: The error message to be displayed when rate limit exceeds. Defaults to `"Too many requests"`. | ||
| Note that when neither `id` or `keys` is provided, the default is to use the user's IP. For that, the default behaviour is to use | ||
| `c:Rajska.Authorization.get_ip/1` to fetch the IP from the absinthe context. That means you need to manually insert the user's IP in the | ||
| absinthe context before using it as an identifier. See the [absinthe docs](https://hexdocs.pm/absinthe/context-and-authentication.html#content) | ||
| for more information. | ||
| """ | ||
| @behaviour Absinthe.Middleware | ||
|
|
||
| alias Absinthe.Resolution | ||
|
|
||
| def call(%Resolution{state: :resolved} = resolution, _config), do: resolution | ||
|
|
||
| def call(%Resolution{} = resolution, config) do | ||
| scale_ms = Keyword.get(config, :scale_ms, 60_000) | ||
| limit = Keyword.get(config, :limit, 10) | ||
| identifier = get_identifier(resolution, config[:keys], config[:id]) | ||
| error_msg = Keyword.get(config, :error_msg, "Too many requests") | ||
|
|
||
| case Hammer.check_rate("query:#{identifier}", scale_ms, limit) do | ||
| {:allow, _count} -> resolution | ||
| {:deny, _limit} -> Resolution.put_result(resolution, {:error, error_msg}) | ||
| end | ||
| end | ||
|
|
||
| defp get_identifier(%Resolution{context: context}, nil, nil), | ||
| do: Rajska.apply_auth_mod(context, :get_ip, [context]) | ||
|
|
||
| defp get_identifier(%Resolution{arguments: arguments}, keys, nil), | ||
| do: get_in(arguments, List.wrap(keys)) || raise "Invalid configuration in Rate Limiter. Key not found in arguments." | ||
|
|
||
| defp get_identifier(%Resolution{}, nil, id), do: id | ||
|
|
||
| defp get_identifier(%Resolution{}, _keys, _id), do: raise "Invalid configuration in Rate Limiter. If key is defined, then id must not be defined" | ||
| end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,133 @@ | ||
| defmodule Rajska.RateLimiterTest do | ||
| use ExUnit.Case, async: false | ||
|
|
||
| import Mock | ||
|
|
||
| defmodule Authorization do | ||
| use Rajska, | ||
| valid_roles: [:user, :admin], | ||
| super_role: :admin | ||
| end | ||
|
|
||
| defmodule Schema do | ||
| use Absinthe.Schema | ||
|
|
||
| def context(ctx), do: Map.put(ctx, :authorization, Authorization) | ||
|
|
||
| input_object :keys_params do | ||
| field :id, :string | ||
| end | ||
|
|
||
| query do | ||
| field :default_config, :string do | ||
| middleware Rajska.RateLimiter | ||
| resolve fn _, _ -> {:ok, "ok"} end | ||
| end | ||
|
|
||
| field :scale_limit, :string do | ||
| middleware Rajska.RateLimiter, scale_ms: 30_000, limit: 5 | ||
| resolve fn _, _ -> {:ok, "ok"} end | ||
| end | ||
|
|
||
| field :id, :string do | ||
| middleware Rajska.RateLimiter, id: :custom_id | ||
| resolve fn _, _ -> {:ok, "ok"} end | ||
| end | ||
|
|
||
| field :key, :string do | ||
| arg :id, :string | ||
|
|
||
| middleware Rajska.RateLimiter, keys: :id | ||
| resolve fn _, _ -> {:ok, "ok"} end | ||
| end | ||
|
|
||
| field :keys, :string do | ||
| arg :params, :keys_params | ||
| middleware Rajska.RateLimiter, keys: [:params, :id] | ||
| resolve fn _, _ -> {:ok, "ok"} end | ||
| end | ||
|
|
||
| field :id_and_key, :string do | ||
| middleware Rajska.RateLimiter, id: :id, keys: :keys | ||
| resolve fn _, _ -> {:ok, "ok"} end | ||
| end | ||
|
|
||
| field :error_msg, :string do | ||
| middleware Rajska.RateLimiter, error_msg: "Rate limit exceeded" | ||
| resolve fn _, _ -> {:ok, "ok"} end | ||
| end | ||
| end | ||
| end | ||
|
|
||
| @default_context [context: %{ip: "ip"}] | ||
|
|
||
| setup_with_mocks([{Hammer, [], [check_rate: fn _a, _b, _c -> {:allow, 1} end]}]) do | ||
| :ok | ||
| end | ||
|
|
||
| test "works with default configs" do | ||
| {:ok, _} = Absinthe.run(query(:default_config), __MODULE__.Schema, @default_context) | ||
| assert_called Hammer.check_rate("query:ip", 60_000, 10) | ||
| end | ||
|
|
||
| test "accepts scale and limit configuration" do | ||
| {:ok, _} = Absinthe.run(query(:scale_limit), __MODULE__.Schema, @default_context) | ||
| assert_called Hammer.check_rate("query:ip", 30_000, 5) | ||
| end | ||
|
|
||
| test "accepts id configuration" do | ||
| {:ok, _} = Absinthe.run(query(:id), __MODULE__.Schema, @default_context) | ||
| assert_called Hammer.check_rate("query:custom_id", 60_000, 10) | ||
| end | ||
|
|
||
| test "accepts key configuration" do | ||
| {:ok, _} = Absinthe.run(query(:key, :id, "id_key"), __MODULE__.Schema, @default_context) | ||
| assert_called Hammer.check_rate("query:id_key", 60_000, 10) | ||
| end | ||
|
|
||
| test "throws error if key is not present" do | ||
| assert_raise RuntimeError, ~r/Invalid configuration in Rate Limiter. Key not found in arguments./, fn -> | ||
| Absinthe.run(query(:key), __MODULE__.Schema, @default_context) | ||
| end | ||
| end | ||
|
|
||
| test "accepts key configuration for nested parameters" do | ||
| {:ok, _} = Absinthe.run(query(:keys, :params, %{id: "id_key"}), __MODULE__.Schema, @default_context) | ||
| assert_called Hammer.check_rate("query:id_key", 60_000, 10) | ||
| end | ||
|
|
||
| test "throws error when id and key are provided as configuration" do | ||
| assert_raise RuntimeError, ~r/Invalid configuration in Rate Limiter. If key is defined, then id must not be defined/, fn -> | ||
| Absinthe.run(query(:id_and_key), __MODULE__.Schema, @default_context) | ||
| end | ||
| end | ||
|
|
||
| test "accepts error msg configuration" do | ||
| with_mock Hammer, [check_rate: fn _a, _b, _c -> {:deny, 1} end] do | ||
| assert {:ok, %{errors: errors}} = Absinthe.run(query(:error_msg), __MODULE__.Schema, @default_context) | ||
| assert [ | ||
| %{ | ||
| locations: [%{column: 0, line: 1}], | ||
| message: "Rate limit exceeded", | ||
| path: ["error_msg"] | ||
| } | ||
| ] == errors | ||
| end | ||
| end | ||
|
|
||
| test "does not apply when resolution is already resolved" do | ||
| resolution = %Absinthe.Resolution{state: :resolved} | ||
| assert resolution == Rajska.RateLimiter.call(resolution, []) | ||
| end | ||
|
|
||
| defp query(name), do: "{ #{name} }" | ||
| defp query(name, key, value) when is_binary(value), do: "{ #{name}(#{key}: \"#{value}\") }" | ||
| defp query(name, key, %{} = value), do: "{ #{name}(#{key}: {#{build_arguments(value)}}) }" | ||
|
|
||
| defp build_arguments(arguments) do | ||
| arguments | ||
| |> Enum.map(fn {k, v} -> if is_nil(v), do: nil, else: "#{k}: #{inspect(v, [charlists: :as_lists])}" end) | ||
| |> Enum.reject(&is_nil/1) | ||
| |> Enum.join(", ") | ||
| end | ||
| end |