Diff: WebSocketServer.cpp

Revision:

5:ce3f1dd90068

Parent:

1:84af7a219830

Child:

6:79ecd4e53456

--- a/WebSocketServer.cpp	Fri Mar 13 06:57:04 2015 +0000
+++ b/WebSocketServer.cpp	Fri Mar 13 20:27:06 2015 +0900
@@ -1,4 +1,5 @@
 #include "WebSocketServer.h"
+// #include "sha1.h"
 
 WebSocketServer::WebSocketServer() {
 }
@@ -34,18 +35,20 @@
 			continue;
 		}
 		printf("New connection\r\n");
+		connection.set_blocking(true);
+
 		while (connection.is_connected()) {
 			ret = connection.receive(buf, sizeof(buf) - 1);
-			if (ret <= 0) {
-				if (ret == 0) {
-					printf("closed\r\n");
-				} else {
-					printf("ERROR: Failed to receive %d\r\n", ret);
-				}
+			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)) {
+				if (this->handleHTTP(buf, ret, connection)) {
 					isWebSocket = true;
 				} else {
 					printf("ERROR: Non websocket\r\n");
@@ -57,6 +60,7 @@
 				}
 			}
 		}
+		printf("closed\r\n");
 		connection.close();
 	}
 }
@@ -65,8 +69,15 @@
 	return true;
 }
 
-bool WebSocketServer::handleHTTP(char* buf, int size) {
+#define UPGRADE_WEBSOCKET	"Upgrade: websocket"
+#define SEC_WEBSOCKET_KEY	"Sec-WebSocket-Key:"
+#define MAGIC_NUMBER		"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
+
+bool WebSocketServer::handleHTTP(char* buf, int size, TCPSocketConnection& connection) {
 	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') {
@@ -75,14 +86,116 @@
 				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);
+
+			} else if (strncmp(line, UPGRADE_WEBSOCKET, strlen(UPGRADE_WEBSOCKET)) == 0) {
+				isUpgradeWebSocket = true;
+				printf("%s found\r\n", UPGRADE_WEBSOCKET);
+			} else if (strncmp(line, SEC_WEBSOCKET_KEY, strlen(SEC_WEBSOCKET_KEY)) == 0) {
+				printf("%s found\r\n", SEC_WEBSOCKET_KEY);
+				isSecWebSocketKeyFound = true;
+				char* ptr = line + strlen(SEC_WEBSOCKET_KEY);
+				while (*ptr == ' ') ++ptr;
+				strcpy(ptr, key);
+				printf("key=%s\r\n", key);
+			}
 			i += 2;
 			line = &buf[i];
 		}
 	}
 
+	if (isUpgradeWebSocket && isSecWebSocketKeyFound) {
+		this->sendUpgradeResponse(key, connection);
+		return true;
+	}
+
+	return false;
+}
+
+bool WebSocketServer::handleWebSocket(char* buf, int size) {
+	buf[size] = '\0';
+	printf("%s\r\n", buf);
 	return true;
 }
 
-bool WebSocketServer::handleWebSocket(char* buf, int size) {
-	return true;
+static char encoding_table[] = {'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', '+', '/'};
+static int mod_table[] = {0, 2, 1};
+
+char *base64_encode(const unsigned char *data,
+                    size_t input_length,
+                    // size_t *output_length
+                    char* encoded_data,
+                    size_t buffer_length) {
+
+    // *output_length = 4 * ((input_length + 2) / 3);
+    size_t output_length = 4 * ((input_length + 2) / 3);
+    if (buffer_length - 1 < output_length) {
+    	return NULL;
+    }
+
+    // char *encoded_data = malloc(*output_length);
+    // if (encoded_data == NULL) return NULL;
+
+    for (int i = 0, j = 0; i < input_length;) {
+
+        uint32_t octet_a = i < input_length ? (unsigned char)data[i++] : 0;
+        uint32_t octet_b = i < input_length ? (unsigned char)data[i++] : 0;
+        uint32_t octet_c = i < input_length ? (unsigned char)data[i++] : 0;
+
+        uint32_t triple = (octet_a << 0x10) + (octet_b << 0x08) + octet_c;
+
+        encoded_data[j++] = encoding_table[(triple >> 3 * 6) & 0x3F];
+        encoded_data[j++] = encoding_table[(triple >> 2 * 6) & 0x3F];
+        encoded_data[j++] = encoding_table[(triple >> 1 * 6) & 0x3F];
+        encoded_data[j++] = encoding_table[(triple >> 0 * 6) & 0x3F];
+    }
+
+    for (int i = 0; i < mod_table[input_length % 3]; i++)
+        // encoded_data[*output_length - 1 - i] = '=';
+        encoded_data[output_length - 1 - i] = '=';
+    encoded_data[output_length] = '\0';
+
+    return encoded_data;
 }
+
+bool WebSocketServer::sendUpgradeResponse(char* key, TCPSocketConnection& connection) {
+	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((unsigned char*)buf, strlen(buf), hash);
+    base64_encode(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 = connection.send_all(resp, strlen(resp));
+    if (ret < 0) {
+    	printf("ERROR: Failed to send response\r\n");
+    	return false;
+    }
+
+    return true;
+}