diff --git a/src/main.c b/src/main.c index d6e7812..84fedcd 100644 --- a/src/main.c +++ b/src/main.c @@ -33,6 +33,7 @@ struct { /************/ int main(int argc, char** argv) { + srand(time(NULL)); FILE* fp; switch(argc) { diff --git a/src/wsock.c b/src/wsock.c index facff13..10a226f 100644 --- a/src/wsock.c +++ b/src/wsock.c @@ -67,6 +67,9 @@ void buffer_append_str(buffer_t* buf, const char* data) { } int _buffer_read(buffer_t* buf, char* out, int n) { + if(buffer_length(buf) == 0) + return 0; + int i = 0; while(buf != NULL && (!n || (n && i < n))) { if(!n || buf->length <= n - i) { @@ -88,13 +91,17 @@ int buffer_peek(buffer_t* buf, char* out, int length) { int buffer_peek_str(buffer_t* buf, char* out, int length) { int read = _buffer_read(buf, out, length); - out[read] = '\0'; + if(read) + out[read] = '\0'; return read; } int buffer_read(buffer_t* buf, char** out) { int length = buffer_length(buf); - *out = malloc(length); + if(length) + *out = malloc(length); + else + *out = NULL; _buffer_read(buf, *out, 0); return length; @@ -102,8 +109,11 @@ int buffer_read(buffer_t* buf, char** out) { int buffer_read_str(buffer_t* buf, char** out) { int length = buffer_length(buf); - *out = malloc(length + 1); - (*out)[length] = '\0'; + if(length) { + *out = malloc(length + 1); + (*out)[length] = '\0'; + } else + *out = NULL; _buffer_read(buf, *out, 0); return length; @@ -160,19 +170,24 @@ 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) + // garbage in check + if((head[0] & 0x70) != 0 + || ((head[0] & 0x0F) > 0x2 && (head[0] & 0x0F) < 0x8) + || ((head[0] & 0x0F) > 0xA && (head[0] & 0x0F) <= 0xF)) return -1; - if((head[0] & 0x0F) > 0x2 && (head[0] & 0x0F) < 0x8) - return -1; - if((head[0] & 0x0F) > 0xA && (head[0] & 0x0F) <= 0xF) + // invalid parameters check + if((out->opcode == WS_CONT && + out->first_opcode == WS_OP_UNDEF) + || (out->opcode != WS_CONT && + !WS_IS_OP_CTRL(out->opcode) && + out->first_opcode != WS_OP_UNDEF) + || (WS_IS_OP_CTRL(out->opcode) && !out->fin)) return -1; + // cache first opcode of fragmented data set + if(out->opcode != WS_CONT && !WS_OP_IS_CTRL(out->opcode)) + out->first_opcode = out->opcode; int mask_start; if((head[1] & 0x7F) < 0x7E) { @@ -189,6 +204,16 @@ int _wsock_read_header(char* head, _ws_head_t* out) { memcpy(out->mask, head + mask_start, 4); } +void _wsock_mask(char* mask, char* msg, int length) { + for(int i = 0; i < length; ++i) + msg[i] = msg[i] ^ mask[i % 4]; +} + +void _wsock_gen_mask(char* mask) { + for(int i = 0; i < 4; ++i) + mask[i] = rand() % 0xFF; +} + wsock_t* wsock_open (const char* host, const char* path, uint16_t port) { @@ -264,12 +289,11 @@ wsock_t* wsock_open wsock_t* wsock = malloc(sizeof(wsock_t)); wsock->sock = sock; - wsock->mode = WS_BLOCK; + wsock->mode = WS_BLOCK | WS_FULL_RECV; wsock->head_buf.first_opcode = WS_OP_UNDEF; wsock->head_buf.opcode = WS_OP_UNDEF; wsock->recv_buf = buffer_create(); wsock->frag_buf = buffer_create(); - wsock->ctrl_buf = buffer_create(); return wsock; } @@ -337,62 +361,83 @@ int wsock_recv(wsock_t* ws, char** out) { // 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); + int buflen = buffer_length(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); + buffer_append(ws->recv_buf, tmp, got); - if(got + buflen < ws->head_buf.body_length && got >= 0) { + if(got + buflen < ws->head_buf.body_length) { got = 0; continue; } + } - + char* body; + int body_length = buffer_flush(ws->recv_buf, &packet); + if(body != NULL) { + if(ws->head_buf.masked) + _ws_mask(body, ws->head_buf.mask, body_length); + if(!WS_PACKET_IS_CTRL(ws->head_buf.opcode)) { + buffer_append(ws->frag_buf, body, body_length); + free(body); + } } switch(ws->head_buf.opcode) { case WS_CONT: case WS_TEXT: case WS_BIN: - + if(ws->head_buf.fin) { + if(ws->head_buf.first_opcode == WS_TEXT + && !(ws->mode & WS_RECV_ALL_BIN)) + got = buffer_flush_str(ws->frag_buf, out); + else + got = buffer_flush(ws->frag_buf, out); + ws->head_buf.first_opcode = WS_OP_UNDEF; + } break; case WS_CLOSE: wsock_close(ws); return -1; case WS_PING: - + wsock_send_raw(ws, WS_PONG, body, body_length); 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; - } + if(WS_OP_IS_CTRL(ws->head_buf.opcode)) + free(body); 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)) + if(WS_OP_IS_CTRL(opcode_was) || !ws->head_buf.fin) { + got = 0; continue; - else + } else break; } while(ws->mode & WS_FULL_RECV); - - return got; + + // catches a weird potential case where the server sends + // empty data frames which will exit the main loop with + // the same condition as a nonblocking call that would + // block. to prevent confusion and the proper functioning + // of the WS_FULL_RECV flag, call recursively in this case + if(ws->mode & WS_FULL_RECV && got == 0) + return wsock_recv(ws, out); + else + return got; +} + +int wsock_send_raw + (wsock_t* ws, int opcode, char* body, int length) +{ + if(!wsock_is_open(ws)) + return -1; + + } int wsock_is_open(wsock_t* ws) { @@ -413,7 +458,6 @@ void wsock_free(wsock_t* ws) { wsock_close(ws); buffer_free(ws->recv_buf); buffer_free(ws->frag_buf); - buffer_free(ws->ctrl_buf); free(ws); } diff --git a/src/wsock.h b/src/wsock.h index 6c41b3d..b8d160c 100644 --- a/src/wsock.h +++ b/src/wsock.h @@ -109,7 +109,7 @@ typedef enum { typedef struct { int sock, mode; _ws_head_t head_buf; - buffer_t *recv_buf, *frag_buf, *ctrl_buf; + buffer_t *recv_buf, *frag_buf; } wsock_t; wsock_t* wsock_open(const char*, const char*, uint16_t); @@ -117,6 +117,7 @@ void wsock_mode_set(wsock_t*, int); int wsock_mode_get(wsock_t*); int wsock_recv(wsock_t*, char**); +int wsock_send_raw(wsock_t*, int, char*, int); int wsock_send(wsock_t*, char*, int); int wsock_send_str(wsock_t*, char*);