! Copyright (C) 2021 Doug Coleman.
! See http://factorcode.org/license.txt for BSD license.
USING: accessors assocs base64 byte-arrays combinators
crypto.xor http http.client io io.binary io.encodings.string
io.encodings.utf8 kernel math math.bitwise multiline namespaces
prettyprint random sequences strings tools.hexdump ;
IN: http.websockets

! TODO: multiplexing

CONSTANT: websocket-version "13"

: random-websocket-key ( -- base64 )
    16 random-bytes >base64 >string ;

: add-websocket-headers ( request -- request )
    "connection" over header>> delete-at
    "Upgrade" "Connection" set-header
    "no-cache" "Pragma" set-header
    "no-cache" "Cache-Control" set-header
    "websocket" "Upgrade" set-header
    "http://www.websocket.org" "Origin" set-header
    websocket-version "Sec-WebSocket-Version" set-header
    random-websocket-key "Sec-WebSocket-Key" set-header
    "permessage-deflate; client_max_window_bits" "Sec-WebSocket-Extensions" set-header
    "gzip, deflate" "Accept-Encoding" set-header
    "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 ;

CONSTANT: websocket-opcode-continue-frame 0
CONSTANT: websocket-opcode-text-frame 1
CONSTANT: websocket-opcode-binary-frame 2
CONSTANT: websocket-opcode-connection-close-frame 8
CONSTANT: websocket-opcode-ping-frame 9
CONSTANT: websocket-opcode-pong-frame 0xa

: get-read-payload-length ( -- length masked? )
    read1 [
        {
            { [ dup 125 <= ] [ ] }
            { [ dup 126 = ] [ drop 2 read be> ] }
            { [ dup 127 = ] [ drop 8 read be> ] }
        } cond
    ] [
        0x80 mask?
    ] bi ;

: get-write-payload-length ( bytes -- length-byte length-bytes/f )
    length {
        { [ dup 125 <= ] [ f ] }
        { [ dup 0xffff <= ] [ [ drop 126 ] [ 2 >be ] bi ] }
        [ [ drop 127 ] [ 8 >be ] bi ]
    } cond ;

! : send-websocket-fragmented ( bytes opcode -- ) 0b10000000 bitor

: send-websocket-bytes ( bytes mask? opcode final? -- )
    0b10000000 0b0 ? bitor write1
    [
        [
            get-write-payload-length [ 0x80 bitor ] dip
            [ write1 ] [ [ write ] when* ] bi*
        ] [
            4 random-bytes
            [ write drop ]
            [ xor-crypt [ write ] when* ] 2bi
        ] bi
    ] [
        [ get-write-payload-length [ write1 ] [ [ write ] when* ] bi* ]
        [ [ write ] when* ] bi
    ] if flush ;

: send-websocket-text ( bytes mask? opcode fin? -- )
    [ utf8 encode ] 3dip send-websocket-bytes ;

: read-payload ( -- payload )
    get-read-payload-length [ [ 4 read ] dip read xor-crypt ] [ read ] if ;

: send-pong ( payload -- )
    t 0xa t send-websocket-bytes ;

ERROR: unsupported-opcode n ;
: read-websocket ( -- loop? obj opcode )
    read1 [
        [ 0x80 mask? drop ] [ 7 clear-bit ] bi
        [
            {
                { f [ f "disconnected" ] }
                { 0 [ t unsupported-opcode ] }
                { 1 [ t read-payload ] }
                { 2 [ t read-payload utf8 decode ] }
                { 8 [ f read1 ] }
                { 9 [ t read-payload [ send-pong ] keep ] }
                { 0xa [ t read-payload ] }
                [ unsupported-opcode ]
            } case
        ] keep
    ] [
        f f f
    ] if* ;

: read-websocket-loop ( quot -- )
    '[
        read-websocket @
    ] loop ; inline

: handle-websocket ( obj opcode -- )
    {
        { f [ [ drop "closed with error" . ] with-global ] }
        { 1 [ [ [ hexdump. ] with-global ] when* ] }
        { 2 [ [ [ hexdump. ] with-global ] when* ] }
        { 8 [ [ [ hexdump. ] when* "ping received" . flush ] with-global ] }
        { 9 [ [ [ hexdump. ] with-global ] when* ] }
        [ 2drop ]
    } case ;



![[
"echo.websocket.org/?encoding=text" <get-request> add-websocket-headers
start-websocket
stream>> [ in>> stream>> ] [ out>> stream>> ] bi
[
    [ handle-websocket ] read-websocket-loop
] with-streams
]]