;;; guile-websocket --- WebSocket client/server
;;; Copyright © 2016 David Thompson <davet@gnu.org>
;;; Copyright © 2017 Christopher Allan Webber <cwebber@dustycloud.org>
;;; Copyright © 2019, 2020 Jan (janneke) Nieuwenhuizen <janneke@gnu.org>
;;;
;;; 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 client.
;;
;;; Code:

(define-module (8sync systems websocket client)
  #:use-module (ice-9 match)
  #:use-module (srfi srfi-26)
  #:use-module (rnrs bytevectors)
  #:use-module (rnrs io ports)
  #:use-module (web request)
  #:use-module (web response)
  #:use-module (web uri)
  #:use-module (oop goops)
  #:use-module (8sync)
  #:use-module (8sync ports)
  #:use-module (8sync contrib base64)
  #:use-module (8sync systems websocket frame)
  #:use-module (8sync systems websocket utils)
  #:export (<websocket>
            .on-close
            .on-error
            .on-message
            .on-open
            .socket
            .state
            .url

            websocket-closed?
            websocket-closing?
            websocket-connect
            websocket-connecting?
            websocket-open?

            websocket-close
            websocket-loop
            websocket-send))

(define no-op (const #f))

(define-actor <websocket> (<actor>)
  ((*init* websocket-init)
   (close websocket-close)
   (open websocket-open)
   (send websocket-send))

  (state #:accessor .state #:init-value 'closed #:init-keyword #:state)
  (socket #:accessor .socket #:init-value #f #:init-keyword #:socket)
  (url #:getter .url #:init-value #f #:init-keyword #:url)
  (uri #:accessor .uri #:init-value #f #:init-keyword #:uri)
  (entropy-port #:accessor .entropy-port #:init-form (open-entropy-port))

  (on-close #:init-keyword #:on-close
                 #:init-value no-op
                 #:accessor .on-close)
  (on-error #:init-keyword #:on-error
            #:init-value no-op
            #:accessor .on-error)
  (on-message #:init-keyword #:on-message
              #:accessor .on-message)
  (on-open #:init-keyword #:on-open
                #:init-value no-op
                #:accessor .on-open))

(define-method (websocket-close (websocket <websocket>) message)
  (when (websocket-open? websocket)
    (false-if-exception (close-port (.socket websocket)))
    (set! (.state websocket) 'closed)
    (false-if-exception ((.on-close websocket) websocket))
    (set! (.socket websocket) #f)))

(define-method (websocket-open (websocket <websocket>) message uri-or-string)
  (if (websocket-closed? websocket)
      (let ((uri (match uri-or-string
                   ((? uri? uri) uri)
                   ((? string? str) (string->uri str)))))
        (if (websocket-uri? uri)
            (catch 'system-error
              (lambda _
                (set! (.uri websocket) uri)
                (let ((sock (make-client-socket uri)))
                  (set! (.socket websocket) sock)
                  (handshake websocket)
                  (websocket-loop websocket message)))
              (lambda (key . args)
                ((.on-error websocket) websocket (format #f "open failed: ~s: ~s" uri-or-string args))))
            ((.on-error websocket) websocket (format #f "not a websocket uri: ~s" uri-or-string))))
      ((.on-error websocket) websocket (format #f "cannot open websocket in state: ~s" (.state websocket)))))

(define (subbytevector bv start end)
  (if (= (bytevector-length bv) end) bv
      (let* ((length (- end start))
             (sub (make-bytevector length)))
        (bytevector-copy! bv start sub 0 length)
        sub)))

(define* (make-fragmented-frames data #:key (fragment-size (expt 2 15)))
  (let ((length (if (string? data) (string-length data)
                    (bytevector-length data))))
    (let loop ((offset 0))
      (let* ((size (min fragment-size (- length offset)))
             (end (+ offset size))
             (final? (= end length))
             (continuation? (not (zero? offset)))
             (frame (if (string? data) (make-text-frame (substring data offset end) #:final? final? #:continuation? continuation?)
                        (make-binary-frame (subbytevector data offset end) #:final? final? #:continuation? continuation?))))
        (if final? (list frame)
            (cons frame (loop end)))))))

(define-method (websocket-send (websocket <websocket>) message data)
  (catch #t           ; expect: wrong-type-arg (open port), system-error
    (lambda _
      (let* ((frames (make-fragmented-frames data)))
        (let loop ((frames frames) (written '(nothing)))
          (when (pair? frames)
            (write-frame (car frames) (.socket websocket))
            (loop (cdr frames) (cons (car frames) written))))))
    (lambda (key . args)
      (let ((message (format #f "~a: ~s" key args)))
        ((.on-error websocket) websocket (format #f "send failed: ~s ~a\n" websocket message))
        (websocket-close websocket message)))))

(define-method (websocket-init (websocket <websocket>) message)
  (and=> (.url websocket) (cut websocket-open websocket message <>)))

(define-method (websocket-loop (websocket <websocket>) message)

  (define (handle-data-frame type data)
    ((.on-message websocket)
     websocket
     (match type
       ('text   (utf8->string data))
       ('binary data))))

  (define (read-frame-maybe)
    (and (not (eof-object? (lookahead-u8 (.socket websocket))))
         (read-frame (.socket websocket))))

  (define (close-down)
    (websocket-close websocket message))

  ((.on-open websocket) websocket)

  (let loop ((fragments '())
             (type #f))
    (let* ((socket (.socket websocket))
           (frame (and (websocket-open? websocket)
                       (read-frame-maybe))))
      (cond
       ;; EOF - port is closed.
       ;; @@: Sometimes the eof object appears here as opposed to
       ;;   at lookahead, but I'm not sure why
       ((or (not frame) (eof-object? frame))
        (close-down))
       ;; Per section 5.4, control frames may appear interspersed
       ;; along with a fragmented message.
       ((close-frame? frame)
        ;; Per section 5.5.1, echo the close frame back to the
        ;; socket before closing the socket.  The socket may no
        ;; longer be listening.
        (false-if-exception
         (write-frame (make-close-frame (frame-data frame)) socket))
        (close-down))
       ((ping-frame? frame)
        ;; Per section 5.5.3, a pong frame must include the exact
        ;; same data as the ping frame.
        (write-frame (make-pong-frame (frame-data frame)) socket)
        (loop fragments type))
       ((pong-frame? frame)           ; silently ignore pongs
        (loop fragments type))
       ((first-fragment-frame? frame) ; begin accumulating fragments
        (loop (list frame) (frame-type frame)))
       ((final-fragment-frame? frame) ; concatenate all fragments
        (handle-data-frame type (frame-concatenate
                                 (reverse (cons frame fragments))))
        (loop '() #f))
       ((fragment-frame? frame)       ; add a fragment
        (loop (cons frame fragments) type))
       ((data-frame? frame)           ; unfragmented data frame
        (handle-data-frame (frame-type frame) (frame-data frame))
        (loop '() #f))))))

;; See Section 3 - WebSocket URIs
(define (encrypted-websocket-scheme? uri)
  "Return #t if the scheme for URI is 'wss', the secure WebSocket
scheme."
  (eq? (uri-scheme uri) 'wss))

(define (unencrypted-websocket-scheme? uri)
  "Return #t if the scheme for URI is 'ws', the insecure WebSocket
scheme."
  (eq? (uri-scheme uri) 'ws))

(define (websocket-uri? uri)
  "Return #t if URI is a valid WebSocket URI."
  (and (or (encrypted-websocket-scheme? uri)
           (unencrypted-websocket-scheme? uri))
       (not (uri-fragment uri))))

(define (set-nonblocking! port)
  (fcntl port F_SETFL (logior O_NONBLOCK (fcntl port F_GETFL)))
  (setvbuf port 'block 1024))

(define (make-client-socket uri)
  "Connect a socket to the remote resource described by URI."
  (let* ((port (uri-port uri))
         (info (car (getaddrinfo (uri-host uri)
                                 (if port
                                     (number->string port)
                                     (symbol->string (uri-scheme uri)))
                                 (if port
                                     AI_NUMERICSERV
                                     0))))
         (sock (with-fluids ((%default-port-encoding #f))
                 (socket (addrinfo:fam info) SOCK_STREAM IPPROTO_IP))))

    (set-nonblocking! sock)
    ;; Disable buffering for websockets
    (setvbuf sock 'none)

    ;; TODO: Configure I/O buffering?
    (connect sock (addrinfo:addr info))
    sock))

(define-method (write (o <websocket>) port)
   (format port "#<websocket ~a ~a>"
           (.url o)
           (.state o)))

(define-method (websocket-connecting? (websocket <websocket>))
  "Return #t if WEBSOCKET is in the connecting state."
  (eq? (.state websocket) 'connecting))

(define-method (websocket-open? (websocket <websocket>))
  "Return #t if WEBSOCKET is in the open state."
  (eq? (.state websocket) 'open))

(define-method (websocket-closing? (websocket <websocket>))
  "Return #t if WEBSOCKET is in the closing state."
  (eq? (.state websocket) 'closing))

(define-method (websocket-closed? (websocket <websocket>))
  "Return #t if WEBSOCKET is in the closed state."
  (eq? (.state websocket) 'closed))

(define-method (generate-client-key (websocket <websocket>))
  "Return a random, base64 encoded nonce using the entropy source of
WEBSOCKET."
  (base64-encode
   (get-bytevector-n (.entropy-port websocket) 16)))

;; See Section 4.1 - Client Requirements
(define (make-handshake-request uri key)
  "Create an HTTP request for initiating a WebSocket connection with
the remote resource described by URI, using a randomly generated nonce
KEY."
  (let ((headers `((host . (,(uri-host uri) . #f))
                   (upgrade . ("WebSocket"))
                   (connection . (upgrade))
                   (sec-websocket-key . ,key)
                   (sec-websocket-version . "13"))))
    (build-request uri #:method 'GET #:headers headers)))

(define-method (handshake (websocket <websocket>))
  "Perform the WebSocket handshake for the client WEBSOCKET."
  (let ((key (generate-client-key websocket)))
    (write-request (make-handshake-request (.uri websocket) key)
                   (.socket websocket))
    (let* ((response (read-response (.socket websocket)))
           (headers (response-headers response))
           (upgrade (assoc-ref headers 'upgrade))
           (connection (assoc-ref headers 'connection))
           (accept (assoc-ref headers 'sec-websocket-accept)))
      ;; Validate the handshake.
      (if (and (= (response-code response) 101)
               (string-ci=? (car upgrade) "websocket")
               (equal? connection '(upgrade))
               (string=? (string-trim-both accept) (make-accept-key key)))
          (set! (.state websocket) 'open)
          (begin
            (websocket-close websocket)
            ((.on-error websocket) websocket
             (format #f "websocket handshake failed: ~s"
                     (uri->string (.uri websocket)))))))))

(define (open-entropy-port)
  "Return an open input port to a reliable source of entropy for the
current system."
  ;; XXX: This works on GNU/Linux and OS X systems, but this isn't
  ;; exactly portable.
  (open-input-file "/dev/urandom"))

(define-method (websocket-close (websocket <websocket>))
  "Close the WebSocket connection for the client WEBSOCKET."
  (let ((socket (.socket websocket)))
    (set! (.state websocket) 'closing)
    (write-frame (make-close-frame (make-bytevector 0)) socket)
    ;; Per section 5.5.1 , wait for the server to close the connection
    ;; for a reasonable amount of time.
    (let loop ()
      (match (select #() (vector socket) #() 1) ; 1 second timeout
        ((#() #(socket) #()) ; there is output to read
         (unless (port-eof? socket)
           (read-frame socket) ; throw it away
           (loop)))))
    (close-port socket)
    (close-port (.entropy-port websocket))
    (set! (.state websocket) 'closed)))