-
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.
- Loading branch information
Showing
6 changed files
with
170 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,31 @@ | ||
| defmodule Rajska.RateLimiter do | ||
| @moduledoc """ | ||
| Rate limiter absinthe middleware. | ||
| """ | ||
| @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, arguments: arguments}, keys, id) do | ||
| case {keys, id} do | ||
| {nil, nil} -> Rajska.apply_auth_mod(context, :get_ip, [context]) | ||
| {keys, nil} -> get_in(arguments, List.wrap(keys)) || raise "Invalid configuration in Rate Limiter. Key not found in arguments." | ||
| {nil, id} -> id | ||
| {_key, _id} -> raise "Invalid configuration in Rate Limiter. If key is defined, then id must not be defined" | ||
| end | ||
| end | ||
| 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,128 @@ | ||
| 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 | ||
|
|
||
| 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 |