From 5e91d919e0554bfb84f407c536396de07a25184a Mon Sep 17 00:00:00 2001
From: Sergey Lyubka <lsm@cesanta.com>
Date: Wed, 9 Nov 2016 17:05:26 +0000
Subject: [PATCH] Make mg_register_http_endpoint work for websocket

PUBLISHED_FROM=5bf5a007929b3b6275121f3cf949526e881fb482
---
 mongoose.c | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/mongoose.c b/mongoose.c
index d14edbce..a3add397 100644
--- a/mongoose.c
+++ b/mongoose.c
@@ -4992,11 +4992,23 @@ void mg_http_handler(struct mg_connection *nc, int ev, void *ev_data) {
       mg_ws_handler(nc, MG_EV_RECV, ev_data);
     } else if (nc->listener != NULL &&
                (vec = mg_get_http_header(hm, "Sec-WebSocket-Key")) != NULL) {
+      mg_event_handler_t handler;
+
       /* This is a websocket request. Switch protocol handlers. */
       mbuf_remove(io, req_len);
       nc->proto_handler = mg_ws_handler;
       nc->flags |= MG_F_IS_WEBSOCKET;
 
+      /*
+       * If we have a handler set up with mg_register_http_endpoint(),
+       * deliver subsequent websocket events to this handler after the
+       * protocol switch.
+       */
+      handler = mg_http_get_endpoint_handler(nc->listener, &hm->uri);
+      if (handler != NULL) {
+        nc->handler = handler;
+      }
+
       /* Send handshake */
       mg_call(nc, nc->handler, MG_EV_WEBSOCKET_HANDSHAKE_REQUEST, hm);
       if (!(nc->flags & MG_F_CLOSE_IMMEDIATELY)) {