ratelimit.clj

(ns x.ratelimit
  (:require
   [taoensso.carmine :as car :refer [wcar]]
   [ring.util.response :as response]
   [ring.mock.request :as mock])
  (:import
   [java.util Calendar]))

;; Util ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(defn seconds-since-midnight []
  (let [c (Calendar/getInstance)
        now-ms (.getTimeInMillis c)]
    (.set c Calendar/HOUR_OF_DAY 0)
    (.set c Calendar/MINUTE 0)
    (.set c Calendar/SECOND 0)
    (.set c Calendar/MILLISECOND 0)
    (let [passed-ms (- now-ms (.getTimeInMillis c))
          seconds-passed (long (/ passed-ms 1000))]
      seconds-passed)))

;; IntSeconds -> Int
(defn calc-curr-window
  "For example: Returns an int in range [1, 96] when window-duration is 900 (15min)
   since there are 96x 15min periods in a day."
  [window-duration]
  (inc (quot (dec (seconds-since-midnight)) window-duration)))

;; Backend impl ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(defprotocol Backend
  (available? [this])
  (reset-limits! [this curr-window])
  (get-window [this])
  (get-limit [self key]))

(deftype RedisBackend [hash-key window-key] Backend
  (available? [_]
    (= "PONG" (try (wcar {} (car/ping)) (catch Throwable _))))
  (reset-limits! [this curr-window]
    (wcar {}
      (car/del hash-key)
      (car/set window-key curr-window)))
  (get-window [_]
    (when-let [n (wcar {} (car/get window-key))]
      (Integer/parseInt n)))
  (get-limit [_ field-key]
    (wcar {} (car/hincrby hash-key field-key 1))))

(defn make-redis-backend [hash-key]
  (let [window-key (str hash-key ":window")]
    (RedisBackend. hash-key window-key)))

;; Middleware ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(defn wrap-ratelimit [handler backend {:keys [window-duration window-limit]}]
  (fn [request]
    (if-not (available? backend)
      (handler request)
      (let [curr-window (calc-curr-window window-duration)]
        (when (not= (get-window backend) curr-window)
          (reset-limits! backend curr-window))
        (let [ip (:remote-addr request)
              curr-limit (get-limit backend ip)
              limit-remaining (max 0 (- window-limit curr-limit))
              secs-til-reset (- (* curr-window window-duration)
                                (seconds-since-midnight))]
          (-> (if (> curr-limit window-limit)
                {:status 429, :body "Too many requests"}
                (handler request))
              (response/header "X-Rate-Limit-Limit" window-limit)
              (response/header "X-Rate-Limit-Remaining" limit-remaining)
              (response/header "X-Rate-Limit-Reset" secs-til-reset)))))))

;; Demo ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(defn mock-handler [request]
  {:status 200, :body "Hello world", :headers {}})

(let [backend (make-redis-backend "ratelimits")
      handler (wrap-ratelimit mock-handler backend {:window-duration 900
                                                    :window-limit 5})]
  (handler (mock/request :get "/")))

;; Responses

{:status 200, :body "Hello world", :headers {"X-Rate-Limit-Reset" "735", "X-Rate-Limit-Remaining" "4", "X-Rate-Limit-Limit" "5"}}
{:status 200, :body "Hello world", :headers {"X-Rate-Limit-Reset" "732", "X-Rate-Limit-Remaining" "3", "X-Rate-Limit-Limit" "5"}}
{:status 200, :body "Hello world", :headers {"X-Rate-Limit-Reset" "731", "X-Rate-Limit-Remaining" "2", "X-Rate-Limit-Limit" "5"}}
{:status 200, :body "Hello world", :headers {"X-Rate-Limit-Reset" "729", "X-Rate-Limit-Remaining" "1", "X-Rate-Limit-Limit" "5"}}
{:status 200, :body "Hello world", :headers {"X-Rate-Limit-Reset" "728", "X-Rate-Limit-Remaining" "0", "X-Rate-Limit-Limit" "5"}}
{:status 429, :body "Too many requests", :headers {"X-Rate-Limit-Reset" "726", "X-Rate-Limit-Remaining" "0", "X-Rate-Limit-Limit" "5"}}
{:status 429, :body "Too many requests", :headers {"X-Rate-Limit-Reset" "705", "X-Rate-Limit-Remaining" "0", "X-Rate-Limit-Limit" "5"}}
{:status 200, :body "Hello world", :headers {"X-Rate-Limit-Reset" "887", "X-Rate-Limit-Remaining" "4", "X-Rate-Limit-Limit" "5"}}