;;; Copyright © 2019 - 2023 Jakob L. Kreuze <zerodaysfordays@sdf.org>
;;;
;;; This program is free software; you can redistribute it and/or
;;; modify it under the terms of the GNU General Public License as
;;; published by the Free Software Foundation; either version 3 of the
;;; License, or (at your option) any later version.
;;;
;;; This program is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
;;; General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with this program. If not, see
;;; <http://www.gnu.org/licenses/>.

(define-module (jakob dynamic rate-limiter)
  #:use-module (jakob dynamic errors)
  #:use-module (jakob dynamic util)
  #:use-module (json)
  #:use-module (rnrs conditions)
  #:use-module (rnrs exceptions)
  #:use-module (srfi-197)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:use-module (web request)
  #:use-module (web response)
  #:export (rate-limit-wrap))

(define-record-type <requester-state>
  (make-requester-state time request-bins)
  requester-state?
  (time         requester-state-time)
  (request-bins requester-state-bins))

(define active-rate-limits (make-hash-table))

(define (rate-limit-for-endpoint name)
  (case name
    ((put-reaction)   1)
    ((get-event-rsvp) 1)
    ((get-event-info) 8)
    ((get-image)      8)
    ((get-gallery)    8)
    ((put-comment)    8)
    ((get-comments)   1024)
    (else 32)))

(define (increment-key! hash-table key)
  (let ((new-value (if (hash-ref hash-table key)
                       (+ 1 (hash-ref hash-table key))
                       1)))
    (hash-set! hash-table key new-value)))

(define (rate-limit-wrap proc)
  (lambda (request body)
    (unless (assoc-ref (request-headers request) 'x-forwarded-for)
      (panic "X-Forwarded-For header not provided"))
    (let ((endpoint-name (procedure-name proc))
          (requester (chain (assoc-ref (request-headers request) 'x-forwarded-for)
                            (string-split _ #\,)
                            (first _))))
      (unless (hash-ref active-rate-limits requester)
        (hash-set! active-rate-limits
                   requester
                   (make-requester-state (current-time) (make-hash-table))))
      (increment-key! (requester-state-bins (hash-ref active-rate-limits requester)) endpoint-name)
      ;; TODO: The `when' body is copy/pasted from above. I think this condition
      ;; (time-based expiry) could be refactored.
      (when (>= (current-time)
                (+ (* 60 60) (requester-state-time (hash-ref active-rate-limits requester))))
        (hash-set! active-rate-limits
                   requester
                   (make-requester-state (current-time) (make-hash-table))))
      (when (and (> (hash-ref (requester-state-bins (hash-ref active-rate-limits requester)) endpoint-name)
                    (rate-limit-for-endpoint endpoint-name)))
        (panic "Your IP address is currently being rate-limited." #:code 429))
      (proc request body))))