websocket.scm

(use rfc.uri)
(use rfc.base64)
(use rfc.sha)
(use rfc.tls)
(use math.mt-random)
(use binary.io)
(use binary.pack)
(use gauche.uvector)
(use gauche.collection)
(use gauche.sequence)
(use gauche.net)
(use gauche.threads)
(use srfi-13)
(use srfi-60)

(debug-print-width #f)
(define guid "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")

(define nonce
  (let ([random-source (make <mersenne-twister>
                         :seed (* (sys-time) (sys-getpid)))]
        [v (make-u32vector 4)])
    (lambda ()
      (mt-random-fill-u32vector! random-source v)
      (base64-encode-string (u8vector->string (uvector-alias <u8vector> v))))))

(define gen-mask
  (let ([random-source (make <mersenne-twister>
                         :seed (* (sys-time) (sys-getpid)))]
        [v (make-u32vector 1)])
    (lambda ()
      (mt-random-fill-u32vector! random-source v)
      (uvector-alias <u8vector> v))))

(define (do-nothing . arg) #f)

(define (accept-key key)
  (base64-encode-string (sha1-digest-string (string-append key guid))))

(define-class <websocket> ()
  ((on-message :init-keyword :on-message :init-value do-nothing)
   (on-open :init-keyword :on-open :init-value do-nothing)
   (on-error :init-keyword :on-error :init-value do-nothing)
   (on-close :init-keyword :on-close :init-value do-nothing)
   (socket :init-value #f)
   (secure-agent :init-value #f)
   (receiver)
   (in-port)
   (out-port)
   (header)
   (buffer :init-value #f)
   (status :init-value 'closed)))

(define-method handshaked? ((self <websocket>))
  (eq? 'open (~ self 'status)))

(define (scheme->default-port scheme)
  (cond ((string=? scheme "ws") 80)
        ((string=? scheme "wss") 443)
        (else (error "unsupported scheme:" scheme))))

(define (url-split url)
  (receive (scheme user host port path query fragment)
      (uri-parse url)
    (let ((path (if path path "/"))
          (query (if query (string-append "?" query) "")))
      (values scheme
            user
            host
            (if port port (scheme->default-port scheme))
            (string-append path query)))))

(define (request-line method path)
  #`",|method| ,|path| HTTP/1.1\r\n")

(define (header-field alst)
  (string-concatenate
   (fold-right
    (^(x s)(cons #`",(x->string(car x)): ,(cdr x)\r\n" s))
    '() alst)))

(define (parse-status-line line)
  (if-let1 m (rxmatch #/^HTTP\/1.1 (\d+) +(.+)$/ line)
    (values (m 1) (m 2))
    (error "Invalid status line:" line)))

(define (request-header host key scheme origin)
  (header-field
   `((Host . ,host)
     (Upgrade . "websocket")
     (Connection . "Upgrade")
     (Sec-WebSocket-Key . ,key)
     (Origin . ,origin)
     (Sec-WebSocket-Version . "13"))))

(define (parse-field line)
  (if-let1 m (rxmatch #/^(\S+): (.*)$/ line)
    (cons (m 1) (m 2))
    (error "Invalid header-field:" line)))

(define (read-header :optional (port (current-input-port)))
  (generator-fold-right
   cons
   '()
   (^() (let1 e (read-line port)
          (if (string-null? e) (eof-object) (parse-field e))))))

(define-method websocket-connect ((self <websocket>) (url <string>))
  (slot-set! self 'status 'connecting)
  (receive (scheme user host port path)
      (url-split url)
    (let* ((sock (make-client-socket 'inet host port))
           (secure-agent (if (equal? "wss" scheme)
                             (rlet1 tls (make-tls)
                               (tls-connect tls (socket-fd sock)))
                             #f))
           (in (if secure-agent
                   (tls-input-port secure-agent)
                   (socket-input-port sock)))
           (out (if secure-agent
                    (tls-output-port secure-agent)
                    (socket-output-port sock :buffering :full)))
           (key (nonce))
           (origin (string-append (cond 
                                   ((string=? "ws" scheme) "http://")
                                   ((string=? "wss" scheme) "https://"))
                                  host)))
      (display (string-append
                (request-line "GET" path)
                (request-header host key scheme origin)
                "\r\n")
               out)
      (flush out)
      (receive (code mes)
          (parse-status-line (read-line in))
        (unless (equal? code "101") (error "fail connect")))
      (let1 header #?=(read-header in)
        (if-let1 r (assoc-ref header "sec-websocket-origin" #f string-ci=?)
          (unless (string-ci=? r origin)
            (slot-set! self 'status 'closed)
            (slot-set! self 'socket #f)
            (error "origin doesn't match:" r)))
        (if-let1 r (assoc-ref header "Sec-WebSocket-Accept" #f string-ci=?)
          (string-ci=? r (accept-key key))
          (begin
            (slot-set! self 'status 'closed)
            (slot-set! self 'socket #f)
            (error "security digest doesn't match:" r)))
        (slot-set! self 'socket sock)
        (slot-set! self 'in-port in)
        (slot-set! self 'out-port out)
        (slot-set! self 'header header)
        (slot-set! self 'secure-agent secure-agent)
        ;; (slot-set! self 'receiver
        ;;            (make-thread
        ;;             (lambda()(until (eq? (websocket-receive self) 'close)
        ;;                        (display 'loop)))))
        (slot-set! self 'status 'open)
        ;; (thread-start! (~ self 'receiver))
        #t))))

(define-method websocket-close ((self <websocket>))
  (unless (handshaked? self)
    (error "websocket hasn't handshake yet."))
  (let1 out (~ self 'out-port)
    (write-u8 (bitwise-ior #x80 8) out)
    (write-u8 0 out)
    (flush out)
    )
  (slot-set! self 'status 'closing))

(define-method websocket-pong ((self <websocket>))
  (unless (handshaked? self)
    (error "websocket hasn't handshake yet."))
  (let1 out (~ self 'out-port)
    (write-u8 (bitwise-ior #x80 10) out)
    (write-u8 0 out)
    (flush out)
    ))

(define-method websocket-ping ((self <websocket>))
  (unless (handshaked? self)
    (error "websocket hasn't handshake yet."))
  (let1 out (~ self 'out-port)
    (write-u8 (bitwise-ior #x80 9) out)
    (write-u8 0 out)
    (flush out)
    ))

(define-method write-payload-length
    ((self <websocket>) (mask <u8vector>) (num <integer>))
  (let ((out (~ self 'out-port))
        (maskbit (if mask #x80 0)))
    (cond ((<= 0 num 125)
           (write-u8 (bitwise-ior num maskbit) out))
          ((<= 126 num #x8000)
           (write-u8 (bitwise-ior 126 maskbit) out)
           (write-u16 num out 'big-endian))
          ((<= #x10000 num #x4000000000000000)
           (write-u8 (bitwise-ior 127 maskbit) out)
           (write-u64 num out 'big-endian))
          (else (error "Payload is too large")))
    (when mask (write-block mask out))
    (flush out)
    ))

(define-method apply-mask-vec ((bvec <u8vector>) (mask <u8vector>))
  (map-to-with-index <u8vector>
   (lambda(i x) (bitwise-xor x (u8vector-ref mask (modulo i 4))))
   bvec))

(define-method apply-mask-str ((str <string>) (mask <u8vector>))
  (apply-mask-vec (string->u8vector str) mask))

(define-method websocket-send ((self <websocket>) (data <u8vector>))
  (unless (handshaked? self) (error "websocket hasn't handshake yet."))
  (let ((out (~ self 'out-port))
        (len (u8vector-length data))
        (mask (gen-mask)))
    (write-u8 (bitwise-ior #x80 2) out)
    (write-payload-length self mask len)
    (write-block (apply-mask-vec data mask) out)
    (flush out)
    ))

(define-method websocket-send ((self <websocket>) (data <string>))
  (unless (handshaked? self) (error "websocket hasn't handshake yet."))
  (let ((out (~ self 'out-port))
        (len (string-size data))
        (mask (gen-mask)))
    (write-u8 (bitwise-ior #x80 1) out)
    (write-payload-length self mask len)
    (write-block (apply-mask-str data mask) out)
    (flush out)
    ))

(define-method websocket-receive ((self <websocket>))
  (unless (handshaked? self) (error "websocket hasn't handshake yet."))
  (let* ((in (~ self 'in-port))
         (b1 (read-byte in))
         (fin (not (zero? (bitwise-and #x80 b1))))
         (op (bitwise-and #x0f b1))
         (b2 (read-byte in))
         (mask-flag (not (zero? (bitwise-and #x80 b2))))
         (plength (bitwise-and #x7f b2)))
    (let* ((plength (case plength
                      ((126) (read-u16 in 'big-endian))
                      ((127) (read-u64 in 'big-endian))
                      (else plength)))
           (mask (make-u8vector 4))
           (payload (make-u8vector plength)))
      (when mask-flag (read-block! mask in))
      (read-block! payload in)
      (case op
        ((1) ((~ self 'on-message)
              (u8vector->string
               (if mask-flag (apply-mask-vec payload mask) payload))
               fin)
             'text)
        ((2) ((~ self 'on-message)
              (if mask-flag (apply-mask-vec payload mask) payload)
              fin)
             'binary)
        ((8) (if (eq? 'closing (~ self 'status))
                 (begin
                   (if-let1 s (~ self 'secure-agent)
                     (tls-close s))
                   (socket-close (~ self 'socket))
                   (slot-set! self 'status 'closed)
                   (slot-set! self 'socket #f))
                 (begin
                   (websocket-close self)
                   (if-let1 s (~ self 'secure-agent)
                     (tls-close s))
                   (socket-close (~ self 'socket))
                   (slot-set! self 'status 'closed)
                   (slot-set! self 'socket #f)))
             'close)
        ((9) (websocket-pong self) 'ping)
        ((10) 'pong)
        (else "unknown opcode"))
      )))



(define websocket (make <websocket>
                    :on-message (lambda(data fin) (display data))))
(websocket-connect websocket "wss://echo.websocket.org")

(websocket-send websocket "365")
(websocket-receive websocket)
(websocket-ping websocket)
(websocket-receive websocket)
(websocket-close websocket)