#lang racket/base

(require crypto
         crypto/argon2
         net/ip
         racket/contract
         racket/list
         racket/random
         racket/string
         web-server/http
         web-server/http/cookie-parse)

(provide make-password-hash
         password-hash-valid?
         make-auth-manager
         auth-manager?
         auth-enabled?
         auth-request-user
         auth-login!
         auth-logout!
         auth-session-cookie
         auth-renewal-cookie
         auth-expired-cookie)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Internal data
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;;; Holds one authenticated browser session and its two activity timestamps.
(struct session
  (username [last-seen #:mutable] [last-cookie-renewal #:mutable])
  #:transparent)

;;; Holds the failed-login count and start time for one client address.
(struct failures ([attempts #:mutable] [started #:mutable]) #:transparent)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; goal   : Hold configured users, trusted proxies, and volatile auth state.
; pre    : Constructor fields contain normalized and parsed internal values.
; post   : Creating or recognizing a value does not change external state.
; result : auth-manager? recognizes values used by the authentication API.
; internals: users and trusted-proxies are immutable configuration references;
;            sessions and failed hold mutable login state protected by lock.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(struct auth-manager
  (users trusted-proxies session-seconds sessions failed lock)
  #:transparent)

(define password-kdf
  (or (get-kdf 'argon2id argon2-factory)
      (error 'rkt-web-player/users "Argon2id is unavailable")))

(define password-parameters
  '((m 19456) (t 2) (p 1)))

;; Used to make an unknown username take the same expensive verification path.
(define dummy-password-hash
  "$argon2id$v=19$m=19456,t=2,p=1$WrJi0t7NsbD3adX8kxrT/g$yCIf2Ork8G8PRdIcA0bEIYIdMwzunrTou/BKcM4cO/0")

(define session-cookie-name "rkt-web-player-session")
(define failure-window-seconds 300)
(define maximum-failures 5)

(define ipv4-mapped-prefix
  #"\0\0\0\0\0\0\0\0\0\0\377\377")

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Provided password functions
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; goal   : Create an Argon2id password hash for configuration storage.
; pre    : Password is a string containing at least twelve characters.
; post   : No module state is changed.
; result : A salted Argon2id hash encoded as a string.
; internals: pwhash uses password-kdf with password-parameters to generate the
;            encoded hash, including its random salt and Argon2 parameters.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define/contract (make-password-hash password)
  (-> string? string?)
  (unless (and (string? password)
               (>= (string-length password) 12))
    (raise-argument-error
     'make-password-hash
     "string containing at least 12 characters"
     password))
  (pwhash password-kdf
          (string->bytes/utf-8 password)
          password-parameters))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; goal   : Verify a password against an encoded Argon2id hash.
; pre    : Password and encoded are arbitrary values.
; post   : No module state is changed.
; result : #t only when both values are strings and the password matches.
; internals: pwhash-verify checks the encoded Argon2id value. Malformed hashes
;            are treated as a failed match rather than escaping as exceptions.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define/contract (password-hash-valid? password encoded)
  (-> any/c any/c boolean?)
  (if (and (string? password) (string? encoded))
      (with-handlers ((exn:fail? (λ (_) #f)))
        (pwhash-verify password-kdf
                       (string->bytes/utf-8 password)
                       encoded))
      #f))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Supporting functions
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;;; Parses an address and normalizes an IPv4-mapped IPv6 value to IPv4.
(define (normal-ip-address value)
  (let* ((address (make-ip-address value))
         (raw (ip-address->bytes address))
         (ipv4-mapped?
          (and (= (bytes-length raw) 16)
               (bytes=? (subbytes raw 0 12) ipv4-mapped-prefix))))
    (if ipv4-mapped?
        (bytes->ipv4-address (subbytes raw 12))
        address)))

;;; Parses one configured IP address or CIDR value as a public net/ip network.
(define (parse-network value)
  (let ((parts (string-split (string-trim value) "/")))
    (unless (member (length parts) '(1 2))
      (raise-argument-error
       'make-auth-manager
       "IP address or CIDR network"
       value))
    (let* ((address
            (with-handlers ((exn:fail?
                             (λ (_)
                               (raise-argument-error
                                'make-auth-manager
                                "IP address or CIDR network"
                                value))))
              (normal-ip-address (car parts))))
           (maximum (ip-address-size address))
           (prefix
            (if (= (length parts) 2)
                (string->number (cadr parts))
                maximum)))
      (unless (and (exact-nonnegative-integer? prefix)
                   (<= prefix maximum))
        (raise-argument-error
         'make-auth-manager
         "IP address or CIDR network"
         value))
      (make-network address prefix))))

;;; Checks whether an address belongs to one configured trusted network.
(define (network-contains? network address-string)
  (with-handlers ((exn:fail? (λ (_) #f)))
    (network-member network (normal-ip-address address-string))))

;;; Reads one request header as a UTF-8 string when present.
(define (header-string request name)
  (let ((value (headers-assq* name (request-headers/raw request))))
    (and value
         (bytes->string/utf-8 (header-value value)))))

;;; Checks whether an address belongs to any configured trusted proxy network.
(define (trusted-proxy? manager address)
  (ormap (λ (network) (network-contains? network address))
         (auth-manager-trusted-proxies manager)))

;;; Resolves the effective client address, honoring only a trusted proxy header.
(define (request-address manager request)
  (let* ((peer (request-client-ip request))
         (forwarded
          (and (trusted-proxy? manager peer)
               (header-string request #"X-Forwarded-For"))))
    (if forwarded
        ;; A trusted reverse proxy appends the address it observed. Earlier
        ;; values can have been supplied by the untrusted client.
        (string-trim (last (string-split forwarded ",")))
        peer)))

;;; Extracts the session token from the request cookies when present.
(define (request-session-token request)
  (let ((cookie
         (findf
          (λ (value)
            (string=? (client-cookie-name value) session-cookie-name))
          (request-cookies request))))
    (if cookie
        (client-cookie-value cookie)
        #f)))

;;; Removes every browser session whose idle lifetime has elapsed.
(define (prune-sessions! manager now)
  (for-each
   (λ (token)
     (let ((value (hash-ref (auth-manager-sessions manager) token)))
       (when (> (- now (session-last-seen value))
                (auth-manager-session-seconds manager))
         (hash-remove! (auth-manager-sessions manager) token))))
   (hash-keys (auth-manager-sessions manager))))

;;; Checks and, when necessary, resets the failure window for one address.
(define (failure-blocked? manager address now)
  (let ((value (hash-ref (auth-manager-failed manager) address #f)))
    (cond
      ((eq? value #f) #f)
      ((> (- now (failures-started value)) failure-window-seconds)
       (hash-remove! (auth-manager-failed manager) address)
       #f)
      (else
       (>= (failures-attempts value) maximum-failures)))))

;;; Adds one failed login to the current address window or starts a new window.
(define (record-failure! manager address now)
  (let ((value (hash-ref (auth-manager-failed manager) address #f)))
    (if (and value
             (<= (- now (failures-started value)) failure-window-seconds))
        (set-failures-attempts! value (add1 (failures-attempts value)))
        (hash-set! (auth-manager-failed manager)
                   address
                   (failures 1 now)))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Provided authentication functions
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; goal   : Report whether browser authentication is configured.
; pre    : Manager is an auth-manager.
; post   : Manager remains unchanged.
; result : #t when at least one configured user can log in, otherwise #f.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define/contract (auth-enabled? manager)
  (-> auth-manager? boolean?)
  (positive? (hash-count (auth-manager-users manager))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; goal   : Resolve the browser user represented by a request cookie.
; pre    : Manager is an auth-manager and request is an HTTP request.
; post   : Expired sessions are removed and a valid session's last-seen time
;          is updated.
; result : "anonymous" when authentication is disabled, the normalized
;          username for a valid session, or #f when login is required.
; internals: request-session-token finds the cookie. The manager lock protects
;            prune-sessions! and the session lookup; a valid lookup updates its
;            idle timestamp before returning the stored username.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define/contract (auth-request-user manager request)
  (-> auth-manager? request? (or/c #f string?))
  (if (not (auth-enabled? manager))
      "anonymous"
      (let ((token (request-session-token request))
            (now (current-seconds)))
        (if (eq? token #f)
            #f
            (call-with-semaphore
             (auth-manager-lock manager)
             (λ ()
               (prune-sessions! manager now)
               (let ((value (hash-ref (auth-manager-sessions manager)
                                      token
                                      #f)))
                 (if value
                     (begin
                       (set-session-last-seen! value now)
                       (session-username value))
                     #f))))))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; goal   : Authenticate credentials and start a browser session.
; pre    : Manager is an auth-manager, request is an HTTP request, and
;          username and password are strings.
; post   : A valid login creates a new session; a failed login updates the
;          rate-limit state for the effective client address.
; result : A new opaque token, #f for invalid credentials, or 'rate-limited.
; internals: request-address selects the rate-limit key and failure-blocked?
;            checks its window while the manager lock is held. Unknown users
;            verify against dummy-password-hash to reduce username-dependent
;            timing differences. Success creates a session; failure delegates
;            to record-failure!.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define/contract (auth-login! manager request username password)
  (-> auth-manager?
      request?
      string?
      string?
      (or/c #f 'rate-limited string?))
  (let ((address (request-address manager request))
        (now (current-seconds))
        (normalized (string-downcase (string-trim username))))
    (call-with-semaphore
     (auth-manager-lock manager)
     (λ ()
       (if (failure-blocked? manager address now)
           'rate-limited
           (let* ((stored (hash-ref (auth-manager-users manager)
                                    normalized
                                    #f))
                  (valid?
                   (password-hash-valid?
                    password
                    (or stored dummy-password-hash))))
             (if (and stored valid?)
                 (let ((token
                        (bytes->hex-string (crypto-random-bytes 32))))
                   (hash-remove! (auth-manager-failed manager) address)
                   (prune-sessions! manager now)
                   (hash-set! (auth-manager-sessions manager)
                              token
                              (session normalized now now))
                   token)
                 (begin
                   (record-failure! manager address now)
                   #f))))))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; goal   : End the browser session named by the request cookie.
; pre    : Manager is an auth-manager and request is an HTTP request.
; post   : The matching server-side session is removed when it exists.
; result : Void.
; internals: request-session-token finds the cookie and the manager lock
;            protects removal from the shared session hash.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define/contract (auth-logout! manager request)
  (-> auth-manager? request? void?)
  (let ((token (request-session-token request)))
    (when token
      (call-with-semaphore
       (auth-manager-lock manager)
       (λ ()
         (hash-remove! (auth-manager-sessions manager) token))))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; goal   : Encode an authenticated session token as a browser cookie.
; pre    : Manager is an auth-manager and token is a session token string.
; post   : Manager remains unchanged.
; result : A Secure, HttpOnly, SameSite=Strict Set-Cookie value whose Max-Age
;          equals the configured session lifetime.
; internals: format combines session-cookie-name, token and the manager's
;            configured lifetime into the complete Set-Cookie header value.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define/contract (auth-session-cookie manager token)
  (-> auth-manager? string? bytes?)
  (string->bytes/utf-8
   (format
    "~a=~a; Path=/; Max-Age=~a; Secure; HttpOnly; SameSite=Strict"
    session-cookie-name
    token
    (auth-manager-session-seconds manager))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; goal   : Renew an actively used browser cookie at a bounded frequency.
; pre    : Manager is an auth-manager and request is an HTTP request.
; post   : Expired sessions are removed. When renewal is due, the session's
;          last-cookie-renewal time is advanced.
; result : A fresh Set-Cookie value after half the configured lifetime has
;          elapsed, otherwise #f.
; internals: request-session-token identifies the session. The manager lock
;            protects prune-sessions! and the renewal timestamp. A half-life
;            threshold prevents the one-second player poll from returning a
;            new cookie every second.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define/contract (auth-renewal-cookie manager request)
  (-> auth-manager? request? (or/c #f bytes?))
  (if (not (auth-enabled? manager))
      #f
      (let ((token (request-session-token request))
            (now (current-seconds)))
        (if (eq? token #f)
            #f
            (call-with-semaphore
             (auth-manager-lock manager)
             (λ ()
               (prune-sessions! manager now)
               (let ((value
                      (hash-ref (auth-manager-sessions manager) token #f)))
                 (if value
                     (let* ((elapsed
                             (- now
                                (session-last-cookie-renewal value)))
                            (renewal-interval
                             (max 1
                                  (quotient
                                   (auth-manager-session-seconds manager)
                                   2))))
                       (if (< elapsed renewal-interval)
                           #f
                           (begin
                             (set-session-last-cookie-renewal! value now)
                             (auth-session-cookie manager token))))
                     #f))))))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; goal   : Encode deletion of the browser session cookie.
; pre    : None.
; post   : No module state is changed.
; result : A Secure, HttpOnly, SameSite=Strict Set-Cookie value with Max-Age 0.
; internals: format uses session-cookie-name and an empty value to instruct the
;            browser to remove the cookie immediately.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define/contract (auth-expired-cookie)
  (-> bytes?)
  (string->bytes/utf-8
   (format
    "~a=; Path=/; Max-Age=0; Secure; HttpOnly; SameSite=Strict"
    session-cookie-name)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; goal   : Create the authentication and session manager used by the server.
; pre    : User-pairs contains username and Argon2id-hash pairs, trusted proxy
;          values are IP addresses or CIDR networks, and session-seconds is a
;          positive exact integer.
; post   : No external state is changed; session and rate-limit tables start
;          empty.
; result : A new auth-manager with normalized usernames and parsed networks.
; internals: Each user entry is validated and copied into a case-insensitive
;            hash. parse-network converts trusted-proxy-values through the
;            public net/ip API; fresh hashes and a semaphore protect sessions
;            and failed-login windows.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define/contract (make-auth-manager
                  user-pairs
                  #:trusted-proxies
                  [trusted-proxy-values '("127.0.0.0/8" "::1/128")]
                  #:session-seconds [session-seconds 604800])
  (->* ((listof (cons/c string? string?)))
       (#:trusted-proxies (listof string?)
        #:session-seconds exact-positive-integer?)
       auth-manager?)
  (unless (exact-positive-integer? session-seconds)
    (raise-argument-error 'make-auth-manager "exact-positive-integer?"
                          session-seconds))
  (let ((users (make-hash)))
    (for-each
     (λ (entry)
       (let ((username (string-trim (car entry)))
             (password-hash (cdr entry)))
         (when (string=? username "")
           (raise-arguments-error
            'make-auth-manager
            "username must not be empty"
            "username" (car entry)))
         (unless (regexp-match? #px"^[$]argon2id[$]" password-hash)
           (raise-arguments-error
            'make-auth-manager
            "user password is not an Argon2id hash"
            "username" (car entry)))
         (hash-set! users (string-downcase username) password-hash)))
     user-pairs)
    (auth-manager users
                  (map parse-network trusted-proxy-values)
                  session-seconds
                  (make-hash)
                  (make-hash)
                  (make-semaphore 1))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Tests for module users.rkt
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(module+ test
  (require net/url
           rackunit
           racket/promise
           web-server/http/request-structs)

  (define test-hash (make-password-hash "correct horse battery staple"))
  (check-true (password-hash-valid? "correct horse battery staple" test-hash))
  (check-false (password-hash-valid? "incorrect password" test-hash))

  (define manager
    (make-auth-manager
     (list (cons "Hans" test-hash))
     #:trusted-proxies '("127.0.0.1/32")))

  (check-true
   (network-contains? (parse-network "192.0.2.0/24") "192.0.2.18"))
  (check-false
   (network-contains? (parse-network "192.0.2.0/24") "192.0.3.18"))
  (check-true
   (network-contains? (parse-network "2001:db8::/32") "2001:db8::12"))
  (check-true
   (network-contains? (parse-network "127.0.0.0/8")
                      "0:0:0:0:0:ffff:7f00:1"))
  (check-exn exn:fail:contract?
             (λ () (parse-network "192.0.2.1/33")))
  (check-exn exn:fail:contract?
             (λ () (parse-network "not-an-address")))

  (define (test-request peer [headers '()])
    (request #"GET" (string->url "http://example.test/api/state")
             headers (delay '()) #f "127.0.0.1" 80 peer))

  (define remote (test-request "198.51.100.2"))
  ;; Local and remote browser requests follow the same login path.
  (check-false (auth-request-user manager (test-request "127.0.0.1")))
  (check-equal?
   (request-address
    manager
    (test-request
     "127.0.0.1"
     (list (header #"X-Forwarded-For"
                   #"198.51.100.8, 203.0.113.9"))))
   "203.0.113.9")
  (check-equal?
   (request-address
    manager
    (test-request
     "198.51.100.2"
     (list (header #"X-Forwarded-For" #"203.0.113.9"))))
   "198.51.100.2")
  (check-equal?
   (request-address
    manager
    (test-request
     "0:0:0:0:0:ffff:7f00:1"
     (list (header #"X-Forwarded-For" #"203.0.113.9"))))
   "203.0.113.9")
  (define token
    (auth-login! manager remote "hans" "correct horse battery staple"))
  (check-true (string? token))
  (check-true
   (regexp-match? #rx#"Max-Age=604800" (auth-session-cookie manager token)))
  (define authenticated
    (test-request
     "198.51.100.2"
     (list
      (header #"Cookie"
              (string->bytes/utf-8
               (format "~a=~a" session-cookie-name token))))))
  (check-equal? (auth-request-user manager authenticated) "hans")
  (define stored-session (hash-ref (auth-manager-sessions manager) token))
  (set-session-last-cookie-renewal! stored-session 0)
  (check-true (bytes? (auth-renewal-cookie manager authenticated)))
  (check-false (auth-renewal-cookie manager authenticated))
  (auth-logout! manager authenticated)
  (check-false (auth-request-user manager authenticated))
  (check-false (auth-login! manager remote "hans" "wrong password"))

  (let ((limited-manager (make-auth-manager '())))
    (for-each
     (λ (_) (record-failure! limited-manager "192.0.2.1" 100))
     (range maximum-failures))
    (check-true (failure-blocked? limited-manager "192.0.2.1" 100))
    (check-false
     (failure-blocked? limited-manager
                       "192.0.2.1"
                       (+ 101 failure-window-seconds))))

  (let ((session-manager (make-auth-manager '() #:session-seconds 10)))
    (hash-set! (auth-manager-sessions session-manager)
               "expired"
               (session "hans" 0 0))
    (prune-sessions! session-manager 11)
    (check-false
     (hash-has-key? (auth-manager-sessions session-manager) "expired"))))