clii/src/wsock.c

#include "wsock.h"

/**  UTILITIES  **/
/*****************/

#define WS_REORDER_BUFLEN (sizeof(long double))

static int _sys_is_net_order() {
	const uint16_t test = 0xB00B;
	return ((char*)&test)[0] == 0xB0;
}

void swap_order(char* in, int length) {
	if(_sys_is_net_order() || length > WS_REORDER_BUFLEN)
		return;

	char buf[WS_REORDER_BUFLEN];
	for(int i = 0; i < length; ++i)
		buf[i] = in[length - i - 1];
	memcpy(in, buf, length);
}

void swap_order_copy(char* in, char* out, int length) {
	memcpy(out, in, length);
	swap_order(out, length);
}


/**  DATA BUFFER  **/
/*******************/

buffer_t* buffer_create() {
	buffer_t* buf = malloc(sizeof(buffer_t));
	buf->data = NULL;
	buf->next = NULL;
	buf->length = 0;

	return buf;
}

int buffer_is_empty(buffer_t* buf) {
	return buf->data == NULL;
}

int buffer_length(buffer_t* buf) {
	if(buf == NULL)
		return 0;

	return buf->length + buffer_length(buf->next);
}

void buffer_append
	(buffer_t* buf, const char* data, int length)
{
	if(buf->data != NULL) {
		if(buf->next == NULL)
			buf->next = buffer_create();
		buffer_append(buf->next, data, length);
	} else {
		buf->data = malloc(length);
		strncpy(buf->data, data, length);
		buf->length = length;
	}
}

void buffer_append_str(buffer_t* buf, const char* data) {
	buffer_append(buf, data, strlen(data));
}

static 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) {
			strncpy(out + i, buf->data, buf->length);
			i += buf->length;
			buf = buf->next;
		} else if(n && buf->length > n - i) {
			strncpy(out + i, buf->data, n - i);
			i += n - i;
		}
	}

	return i;
}

int buffer_peek(buffer_t* buf, char* out, int length) {
	return _buffer_read(buf, out, length);
}

int buffer_peek_str(buffer_t* buf, char* out, int length) {
	int read = _buffer_read(buf, out, length);
	if(read)
		out[read] = '\0';
	return read;
}

int buffer_read(buffer_t* buf, char** out) {
	int length = buffer_length(buf);
	if(length)
		*out = malloc(length);
	else
		*out = NULL;

	_buffer_read(buf, *out, 0);
	return length;
}

int buffer_read_str(buffer_t* buf, char** out) {
	int length = buffer_length(buf);
	if(length) {
		*out = malloc(length + 1);
		(*out)[length] = '\0';
	} else
		*out = NULL;

	_buffer_read(buf, *out, 0);
	return length;
}

int buffer_flush(buffer_t* buf, char** out) {
	int length = buffer_read(buf, out);
	buffer_clear(buf);
	return length;
}

int buffer_flush_str(buffer_t* buf, char** out) {
	int length = buffer_read_str(buf, out);
	buffer_clear(buf);
	return length;
}

static void _buffer_clear(buffer_t* buf, int root) {
	if(buf == NULL)
		return;
	_buffer_clear(buf->next, 0);

	free(buf->data);
	if(!root)
		free(buf);
	else {
		buf->data = NULL;
		buf->next = NULL;
		buf->length = 0;
	}
}

void buffer_clear(buffer_t* buf) {
	_buffer_clear(buf, 1);
}

void buffer_free(buffer_t* buf) {
	buffer_clear(buf);
	free(buf);
}


/**  WEB SOCKET  **/
/******************/

static int _wsock_header_length(char* head) {
	return 2 
		+ ((head[1] & 0x80) == 0 ? 0 : 4)
		+ ((head[1] & 0x7F) < 0x7E ? 0
			: ((head[1] & 0x7F) == 0x7E ? 2 : 8));
}

static 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];
	out->masked = 0x80 & head[1];

	// 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 0;
	// invalid parameters check
	if((out->opcode == WS_CONT && 
				out->first_opcode == WS_OP_UNDEF)
			|| (out->opcode != WS_CONT &&
				!WS_OP_IS_CTRL(out->opcode) && 
				out->first_opcode != WS_OP_UNDEF)
			|| (WS_OP_IS_CTRL(out->opcode) && !out->fin))
		return 0;
	// 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) {
		out->body_length = head[1] & 0x7F;
		mask_start = 2;
	} else if((head[1] & 0x7F) == 0x7E) {
		swap_order_copy(head + 2, (char*)&out->body_length, 2);
		mask_start = 4;
	} else {
		swap_order_copy(head + 2, (char*)&out->body_length, 8);
		mask_start = 10;
	}
	if(out->masked)
		memcpy(out->mask, head + mask_start, 4);

	return 1;
}

static void _wsock_mask(char* mask, char* msg, int length) {
	for(int i = 0; i < length; ++i)
		msg[i] = msg[i] ^ mask[i % 4];
}

