├── .formatter.exs ├── .github └── workflows │ └── elixir.yml ├── .gitignore ├── LICENSE ├── README.md ├── config ├── config.exs ├── dev.exs └── test.exs ├── lib ├── minirate.ex └── minirate │ ├── counter.ex │ ├── plug.ex │ └── worker.ex ├── mix.exs ├── mix.lock └── test ├── minirate ├── counter_test.exs ├── plug_test.exs └── worker_test.exs ├── minirate_test.exs └── test_helper.exs /.formatter.exs: -------------------------------------------------------------------------------- 1 | # Used by "mix format" 2 | [ 3 | inputs: ["{mix,.formatter}.exs", "{config,lib,test}/**/*.{ex,exs}"] 4 | ] 5 | -------------------------------------------------------------------------------- /.github/workflows/elixir.yml: -------------------------------------------------------------------------------- 1 | # This workflow uses actions that are not certified by GitHub. 2 | # They are provided by a third-party and are governed by 3 | # separate terms of service, privacy policy, and support 4 | # documentation. 5 | 6 | name: Elixir CI 7 | 8 | on: 9 | push: 10 | branches: [ "master" ] 11 | pull_request: 12 | branches: [ "master" ] 13 | 14 | permissions: 15 | contents: read 16 | 17 | jobs: 18 | build: 19 | 20 | name: Build and test 21 | runs-on: ubuntu-latest 22 | 23 | steps: 24 | - uses: actions/checkout@v4 25 | - name: Set up Elixir 26 | uses: erlef/setup-beam@61e01a43a562a89bfc54c7f9a378ff67b03e4a21 # v1.16.0 27 | with: 28 | elixir-version: '1.17.2' # [Required] Define the Elixir version 29 | otp-version: '27.0.1' # [Required] Define the Erlang/OTP version 30 | - name: Restore dependencies cache 31 | uses: actions/cache@v3 32 | with: 33 | path: deps 34 | key: ${{ runner.os }}-mix-${{ hashFiles('**/mix.lock') }} 35 | restore-keys: ${{ runner.os }}-mix- 36 | - name: Install dependencies 37 | run: mix deps.get 38 | - name: Run tests 39 | run: mix test 40 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # The directory Mix will write compiled artifacts to. 2 | /_build/ 3 | 4 | # If you run "mix test --cover", coverage assets end up here. 5 | /cover/ 6 | 7 | # The directory Mix downloads your dependencies sources to. 8 | /deps/ 9 | 10 | # Where third-party dependencies like ExDoc output generated docs. 11 | /doc/ 12 | 13 | # Ignore .fetch files in case you like to edit your project deps locally. 14 | /.fetch 15 | 16 | # If the VM crashes, it generates a dump, let's ignore it too. 17 | erl_crash.dump 18 | 19 | # Also ignore archive artifacts (built via "mix archive.build"). 20 | *.ez 21 | 22 | # Ignore package tarball (built via "mix hex.build"). 23 | minirate-*.tar 24 | 25 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Jose Daniel Gil Climent 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ![Elixir CI](https://github.com/4xposed/minirate/workflows/Elixir%20CI/badge.svg?event=push) 2 | 3 | # Minirate 4 | 5 | A dead simple distributed rate limiting library in Elixir using Mnesia. 6 | 7 | # What is it? 8 | 9 | A distributed rate limiter with a focus on readable and well tested code. 10 | 11 | The counter is syncronized over all connected nodes 12 | 13 | ```elixir 14 | iex(test2@127.0.0.1)19> Minirate.check_limit("download", "user_1", 100) 15 | {:allow, 1} 16 | ``` 17 | ```elixir 18 | iex(test1@127.0.0.1)14> Minirate.check_limit("download", "user_1", 100) 19 | {:allow, 2} 20 | ``` 21 | 22 | ## Installation 23 | 24 | Minirate is availabe as a package in Hex, just add it to your `mix.exs` file: 25 | 26 | ```elixir 27 | def deps 28 | [{:minirate, "~> 0.1"}] 29 | end 30 | ``` 31 | 32 | and add it to your extra applications: 33 | ```elixir 34 | def applications do 35 | [ 36 | extra_applications: [:minirate] 37 | ] 38 | ``` 39 | 40 | ## Configuration 41 | 42 | Minirate needs to be configured using Mix.Config. 43 | 44 | For example, in `config/config.exs`: 45 | 46 | ``` 47 | config :minirate, 48 | mnesia_table: :rate_limiter, 49 | expiry_ms: 60_000 50 | cleanup_period_ms: 10_000 51 | ``` 52 | 53 | `mnesia_table` specifies which table will Mnesia use to write the counters. 54 | `expiry_ms` specifies the counter life in millisecconds (for example to have rates like x request every 10 seconds, you would set `expiry_ms` to 10_000) 55 | `cleanup_period_ms` specifies how often minirate will clean expired counters from the mnesia database 56 | 57 | 58 | ## Usage 59 | 60 | With Minirate you can rate limit any action on your application. 61 | 62 | The module `Minirate` the function `check_limit(action_name, identifier, limit)` 63 | 64 | An Example: 65 | 66 | ```elixir 67 | @download_limit 1_000 68 | 69 | def download_file(file, user_id) do 70 | case Minirate.check_limit("download_file", user_id, @download_limit) do 71 | {:allow, _count} -> 72 | # Logic to download the file 73 | 74 | {:block, _reason} -> 75 | # Logic when the limit has been reached 76 | 77 | {:skip, _reason} -> 78 | # Skip will only happen if there's a problem with your nodes or mnesia setup and a count cannot be determined. 79 | end 80 | ``` 81 | 82 | ## Using Minirate.Plug 83 | 84 | `Minirate.Plug` can rate-limit actions in your web application using the ip address of the requester. 85 | 86 | You can just put in the pipeline of your web application something like this: 87 | 88 | ```elixir 89 | plug Minirate.Plug, [action: action, limit: 10_000] 90 | ``` 91 | 92 | or for more flexibilty: 93 | ```elixir 94 | plug Minirate.Plug, [action: "custom_action", limit: 10_000] when action == :update or action == :create 95 | ``` 96 | -------------------------------------------------------------------------------- /config/config.exs: -------------------------------------------------------------------------------- 1 | import Config 2 | 3 | config :minirate, 4 | mnesia_table: :rate_limiter, 5 | expiry_ms: 60_000, 6 | cleanup_period_ms: 10_000 7 | 8 | import_config "#{Mix.env()}.exs" 9 | -------------------------------------------------------------------------------- /config/dev.exs: -------------------------------------------------------------------------------- 1 | import Config 2 | 3 | config :minirate, 4 | mnesia_table: :rate_limiter, 5 | expiry_ms: 60_000, 6 | cleanup_period_ms: 10_000 7 | -------------------------------------------------------------------------------- /config/test.exs: -------------------------------------------------------------------------------- 1 | import Config 2 | 3 | config :minirate, 4 | mnesia_table: :rate_limiter, 5 | expiry_ms: 300, 6 | cleanup_period_ms: 50 7 | -------------------------------------------------------------------------------- /lib/minirate.ex: -------------------------------------------------------------------------------- 1 | defmodule Minirate do 2 | @moduledoc false 3 | 4 | use Application 5 | 6 | def start(_type, _args) do 7 | children = [ 8 | %{ 9 | id: Minirate.Worker, 10 | start: {Minirate.Worker, :start_link, [mnesia_table(), expiry_ms(), cleanup_period_ms()]} 11 | } 12 | ] 13 | 14 | Supervisor.start_link(children, strategy: :one_for_one) 15 | end 16 | 17 | def check_limit(action, id, limit) do 18 | Minirate.Counter.check_limit(mnesia_table(), {action, id, limit, now()}) 19 | end 20 | 21 | def check_limit(action, id, limit, increment) do 22 | Minirate.Counter.check_limit(mnesia_table(), {action, id, limit, now()}, increment) 23 | end 24 | 25 | defp expiry_ms, do: get_config(:expiry_ms) 26 | 27 | defp mnesia_table, do: get_config(:mnesia_table) 28 | 29 | defp cleanup_period_ms, do: get_config(:cleanup_period_ms) 30 | 31 | defp get_config(key) do 32 | case Application.get_env(:minirate, key) do 33 | nil -> 34 | Kernel.raise( 35 | ArgumentError, 36 | "the configuration parameter #{Kernel.inspect(key)} is not set" 37 | ) 38 | 39 | value -> 40 | value 41 | end 42 | end 43 | 44 | defp now do 45 | DateTime.utc_now() |> DateTime.to_unix(:millisecond) 46 | end 47 | end 48 | -------------------------------------------------------------------------------- /lib/minirate/counter.ex: -------------------------------------------------------------------------------- 1 | defmodule Minirate.Counter do 2 | @moduledoc false 3 | 4 | alias :mnesia, as: Mnesia 5 | 6 | @default_increment 1 7 | 8 | def check_limit(table, {action, id, limit, now}) do 9 | check_limit(table, {action, id, limit, now}, @default_increment) 10 | end 11 | 12 | def check_limit(table, {action, id, limit, now}, increment) do 13 | with {:ok, count} <- get_count(table, {action, id, now}, increment), 14 | true <- count <= limit do 15 | {:allow, count} 16 | else 17 | false -> {:block, :limit_exceeded} 18 | {:error, reason} -> {:skip, reason} 19 | end 20 | end 21 | 22 | def get_count(table, {action, id, now}) do 23 | get_count(table, {action, id, now}, @default_increment) 24 | end 25 | 26 | def get_count(table, {action, id, now}, increment) do 27 | key = "#{action}_#{id}" 28 | 29 | transac_fn = fn -> 30 | case Mnesia.read({table, key}) do 31 | # No existing count found 32 | [] -> 33 | # Create count entry 34 | Mnesia.write({table, key, increment, now}) 35 | {:ok, increment} 36 | 37 | # Count entry has been found 38 | [{^table, ^key, count, timestamp}] -> 39 | # Increment count 40 | current_count = count + increment 41 | Mnesia.write({table, key, current_count, timestamp}) 42 | {:ok, current_count} 43 | end 44 | end 45 | 46 | case Mnesia.transaction(transac_fn) do 47 | {:atomic, {:ok, count}} -> {:ok, count} 48 | {:aborted, reason} -> {:error, reason} 49 | end 50 | end 51 | 52 | def expire_keys(table, time_of_expiration) do 53 | transac_fn = fn -> 54 | match = {table, :"$1", :_, :"$2"} 55 | filter = [{:or, {:<, :"$2", time_of_expiration}, {:==, :"$2", time_of_expiration}}] 56 | result = [:"$1"] 57 | 58 | search_term = {match, filter, result} 59 | 60 | Mnesia.select(table, [search_term]) 61 | |> Enum.each(fn key -> 62 | Mnesia.delete(table, key, :write) 63 | end) 64 | end 65 | 66 | case Mnesia.transaction(transac_fn) do 67 | {:atomic, :ok} -> :ok 68 | {:aborted, _reason} -> :error 69 | end 70 | end 71 | 72 | def create_mnesia_table(table) do 73 | Mnesia.create_table(table, attributes: [:key, :count, :timestamp]) 74 | Mnesia.add_table_index(table, :timestamp) 75 | end 76 | end 77 | -------------------------------------------------------------------------------- /lib/minirate/plug.ex: -------------------------------------------------------------------------------- 1 | defmodule Minirate.Plug do 2 | @moduledoc false 3 | 4 | import Plug.Conn 5 | 6 | alias Minirate 7 | 8 | def init, do: [] 9 | 10 | def init(opts) do 11 | Enum.each([:action, :limit], fn key -> 12 | option = Keyword.get(opts, key) 13 | 14 | if Kernel.is_nil(option) do 15 | raise( 16 | ArgumentError, 17 | "Minirate.Plug requiers the option #{Kernel.inspect(option)} to be set." 18 | ) 19 | end 20 | end) 21 | 22 | opts 23 | end 24 | 25 | def call(conn, action: action, limit: limit) do 26 | ip = fetch_ip(conn) 27 | process(conn, action, limit, ip) 28 | end 29 | 30 | def call(conn, _), do: conn 31 | 32 | def process(conn, _action, _limit, ip) when is_nil(ip) do 33 | conn 34 | end 35 | 36 | def process(conn, action, limit, ip) do 37 | case Minirate.check_limit(action, ip, limit) do 38 | {:allow, _} -> conn 39 | {:skip, _} -> conn 40 | {:block, :limit_exceeded} -> block_request(conn) 41 | end 42 | end 43 | 44 | defp block_request(conn) do 45 | conn 46 | |> send_resp(429, "Too Many Requests") 47 | |> halt() 48 | end 49 | 50 | defp fetch_ip(conn) do 51 | conn.remote_ip 52 | |> Tuple.to_list() 53 | |> Enum.join(".") 54 | end 55 | end 56 | -------------------------------------------------------------------------------- /lib/minirate/worker.ex: -------------------------------------------------------------------------------- 1 | defmodule Minirate.Worker do 2 | @moduledoc false 3 | 4 | use GenServer 5 | 6 | alias Minirate.Counter 7 | 8 | # Public API 9 | 10 | def start_link(mnesia_table, expiry_ms, cleanup_period_ms) do 11 | args = %{ 12 | mnesia_table: mnesia_table, 13 | expiry_ms: expiry_ms, 14 | cleanup_period_ms: cleanup_period_ms 15 | } 16 | 17 | GenServer.start_link(__MODULE__, args, name: __MODULE__) 18 | end 19 | 20 | # GenServer Callbacks 21 | 22 | def init(args) do 23 | Process.send_after(self(), :create_table, 500) 24 | :timer.send_interval(args.cleanup_period_ms, :expire) 25 | 26 | {:ok, args} 27 | end 28 | 29 | def handle_info(:expire, state) do 30 | expiration = now() - state.expiry_ms 31 | Counter.expire_keys(state.mnesia_table, expiration) 32 | 33 | {:noreply, state} 34 | end 35 | 36 | def handle_info(:create_table, state) do 37 | Counter.create_mnesia_table(state.mnesia_table) 38 | {:noreply, state} 39 | end 40 | 41 | defp now do 42 | DateTime.utc_now() |> DateTime.to_unix(:millisecond) 43 | end 44 | end 45 | -------------------------------------------------------------------------------- /mix.exs: -------------------------------------------------------------------------------- 1 | defmodule Minirate.MixProject do 2 | use Mix.Project 3 | 4 | def project do 5 | [ 6 | app: :minirate, 7 | description: "A dead simple distributed rate limiter using Mnesia", 8 | source_url: "https://github.com/4xposed/minirate", 9 | docs: docs(), 10 | package: package(), 11 | version: "0.1.2", 12 | elixir: "~> 1.15", 13 | build_embedded: Mix.env() == :prod, 14 | start_permanent: Mix.env() == :prod, 15 | deps: deps() 16 | ] 17 | end 18 | 19 | def application do 20 | [ 21 | mod: {Minirate, []}, 22 | extra_applications: [:logger, :mnesia] 23 | ] 24 | end 25 | 26 | def docs do 27 | [ 28 | main: "Minirate", 29 | extras: ["README.md"] 30 | ] 31 | end 32 | 33 | defp deps do 34 | [ 35 | {:plug, "~> 1.16"}, 36 | {:credo, "~> 1.7", only: [:dev, :test], runtime: false}, 37 | {:ex_doc, ">= 0.34.0", only: :dev, runtime: false} 38 | ] 39 | end 40 | 41 | defp package do 42 | [ 43 | name: :minirate, 44 | maintainers: ["Daniel Climent"], 45 | licenses: ["MIT"], 46 | links: %{"GitHub" => "https://github.com/4xposed/minirate"} 47 | ] 48 | end 49 | end 50 | -------------------------------------------------------------------------------- /mix.lock: -------------------------------------------------------------------------------- 1 | %{ 2 | "bunt": {:hex, :bunt, "1.0.0", "081c2c665f086849e6d57900292b3a161727ab40431219529f13c4ddcf3e7a44", [:mix], [], "hexpm", "dc5f86aa08a5f6fa6b8096f0735c4e76d54ae5c9fa2c143e5a1fc7c1cd9bb6b5"}, 3 | "credo": {:hex, :credo, "1.7.7", "771445037228f763f9b2afd612b6aa2fd8e28432a95dbbc60d8e03ce71ba4446", [:mix], [{:bunt, "~> 0.2.1 or ~> 1.0", [hex: :bunt, repo: "hexpm", optional: false]}, {:file_system, "~> 0.2 or ~> 1.0", [hex: :file_system, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "8bc87496c9aaacdc3f90f01b7b0582467b69b4bd2441fe8aae3109d843cc2f2e"}, 4 | "earmark": {:hex, :earmark, "1.4.3", "364ca2e9710f6bff494117dbbd53880d84bebb692dafc3a78eb50aa3183f2bfd", [:mix], [], "hexpm", "8cf8a291ebf1c7b9539e3cddb19e9cef066c2441b1640f13c34c1d3cfc825fec"}, 5 | "earmark_parser": {:hex, :earmark_parser, "1.4.41", "ab34711c9dc6212dda44fcd20ecb87ac3f3fce6f0ca2f28d4a00e4154f8cd599", [:mix], [], "hexpm", "a81a04c7e34b6617c2792e291b5a2e57ab316365c2644ddc553bb9ed863ebefa"}, 6 | "ex_doc": {:hex, :ex_doc, "0.34.2", "13eedf3844ccdce25cfd837b99bea9ad92c4e511233199440488d217c92571e8", [:mix], [{:earmark_parser, "~> 1.4.39", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "5ce5f16b41208a50106afed3de6a2ed34f4acfd65715b82a0b84b49d995f95c1"}, 7 | "file_system": {:hex, :file_system, "1.0.1", "79e8ceaddb0416f8b8cd02a0127bdbababe7bf4a23d2a395b983c1f8b3f73edd", [:mix], [], "hexpm", "4414d1f38863ddf9120720cd976fce5bdde8e91d8283353f0e31850fa89feb9e"}, 8 | "jason": {:hex, :jason, "1.4.4", "b9226785a9aa77b6857ca22832cffa5d5011a667207eb2a0ad56adb5db443b8a", [:mix], [{:decimal, "~> 1.0 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "c5eb0cab91f094599f94d55bc63409236a8ec69a21a67814529e8d5f6cc90b3b"}, 9 | "makeup": {:hex, :makeup, "1.1.2", "9ba8837913bdf757787e71c1581c21f9d2455f4dd04cfca785c70bbfff1a76a3", [:mix], [{:nimble_parsec, "~> 1.2.2 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "cce1566b81fbcbd21eca8ffe808f33b221f9eee2cbc7a1706fc3da9ff18e6cac"}, 10 | "makeup_elixir": {:hex, :makeup_elixir, "0.16.2", "627e84b8e8bf22e60a2579dad15067c755531fea049ae26ef1020cad58fe9578", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "41193978704763f6bbe6cc2758b84909e62984c7752b3784bd3c218bb341706b"}, 11 | "makeup_erlang": {:hex, :makeup_erlang, "1.0.1", "c7f58c120b2b5aa5fd80d540a89fdf866ed42f1f3994e4fe189abebeab610839", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "8a89a1eeccc2d798d6ea15496a6e4870b75e014d1af514b1b71fa33134f57814"}, 12 | "mime": {:hex, :mime, "2.0.6", "8f18486773d9b15f95f4f4f1e39b710045fa1de891fada4516559967276e4dc2", [:mix], [], "hexpm", "c9945363a6b26d747389aac3643f8e0e09d30499a138ad64fe8fd1d13d9b153e"}, 13 | "nimble_parsec": {:hex, :nimble_parsec, "1.4.0", "51f9b613ea62cfa97b25ccc2c1b4216e81df970acd8e16e8d1bdc58fef21370d", [:mix], [], "hexpm", "9c565862810fb383e9838c1dd2d7d2c437b3d13b267414ba6af33e50d2d1cf28"}, 14 | "plug": {:hex, :plug, "1.16.1", "40c74619c12f82736d2214557dedec2e9762029b2438d6d175c5074c933edc9d", [:mix], [{:mime, "~> 1.0 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:plug_crypto, "~> 1.1.1 or ~> 1.2 or ~> 2.0", [hex: :plug_crypto, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.3 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "a13ff6b9006b03d7e33874945b2755253841b238c34071ed85b0e86057f8cddc"}, 15 | "plug_crypto": {:hex, :plug_crypto, "2.1.0", "f44309c2b06d249c27c8d3f65cfe08158ade08418cf540fd4f72d4d6863abb7b", [:mix], [], "hexpm", "131216a4b030b8f8ce0f26038bc4421ae60e4bb95c5cf5395e1421437824c4fa"}, 16 | "telemetry": {:hex, :telemetry, "1.3.0", "fedebbae410d715cf8e7062c96a1ef32ec22e764197f70cda73d82778d61e7a2", [:rebar3], [], "hexpm", "7015fc8919dbe63764f4b4b87a95b7c0996bd539e0d499be6ec9d7f3875b79e6"}, 17 | } 18 | -------------------------------------------------------------------------------- /test/minirate/counter_test.exs: -------------------------------------------------------------------------------- 1 | defmodule Minirate.CounterTest do 2 | use ExUnit.Case 3 | 4 | alias Minirate.Counter 5 | 6 | setup do 7 | :mnesia.create_table(:rate_limiter, attributes: [:key, :count, :timestamp]) 8 | 9 | on_exit(fn -> :mnesia.delete_table(:rate_limiter) end) 10 | end 11 | 12 | describe ".get_count/3 - increments by 1" do 13 | test "inserts entry when it's called for the first time" do 14 | assert [] == :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 15 | 16 | fake_now = 158_037_147_724 17 | 18 | assert {:ok, 1} == Counter.get_count(:rate_limiter, {"test", 1, fake_now}) 19 | 20 | assert [{:rate_limiter, "test_1", 1, fake_now}] == 21 | :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 22 | end 23 | 24 | test "updates only the count when it's called more than once" do 25 | fake_now = 158_037_147_724 26 | :mnesia.dirty_write({:rate_limiter, "test_1", 1, fake_now}) 27 | 28 | assert [{:rate_limiter, "test_1", 1, fake_now}] == 29 | :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 30 | 31 | fake_future_time = fake_now + 1_000 32 | 33 | assert {:ok, 2} == Counter.get_count(:rate_limiter, {"test", 1, fake_future_time}) 34 | 35 | assert [{:rate_limiter, "test_1", 2, fake_now}] == 36 | :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 37 | end 38 | end 39 | 40 | describe ".get_count/4 - increments by n" do 41 | test "inserts entry when it's called for the first time" do 42 | assert [] == :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 43 | 44 | fake_now = 158_037_147_724 45 | increment = 2 46 | 47 | assert {:ok, 2} == Counter.get_count(:rate_limiter, {"test", 1, fake_now}, increment) 48 | 49 | assert [{:rate_limiter, "test_1", 2, fake_now}] == 50 | :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 51 | end 52 | 53 | test "updates only the count when it's called more than once" do 54 | fake_now = 158_037_147_724 55 | :mnesia.dirty_write({:rate_limiter, "test_1", 2, fake_now}) 56 | 57 | assert [{:rate_limiter, "test_1", 2, fake_now}] == 58 | :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 59 | 60 | fake_future_time = fake_now + 1_000 61 | increment = 2 62 | 63 | assert {:ok, 4} == 64 | Counter.get_count(:rate_limiter, {"test", 1, fake_future_time}, increment) 65 | 66 | assert [{:rate_limiter, "test_1", 4, fake_now}] == 67 | :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 68 | end 69 | end 70 | 71 | describe ".check_limit/3 - increments by 1 and checks if the limit has exceeded" do 72 | test "when there's no count yet" do 73 | assert [] == :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 74 | 75 | fake_now = 158_037_147_724 76 | limit = 10 77 | 78 | assert {:allow, 1} == Counter.check_limit(:rate_limiter, {"test", 1, limit, fake_now}) 79 | end 80 | 81 | test "when the limit is lower than the count" do 82 | fake_now = 158_037_147_724 83 | :mnesia.dirty_write({:rate_limiter, "test_1", 15, fake_now}) 84 | 85 | limit = 10 86 | fake_future_time = fake_now + 1_000 87 | 88 | assert {:block, :limit_exceeded} == 89 | Counter.check_limit(:rate_limiter, {"test", 1, limit, fake_future_time}) 90 | end 91 | 92 | test "when the limit is higher than the count" do 93 | fake_now = 158_037_147_724 94 | :mnesia.dirty_write({:rate_limiter, "test_1", 15, fake_now}) 95 | 96 | limit = 20 97 | fake_future_time = fake_now + 1_000 98 | 99 | assert {:allow, 16} == 100 | Counter.check_limit(:rate_limiter, {"test", 1, limit, fake_future_time}) 101 | end 102 | end 103 | 104 | describe ".check_limit/4 - increments by n and checks if the limit has exceeded" do 105 | test "when there's no count yet" do 106 | assert [] == :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 107 | 108 | fake_now = 158_037_147_724 109 | increment = 5 110 | limit = 10 111 | 112 | assert {:allow, 5} == 113 | Counter.check_limit(:rate_limiter, {"test", 1, limit, fake_now}, increment) 114 | end 115 | 116 | test "when the limit is lower than the count" do 117 | fake_now = 158_037_147_724 118 | :mnesia.dirty_write({:rate_limiter, "test_1", 15, fake_now}) 119 | 120 | increment = 6 121 | limit = 20 122 | fake_future_time = fake_now + 1_000 123 | 124 | assert {:block, :limit_exceeded} == 125 | Counter.check_limit(:rate_limiter, {"test", 1, limit, fake_future_time}, increment) 126 | end 127 | 128 | test "when the limit is higher than the count" do 129 | fake_now = 158_037_147_724 130 | :mnesia.dirty_write({:rate_limiter, "test_1", 15, fake_now}) 131 | 132 | increment = 5 133 | limit = 20 134 | fake_future_time = fake_now + 1_000 135 | 136 | assert {:allow, 20} == 137 | Counter.check_limit(:rate_limiter, {"test", 1, limit, fake_future_time}, increment) 138 | end 139 | end 140 | 141 | describe ".expire_keys/1 - deletes the entries older than the expiration set" do 142 | test "record not older" do 143 | fake_now = 158_037_147_724 144 | :mnesia.dirty_write({:rate_limiter, "test_1", 1, fake_now}) 145 | 146 | time_of_expiration = fake_now - 10_000 147 | 148 | Counter.expire_keys(:rate_limiter, time_of_expiration) 149 | 150 | assert [{:rate_limiter, "test_1", 1, fake_now}] == 151 | :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 152 | end 153 | 154 | test "record older than expiration time " do 155 | fake_now = 158_037_147_724 156 | :mnesia.dirty_write({:rate_limiter, "test_1", 1, fake_now}) 157 | 158 | time_of_expiration = fake_now + 10_000 159 | 160 | Counter.expire_keys(:rate_limiter, time_of_expiration) 161 | 162 | assert [] == :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 163 | end 164 | 165 | test "record exactly on expiration time " do 166 | fake_now = 158_037_147_724 167 | :mnesia.dirty_write({:rate_limiter, "test_1", 1, fake_now}) 168 | 169 | time_of_expiration = fake_now 170 | 171 | Counter.expire_keys(:rate_limiter, time_of_expiration) 172 | 173 | assert [] == :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 174 | end 175 | end 176 | end 177 | -------------------------------------------------------------------------------- /test/minirate/plug_test.exs: -------------------------------------------------------------------------------- 1 | defmodule Minirate.PlugTest do 2 | use ExUnit.Case 3 | use Plug.Test 4 | 5 | alias Minirate.Plug 6 | 7 | setup do 8 | :mnesia.create_table(:rate_limiter, attributes: [:key, :count, :timestamp]) 9 | 10 | on_exit(fn -> :mnesia.delete_table(:rate_limiter) end) 11 | end 12 | 13 | test "rate limits requests" do 14 | req_conn = 15 | conn(:get, "/", []) 16 | |> Map.put(:remote_ip, {0, 0, 0, 0}) 17 | 18 | conn = Plug.call(req_conn, action: "device_request", limit: 2) 19 | refute conn.halted 20 | 21 | conn = Plug.call(req_conn, action: "device_request", limit: 2) 22 | refute conn.halted 23 | 24 | conn = Plug.call(req_conn, action: "device_request", limit: 2) 25 | assert conn.halted 26 | assert conn.status == 429 27 | end 28 | 29 | test "resets limit with time" do 30 | req_conn = 31 | conn(:get, "/", []) 32 | |> Map.put(:remote_ip, {0, 0, 0, 0}) 33 | 34 | conn = Plug.call(req_conn, action: "device_request", limit: 2) 35 | refute conn.halted 36 | 37 | conn = Plug.call(req_conn, action: "device_request", limit: 2) 38 | refute conn.halted 39 | 40 | conn = Plug.call(req_conn, action: "device_request", limit: 2) 41 | assert conn.halted 42 | assert conn.status == 429 43 | 44 | Process.sleep(400) 45 | 46 | conn = Plug.call(req_conn, action: "device_request", limit: 2) 47 | refute conn.halted 48 | end 49 | end 50 | -------------------------------------------------------------------------------- /test/minirate/worker_test.exs: -------------------------------------------------------------------------------- 1 | defmodule Minirate.WorkerTest do 2 | use ExUnit.Case 3 | 4 | alias Minirate.Worker 5 | 6 | setup do 7 | :mnesia.create_table(:rate_limiter, attributes: [:key, :count, :timestamp]) 8 | 9 | on_exit(fn -> :mnesia.delete_table(:rate_limiter) end) 10 | end 11 | 12 | describe ".expire" do 13 | test "removes from the mnesia table the expired counters" do 14 | time_in_ms = DateTime.utc_now() |> DateTime.to_unix(:millisecond) 15 | old_timestamp = time_in_ms - 100_000 16 | 17 | :mnesia.dirty_write({:rate_limiter, "test_1", 2, old_timestamp}) 18 | :mnesia.dirty_write({:rate_limiter, "test_2", 4, time_in_ms}) 19 | 20 | Task.async(fn -> 21 | Kernel.send(Worker, :expire) 22 | Process.sleep(10) 23 | end) 24 | |> Task.await(:infinity) 25 | 26 | assert [] == :mnesia.dirty_match_object({:rate_limiter, "test_1", :_, :_}) 27 | 28 | assert [{:rate_limiter, "test_2", 4, time_in_ms}] == 29 | :mnesia.dirty_match_object({:rate_limiter, "test_2", :_, :_}) 30 | end 31 | end 32 | end 33 | -------------------------------------------------------------------------------- /test/minirate_test.exs: -------------------------------------------------------------------------------- 1 | defmodule MinirateTest do 2 | use ExUnit.Case 3 | 4 | setup do 5 | :mnesia.create_table(:rate_limiter, attributes: [:key, :count, :timestamp]) 6 | 7 | on_exit(fn -> :mnesia.delete_table(:rate_limiter) end) 8 | end 9 | 10 | describe ".check_limit/3" do 11 | test "increments the counter" do 12 | assert {:allow, 1} == Minirate.check_limit("test", "user_1", 10) 13 | assert {:allow, 2} == Minirate.check_limit("test", "user_1", 10) 14 | assert {:allow, 1} == Minirate.check_limit("test", "user_2", 10) 15 | end 16 | 17 | test "returns a a tuple with value {:block, :limit_exceeded} when the limit is reached" do 18 | assert {:allow, 1} == Minirate.check_limit("test", "user_1", 2) 19 | assert {:allow, 2} == Minirate.check_limit("test", "user_1", 2) 20 | assert {:allow, 1} == Minirate.check_limit("test", "user_2", 2) 21 | assert {:block, :limit_exceeded} == Minirate.check_limit("test", "user_1", 2) 22 | end 23 | 24 | test "counter is reset after expiration" do 25 | assert {:allow, 1} == Minirate.check_limit("test", "user_1", 2) 26 | assert {:allow, 2} == Minirate.check_limit("test", "user_1", 2) 27 | assert {:block, :limit_exceeded} == Minirate.check_limit("test", "user_1", 2) 28 | 29 | Process.sleep(400) 30 | 31 | assert {:allow, 1} == Minirate.check_limit("test", "user_1", 2) 32 | end 33 | end 34 | 35 | describe ".check_limit/4" do 36 | test "increments the counter with a custom increment" do 37 | assert {:allow, 5} == Minirate.check_limit("test", "user_1", 10, 5) 38 | assert {:allow, 8} == Minirate.check_limit("test", "user_1", 10, 3) 39 | assert {:allow, 9} == Minirate.check_limit("test", "user_2", 10, 9) 40 | end 41 | 42 | test "returns a a tuple with value {:block, :limit_exceeded} when the limit is reached" do 43 | assert {:allow, 2} == Minirate.check_limit("test", "user_1", 2, 2) 44 | assert {:allow, 1} == Minirate.check_limit("test", "user_2", 2, 1) 45 | assert {:block, :limit_exceeded} == Minirate.check_limit("test", "user_1", 2, 2) 46 | end 47 | 48 | test "counter is reset after expiration" do 49 | assert {:allow, 2} == Minirate.check_limit("test", "user_1", 2, 2) 50 | assert {:block, :limit_exceeded} == Minirate.check_limit("test", "user_1", 2, 2) 51 | 52 | Process.sleep(400) 53 | 54 | assert {:allow, 2} == Minirate.check_limit("test", "user_1", 2, 2) 55 | end 56 | end 57 | end 58 | -------------------------------------------------------------------------------- /test/test_helper.exs: -------------------------------------------------------------------------------- 1 | ExUnit.start() 2 | :mnesia.create_schema([node()]) 3 | :mnesia.start() 4 | --------------------------------------------------------------------------------