#lang racket/base

(require crypto
         crypto/argon2
         net/private/ip
         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)

(struct ip-network (address prefix) #:transparent)
(struct session
  (username [last-seen #:mutable] [last-cookie-renewal #:mutable])
  #:transparent)
(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.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(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)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; 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.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define (make-password-hash password)
  (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.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define (password-hash-valid? password encoded)
  (and (string? password)
       (string? encoded)
       (with-handlers ((exn:fail? (lambda (_) #f)))
         (pwhash-verify password-kdf
                        (string->bytes/utf-8 password)
                        encoded))))

(define (normal-ip-bytes value)
  (define raw
    (ip-address->bytes (make-ip-address value)))
  ;; Normalize IPv4-mapped IPv6 addresses to four bytes.
  (if (and (= (bytes-length raw) 16)
           (for/and ((index (in-range 10)))
             (zero? (bytes-ref raw index)))
           (= (bytes-ref raw 10) #xff)
           (= (bytes-ref raw 11) #xff))
      (subbytes raw 12)
      raw))

(define (parse-network value)
  (define parts (string-split (string-trim value) "/"))
  (unless (member (length parts) '(1 2))
    (raise-argument-error 'make-auth-manager "IP address or CIDR network" value))
  (define address
    (with-handlers ((exn:fail?
                     (lambda (_)
                       (raise-argument-error
                        'make-auth-manager
                        "IP address or CIDR network"
                        value))))
      (normal-ip-bytes (car parts))))
  (define maximum (* 8 (bytes-length address)))
  (define 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))
  (ip-network address prefix))

(define (network-contains? network address-string)
  (with-handlers ((exn:fail? (lambda (_) #f)))
    (define candidate (normal-ip-bytes address-string))
    (define expected (ip-network-address network))
    (and (= (bytes-length candidate) (bytes-length expected))
         (let-values (((whole remainder)
                       (quotient/remainder (ip-network-prefix network) 8)))
           (and (for/and ((index (in-range whole)))
                  (= (bytes-ref candidate index)
                     (bytes-ref expected index)))
                (or (zero? remainder)
                    (let ((mask
                           (bitwise-and #xff
                                        (arithmetic-shift #xff (- remainder 8)))))
                      (= (bitwise-and (bytes-ref candidate whole) mask)
                         (bitwise-and (bytes-ref expected whole) mask)))))))))

(define (header-string request name)
  (let ((value (headers-assq* name (request-headers/raw request))))
    (and value
         (bytes->string/utf-8 (header-value value)))))

(define (trusted-proxy? manager address)
  (ormap (lambda (network) (network-contains? network address))
         (auth-manager-trusted-proxies manager)))

(define (request-address manager request)
  (define peer (request-client-ip request))
  (define 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))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; 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 (auth-enabled? manager)
  (positive? (hash-count (auth-manager-users manager))))

(define (request-session-token request)
  (for/or ((cookie (in-list (request-cookies request))))
    (and (string=? (client-cookie-name cookie) session-cookie-name)
         (client-cookie-value cookie))))

(define (prune-sessions! manager now)
  (for ((token (in-list (hash-keys (auth-manager-sessions manager)))))
    (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)))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; 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.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define (auth-request-user manager request)
  (cond
    ((not (auth-enabled? manager)) "anonymous")
    (else
     (let ((token (request-session-token request))
           (now (current-seconds)))
       (and token
            (call-with-semaphore
             (auth-manager-lock manager)
             (lambda ()
               (prune-sessions! manager now)
               (let ((value (hash-ref (auth-manager-sessions manager)
                                      token
                                      #f)))
                 (and value
                      (begin
                        (set-session-last-seen! value now)
                        (session-username value)))))))))))

(define (failure-blocked? manager address now)
  (define value (hash-ref (auth-manager-failed manager) address #f))
  (and value
       (if (> (- now (failures-started value)) failure-window-seconds)
           (begin
             (hash-remove! (auth-manager-failed manager) address)
             #f)
           (>= (failures-attempts value) maximum-failures))))

(define (record-failure! manager address now)
  (define value (hash-ref (auth-manager-failed manager) address #f))
  (if (and value
           (<= (- now (failures-started value)) failure-window-seconds))
      (set-failures-attempts! value (+ 1 (failures-attempts value)))
      (hash-set! (auth-manager-failed manager)
                 address
                 (failures 1 now))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; 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: Unknown users follow the same Argon2id verification path as known
;            users to reduce username-dependent timing differences.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define (auth-login! manager request username password)
  (define address (request-address manager request))
  (define now (current-seconds))
  (define normalized (string-downcase (string-trim username)))
  (call-with-semaphore
   (auth-manager-lock manager)
   (lambda ()
     (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.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define (auth-logout! manager request)
  (let ((token (request-session-token request)))
    (when token
      (call-with-semaphore
       (auth-manager-lock manager)
       (lambda ()
         (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.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define (auth-session-cookie manager token)
  (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: The server idle timer moves on every authenticated request, while
;            this half-life threshold prevents the one-second player poll from
;            returning Set-Cookie every second.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define (auth-renewal-cookie manager request)
  (and (auth-enabled? manager)
       (let ((token (request-session-token request))
             (now (current-seconds)))
         (and token
              (call-with-semaphore
               (auth-manager-lock manager)
               (λ ()
                 (prune-sessions! manager now)
                 (define value
                   (hash-ref (auth-manager-sessions manager) token #f))
                 (and value
                      (>= (- now (session-last-cookie-renewal value))
                          (max 1
                               (quotient
                                (auth-manager-session-seconds manager)
                                2)))
                      (begin
                        (set-session-last-cookie-renewal! value now)
                        (auth-session-cookie manager token)))))))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; 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.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define (auth-expired-cookie)
  (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.
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(define (make-auth-manager user-pairs
                           #:trusted-proxies
                           [trusted-proxy-values '("127.0.0.0/8" "::1/128")]
                           #:session-seconds [session-seconds 604800])
  (unless (exact-positive-integer? session-seconds)
    (raise-argument-error 'make-auth-manager "exact-positive-integer?"
                          session-seconds))
  (define users (make-hash))
  (for ((entry (in-list user-pairs)))
    (unless (and (pair? entry)
                 (string? (car entry))
                 (string? (cdr entry)))
      (raise-argument-error
       'make-auth-manager
       "(listof (cons/c string? string?))"
       user-pairs))
    (when (string=? (string-trim (car entry)) "")
      (raise-arguments-error
       'make-auth-manager
       "username must not be empty"
       "username" (car entry)))
    (unless (regexp-match? #px"^[$]argon2id[$]" (cdr entry))
      (raise-arguments-error
       'make-auth-manager
       "user password is not an Argon2id hash"
       "username" (car entry)))
    (hash-set! users (string-downcase (string-trim (car entry)))
               (cdr entry)))
  (auth-manager users
                (map parse-network trusted-proxy-values)
                session-seconds
                (make-hash)
                (make-hash)
                (make-semaphore 1)))

(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")))

  (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")
  (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")))