Add handling of 429 Too Many Requests

lib/liberator/evaluator.ex

@@ -11,7 +11,8 @@ defmodule Liberator.Evaluator do
     method_allowed?: {:malformed?, :handle_method_not_allowed},
     malformed?: {:handle_malformed, :authorized?},
     authorized?: {:allowed?, :handle_unauthorized},
-    allowed?: {:valid_content_header?, :handle_forbidden},
+    allowed?: {:too_many_requests?, :handle_forbidden},
+    too_many_requests?: {:handle_too_many_requests, :valid_content_header?},
     valid_content_header?: {:known_content_type?, :handle_not_implemented},
     known_content_type?: {:valid_entity_length?, :handle_unsupported_media_type},
     valid_entity_length?: {:is_options?, :handle_request_entity_too_large},
@@ -99,6 +100,7 @@ defmodule Liberator.Evaluator do
     handle_uri_too_long: 414,
     handle_unsupported_media_type: 415,
     handle_unprocessable_entity: 422,
+    handle_too_many_requests: 429,
     handle_unavailable_for_legal_reasons: 451,
     handle_unknown_method: 501,
     handle_not_implemented: 501,
@@ -155,6 +157,27 @@ defmodule Liberator.Evaluator do
           conn
         end
 
+        conn =
+          if retry_after = Map.get(conn.assigns, :retry_after) do
+            retry_after_value =
+              cond do
+                Timex.is_valid?(retry_after) ->
+                  Timex.format!(retry_after, "%a, %d %b %Y %H:%M:%S GMT", :strftime)
+                is_integer(retry_after)->
+                  Integer.to_string(retry_after)
+                String.valid?(retry_after) ->
+                  retry_after
+                true ->
+                  raise "Value for :retry_after was not a valid DateTime, integer, or String, but was #{inspect retry_after}. " <>
+                    "Make sure the too_many_requests/1 function of #{inspect module} is setting that key to one of those types. " <>
+                    "Remember that you can also just return true or false."
+              end
+
+            put_resp_header(conn, "retry-after", retry_after_value)
+          else
+            conn
+          end
+
         status = @handlers[next_step]
         body = apply(module, next_step, [conn])
 

lib/liberator/resource.ex

@@ -143,6 +143,7 @@ defmodule Liberator.Resource do
   | `c:handle_uri_too_long/1`                   | 414 |
   | `c:handle_unsupported_media_type/1`         | 415 |
   | `c:handle_unprocessable_entity/1`           | 422 |
+  | `c:handle_too_many_requests/1`              | 429 |
   | `c:handle_unavailable_for_legal_reasons/1`  | 451 |
   | `c:handle_unknown_method/1`                 | 501 |
   | `c:handle_not_implemented/1`                | 501 |
@@ -192,6 +193,7 @@ defmodule Liberator.Resource do
   | `c:post_redirect?/1`        | Should the response redirect after a `POST`?                                        | false |
   | `c:put_to_different_url?/1` | Should the `PUT` request be made to a different URL?                                | false |
   | `c:processable?/1`          | Is the request body processable?                                                    | true |
+  | `c:too_many_requests?/1`    | Has the client or user issued too many requests in a period of time?                | false |
   | `c:service_available?/1`    | Is the service available?                                                           | true |
   | `c:unavailable_for_legal_reasons?/1` | Is the resource not available, for legal reasons?                          | false |
   | `c:uri_too_long?/1`         | Is the request URI too long?                                                        | false |
@@ -395,6 +397,27 @@ defmodule Liberator.Resource do
   @doc since: "1.0"
   @callback allowed?(Plug.Conn.t) :: true | false
 
