pmd: handle case we are already on drain list

Provide internal helper for adding to list that takes care of the
case we are already on the list.

https://github.com/warmcat/libwebsockets/issues/847

lib/client-parser.c

@@ -23,7 +23,6 @@
 
 int lws_client_rx_sm(struct lws *wsi, unsigned char c)
 {
-	struct lws_context_per_thread *pt = &wsi->context->pt[(int)wsi->tsi];
 	int callback_action = LWS_CALLBACK_CLIENT_RECEIVE;
 	int handled, n, m, rx_draining_ext = 0;
 	unsigned short close_code;
@@ -31,21 +30,10 @@ int lws_client_rx_sm(struct lws *wsi, unsigned char c)
 	unsigned char *pp;
 
 	if (wsi->u.ws.rx_draining_ext) {
-		struct lws **w = &pt->rx_draining_ext_list;
-		lwsl_ext("%s: RX EXT DRAINING: Removing from list\n", __func__, c);
 		assert(!c);
 		eff_buf.token = NULL;
 		eff_buf.token_len = 0;
-		wsi->u.ws.rx_draining_ext = 0;
-		/* remove us from context draining ext list */
-		while (*w) {
-			if (*w == wsi) {
-				*w = wsi->u.ws.rx_draining_ext_list;
-				break;
-			}
-			w = &((*w)->u.ws.rx_draining_ext_list);
-		}
-		wsi->u.ws.rx_draining_ext_list = NULL;
+		lws_remove_wsi_from_draining_ext_list(wsi);
 		rx_draining_ext = 1;
 
 		goto drain_extension;
@@ -523,17 +511,16 @@ utf8_fail:			lwsl_info("utf8 error\n");
 		if (callback_action == LWS_CALLBACK_CLIENT_RECEIVE_PONG)
 			lwsl_info("Client doing pong callback\n");
 
-		if (n && eff_buf.token_len) {
+		if (n && eff_buf.token_len)
 			/* extension had more... main loop will come back
 			 * we want callback to be done with this set, if so,
 			 * because lws_is_final() hides it was final until the
 			 * last chunk
 			 */
-			wsi->u.ws.rx_draining_ext = 1;
-			wsi->u.ws.rx_draining_ext_list = pt->rx_draining_ext_list;
-			pt->rx_draining_ext_list = wsi;
-			lwsl_ext("%s: RX EXT DRAINING: Adding to list\n", __func__);
-		}
+			lws_add_wsi_to_draining_ext_list(wsi);
+		else
+			lws_remove_wsi_from_draining_ext_list(wsi);
+
 		if (wsi->state == LWSS_RETURNED_CLOSE_ALREADY ||
 		    wsi->state == LWSS_AWAITING_CLOSE_ACK)
 			goto already_done;

lib/parsers.c

@@ -986,16 +986,39 @@ LWS_VISIBLE int lws_frame_is_binary(struct lws *wsi)
 {
 	return wsi->u.ws.frame_is_binary;
 }
-static void
+
+void
+lws_add_wsi_to_draining_ext_list(struct lws *wsi)
+{
+	struct lws_context_per_thread *pt = &wsi->context->pt[(int)wsi->tsi];
+
+	if (wsi->u.ws.rx_draining_ext)
+		return;
+
+	lwsl_ext("%s: RX EXT DRAINING: Adding to list\n", __func__);
+
+	wsi->u.ws.rx_draining_ext = 1;
+	wsi->u.ws.rx_draining_ext_list = pt->rx_draining_ext_list;
+	pt->rx_draining_ext_list = wsi;
+}
+
+void
 lws_remove_wsi_from_draining_ext_list(struct lws *wsi)
 {
 	struct lws_context_per_thread *pt = &wsi->context->pt[(int)wsi->tsi];
 	struct lws **w = &pt->rx_draining_ext_list;
 
+	if (!wsi->u.ws.rx_draining_ext)
+		return;
+
+	lwsl_ext("%s: RX EXT DRAINING: Removing from list\n", __func__);
+
 	wsi->u.ws.rx_draining_ext = 0;
+
 	/* remove us from context draining ext list */
 	while (*w) {
 		if (*w == wsi) {
+			/* if us, point it instead to who we were pointing to */
 			*w = wsi->u.ws.rx_draining_ext_list;
 			break;
 		}
@@ -1008,7 +1031,6 @@ lws_remove_wsi_from_draining_ext_list(struct lws *wsi)
 int
 lws_rx_sm(struct lws *wsi, unsigned char c)
 {
-	struct lws_context_per_thread *pt = &wsi->context->pt[(int)wsi->tsi];
 	int callback_action = LWS_CALLBACK_RECEIVE;
 	int ret = 0, n, rx_draining_ext = 0;
 	struct lws_tokens eff_buf;
@@ -1458,15 +1480,11 @@ drain_extension:
 		if (rx_draining_ext && eff_buf.token_len == 0)
 			goto already_done;
 
-		if (n && eff_buf.token_len) {
+		if (n && eff_buf.token_len)
 			/* extension had more... main loop will come back */
-			// lwsl_notice("ext has stuff to drain\n");
-			wsi->u.ws.rx_draining_ext = 1;
-			wsi->u.ws.rx_draining_ext_list = pt->rx_draining_ext_list;
-			pt->rx_draining_ext_list = wsi;
-		} else {
+			lws_add_wsi_to_draining_ext_list(wsi);
+		else
 			lws_remove_wsi_from_draining_ext_list(wsi);
-		}
 
 		if (eff_buf.token_len > 0 ||
 		    callback_action == LWS_CALLBACK_RECEIVE_PONG) {

lib/private-libwebsockets.h

@@ -1775,6 +1775,10 @@ lws_plat_service_periodic(struct lws_context *context);
 LWS_EXTERN int
 lws_plat_change_pollfd(struct lws_context *context, struct lws *wsi,
 		       struct lws_pollfd *pfd);
+LWS_EXTERN void
+lws_add_wsi_to_draining_ext_list(struct lws *wsi);
+LWS_EXTERN void
+lws_remove_wsi_from_draining_ext_list(struct lws *wsi);
 LWS_EXTERN int
 lws_plat_context_early_init(void);
 LWS_EXTERN void

lib/service.c

@@ -445,9 +445,10 @@ lws_service_flag_pending(struct lws_context *context, int tsi)
 	while (wsi) {
 		pt->fds[wsi->position_in_fds_table].revents |=
 			pt->fds[wsi->position_in_fds_table].events & LWS_POLLIN;
-		if (pt->fds[wsi->position_in_fds_table].revents &
-		    LWS_POLLIN)
+		if (pt->fds[wsi->position_in_fds_table].revents & LWS_POLLIN) {
 			forced = 1;
+			break;
+		}
 		wsi = wsi->u.ws.rx_draining_ext_list;
 	}