Christoph's web socket close fix

jks-prv · Aug 29, 2018 · b7568de · b7568de
1 parent b505925
commit b7568de
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 19 deletions.
diff --git a/pkgs/mongoose/mongoose.cpp b/pkgs/mongoose/mongoose.cpp
@@ -2621,6 +2621,12 @@ int mg_websocket_write(struct mg_connection* conn, int opcode,
     }
     free(copy);
 
+    // If we send closing frame, schedule a connection to be closed after
+    // data is drained to the client.
+    if (opcode == 0x08) {
+        MG_CONN_2_CONN(conn)->ns_conn->flags |= NSF_FINISHED_SENDING_DATA;
+    }
+
     return retval;
 }
 

diff --git a/web/web.cpp b/web/web.cpp
@@ -344,6 +344,14 @@ int web_request(struct mg_connection *mc, enum mg_event evt) {
 	if (mc->is_websocket) {
 		// This handler is called for each incoming websocket frame, one or more
 		// times for connection lifetime.
+
+        if ((mc->wsbits & 0x0F) == WS_OPCODE_CLOSE) { // close request from client
+            // respond with a close request to the client
+            // after this close request, the connection is scheduled to be closed
+            mg_websocket_write(mc, WS_OPCODE_CLOSE, mc->content, mc->content_len);
+            return MG_TRUE;
+        }
+
 		char *s = mc->content;
 		int sl = mc->content_len;
 		//printf("WEBSOCKET: len %d uri <%s>\n", sl, mc->uri);

diff --git a/web/web.h b/web/web.h
@@ -40,6 +40,7 @@ Boston, MA  02110-1301, USA.
 
 #define	WS_OPCODE_TEXT		1
 #define	WS_OPCODE_BINARY	2
+#define	WS_OPCODE_CLOSE     8
 
 #define NREQ_BUF (16*1024)		// the dx list can easily get longer than 1K
 
@@ -56,11 +57,11 @@ typedef struct {
 
 extern rx_stream_t streams[];
 
-#define	N_ADMIN			8
-#define N_CONN_SND_WF	2
+#define	N_ADMIN             8
+#define N_CONN_SND_WF_EXT   3
 
-// N_EXT below because it's possible that a user could have loaded, and idled all possible extensions
-#define	N_CONNS	(MAX_RX_CHANS * (N_CONN_SND_WF + N_EXT) + N_ADMIN)
+// +1 below for safety
+#define	N_CONNS	(MAX_RX_CHANS * (N_CONN_SND_WF_EXT + 1) + N_ADMIN)
 
 extern char *web_server_hdr;
 extern u4_t mtime_obj_keep_edata_always_o;

diff --git a/web/web_server.cpp b/web/web_server.cpp
@@ -160,23 +160,23 @@ static int ev_handler(struct mg_connection *mc, enum mg_event evt) {
 
 	//printf("ev_handler %d:%d len %d\n", mc->local_port, mc->remote_port, (int) mc->content_len);
 	//printf("MG_REQUEST: URI:%s query:%s\n", mc->uri, mc->query_string);
-	if (evt == MG_REQUEST || evt == MG_CACHE_INFO || evt == MG_CACHE_RESULT) {
-		int r = web_request(mc, evt);
-		return r;
-	} else
-	if (evt == MG_CLOSE) {
+    switch (evt) {
+    case MG_REQUEST:
+    case MG_CACHE_INFO:
+    case MG_CACHE_RESULT:
+		return web_request(mc, evt);
+    case MG_CLOSE:
 		//printf("MG_CLOSE\n");
-		rx_server_websocket(WS_MODE_CLOSE, mc);
-		mc->connection_param = NULL;
-		return MG_TRUE;
-	} else
-	if (evt == MG_AUTH) {
+        rx_server_websocket(WS_MODE_CLOSE, mc);
+        mc->connection_param = NULL;
+        return MG_TRUE;
+    case MG_AUTH:
 		//printf("MG_AUTH\n");
-		return MG_TRUE;
-	} else {
-		//printf("MG_OTHER\n");
-		return MG_FALSE;
-	}
+        return MG_TRUE;
+    default:
+		//printf("MG_OTHER evt=%d\n", evt);
+        return MG_FALSE;
+    }
 }
 
 // polled send of data _to_ web server