From 29c29a8fc148e25abaf626a5f2f9d5698d1a1eba Mon Sep 17 00:00:00 2001
From: Jiri Sedlacek <jiri.sedlacek@futu.cz>
Date: Mon, 23 Dec 2024 18:52:36 +0100
Subject: [PATCH] Add on_error handler (#35)

# Conflicts:
#	lib/hammer_plug.ex
---
 lib/hammer_plug.ex        | 45 +++++++++++++++++++++++++++++++++------
 test/hammer_plug_test.exs | 29 +++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 7 deletions(-)

diff --git a/lib/hammer_plug.ex b/lib/hammer_plug.ex
index 3af9238..6c8c440 100644
--- a/lib/hammer_plug.ex
+++ b/lib/hammer_plug.ex
@@ -12,7 +12,8 @@ defmodule Hammer.Plug do
         rate_limit: {"chat:message:post", 60_000, 20},
         by: {:session, :user, &Helpers.get_user_id/2},
         when_nil: :raise,
-        on_deny: &Helpers.handle_deny/2
+        on_deny: &Helpers.handle_deny/2,
+        on_error: &Helpers.handle_error/2
       ] when action == :post_chat_message
   ## Options
   ### :rate_limit (`{prefix::string, time_scale::int, limit::int}`)
@@ -78,6 +79,11 @@ defmodule Hammer.Plug do
       #     def handle_rate_limit_deny(conn, _opts) do
       #       ...
       #     end
+
+  ### :on_error
+  A plug function to be invoked when there is an error when
+  checking the limit.
+
   """
   @behaviour Plug
   import Plug.Conn
@@ -113,13 +119,21 @@ defmodule Hammer.Plug do
         &Hammer.Plug.default_on_deny_handler/2
       )
 
+    on_error_handler =
+      Keyword.get(
+        opts,
+        :on_error,
+        &Hammer.Plug.default_on_deny_handler/2
+      )
+
     config = %{
       id_prefix: id_prefix,
       scale: scale,
       limit: limit,
       by: by,
       when_nil: when_nil,
-      on_deny: on_deny_handler
+      on_deny: on_deny_handler,
+      on_error: on_error_handler
     }
 
     plug_name = plug_name(id_prefix)
@@ -142,7 +156,8 @@ defmodule Hammer.Plug do
       limit: limit,
       by: by,
       when_nil: when_nil,
-      on_deny: on_deny_handler
+      on_deny: on_deny_handler,
+      on_error: on_error_handler
     } = config
 
     request_identifier = get_request_identifier(conn, by)
@@ -152,7 +167,15 @@ defmodule Hammer.Plug do
         case when_nil do
           # Proceed
           :use_nil ->
-            do_rate_limit_check(conn, id_prefix, nil, scale, limit, on_deny_handler)
+            do_rate_limit_check(
+              conn,
+              id_prefix,
+              nil,
+              scale,
+              limit,
+              on_deny_handler,
+              on_error_handler
+            )
 
           :raise ->
             raise Hammer.Plug.NilError
@@ -163,7 +186,7 @@ defmodule Hammer.Plug do
         end
 
       id ->
-        do_rate_limit_check(conn, id_prefix, id, scale, limit, on_deny_handler)
+        do_rate_limit_check(conn, id_prefix, id, scale, limit, on_deny_handler, on_error_handler)
     end
   end
 
@@ -200,7 +223,15 @@ defmodule Hammer.Plug do
     end
   end
 
-  defp do_rate_limit_check(conn, id_prefix, request_id, scale, limit, on_deny_handler) do
+  defp do_rate_limit_check(
+         conn,
+         id_prefix,
+         request_id,
+         scale,
+         limit,
+         on_deny_handler,
+         on_error_handler
+       ) do
     full_id = "#{id_prefix}:#{request_id}"
 
     case Hammer.check_rate(full_id, scale, limit) do
@@ -211,7 +242,7 @@ defmodule Hammer.Plug do
         on_deny_handler.(conn, [])
 
       {:error, _reason} ->
-        on_deny_handler.(conn, [])
+        on_error_handler.(conn, [])
     end
   end
 
diff --git a/test/hammer_plug_test.exs b/test/hammer_plug_test.exs
index 3d65293..87b468a 100644
--- a/test/hammer_plug_test.exs
+++ b/test/hammer_plug_test.exs
@@ -15,6 +15,13 @@ defmodule Helpers do
     |> send_resp(404, "Not Found")
     |> halt()
   end
+
+  def on_error(conn, _opts) do
+    conn
+    |> put_resp_header("x-hammer-test", "yes")
+    |> send_resp(500, "Internal Server Error")
+    |> halt()
+  end
 end
 
 defmodule Hammer.PlugTest do
@@ -87,6 +94,28 @@ defmodule Hammer.PlugTest do
     end
   end
 
+  describe "with custom on_error handler" do
+    @opts Hammer.Plug.init(
+            rate_limit: {"test", 1_000, 3},
+            by: :ip,
+            on_error: &Helpers.on_error/2
+          )
+
+    test "halts the conn and sends a 500 (non default) on failure" do
+      with_mock Hammer, check_rate: fn _a, _b, _c -> {:error, true} end do
+        conn =
+          :get
+          |> conn("/hello")
+          |> Hammer.Plug.call(@opts)
+
+        assert conn.status == 500
+        assert get_resp_header(conn, "x-hammer-test") == ["yes"]
+        assert conn.halted == true
+        assert called(Hammer.check_rate("test:127.0.0.1", 1_000, 3))
+      end
+    end
+  end
+
   describe "by session" do
     @opts Hammer.Plug.init(
             rate_limit: {"test", 1_000, 3},