From 22877d8f2fbe25c0c64059d8fc0026205b63dee9 Mon Sep 17 00:00:00 2001
From: tnoho <teaotya@gmail.com>
Date: Fri, 16 Oct 2015 16:30:49 +0900
Subject: [PATCH] Fix masking method

---
 examples/websocket_client/Makefile           |  2 +-
 examples/websocket_client/websocket_client.c | 11 +++---
 src/civetweb.c                               | 38 ++++++++------------
 3 files changed, 22 insertions(+), 29 deletions(-)

diff --git a/examples/websocket_client/Makefile b/examples/websocket_client/Makefile
index 3c70bf60..e3ef1346 100644
--- a/examples/websocket_client/Makefile
+++ b/examples/websocket_client/Makefile
@@ -28,7 +28,7 @@ $(PROG): $(CIVETWEB_LIB) $(SRC)
 	$(CC) -o $@ $(CFLAGS) $(LDFLAGS) $(SRC) $(CIVETWEB_LIB) $(LIBS)
 
 $(CIVETWEB_LIB):
-	$(MAKE) -C $(TOP) clean lib
+	$(MAKE) -C $(TOP) clean lib WITH_WEBSOCKET=1
 	cp $(TOP)/$(CIVETWEB_LIB) .
 
 clean:
diff --git a/examples/websocket_client/websocket_client.c b/examples/websocket_client/websocket_client.c
index c42cda77..f9756a84 100644
--- a/examples/websocket_client/websocket_client.c
+++ b/examples/websocket_client/websocket_client.c
@@ -82,6 +82,9 @@ websocket_server_data(struct mg_connection *conn,
 #endif
 {
 	printf("Server: Got %u bytes from the client\n", data_len);
+	printf("Server received data from client: ");
+	fwrite(data, 1, data_len, stdout);
+	printf("\n");
 
 	if (data_len < 3 || 0 != memcmp(data, "bye", 3)) {
 		/* Send websocket acknowledge message */
@@ -258,7 +261,7 @@ main(int argc, char *argv[])
 	client1_data.data = NULL;
 	client1_data.len = 0;
 
-	mg_websocket_write(newconn1, WEBSOCKET_OPCODE_TEXT, "data1", 5);
+	mg_websocket_client_write(newconn1, WEBSOCKET_OPCODE_TEXT, "data1", 5);
 
 	sleep(1); /* Should get the acknowledge message */
 	assert(client1_data.closed == 0);
@@ -305,7 +308,7 @@ main(int argc, char *argv[])
 	client2_data.data = NULL;
 	client2_data.len = 0;
 
-	mg_websocket_write(newconn1, WEBSOCKET_OPCODE_TEXT, "data2", 5);
+	mg_websocket_client_write(newconn1, WEBSOCKET_OPCODE_TEXT, "data2", 5);
 
 	sleep(1); /* Should get the acknowledge message */
 	assert(client1_data.closed == 0);
@@ -321,7 +324,7 @@ main(int argc, char *argv[])
 	client1_data.data = NULL;
 	client1_data.len = 0;
 
-	mg_websocket_write(newconn1, WEBSOCKET_OPCODE_TEXT, "bye", 3);
+	mg_websocket_client_write(newconn1, WEBSOCKET_OPCODE_TEXT, "bye", 3);
 
 	sleep(1); /* Should get the goodbye message */
 	assert(client1_data.closed == 0);
@@ -347,7 +350,7 @@ main(int argc, char *argv[])
 	assert(client2_data.data == NULL);
 	assert(client2_data.len == 0);
 
-	mg_websocket_write(newconn2, WEBSOCKET_OPCODE_TEXT, "bye", 3);
+	mg_websocket_client_write(newconn2, WEBSOCKET_OPCODE_TEXT, "bye", 3);
 
 	sleep(1); /* Should get the goodbye message */
 	assert(client1_data.closed == 1);
diff --git a/src/civetweb.c b/src/civetweb.c
index 30987599..2443e37c 100755
--- a/src/civetweb.c
+++ b/src/civetweb.c
@@ -8081,14 +8081,10 @@ mg_websocket_write_exec(struct mg_connection *conn,
                         int opcode,
                         const char *data,
                         size_t dataLen,
-                        char mask)
+                        uint32_t masking_key)
 {
 	unsigned char header[14];
 	size_t headerLen = 1;
-	uint8_t masking_key[4];
-	unsigned char masked_data[dataLen];
-	const char* send_data;
-	size_t i = 0;
 
 	int retval = -1;
 
@@ -8112,26 +8108,11 @@ mg_websocket_write_exec(struct mg_connection *conn,
 		headerLen = 10;
 	}
 
-	if(mask) {
+	if(masking_key) {
 		/* add mask */
-		srand(time(NULL));
-		masking_key[0] = rand() & 0xff;
-		masking_key[1] = rand() & 0xff;
-		masking_key[2] = rand() & 0xff;
-		masking_key[3] = rand() & 0xff;
-
 		header[1] |= 0x80;
-		header[headerLen] = masking_key[0];
-		header[headerLen + 1] = masking_key[1];
-		header[headerLen + 2] = masking_key[2];
-		header[headerLen + 3] = masking_key[3];
+		*(uint32_t *)(void *)(header + headerLen) = masking_key;
 		headerLen += 4;
-		for (i = 0; i != dataLen; ++i) {
-			masked_data[i] = *data++ ^ masking_key[i&0x3];
-		}
-		send_data = (const char*)masked_data;
-	} else {
-		send_data = data;
 	}
 
 
@@ -8143,7 +8124,7 @@ mg_websocket_write_exec(struct mg_connection *conn,
 	(void)mg_lock_connection(conn);
 	retval = mg_write(conn, header, headerLen);
 	if (dataLen > 0) {
-		retval = mg_write(conn, send_data, dataLen);
+		retval = mg_write(conn, data, dataLen);
 	}
 	mg_unlock_connection(conn);
 
@@ -8165,7 +8146,16 @@ mg_websocket_client_write(struct mg_connection *conn,
                           const char *data,
                           size_t dataLen)
 {
-	return mg_websocket_write_exec(conn, opcode, data, dataLen, 1);
+	int retval = -1;
+	size_t i = 0;
+	uint32_t masking_key = 0x1594DAC0;
+	char* masked_data = (char*)mg_malloc(dataLen + 4);
+	for (i = 0; i < dataLen; i+= 4) {
+		*(uint32_t *)(void *)(masked_data + i) = *(uint32_t *)(void *)(data + i) ^ masking_key;
+	}
+	retval = mg_websocket_write_exec(conn, opcode, masked_data, dataLen, masking_key);
+	mg_free(masked_data);
+	return retval;
 }
 
 static void