Sagittarius で WebSocket 通信をするスクリプト。 Sagittarius 0.4.8.1 では wss に繋ぐことは出来ず…。 (おそらく make-client-tls-socket が Server Name Indication に非対応のため)

websocket.scm

(import (rnrs)
        (rnrs io ports)
        (math random)
        (scheme time)
        (rfc base64)
        (math hash)
        (rfc uri)
        (rfc tls)
        (srfi :13 strings)
        (sagittarius object)
        (sagittarius regex)
        (sagittarius socket)
        (sagittarius control)
        (clos user))

(define-syntax debug
  (syntax-rules ()
    ((_ expr)
     (call-with-values (lambda()expr)
       (lambda args
         (display "debug: " (current-error-port))
         (write args (current-error-port))
         (newline (current-error-port))
         (apply values args))))))

(define-syntax rlet1
  (syntax-rules ()
    ((_ var expr body ...)
     (let ((var expr))
       body ...
       var))))

(define-syntax push!
  (syntax-rules ()
    ((_ var obj)
     (set! var (cons obj var)))))

(define-syntax if-let1
  (syntax-rules ()
    ((_ var expr act alt)
     (let1 var expr
       (if var act alt)))
    ((_ var expr act)
     (if-let1 var expr act #f))))

(define (read-line :optional (port (current-input-port)))
  (call-with-values open-bytevector-output-port
    (lambda(out get)
      (let loop ((ch (get-u8 port)))
        (cond ((eqv? #x0d ch)
               (if (eqv? (get-u8 port) #x0a)
                   (utf8->string (get))
                   (error 'read-line "Invalid character")))
              ((eof-object? ch) (utf8->string (get)))
              (else (put-u8 out ch)
                    (loop (get-u8 port))))))))

(define (assoc-ref alist obj :optional (default #f) (comp equal?))
  (if-let1 r (assoc obj alist comp)
    (cdr r)
    default))

(define (put-u16 port num :optional (endian (endianness big)))
  (rlet1 v (make-bytevector 2)
    (bytevector-u16-set! v 0 num endian)
    (put-bytevector port v)))

(define (put-u64 port num :optional (endian (endianness big)))
  (rlet1 v (make-bytevector 8)
    (bytevector-u64-set! v 0 num endian)
    (put-bytevector port v)))

(define (get-u16 port :optional (endian (endianness big)))
  (rlet1 v (make-bytevector 2)
    (get-bytevector-n! port v 0 2)
    (bytevector-u16-ref v 0 endian)))

(define (get-u64 port :optional (endian (endianness big)))
  (rlet1 v (make-bytevector 8)
    (get-bytevector-n! port v 0 8)
    (bytevector-u64-ref v 0 endian)))

(define guid "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")

(define rand
  (rlet1 prng (secure-random Yarrow)
    (random-seed-set! prng (current-jiffy))))

(define (nonce)
  (utf8->string (base64-encode (read-random-bytes rand 16) :line-width #f)))

(define (gen-mask) (read-random-bytes rand 4))

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

(define (sha1-digest-string str)
  (hash SHA-1 (string->utf8 str)))

(define (accept-key key)
  (utf8->string
   (base64-encode
    (sha1-digest-string (string-append key guid))
    :line-width #f)))

(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)
   (receiver)
   (port)
   (header)
   (buffer :init-value #f)
   (status :init-value 'closed)))

(define-method handshaked? ((self <websocket>))
  (eq? 'open (slot-ref 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)
  (string-append method " " path " HTTP/1.1\r\n"))

(define (header-field alst)
  (string-concatenate
   (fold-right
    (lambda(x s)
      (cons (string-append (->string (car x)) ": " (cdr x) "\r\n") s))
    '() alst)))

(define (parse-status-line line)
  (let* ((r (regex "^HTTP/1.1 (\\d+) +(.+)$"))
         (m (looking-at r line)))
    (if m (values (m 1) (m 2))
        (error 'parse-status-line "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)
  (let* ((r (regex "^(\\S+): (.*)$"))
         (m (looking-at r line)))
    (if m
        (cons (m 1) (m 2))
        (error 'parse-field "Invalid header-field:" line))))

(define (read-header :optional (port (current-input-port)))
  (rlet1 r '()
    (do ((line (read-line port) (read-line port)))
        ((string-null? line))
      (push! r (parse-field line)))))

(define-method websocket-connect ((self <websocket>) (url <string>))
  (debug 'b)
  (slot-set! self 'status 'connecting)
  (receive (scheme user host port path)
      (debug (url-split url))
    (let* ((sock (if (debug (equal? "wss" scheme))
                     (debug (make-client-tls-socket host (debug port)))
                     (make-client-socket host port AF_INET)))
           (port (if (equal? "wss" scheme)
                     (tls-socket-port sock)
                     (socket-port sock)))
           (key (nonce))
           (origin (string-append (cond 
                                   ((string=? "ws" scheme) "http://")
                                   ((string=? "wss" scheme) "https://"))
                                  host)))
      (put-bytevector port
                      (string->utf8 (string-append
                                     (request-line "GET" path)
                                     (request-header host key scheme origin)
                                     "\r\n")))
      (debug 'a)
      (flush-output-port port)
      (receive (code mes)
          (parse-status-line (debug (read-line port)))
        (unless (equal? code "101") (error 'websocket-connect "fail connect")))
      (let1 header (read-header port)
        (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 'port port)
        (slot-set! self 'header header)
        ;; (slot-set! self 'receiver
        ;;            (make-thread
        ;;             (lambda()(until (eq? (websocket-receive self) 'close)
        ;;                        (display 'loop)))))
        (slot-set! self 'status 'open)
        ;; (thread-start! (slot-ref self 'receiver))
        #t))))

(define-method websocket-close ((self <websocket>))
  (unless (handshaked? self)
    (error "websocket hasn't handshake yet."))
  (let1 port (slot-ref self 'port)
    (put-u8 port (fxior #x80 8))
    (put-u8 port #x80)
    (put-bytevector port (gen-mask))
    (flush-output-port port))
  (slot-set! self 'status 'closing))

(define-method websocket-pong ((self <websocket>))
  (unless (handshaked? self)
    (error "websocket hasn't handshake yet."))
  (let1 port (slot-ref self 'port)
    (put-u8 port (fxior #x80 10))
    (put-u8 port #x80)
    (put-bytevector port (gen-mask))
    (flush-output-port port)))

(define-method websocket-ping ((self <websocket>))
  (unless (handshaked? self)
    (error "websocket hasn't handshake yet."))
  (let1 port (slot-ref self 'port)
    (put-u8 port (fxior #x80 9))
    (put-u8 port #x80)
    (put-bytevector port (gen-mask))
    (flush-output-port port)))

(define-method write-payload-length
    ((self <websocket>) (mask <bytevector>) (num <integer>))
  (let ((port (slot-ref self 'port))
        (maskbit (if mask #x80 0)))
    (cond ((<= 0 num 125)
           (put-u8 port (fxior num maskbit)))
          ((<= 126 num #x8000)
           (put-u8 port (fxior 126 maskbit))
           (put-u16 port num))
          ((<= #x10000 num #x4000000000000000)
           (put-u8 port (fxior 127 maskbit))
           (put-u64 port num))
          (else (error "Payload is too large")))
    (when mask (put-bytevector port mask))
    (flush-output-port port)))

(define-method apply-mask ((bvec <bytevector>) (mask <bytevector>))
  (let* ((len (bytevector-length bvec))
         (nvec (make-bytevector len)))
    (do ((i 0 (+ i 1)))
        ((= i len) nvec)
      (bytevector-u8-set! nvec i
        (fxxor (bytevector-u8-ref bvec i)
               (bytevector-u8-ref mask (modulo i 4)))))))

(define-method websocket-send ((self <websocket>) (data <bytevector>))
  (unless (handshaked? self) (error "websocket hasn't handshake yet."))
  (let ((port (slot-ref self 'port))
        (len (bytevector-length data))
        (mask (gen-mask)))
    (put-u8 port (fxior #x80 2))
    (write-payload-length self mask len)
    (put-bytevector port (apply-mask data mask))
    (flush-output-port port)))

(define-method websocket-send ((self <websocket>) (data <string>))
  (unless (handshaked? self) (error "websocket hasn't handshake yet."))
  (let* ((port (slot-ref self 'port))
         (vec (string->utf8 data))
         (len (bytevector-length vec))
         (mask (gen-mask)))
    (put-u8 port (fxior #x80 1))
    (write-payload-length self mask len)
    (put-bytevector port (apply-mask vec mask))
    (flush-output-port port)))

(define-method websocket-receive ((self <websocket>))
  (unless (handshaked? self) (error "websocket hasn't handshake yet."))
  (let* ((port (slot-ref self 'port))
         (b1 (get-u8 port))
         (fin (not (zero? (fxand #x80 b1))))
         (op (fxand #x0f b1))
         (b2 (get-u8 port))
         (mask-flag (not (zero? (fxand #x80 b2))))
         (plength (fxand #x7f b2)))
    (let* ((plength (case plength
                      ((126) (get-u16 in))
                      ((127) (get-u64 in))
                      (else plength)))
           (mask (make-bytevector 4))
           (payload (make-bytevector plength)))
      (when mask-flag (get-bytevector-n! port mask 0 4))
      (get-bytevector-n! port payload 0 plength)
      (case op
        ((1) ((slot-ref self 'on-message)
              (utf8->string
               (if mask-flag (apply-mask payload mask) payload))
               fin)
             'text)
        ((2) ((slot-ref self 'on-message)
              (if mask-flag (apply-mask payload mask) payload)
              fin)
             'binary)
        ((8) (if (eq? 'closing (slot-ref self 'status))
                 (begin
                   (socket-close (slot-ref self 'socket))
                   (slot-set! self 'status 'closed)
                   (slot-set! self 'socket #f))
                 (begin
                   (websocket-close self)
                   (socket-close (slot-ref self 'socket))
                   (slot-set! self 'status 'closed)
                   (slot-set! self 'socket #f)))
             'close)
        ((9) (websocket-pong self) 'ping)
        ((10) (debug 'pong))
        (else "unknown opcode"))
      )))

(define websocket (make <websocket>
                    :on-message (lambda(data fin) (display data))))

(websocket-connect websocket "ws://echo.websocket.org")
(display "connected.\n")
(websocket-send websocket "Hello, world.")
(websocket-receive websocket)
;; (websocket-ping websocket)
;; (websocket-receive websocket)
(websocket-close websocket)

最新のHEADではwssでも動くようになりました。以下は実行ログです。

debug: (b)
debug: ("wss" #f "echo.websocket.org" "443" "/")
debug: (#t)
debug: ("443")
debug: (#<tls-client-socket 301>)
debug: (a)
debug: ("HTTP/1.1 101 Web Socket Protocol Handshake")
connected.
Hello, world.