Diff: WebSocketConnection.cpp

Revision:

10:578778037efb

diff -r 774f408b9740 -r 578778037efb WebSocketConnection.cpp
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/WebSocketConnection.cpp	Mon Mar 16 18:48:50 2015 +0900
@@ -0,0 +1,235 @@
+#include "WebSocketConnection.h"
+#include "WebSocketServer.h"
+#include "sha1.h"
+
+#define UPGRADE_WEBSOCKET	"Upgrade: websocket"
+#define SEC_WEBSOCKET_KEY	"Sec-WebSocket-Key:"
+#define MAGIC_NUMBER		"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
+#define OP_CONT		0x0
+#define OP_TEXT		0x1
+#define OP_BINARY	0x2
+#define OP_CLOSE	0x8
+#define OP_PING		0x9
+#define OP_PONG		0xA
+
+WebSocketConnection::WebSocketConnection(WebSocketServer* server)
+{
+	mServer = server;
+}
+
+WebSocketConnection::~WebSocketConnection()
+{
+}
+
+void WebSocketConnection::run()
+{
+	char buf[1024];
+	bool isWebSocket = false;
+
+	// it doesn't work...
+	// mConnection.set_blocking(true);
+
+	while (mConnection.is_connected()) {
+		int ret = mConnection.receive(buf, sizeof(buf) - 1);
+		if (ret == 0) {
+			// printf("No data to receive\r\n");
+			continue;
+		}
+		if (ret < 0) {
+			printf("ERROR: Failed to receive %d\r\n", ret);
+			break;
+		}
+		if (!isWebSocket) {
+			if (this->handleHTTP(buf, ret)) {
+				isWebSocket = true;
+			} else {
+				printf("ERROR: Non websocket\r\n");
+				break;
+			}
+		} else {
+			if (!this->handleWebSocket(buf, ret)) {
+				break;
+			}
+		}
+	}
+	// printf("Closed\r\n");
+	mConnection.close();
+}
+
+bool WebSocketConnection::handleHTTP(char* buf, int size)
+{
+	char* line = &buf[0];
+	char key[128];
+	bool isUpgradeWebSocket = false;
+	bool isSecWebSocketKeyFound = false;
+
+	for (int i = 0; i < size; i++) {
+		if (buf[i] == '\r' && i+1 < size && buf[i+1] == '\n') {
+			buf[i] = '\0';
+			if (strlen(buf) <= 0) {
+				break;
+			}
+			printf("[%s]\r\n", line);
+			if (line == &buf[0]) {
+				char* method = strtok(buf, " ");
+				char* path = strtok(NULL, " ");
+				char* version = strtok(NULL, " ");
+				// printf("[%s] [%s] [%s]\r\n", method, path, version);
+				mHandler = mServer->getHandler(path);
+				if (!mHandler) {
+					printf("ERROR: Handler not found for %s\r\n", path);
+					return false;
+				}
+			} else if (strncmp(line, UPGRADE_WEBSOCKET, strlen(UPGRADE_WEBSOCKET)) == 0) {
+				isUpgradeWebSocket = true;
+			} else if (strncmp(line, SEC_WEBSOCKET_KEY, strlen(SEC_WEBSOCKET_KEY)) == 0) {
+				isSecWebSocketKeyFound = true;
+				char* ptr = line + strlen(SEC_WEBSOCKET_KEY);
+				while (*ptr == ' ') ++ptr;
+				strcpy(key, ptr);
+			}
+			i += 2;
+			line = &buf[i];
+		}
+	}
+
+	if (isUpgradeWebSocket && isSecWebSocketKeyFound) {
+		this->sendUpgradeResponse(key);
+		if (mHandler) {
+			mHandler->onOpen();
+		}
+		mPrevFin = true;
+		return true;
+	}
+
+	return false;
+}
+
+bool WebSocketConnection::handleWebSocket(char* buf, int size)
+{
+	uint8_t* ptr = (uint8_t*)buf;
+
+	bool fin = (*ptr & 0x80) == 0x80;
+	uint8_t opcode = *ptr & 0xF;
+
+	if (opcode == OP_PING) {
+		*ptr = ((*ptr & 0xF0) | OP_PONG);
+		mConnection.send_all(buf, size);
+		return true;
+	}
+	if (opcode == OP_CLOSE) {
+		if (mHandler) {
+			mHandler->onClose();
+		}
+		return false;
+	}
+	ptr++;
+
+	if (!fin || !mPrevFin) {	
+		printf("WARN: Data consists of multiple frame not supported\r\n");
+		mPrevFin = fin;
+		return true; // not an error, just discard it
+	}
+	mPrevFin = fin;
+
+	bool mask = (*ptr & 0x80) == 0x80;
+	uint8_t len = *ptr & 0x7F;
+	ptr++;
+	
+	if (len > 125) {
+		printf("WARN: Extended payload length not supported\r\n");
+		return true; // not an error, just discard it
+	}
+
+	char* data;
+	if (mask) {
+		char* maskingKey = (char*)ptr;
+		data = (char*)(ptr + 4);
+		for (int i = 0; i < len; i++) {
+        	data[i] = data[i] ^ maskingKey[(i % 4)];
+        }
+	} else {
+		data = (char*)ptr;
+	}
+	if (mHandler) {
+		if (opcode == OP_TEXT) {
+			data[len] = '\0';
+			mHandler->onMessage(data);
+		} else if (opcode == OP_BINARY) {
+			mHandler->onMessage(data, len);
+		}
+	}
+	return true;
+}
+
+char* base64Encode(const uint8_t* data, size_t size,
+                   char* outputBuffer, size_t outputBufferSize)
+{
+	static char encodingTable[] = {'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H',
+	                               'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P',
+	                               'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X',
+	                               'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f',
+	                               'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
+	                               'o', 'p', 'q', 'r', 's', 't', 'u', 'v',
+	                               'w', 'x', 'y', 'z', '0', '1', '2', '3',
+	                               '4', '5', '6', '7', '8', '9', '+', '/'};
+    size_t outputLength = 4 * ((size + 2) / 3);
+    if (outputBufferSize - 1 < outputLength) { // -1 for NUL
+    	return NULL;
+    }
+
+    for (size_t i = 0, j = 0; i < size; /* nothing */) {
+        uint32_t octet1 = i < size ? (unsigned char)data[i++] : 0;
+        uint32_t octet2 = i < size ? (unsigned char)data[i++] : 0;
+        uint32_t octet3 = i < size ? (unsigned char)data[i++] : 0;
+
+        uint32_t triple = (octet1 << 0x10) + (octet2 << 0x08) + octet3;
+
+        outputBuffer[j++] = encodingTable[(triple >> 3 * 6) & 0x3F];
+        outputBuffer[j++] = encodingTable[(triple >> 2 * 6) & 0x3F];
+        outputBuffer[j++] = encodingTable[(triple >> 1 * 6) & 0x3F];
+        outputBuffer[j++] = encodingTable[(triple >> 0 * 6) & 0x3F];
+    }
+
+	static int padTable[] = { 0, 2, 1 };
+	int paddingCount = padTable[size % 3];
+
+    for (int i = 0; i < paddingCount; i++) {
+        outputBuffer[outputLength - 1 - i] = '=';
+    }
+    outputBuffer[outputLength] = '\0'; // NUL
+
+    return outputBuffer;
+}
+
+bool WebSocketConnection::sendUpgradeResponse(char* key)
+{
+	char buf[128];
+
+	if (strlen(key) + sizeof(MAGIC_NUMBER) > sizeof(buf)) {
+		return false;
+	}
+	strcpy(buf, key);
+	strcat(buf, MAGIC_NUMBER);
+
+    unsigned char hash[20];
+	char encoded[30];
+    sha1(buf, strlen(buf), (char*)hash);
+    base64Encode(hash, 20, encoded, sizeof(encoded));
+
+    char resp[] = "HTTP/1.1 101 Switching Protocols\r\n" \
+	    "Upgrade: websocket\r\n" \
+    	"Connection: Upgrade\r\n" \
+    	"Sec-WebSocket-Accept: XXXXXXXXXXXXXXXXXXXXXXXXXXXXX\r\n\r\n";
+    char* ptr = strstr(resp, "XXXXX");
+    strcpy(ptr, encoded);
+    strcpy(ptr+strlen(encoded), "\r\n\r\n");
+
+    int ret = mConnection.send_all(resp, strlen(resp));
+    if (ret < 0) {
+    	printf("ERROR: Failed to send response\r\n");
+    	return false;
+    }
+
+    return true;
+}