;;; guile-websocket --- WebSocket client/server
;;; Copyright © 2015, 2025 David Thompson <dthompson2@worcester.edu>
;;; Copyright © 2021 Jan (janneke) Nieuwenhuizen <janneke@gnu.org>
;;; Copyright © 2023 Andrew Whatson <whatson@tailcall.au>
;;;
;;; This file is part of guile-websocket.
;;;
;;; Guile-websocket is free software; you can redistribute it and/or modify
;;; it under the terms of the GNU Lesser General Public License as
;;; published by the Free Software Foundation; either version 3 of the
;;; License, or (at your option) any later version.
;;;
;;; Guile-websocket 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
;;; Lesser General Public License for more details.
;;;
;;; You should have received a copy of the GNU Lesser General Public
;;; License along with guile-websocket.  If not, see
;;; <http://www.gnu.org/licenses/>.

;;; Commentary:
;;
;; WebSocket server.
;;
;;; Code:

(define-module (web socket server)
  #:use-module (gnutls)
  #:use-module (ice-9 match)
  #:use-module (ice-9 suspendable-ports)
  #:use-module (rnrs bytevectors)
  #:use-module (rnrs io ports)
  #:use-module (web request)
  #:use-module (web response)
  #:use-module (web uri)
  #:use-module (web socket frame)
  #:use-module (web socket utils)
  #:export (websocket-upgrade-request?
            websocket-upgrade-response?
            make-websocket-upgrade-response
            make-server-socket
            run-websocket-server))

(define* (make-server-socket #:key
                             (host #f)
                             (family AF_INET)
                             (addr (if host (inet-pton family host) INADDR_LOOPBACK))
                             (port 8080))
  (let ((sock (socket PF_INET SOCK_STREAM 0)))
    (setsockopt sock SOL_SOCKET SO_REUSEADDR 1)
    (fcntl sock F_SETFD FD_CLOEXEC)
    (bind sock AF_INET addr port)
    sock))

(define (accept-client server-sock)
  (match (accept server-sock O_CLOEXEC)
    ((client-sock . _)
     client-sock)))

(define (make-tls-session cert key)
  (let ((tls-session (make-session connection-end/server))
        (creds (make-certificate-credentials)))
    (set-session-default-priority! tls-session)
    (set-certificate-credentials-x509-key-data! creds cert key
                                                x509-certificate-format/pem)
    (set-session-credentials! tls-session creds)
    tls-session))

(define (tls-retry proc tls-session)
  (let lp ((retries 5))
    (define (handle-error exn)
      (if (zero? retries)
          (raise-exception exn)
          (match (exception-args exn)
            (((? fatal-error?) . _)
             (raise-exception exn))
            ;; For non-fatal errors such as EAGAIN or EINTERRUPTED, give
            ;; control to the current read waiter and try again when it
            ;; returns.
            (_
             ((current-read-waiter) (session-record-port tls-session))
             (lp (1- retries))))))
    (with-exception-handler handle-error
      (lambda ()
        (proc tls-session))
      #:unwind? #t
      #:unwind-for-type 'gnutls-error)))

(define (tls-wrap sock key cert)
  (let* ((tls-session (make-tls-session cert key))
         (tls-port (session-record-port tls-session)))
    (define (close tls-port) (close-port sock))
    (set-session-transport-fd! tls-session (fileno sock))
    (set-session-record-port-close! tls-port close)
    (setvbuf tls-port 'block)
    (tls-retry handshake tls-session)
    tls-port))

(define (websocket-upgrade-request? request)
  "Return @code{#t} if @var{request} is a valid WebSocket upgrade
request."
  (let ((headers (request-headers request)))
    ;; See section 4.2.1.
    (and (eq? (request-method request) 'GET)
         ;; TODO: Validate Host header contents.
         (assq-ref headers 'host)
         (match (assq-ref headers 'upgrade)
           ((upgrade)
            (string-ci=? upgrade "websocket"))
           (_ #f))
         (memq 'upgrade (assq-ref headers 'connection))
         ;; TODO: We could validate, without decoding, that this is
         ;; base64 text that produces a 16 byte value.
         (assq-ref headers 'sec-websocket-key)
         (string=? (assq-ref headers 'sec-websocket-version) "13"))))

(define (make-websocket-upgrade-response request)
  "Return a WebSocket upgrade response corresponding to @var{request}."
  ;; See section 4.2.2.
  (let* ((client-key (assoc-ref (request-headers request) 'sec-websocket-key))
         (accept-key (make-accept-key (string-trim-both client-key))))
    (build-response #:code 101
                    #:headers `((upgrade . ("websocket"))
                                (connection . (upgrade))
                                (sec-websocket-accept . ,accept-key)))))

(define (websocket-upgrade-response? response)
  "Return @code{#t} if @var{response} is a valid WebSocket upgrade
response."
  (and (= (response-code response) 101)
       (let ((headers (response-headers response)))
         (and (equal? (assq-ref headers 'upgrade) '("websocket"))
              (equal? (assq-ref headers 'connection) '(upgrade))))))

(define* (run-websocket-server handle-client #:key
                               (server-socket (make-server-socket))
                               (max-clients 1)
                               (accept-client accept-client)
                               (spawn-client (lambda (thunk) (thunk)))
                               tls-private-key tls-certificate)
  "Run WebSocket server on @var{server-socket} that accepts up to
@var{max-clients} connected clients.  When a client connects,
@var{handle-client} is called with the client socket as its argument.

The default server implementation handles only one client at a time.
To change this, specify @var{accept-client} and @var{spawn-client}.
@var{accept-client} is a procedure that takes the server socket as its
only argument and returns a client socket.  @var{spawn-client} is a
procedure that takes a thunk as its only argument and should apply
thunk to start the client handshake process.  Using these two
procedures, it is possible to do things like spawn a thread for each
client or use suspendable ports or Fibers to create a non-blocking
server and client(s).

@var{tls-private-key} and @var{tls-certificate} are optionally
bytevectors that contain PEM formatted X.509 private key and
certificate data.  When specified, the server will be encrypted with
TLS using the given key."
  (define (serve-client client)
    ;; Perform the HTTP handshake and upgrade to WebSocket protocol.
    (let ((request (read-request client)))
      (if (websocket-upgrade-request? request)
          (let ((response (make-websocket-upgrade-response request)))
            (write-response response client)
            (force-output client)
            (handle-client client))
          (let ((response (build-response #:code 400)))
            (write-response response client)
            (close-port client)))))
  (listen server-socket max-clients)
  (sigaction SIGPIPE SIG_IGN)
  (let loop ()
    (match (accept-client server-socket)
      (#f (values))                     ; terminate
      (#t (loop))                       ; no new client, but continue
      (client
       ;; Disable Nagle's algorithm.
       (setsockopt client IPPROTO_TCP TCP_NODELAY 1)
       (setvbuf client 'block 1024)
       (spawn-client
        (lambda ()
          (serve-client
           (if (and tls-private-key tls-certificate)
               (tls-wrap client tls-private-key tls-certificate)
               client))))
       (loop)))))
