src/server/daemon/websocket.c

/*
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS HEADER.
 *
 * Copyright 2016 Olaf Wintermann. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *   1. Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 *
 *   2. Redistributions in binary form must reproduce the above copyright
 *      notice, this list of conditions and the following disclaimer in the
 *      documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include <stdio.h>
#include <stdlib.h>

#include "websocket.h"

#include "../util/io.h"
#include "../util/pblock.h"
#include "../util/util.h"
#include "../util/strbuf.h"
#include <ucx/string.h>

#define WS_BUFFER_LEN 2048

NSAPI_PUBLIC int http_handle_websocket(Session *sn, Request *rq, WebSocket *websocket) {
    char *connection = pblock_findkeyval(pb_key_connection, rq->headers);
    char *upgrade = pblock_findval("upgrade", rq->headers);
    char *origin = pblock_findval("origin", rq->headers);
    char *wskey = pblock_findval("sec-websocket-key", rq->headers);
    char *wsprot = pblock_findval("sec-websocket-protocol", rq->headers);
    char *wsv = pblock_findval("sec-websocket-version", rq->headers);
    
    if(!connection || !upgrade) {
        return REQ_NOACTION;
    }
    
    if(sstrcasecmp(sstr(connection), S("upgrade"))) {
        return REQ_NOACTION;
    }
    if(sstrcasecmp(sstr(upgrade), S("websocket"))) {
        return REQ_NOACTION;
    }
    
    sstr_t wsaccept = sstrcat(2, sstr(wskey), S("258EAFA5-E914-47DA-95CA-C5AB0DC85B11"));
    unsigned char hash[20];
    SHA1((const unsigned char*)wsaccept.ptr, wsaccept.length, hash);
    char *websocket_accept = util_base64encode((char*)hash, 20);
    
    sbuf_t *response = sbuf_new(512);
    sbuf_append(response, S("HTTP/1.1 101 Switching Protocols\r\n"));
    sbuf_append(response, S("Upgrade: websocket\r\n"));
    sbuf_append(response, S("Connection: Upgrade\r\n"));
    sbuf_append(response, S("Sec-WebSocket-Accept: "));
    sbuf_puts(response, websocket_accept);
    sbuf_append(response, S("\r\n\r\n"));
    
    net_write(sn->csd, response->ptr, response->length);
    sbuf_free(response);
    free(websocket_accept);
    free(wsaccept.ptr);
    
    // start websocket I/O
    WSParser *parser = websocket_parser(sn);
    
    WSFrame frame;
    
    int ret = REQ_PROCEED;
    char *inbuf = pool_malloc(sn->pool, WS_BUFFER_LEN);
    ssize_t r = 0;
    while((r = net_read(sn->csd, inbuf, WS_BUFFER_LEN)) > 0) {
        websocket_input(parser, inbuf, r);
        WSMessage *msg;
        int error;
        while((msg = websocket_get_message(parser, &error)) != NULL) {
            websocket->on_message(websocket, msg);
        }
        if(error) {
            log_ereport(LOG_FAILURE, "websocket protocol error");
            break;
        }
    }
    
    return ret;
}


WSParser* websocket_parser(Session *sn) {
    WSParser *parser = pool_malloc(sn->pool, sizeof(WSParser));
    if(!parser) {
        return NULL;
    }
    ZERO(parser, sizeof(WSParser));
    parser->pool = sn->pool;
    return parser;
}

void websocket_input(WSParser *parser, const char *data, size_t length) {
    parser->inbuf = data;
    parser->length = length;
    parser->pos = 0;
}

WSMessage* websocket_get_message(WSParser *parser, int *error) {
    WSFrame rframe;
    WSMessage *retmsg = NULL;
    
    while(parser->pos < parser->length) {
        const char *inbuf = parser->inbuf + parser->pos;
        size_t length = parser->length - parser->pos;
        
        if(parser->state == 0) {
            WSFrame frame;
            ZERO(&frame, sizeof(WSFrame));
            
            /*
             * small buffer for a websocket frame without payload data
             * I know using so many buffers it not zero copy but
             * it makes things a little bit easier :)
             */
            char frame_data[WS_FRAMEHEADER_BUFLEN];
            size_t flen = 0;
            
            /*
             * when the last call of websocket_get_message didn't completed
             * a frame header, the tmpbuf contains the remaining bytes
             * in this case we combine tmpbuf and inputbuf
             */
            if(parser->tmplen > 0) {
                memcpy(parser->tmpbuf, frame_data, parser->tmplen);
                flen = parser->tmplen;
            }
            size_t cp_remaining = length < WS_FRAMEHEADER_BUFLEN-flen ?
                length : WS_FRAMEHEADER_BUFLEN-flen;
            memcpy(&frame_data[flen], inbuf, cp_remaining);
            flen += cp_remaining;
            
            // ready to parse the frame
            ssize_t frame_hlen = websocket_get_frameheader(
                    &frame,
                    frame_data,
                    flen);
            
            if(frame_hlen == -1) {
                // protocol error, abort
                *error = 1;
                return NULL;
            }
            if(frame_hlen == 0) {
                memcpy(parser->tmpbuf, frame_data, flen);
            } else {
                inbuf += frame_hlen;
                length -= frame_hlen;
                parser->pos += frame_hlen;
                
                // frame complete, create a message object
                if(frame.payload_length > 0) {
                    WSMessage *msg = pool_malloc(parser->pool, sizeof(WSMessage));
                    msg->data = pool_malloc(parser->pool, frame.payload_length);
                    msg->length = frame.payload_length;
                    msg->next = NULL;
                    msg->type = frame.opcode;
                    
                    if(frame.payload_length >= length) {
                        // message complete
                        memcpy(msg->data, inbuf, frame.payload_length);
                        parser->pos += frame.payload_length;
                        
                        rframe = frame;
                        retmsg = msg;
                        break;
                    } else {
                        memcpy(msg->data, inbuf, length);
                        parser->state = 1;
                        parser->current = msg;
                        parser->cur_plen = length;
                        parser->frame = frame;
                        return NULL;
                    }
                }
            }
        } else {
            WSMessage *msg = parser->current;
            if(msg->length >= parser->cur_plen + length) {
                // still incomplete message
                memcpy(msg->data + parser->cur_plen, inbuf, length);
                parser->cur_plen += length;
                return NULL;
            } else {
                size_t cplen = msg->length - parser->cur_plen;
                memcpy(msg->data + parser->cur_plen, inbuf, cplen);
                parser->pos += cplen;
                parser->state = 0;
                parser->current = NULL;
                
                rframe = parser->frame;
                retmsg = msg;
                break;
            }
        }
    }
    
    if(retmsg && rframe.mask) {
        websocket_mask_data(retmsg->data, retmsg->length, rframe.masking_key);
    }
    return retmsg;
}


