basis/http/websockets/websockets.factor

   1 ! Copyright (C) 2021 Doug Coleman.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors alien.enums alien.syntax assocs base64
   4 byte-arrays combinators combinators.short-circuit crypto.xor
   5 http http.client io io.binary io.encodings.string
   6 io.encodings.utf8 kernel math math.bitwise multiline namespaces
   7 prettyprint random sequences strings tools.hexdump ;
   8 IN: http.websockets
   9
  10 ! TODO: multiplexing, fragmented send
  11
  12 CONSTANT: websocket-version "13"
  13
  14 : random-websocket-key ( -- base64 )
  15     16 random-bytes >base64 >string ;
  16
  17 : add-websocket-headers ( request -- request )
  18     "connection" over header>> delete-at
  19     "Upgrade" "Connection" set-header
  20     "no-cache" "Pragma" set-header
  21     "no-cache" "Cache-Control" set-header
  22     "websocket" "Upgrade" set-header
  23     ! "http://www.websocket.org" "Origin" set-header
  24     "https://www.piesocket.com" "Origin" set-header
  25     websocket-version "Sec-WebSocket-Version" set-header
  26     random-websocket-key "Sec-WebSocket-Key" set-header
  27     "permessage-deflate; client_max_window_bits" "Sec-WebSocket-Extensions" set-header
  28     "gzip, deflate" "Accept-Encoding" set-header
  29     "en-US,en;q=0.9,sw-TZ;q=0.8,sw;q=0.7,es-US;q=0.6,es;q=0.5,de-DE;q=0.4,de;q=0.3,fr-FR;q=0.2,fr;q=0.1" "Accept-Language" set-header ;
  30
  31 : add-origin-header ( request origin -- request ) "Origin" set-header ;
  32
  33 ENUM: WEBSOCKET-OPCODE
  34     { WS-CONTINUE 0 }
  35     { WS-TEXT 1 }
  36     { WS-BINARY 2 }
  37     { WS-CONNECTION-CLOSE 8 }
  38     { WS-PING 9 }
  39     { WS-PONG 0xa } ;
  40
  41 : get-read-payload-length ( -- length masked? )
  42     read1 [
  43         {
  44             { [ dup 125 <= ] [ ] }
  45             { [ dup 126 = ] [ drop 2 read be> ] }
  46             { [ dup 127 = ] [ drop 8 read be> ] }
  47         } cond
  48     ] [
  49         0x80 mask?
  50     ] bi ;
  51
  52 : get-write-payload-length ( bytes -- length-byte length-bytes/f )
  53     length {
  54         { [ dup 125 <= ] [ f ] }
  55         { [ dup 0xffff <= ] [ [ drop 126 ] [ 2 >be ] bi ] }
  56         [ [ drop 127 ] [ 8 >be ] bi ]
  57     } cond ;
  58
  59 ! The final packet of a fragmented send has high bit set
  60 ! opcode should be WS-TEXT or WS-binary
  61 ! mask is a random 4 bytes to XOR with the data, optional
  62 : send-websocket-bytes ( bytes mask? opcode final? -- )
  63     0b10000000 0b0 ? bitor write1
  64     [
  65         [
  66             get-write-payload-length [ 0x80 bitor ] dip
  67             [ write1 ] [ [ write ] when* ] bi*
  68         ] [
  69             4 random-bytes
  70             [ write drop ]
  71             [ xor-crypt [ write ] when* ] 2bi
  72         ] bi
  73     ] [
  74         [ get-write-payload-length [ write1 ] [ [ write ] when* ] bi* ]
  75         [ [ write ] when* ] bi
  76     ] if flush ;
  77
  78 : send-websocket-text ( bytes mask? opcode fin? -- )
  79     [ utf8 encode ] 3dip send-websocket-bytes ;
  80
  81 : read-payload ( -- payload )
  82     get-read-payload-length [ [ 4 read ] dip read xor-crypt ] [ read ] if ;
  83
  84 : send-pong ( payload -- )
  85     t 0xa t send-websocket-bytes ;
  86
  87 SYMBOL: websocket-received
  88
  89 ERROR: unsupported-opcode n ;
  90 : read-websocket ( -- loop? obj opcode )
  91     read1 [
  92         [ 0x80 mask? drop ] [ 7 clear-bit ] bi
  93         [
  94             WEBSOCKET-OPCODE number>enum
  95             {
  96                 { f [ f "disconnected" ] }
  97                 ! { WS-CONTINUE [ t websocket-received dup get push ] }
  98                 { WS-TEXT [ t read-payload ] }
  99                 { WS-BINARY [ t read-payload utf8 decode ] }
 100                 { WS-CONNECTION-CLOSE [ f read1 ] }
 101                 { WS-PING [ t read-payload [ send-pong ] keep ] }
 102                 { WS-PONG [ t read-payload ] }
 103                 [ unsupported-opcode ]
 104             } case
 105         ] keep
 106     ] [
 107         f f f
 108     ] if* ;
 109
 110 : read-websocket-loop ( quot -- )
 111     '[
 112         websocket-received V{ } clone [
 113             read-websocket @
 114         ] with-variable
 115     ] loop ; inline
 116
 117
 118 : default-handle-websocket ( obj opcode -- )
 119     WEBSOCKET-OPCODE number>enum
 120     {
 121         { f [ [ drop "closed with error" . ] with-global ] }
 122         ! { WS-CONTINUE [  ] }
 123         { WS-TEXT [ [ [ hexdump. ] with-global ] when* ] }
 124         { WS-BINARY [ [ [ hexdump. ] with-global ] when* ] }
 125         { WS-CONNECTION-CLOSE [ [ [ . ] when* ] with-global ] }
 126         { WS-PING [ [ [ hexdump. ] with-global ] when* ] }
 127         [ 2drop ]
 128     } case ;
 129
 130 : check-websocket-upgraded? ( response -- ? )
 131     {
 132         [ code>> 101 = ]
 133         [
 134             header>> {
 135                 [ "connection" of "upgrade" = ]
 136                 [ "upgrade" of "websocket" = ]
 137             } 1&&
 138         ]
 139     } 1&& ;
 140
 141 : start-websocket ( url -- response )
 142     <get-request> add-websocket-headers http-request* drop
 143     dup check-websocket-upgraded? [ ] [ ] if ;