diff --git a/src/wsock.c b/src/wsock.c index 5ac4b44..facff13 100644 --- a/src/wsock.c +++ b/src/wsock.c @@ -149,21 +149,6 @@ void buffer_free(buffer_t* buf) { /** WEB SOCKET **/ /******************/ -typedef struct { - uint8_t fin, opcode, masked, head_length; - uint8_t mask[4]; - uint64_t body_length; -} _ws_head_t; - -typedef enum { - WS_CONT = 0x0, - WS_TEXT = 0x1, - WS_BIN = 0x2, - WS_CLOSE = 0x8, - WS_PING = 0x9, - WS_PONG = 0xA -} _ws_opcode_t; - int _wsock_header_length(char* head) { return 2 + ((head[1] & 0x80) == 0 ? 0 : 4) @@ -175,6 +160,11 @@ int _wsock_read_header(char* head, _ws_head_t* out) { out->head_length = _wsock_header_length(head); out->fin = 0x80 & head[0]; out->opcode = 0x0F & head[0]; + if(out->opcode == WS_CONT && + out->first->opcode == WS_OP_UNDEF) + return -1; + if(out->opcode != WS_CONT && !WS_OP_IS_CTRL(out->opcode)) + out->first_opcode = out->opcode; out->masked = 0x80 & head[1]; if((head[0] & 0x70) != 0) @@ -275,8 +265,11 @@ wsock_t* wsock_open wsock_t* wsock = malloc(sizeof(wsock_t)); wsock->sock = sock; wsock->mode = WS_BLOCK; + wsock->head_buf.first_opcode = WS_OP_UNDEF; + wsock->head_buf.opcode = WS_OP_UNDEF; wsock->recv_buf = buffer_create(); - wsock->send_buf = buffer_create(); + wsock->frag_buf = buffer_create(); + wsock->ctrl_buf = buffer_create(); return wsock; } @@ -288,13 +281,117 @@ int wsock_mode_get(wsock_t* ws) { return ws->mode; } -int wsock_recv(wsock_t* ws, char** out) { - char buffer[4096]; - int got = 0; - do { - //if((ws->mode & WS_NOBLOCK) != 0) +int _srecv(wsock_t* ws, char* out, int len, int flags) { + int got = recv(ws->sock, out, len, flags); + if(got == EAGAIN || got == EWOULDBLOCK) + return 0; + else if(got <= 0) { + wsock_close(ws); + return -1; + } else + return got; +} - } while((ws->mode & WS_FULL_RECV) != 0); +int wsock_recv(wsock_t* ws, char** out) { + if(!wsock_is_open(ws)) + return -1; + + char tmp[WS_INTERNAL_BUFLEN]; + int got, flags; + do { + if(ws->head_buf.opcode == WS_OP_UNDEF) { + flags = MSG_PEEK | + ((ws->mode & WS_NOBLOCK) ? MSG_DONTWAIT : 0); + got = _srecv(ws, tmp, 2, flags); + if(got < 2 && got >= 0) { + got = 0; + continue; + } else if(got < 0) + return got; + + int head_length = _wsock_header_length(tmp); + got = _srecv(ws, tmp, head_length, flags); + if(got < head_length && got >= 0) { + got = 0; + continue; + } else if(got < 0) + return got; + + flags &= ~MSG_PEEK; + got = _srecv(ws, tmp, head_length, flags); + if(got < 0) + return got; + if(!_wsock_read_header(tmp, &ws->head_buf) + || ws->head_buf.body_length > WS_MAX_LENGTH) { + wsock_close(ws); + return -1; + } + } + + if(ws->head_buf.body_length > 0) { + // todo: complex fetching algorithm that queries + // internal recv buffer if more info is available + // and continues to buffer data in the same call + // if WS_FULL_RECV is not enabled, however i have + // no current use case for this and this will still + // work provided wsock_recv is called enough times + // that the wsock buffer eventually received all + // information off of the recv buffer + int buflen = buffer_length( + WS_OP_IS_CTRL(ws->head_buf.opcode) + ? ws->ctrl_buf : ws->recv_buf); + flags = (ws->mode & WS_NOBLOCK) ? MSG_DONTWAIT : 0; + got = _srecv(ws, tmp, __MIN(sizeof(tmp), ws->head_buf.body_length - buflen), flags); + + if(got < 0) + return got; + else if(got > 0) + buffer_append((WS_OP_IS_CTRL(ws->head_buf.opcode) + ? ws->ctrl_buf : ws->recv_buf), tmp, got); + + if(got + buflen < ws->head_buf.body_length && got >= 0) { + got = 0; + continue; + } + + + } + + switch(ws->head_buf.opcode) { + case WS_CONT: + case WS_TEXT: + case WS_BIN: + + break; + case WS_CLOSE: + wsock_close(ws); + return -1; + case WS_PING: + + break; + } + + if(ws->head_buf.fin) { + if(!WS_OP_IS_CTRL(ws->opcode)) { + buffer_clear(ws->recv_buf); + ws->head_buf.first_opcode = WS_OP_UNDEF; + } else + buffer_clear(ws->ctrl_buf); + } else if(WS_OP_IS_CTRL(ws->opcode)) { + wsock_close(ws); + return -1; + } + + int opcode_was = ws->head_buf.opcode; + ws->head_buf.opcode = WS_OP_UNDEF; + if(WS_OP_IS_CTRL(opcode_was) || + (!WS_OP_IS_CTRL(opcode_was) && + ws->head_buf.body_length == 0)) + continue; + else + break; + } while(ws->mode & WS_FULL_RECV); + return got; } @@ -303,10 +400,20 @@ int wsock_is_open(wsock_t* ws) { } void wsock_close(wsock_t* ws) { - // todo: send close frame + if(!wsock_is_open(ws)) + return; + + const char close_frame[2] = {0x88, 0x00}; + send(ws->sock, close_frame, sizeof(close_frame)); shutdown(ws->sock, SHUT_RDWR); + ws->sock = -1; +} + +void wsock_free(wsock_t* ws) { + wsock_close(ws); buffer_free(ws->recv_buf); - buffer_free(ws->send_buf); + buffer_free(ws->frag_buf); + buffer_free(ws->ctrl_buf); free(ws); } diff --git a/src/wsock.h b/src/wsock.h index 5dd02ce..6c41b3d 100644 --- a/src/wsock.h +++ b/src/wsock.h @@ -11,6 +11,9 @@ /** UTILITIES **/ /*****************/ +#define __MAX(A, B) ((A) > (B) ? (A) : (B)) +#define __MIN(A, B) ((A) < (B) ? (A) : (B)) + void swap_order(char*, int); void swap_order_copy(char*, char*, int); @@ -46,6 +49,11 @@ void buffer_free(buffer_t*); /** WEB SOCKET **/ /******************/ +#define WS_MAX_LENGTH 0xFFFFFF +#define WS_INTERNAL_BUFLEN 4096 +#define WS_OP_UNDEF 0xFF +#define WS_OP_IS_CTRL(A) (A >= WS_CLOSE) + typedef enum { // block on system calls WS_BLOCK = 0, @@ -83,9 +91,25 @@ typedef enum { // has an EXTREME THIRST for partial sends, they can // implement this behavior themselves +typedef struct { + uint8_t fin, opcode, first_opcode, + masked, mask[4], head_length; + uint64_t body_length; +} _ws_head_t; + +typedef enum { + WS_CONT = 0x0, + WS_TEXT = 0x1, + WS_BIN = 0x2, + WS_CLOSE = 0x8, + WS_PING = 0x9, + WS_PONG = 0xA +} _ws_opcode_t; + typedef struct { int sock, mode; - buffer_t *recv_buf, *send_buf; + _ws_head_t head_buf; + buffer_t *recv_buf, *frag_buf, *ctrl_buf; } wsock_t; wsock_t* wsock_open(const char*, const char*, uint16_t); @@ -98,4 +122,4 @@ int wsock_send_str(wsock_t*, char*); int wsock_is_open(wsock_t*); void wsock_close(wsock_t*); - +void wsock_free(wsock_t*);