ssize_t websocket_get_frameheader(WSFrame *frame, const char *buf, size_t len) {
    if(len < 2) {
        return 0; // too small for anything
    }
    
/*
    printf("websocket_get_frameheader: ");
    for(int i=0;i<len;i++) {
        printf("%x ", buf[i]);
        if(len > 15) {
            break;
        }
    }
    printf("\n");
*/
    
    size_t msglen = 2; // minimal length
    
    uint8_t fin = (buf[0] & 0x80) != 0;
    uint8_t opcode = buf[0] & 0xf;
    
    uint8_t mask = (buf[1] & 0x80) != 0;
    uint8_t payload_len = buf[1] & 0x7f;
    
    uint64_t payload_length = payload_len;
    if(payload_len == 126) {
        msglen += 2;
        if(len < msglen) {
            return 0;
        }
        payload_length = *((uint16_t*)(buf+2));
    } else if(payload_len == 127) {
        msglen += 8;
        if(len < msglen) {
            return 0;
        }
        payload_length = *((uint64_t*)(buf+2));
    } else if(payload_len > 127) {
        return -1;
    }
    
    uint32_t masking_key = 0;
    if(mask) {
        msglen += 4;
        if(len < msglen) {
            return 0;
        }
        masking_key = *((uint32_t*)(buf+msglen-4));
    }
    
    frame->header_complete = TRUE;
    frame->fin = fin;
    frame->opcode = opcode;
    frame->mask = mask;
    frame->masking_key = masking_key;
    frame->payload_length = payload_length;
    
    return msglen;
}

void websocket_mask_data(char *buf, size_t len, uint32_t mask) {
    size_t m = len % 4;
    size_t alen = (len - m) / 4;
    
    uint32_t *data = (uint32_t*)buf;
    for(int i=0;i<alen;i++) {
        data[i] = data[i] ^ mask;
    }
    
    int j = 0;
    char *cmask = (char*)&mask;
    for(int i=len-m;i<len;i++) {
        buf[i] = buf[i] ^ cmask[j];
        j++;
    }
}

/* ------------------------------ public API ------------------------------*/

NSAPI_PUBLIC int websocket_send_text(SYS_NETFD csd, char *msg, size_t len) {
    char frame[WS_FRAMEHEADER_BUFLEN];
    frame[0] = 129; // 0b10000001
    size_t hlen;
    if(len < 126) {
        frame[1] = (char)len;
        hlen = 2;
    } else if(len < 65536) {
        frame[1] = 126;
        uint16_t plen = htons(len);
        memcpy(frame + 2, &plen, 2);
        hlen = 4;
    } else {
        frame[1] = 127;
        // TODO
        hlen = 10;
    }
    
    struct iovec iov[2];
    iov[0].iov_base = frame;
    iov[0].iov_len = hlen;
    
    iov[1].iov_base = msg;
    iov[1].iov_len = len;
    
    ssize_t w = net_writev(csd, iov, 2);
    if(w > 0) {
        return 0;
    } else {
        return 1;
    }
}