static 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)
{
	char s_port[6];
	sprintf(s_port, "%i", port);

	int sock;
	struct addrinfo hints, *result, *rp;
	memset(&hints, 0, sizeof(hints));
	hints.ai_family		= AF_UNSPEC;
	hints.ai_socktype	= SOCK_STREAM;
	hints.ai_flags		= 0;
	hints.ai_protocol	= 0;

	if(getaddrinfo(host, s_port, &hints, &result) != 0)
		return NULL;
	for(rp = result; rp != NULL; rp = rp->ai_next) {
		if((sock = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol)) == -1)
			continue;
		if(connect(sock, rp->ai_addr, rp->ai_addrlen) != -1)
			break;
	}
	freeaddrinfo(result);
	if(rp == NULL)
		return NULL;
	
	// todo: randomize wsock key
	char s_shake[4096];
	sprintf(s_shake, 
			"GET %s HTTP/1.1\r\n"
			"Host: %s:%i\r\n"
			"Upgrade: websocket\r\n"
			"Connection: Upgrade\r\n"
			"Sec-WebSocket-Key: AQIDBAUGBwgJCgsMDQ4PEC==\r\n"
			"Sec-WebSocket-Version: 13\r\n\r\n", 
		path,
		host, port
	);
	if(send(sock, s_shake, strlen(s_shake), MSG_NOSIGNAL) <= 0) {
		shutdown(sock, SHUT_RDWR);
		return NULL;
	}
	s_shake[0] = '\0';

	struct pollfd pfd;
	pfd.fd = sock;
	pfd.events = POLLIN;

	int total = 0;
	char* end;
	while((end = strstr(s_shake, "\r\n\r\n")) == NULL) {
		if(total + 1 >= sizeof(s_shake) 
				|| poll(&pfd, 1, 5000) <= 0) 
		{
			shutdown(sock,SHUT_RDWR);
			return NULL;
		}

		total = recv(sock, s_shake, sizeof(s_shake) - 1, MSG_PEEK);
		s_shake[total] = '\0';
	}

	total = (end - s_shake) + 4;
	recv(sock, s_shake, total, 0);
	s_shake[total + 1] = '\0'; 

	const char resp_check[] = "HTTP/1.1 101 Switching Protocols";
	if(strncmp(s_shake, resp_check, sizeof(resp_check) - 1) != 0)
	{
		shutdown(sock, SHUT_RDWR);
		return NULL;
	}

	wsock_t* wsock = malloc(sizeof(wsock_t));
	wsock->sock = sock;
	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();
	return wsock;
}

void wsock_mode_set(wsock_t* ws, int mode) {
	ws->mode = mode | WS_FULL_SEND;
}

int wsock_mode_get(wsock_t* ws) {
	return ws->mode;
}

static 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;
}

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->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->recv_buf, tmp, got);
	
			if(got + buflen < ws->head_buf.body_length) {
				got = 0;
				continue;
			}
		}

		char* body;
		int body_length = buffer_flush(ws->recv_buf, &body);
		if(body != NULL) {
			if(ws->head_buf.masked)
				_wsock_mask((char*)ws->head_buf.mask, 
						body, body_length);
			if(!WS_OP_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;
			case WS_PONG:
				break;
		}

		if(WS_OP_IS_CTRL(ws->head_buf.opcode)) {
			got = 0;
			free(body);
		}

		ws->head_buf.opcode = WS_OP_UNDEF;
		if(!ws->head_buf.fin) {
			got = 0;
			continue;
		} else
			break;
	} while(ws->mode & WS_FULL_RECV);

	return got;
}

// todo: support sending fragments at some point
int wsock_send_raw
	(wsock_t* ws, int opcode, char* body, uint64_t length)
{
	if(!wsock_is_open(ws))
		return -1;

	char header[14], *mask;
	int header_bsize_len = 
		(length < 0x7E) ? 0 :
		((length <= 0xFFFF) ? 2 : 8);
	int header_length = 
		2 + (length > 0 ? 4 : 0) + header_bsize_len;
	header[0] = 0x80 | opcode;
	header[1] = 0x80 | ((length < 0x7E) ? length
		: (length <= 0xFFFF ? 0x7E : 0x7F));

	uint16_t length_u16 = (uint16_t)length;
	if(header_bsize_len > 0)
		swap_order_copy(
				(header_bsize_len == 2 
					? (char*)&length_u16 : (char*)&length),
				header + 2, header_bsize_len);

	if(length > 0) {
		mask = header + 2 + header_bsize_len;
		_wsock_gen_mask(mask);
	}
	
	if(send(ws->sock, header, 
			header_length, MSG_NOSIGNAL) <= 0) 
	{
		wsock_close(ws);
		return -1;
	}

	if(length > 0) {
		_wsock_mask(mask, body, length);
		if(send(ws->sock, body, length, MSG_NOSIGNAL) <= 0) {
			wsock_close(ws);
			return -1;
		}
		_wsock_mask(mask, body, length);	
	}

	return header_length + length;
}

int wsock_send(wsock_t* ws, char* body, int length) {
	return wsock_send_raw(ws, WS_BIN, body, length);
}

int wsock_send_str(wsock_t* ws, char* body) {
	if(ws->mode & WS_SEND_ALL_BIN)
		return wsock_send(ws, body, strlen(body));
	else
		return wsock_send_raw(ws, WS_TEXT, body, strlen(body));
}

int wsock_ping(wsock_t* ws, char* msg) {
	if(msg == NULL)
		return wsock_send_raw(ws, WS_PING, NULL, 0);
	else
		return wsock_send_raw(ws, WS_PING, msg, strlen(msg));
}

int wsock_is_open(wsock_t* ws) {
	return ws->sock >= 0;
}

void wsock_close(wsock_t* ws) {
	if(!wsock_is_open(ws))
		return;

	const char close_frame[2] = {0x88, 0x00};
	send(ws->sock, close_frame, 
			sizeof(close_frame), MSG_NOSIGNAL);
	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->frag_buf);
	free(ws);
}