+  @doc """
+  Check to see if the client has performed too many requests.
+  Used as part of a rate limiting scheme.
+
+  If you return a map containing a `:retry_after` key,
+  then the response's `retry-after` header will be automatically set.
+  The value of this key can be either an Elixir `DateTime` object,
+  a `String` HTTP date, or an integer of seconds.
+  All of these values tell the client when they can attempt their request again.
+  Note that if you provide a `String` for this value,
+  it should be formatted as an HTTP date.
+
+  If you do return map with the key `:retry_after` set,
+  and its value is not a `DateTime`, integer, or valid `String`,
+  then Liberator will raise an exception.
+
+  By default, always returns `false`.
+  """
+  @doc since: "1.2"
+  @callback too_many_requests?(Plug.Conn.t) :: true | false
+
   @doc """
   Check if the Content-Type of the body is valid.
 
@@ -1001,6 +1024,14 @@ defmodule Liberator.Resource do
   @doc since: "1.0"
   @callback handle_unprocessable_entity(Plug.Conn.t) :: Plug.Conn.t
 
+  @doc """
+  Returns content for a `429 Too Many Requests` response.
+
+  For more information on this response type, see [RFC 6585, section 4](https://tools.ietf.org/html/rfc6585#section-4).
+  """
+  @doc since: "1.2"
+  @callback handle_too_many_requests(Plug.Conn.t) :: Plug.Conn.t
+
   @doc """
   Returns content for a `451 Unavailable for Legal Reasons` response.
 
@@ -1094,6 +1125,8 @@ defmodule Liberator.Resource do
       @impl true
       def allowed?(_conn), do: true
       @impl true
+      def too_many_requests?(_conn), do: false
+      @impl true
       def valid_content_header?(_conn), do: true
       @impl true
       def known_content_type?(_conn), do: true
@@ -1422,6 +1455,11 @@ defmodule Liberator.Resource do
         "Unprocessable Entity"
       end
 
+      @impl true
+      def handle_too_many_requests(conn) do
+        "Too Many Requests"
+      end
+
       @impl true
       def handle_unavailable_for_legal_reasons(_conn) do
         "Unavailable for Legal Reasons"

test/liberator/resource_test.exs

@@ -24,6 +24,7 @@ defmodule Liberator.ResourceTest do
       malformed?: false,
       authorized?: true,
       allowed?: true,
+      too_many_requests?: false,
       valid_content_header?: true,
       known_content_type?: true,
       valid_entity_length?: true,
@@ -322,6 +323,88 @@ defmodule Liberator.ResourceTest do
     assert conn.resp_body == "Unprocessable Entity"
   end
 
+  test "returns 429 when too_many_requests? returns true" do
+    defmodule RateLimitedResource do
+      use Liberator.Resource
+      @impl true
+      def too_many_requests?(_conn), do: true
+    end
+
+    conn = conn(:get, "/")
+    conn = RateLimitedResource.call(conn, [])
+
+    assert conn.state == :sent
+    assert conn.status == 429
+    assert conn.resp_body == "Too Many Requests"
+  end
+
+  test "sets retry-after header of resource if too_many_requests returns %{retry_after: %DateTime{}}" do
+    defmodule RateLimitedUntilOctoberResource do
+      use Liberator.Resource
+      @impl true
+      def too_many_requests?(_conn), do: %{retry_after: ~U[2020-10-12 17:06:00Z]}
+    end
+
+    conn = conn(:get, "/")
+    conn = RateLimitedUntilOctoberResource.call(conn, [])
+
+    assert conn.state == :sent
+    assert conn.status == 429
+    assert conn.resp_body == "Too Many Requests"
+    assert "Mon, 12 Oct 2020 17:06:00 GMT" in get_resp_header(conn, "retry-after")
+  end
+
+  test "sets retry-after header of resource if too_many_requests returns %{retry_after: 60}" do
+    defmodule RateLimitedForAMinuteResource do
+      use Liberator.Resource
+      @impl true
+      def too_many_requests?(_conn), do: %{retry_after: 60}
+    end
+
+    conn = conn(:get, "/")
+    conn = RateLimitedForAMinuteResource.call(conn, [])
+
+    assert conn.state == :sent
+    assert conn.status == 429
+    assert conn.resp_body == "Too Many Requests"
+    assert "60" in get_resp_header(conn, "retry-after")
+  end
+
+  test "sets retry-after header of resource if too_many_requests returns %{retry_after: \"whenever, man\"}" do
+    defmodule RateLimitedByLebowskiResource do
+      use Liberator.Resource
+      @impl true
+      def too_many_requests?(_conn), do: %{retry_after: "whenever, man"}
+    end
+
+    conn = conn(:get, "/")
+    conn = RateLimitedByLebowskiResource.call(conn, [])
+
+    assert conn.state == :sent
+    assert conn.status == 429
+    assert conn.resp_body == "Too Many Requests"
+    assert "whenever, man" in get_resp_header(conn, "retry-after")
+  end
+
+  test "raises if the value of :retry_after is not a valid string" do
+    defmodule RateLimitedByUnicodeConsortiumResource do
+      use Liberator.Resource
+      @impl true
+      def too_many_requests?(_conn), do: %{retry_after: <<0xFFFF::16>>}
+    end
+
+    expected_message =
+      "Value for :retry_after was not a valid DateTime, integer, or String, but was <<255, 255>>. " <>
+      "Make sure the too_many_requests/1 function of " <>
+      "Liberator.ResourceTest.RateLimitedByUnicodeConsortiumResource is setting " <>
+      "that key to one of those types. Remember that you can also just return true or false."
+
+    conn = conn(:get, "/")
+    assert_raise RuntimeError, expected_message, fn ->
+      RateLimitedByUnicodeConsortiumResource.call(conn, [])
+    end
+  end
+
   test "returns 451 when unavailable_for_legal_reasons? returns true" do
     defmodule UnavailableForLegalReasonsResource do
       use Liberator.Resource