Fix handling of WS handshake error response

Check response code, make sure it's 101.
Pass http_message to the client to keep it appraised.
This represents a slight change in the API -
in case of an error MG_EV_WEBSOCKET_HANDSHAKE_DONE will now be delivered where previosuly connection would just hang.
Clients that do not examine the argument may for a moment think handshake has succeeded but in fact connection will be closed immediately.

CL: mg: Fix handling of WS handshake error response

PUBLISHED_FROM=645a43d9e5bee216e54411f85827c9b974e9a7d1

docs/c-api/mg_http.h/mg_set_protocol_http_websocket.md

@@ -30,7 +30,9 @@ The user-defined event handler will receive following extra events:
 - MG_EV_WEBSOCKET_HANDSHAKE_REQUEST: server has received the WebSocket
   handshake request. `ev_data` contains parsed HTTP request.
 - MG_EV_WEBSOCKET_HANDSHAKE_DONE: server has completed the WebSocket
-  handshake. `ev_data` is `NULL`.
+  handshake. `ev_data` is a `struct http_message` containing the
+  client's request (server mode) or server's response (client).
+  In client mode handler can examine `resp_code`, which should be 101.
 - MG_EV_WEBSOCKET_FRAME: new WebSocket frame has arrived. `ev_data` is
   `struct websocket_message *`
 

examples/websocket_chat_client/websocket_chat_client.c

@@ -14,7 +14,6 @@ static int s_done = 0;
 static int s_is_connected = 0;
 
 static void ev_handler(struct mg_connection *nc, int ev, void *ev_data) {
-  struct websocket_message *wm = (struct websocket_message *) ev_data;
   (void) nc;
 
   switch (ev) {
@@ -26,8 +25,14 @@ static void ev_handler(struct mg_connection *nc, int ev, void *ev_data) {
       break;
     }
     case MG_EV_WEBSOCKET_HANDSHAKE_DONE: {
-      printf("-- Connected\n");
-      s_is_connected = 1;
+      struct http_message *hm = (struct http_message *) ev_data;
+      if (hm->resp_code == 101) {
+        printf("-- Connected\n");
+        s_is_connected = 1;
+      } else {
+        printf("-- Connection failed! HTTP code %d\n", hm->resp_code);
+        /* Connection will be closed after this. */
+      }
       break;
     }
     case MG_EV_POLL: {
@@ -66,6 +71,7 @@ static void ev_handler(struct mg_connection *nc, int ev, void *ev_data) {
       break;
     }
     case MG_EV_WEBSOCKET_FRAME: {
+      struct websocket_message *wm = (struct websocket_message *) ev_data;
       printf("%.*s\n", (int) wm->size, wm->data);
       break;
     }

mongoose.c

@@ -6649,16 +6649,23 @@ void mg_http_handler(struct mg_connection *nc, int ev,
       /* Do nothing, request is not yet fully buffered */
     }
 #if MG_ENABLE_HTTP_WEBSOCKET
-    else if (nc->listener == NULL &&
-             mg_get_http_header(hm, "Sec-WebSocket-Accept")) {
+    else if (nc->listener == NULL && (nc->flags & MG_F_IS_WEBSOCKET)) {
       /* We're websocket client, got handshake response from server. */
-      /* TODO(lsm): check the validity of accept Sec-WebSocket-Accept */
-      mbuf_remove(io, req_len);
-      nc->proto_handler = mg_ws_handler;
-      nc->flags |= MG_F_IS_WEBSOCKET;
-      mg_call(nc, nc->handler, nc->user_data, MG_EV_WEBSOCKET_HANDSHAKE_DONE,
-              NULL);
-      mg_ws_handler(nc, MG_EV_RECV, ev_data MG_UD_ARG(user_data));
+      DBG(("%p WebSocket upgrade code %d", nc, hm->resp_code));
+      if (hm->resp_code == 101 &&
+          mg_get_http_header(hm, "Sec-WebSocket-Accept")) {
+        /* TODO(lsm): check the validity of accept Sec-WebSocket-Accept */
+        mg_call(nc, nc->handler, nc->user_data, MG_EV_WEBSOCKET_HANDSHAKE_DONE,
+                hm);
+        mbuf_remove(io, req_len);
+        nc->proto_handler = mg_ws_handler;
+        mg_ws_handler(nc, MG_EV_RECV, ev_data MG_UD_ARG(user_data));
+      } else {
+        mg_call(nc, nc->handler, nc->user_data, MG_EV_WEBSOCKET_HANDSHAKE_DONE,
+                hm);
+        nc->flags |= MG_F_CLOSE_IMMEDIATELY;
+        mbuf_remove(io, req_len);
+      }
     } else if (nc->listener != NULL &&
                (vec = mg_get_http_header(hm, "Sec-WebSocket-Key")) != NULL) {
       struct mg_http_endpoint *ep;
@@ -6689,7 +6696,7 @@ void mg_http_handler(struct mg_connection *nc, int ev,
           mg_ws_handshake(nc, vec, hm);
         }
         mg_call(nc, nc->handler, nc->user_data, MG_EV_WEBSOCKET_HANDSHAKE_DONE,
-                NULL);
+                hm);
         mg_ws_handler(nc, MG_EV_RECV, ev_data MG_UD_ARG(user_data));
       }
     }
@@ -10402,6 +10409,8 @@ void mg_send_websocket_handshake3v(struct mg_connection *nc,
   }
   mg_printf(nc, "\r\n");
 
+  nc->flags |= MG_F_IS_WEBSOCKET;
+
   mbuf_free(&auth);
 }
 

mongoose.h

@@ -4806,7 +4806,7 @@ struct mg_ssi_call_ctx {
 
 #if MG_ENABLE_HTTP_WEBSOCKET
 #define MG_EV_WEBSOCKET_HANDSHAKE_REQUEST 111 /* struct http_message * */
-#define MG_EV_WEBSOCKET_HANDSHAKE_DONE 112    /* NULL */
+#define MG_EV_WEBSOCKET_HANDSHAKE_DONE 112    /* struct http_message * */
 #define MG_EV_WEBSOCKET_FRAME 113             /* struct websocket_message * */
 #define MG_EV_WEBSOCKET_CONTROL_FRAME 114     /* struct websocket_message * */
 #endif
@@ -4845,7 +4845,9 @@ struct mg_ssi_call_ctx {
  * - MG_EV_WEBSOCKET_HANDSHAKE_REQUEST: server has received the WebSocket
  *   handshake request. `ev_data` contains parsed HTTP request.
  * - MG_EV_WEBSOCKET_HANDSHAKE_DONE: server has completed the WebSocket
- *   handshake. `ev_data` is `NULL`.
+ *   handshake. `ev_data` is a `struct http_message` containing the
+ *   client's request (server mode) or server's response (client).
+ *   In client mode handler can examine `resp_code`, which should be 101.
  * - MG_EV_WEBSOCKET_FRAME: new WebSocket frame has arrived. `ev_data` is
  *   `struct websocket_message *`
  *

src/mg_http.c

@@ -827,16 +827,23 @@ void mg_http_handler(struct mg_connection *nc, int ev,
       /* Do nothing, request is not yet fully buffered */
     }
 #if MG_ENABLE_HTTP_WEBSOCKET
-    else if (nc->listener == NULL &&
-             mg_get_http_header(hm, "Sec-WebSocket-Accept")) {
+    else if (nc->listener == NULL && (nc->flags & MG_F_IS_WEBSOCKET)) {
       /* We're websocket client, got handshake response from server. */
-      /* TODO(lsm): check the validity of accept Sec-WebSocket-Accept */
-      mbuf_remove(io, req_len);
-      nc->proto_handler = mg_ws_handler;
-      nc->flags |= MG_F_IS_WEBSOCKET;
-      mg_call(nc, nc->handler, nc->user_data, MG_EV_WEBSOCKET_HANDSHAKE_DONE,
-              NULL);
-      mg_ws_handler(nc, MG_EV_RECV, ev_data MG_UD_ARG(user_data));
+      DBG(("%p WebSocket upgrade code %d", nc, hm->resp_code));
+      if (hm->resp_code == 101 &&
+          mg_get_http_header(hm, "Sec-WebSocket-Accept")) {
+        /* TODO(lsm): check the validity of accept Sec-WebSocket-Accept */
+        mg_call(nc, nc->handler, nc->user_data, MG_EV_WEBSOCKET_HANDSHAKE_DONE,
+                hm);
+        mbuf_remove(io, req_len);
+        nc->proto_handler = mg_ws_handler;
+        mg_ws_handler(nc, MG_EV_RECV, ev_data MG_UD_ARG(user_data));
+      } else {
+        mg_call(nc, nc->handler, nc->user_data, MG_EV_WEBSOCKET_HANDSHAKE_DONE,
+                hm);
+        nc->flags |= MG_F_CLOSE_IMMEDIATELY;
+        mbuf_remove(io, req_len);
+      }
     } else if (nc->listener != NULL &&
                (vec = mg_get_http_header(hm, "Sec-WebSocket-Key")) != NULL) {
       struct mg_http_endpoint *ep;
@@ -867,7 +874,7 @@ void mg_http_handler(struct mg_connection *nc, int ev,
           mg_ws_handshake(nc, vec, hm);
         }
         mg_call(nc, nc->handler, nc->user_data, MG_EV_WEBSOCKET_HANDSHAKE_DONE,
-                NULL);
+                hm);
         mg_ws_handler(nc, MG_EV_RECV, ev_data MG_UD_ARG(user_data));
       }
     }

src/mg_http.h

@@ -106,7 +106,7 @@ struct mg_ssi_call_ctx {
 
 #if MG_ENABLE_HTTP_WEBSOCKET
 #define MG_EV_WEBSOCKET_HANDSHAKE_REQUEST 111 /* struct http_message * */
-#define MG_EV_WEBSOCKET_HANDSHAKE_DONE 112    /* NULL */
+#define MG_EV_WEBSOCKET_HANDSHAKE_DONE 112    /* struct http_message * */
 #define MG_EV_WEBSOCKET_FRAME 113             /* struct websocket_message * */
 #define MG_EV_WEBSOCKET_CONTROL_FRAME 114     /* struct websocket_message * */
 #endif
@@ -145,7 +145,9 @@ struct mg_ssi_call_ctx {
  * - MG_EV_WEBSOCKET_HANDSHAKE_REQUEST: server has received the WebSocket
  *   handshake request. `ev_data` contains parsed HTTP request.
  * - MG_EV_WEBSOCKET_HANDSHAKE_DONE: server has completed the WebSocket
- *   handshake. `ev_data` is `NULL`.
+ *   handshake. `ev_data` is a `struct http_message` containing the
+ *   client's request (server mode) or server's response (client).
+ *   In client mode handler can examine `resp_code`, which should be 101.
  * - MG_EV_WEBSOCKET_FRAME: new WebSocket frame has arrived. `ev_data` is
  *   `struct websocket_message *`
  *

src/mg_http_websocket.c

@@ -478,6 +478,8 @@ void mg_send_websocket_handshake3v(struct mg_connection *nc,
   }
   mg_printf(nc, "\r\n");
 
+  nc->flags |= MG_F_IS_WEBSOCKET;
+
   mbuf_free(&auth);
 }
 

test/unit_test.c

@@ -2540,32 +2540,6 @@ static const char *test_http_range(void) {
   return NULL;
 }
 
-static void cb3(struct mg_connection *nc, int ev, void *ev_data) {
-  struct websocket_message *wm = (struct websocket_message *) ev_data;
-
-  if (ev == MG_EV_WEBSOCKET_FRAME) {
-    const char *reply = wm->size == 2 && !memcmp(wm->data, "hi", 2) ? "A" : "B";
-    mg_printf_websocket_frame(nc, WEBSOCKET_OP_TEXT, "%s", reply);
-  }
-}
-
-static void cb4(struct mg_connection *nc, int ev, void *ev_data) {
-  struct websocket_message *wm = (struct websocket_message *) ev_data;
-
-  if (ev == MG_EV_WEBSOCKET_FRAME) {
-    memcpy(nc->user_data, wm->data, wm->size);
-    mg_send_websocket_frame(nc, WEBSOCKET_OP_CLOSE, NULL, 0);
-  } else if (ev == MG_EV_WEBSOCKET_HANDSHAKE_DONE) {
-    /* Send "hi" to server. server must reply "A". */
-    struct mg_str h[2];
-    h[0].p = "h";
-    h[0].len = 1;
-    h[1].p = "i";
-    h[1].len = 1;
-    mg_send_websocket_framev(nc, WEBSOCKET_OP_TEXT, h, 2);
-  }
-}
-
 static void cb_ws_server(struct mg_connection *nc, int ev, void *ev_data) {
   struct websocket_message *wm = (struct websocket_message *) ev_data;
 
@@ -2780,6 +2754,38 @@ static const char *test_websocket(void) {
   return NULL;
 }
 
+static void cb3(struct mg_connection *nc, int ev, void *ev_data) {
+  struct websocket_message *wm = (struct websocket_message *) ev_data;
+  if (ev != MG_EV_WEBSOCKET_FRAME) return;
+  DBG(("server data '%.*s'", (int) wm->size, wm->data));
+  const char *reply = wm->size == 2 && !memcmp(wm->data, "hi", 2) ? "A" : "B";
+  mg_printf_websocket_frame(nc, WEBSOCKET_OP_TEXT, "%s", reply);
+}
+
+static void cb4(struct mg_connection *nc, int ev, void *ev_data) {
+  char *buf = (char *) nc->user_data;
+  if (ev == MG_EV_WEBSOCKET_FRAME) {
+    struct websocket_message *wm = (struct websocket_message *) ev_data;
+    DBG(("client data '%.*s'", (int) wm->size, wm->data));
+    memcpy(buf, wm->data, wm->size);
+    mg_send_websocket_frame(nc, WEBSOCKET_OP_CLOSE, NULL, 0);
+  } else if (ev == MG_EV_WEBSOCKET_HANDSHAKE_DONE) {
+    struct http_message *hm = (struct http_message *) ev_data;
+    DBG(("code %d", hm->resp_code));
+    if (hm->resp_code == 101) {
+      /* Send "hi" to server. server must reply "A". */
+      struct mg_str h[2];
+      h[0].p = "h";
+      h[0].len = 1;
+      h[1].p = "i";
+      h[1].len = 1;
+      mg_send_websocket_framev(nc, WEBSOCKET_OP_TEXT, h, 2);
+    } else {
+      snprintf(buf, 20, "code %d", hm->resp_code);
+    }
+  }
+}
+
 static void cbwep(struct mg_connection *c, int ev, void *ev_data) {
   struct websocket_message *wm = (struct websocket_message *) ev_data;
   char *buf = (char *) c->user_data;
@@ -2787,6 +2793,9 @@ static void cbwep(struct mg_connection *c, int ev, void *ev_data) {
   switch (ev) {
     case MG_EV_WEBSOCKET_HANDSHAKE_REQUEST:
       strcat(buf, "R");
+      if (buf[0] != '0') {
+        mg_http_send_error(c, 403, "I don't like you");
+      }
       break;
     case MG_EV_WEBSOCKET_HANDSHAKE_DONE:
       strcat(buf, "D");
@@ -2803,7 +2812,7 @@ static const char *test_websocket_endpoint(void) {
   struct mg_mgr mgr;
   struct mg_connection *nc;
   const char *local_addr = "127.0.0.1:7798";
-  char buf[20] = "", buf2[20] = "";
+  char buf[20] = "", buf2[20] = "0";
 
   mg_mgr_init(&mgr, NULL);
   /* mgr.hexdump_file = "-"; */
@@ -2818,10 +2827,21 @@ static const char *test_websocket_endpoint(void) {
   nc->user_data = buf;
   mg_send_websocket_handshake(nc, "/boo", NULL);
   poll_until(&mgr, 1, c_str_ne, buf, (void *) "");
-  mg_mgr_free(&mgr);
-
   /* Check that test buffer has been filled by the callback properly. */
-  ASSERT_STREQ(buf, "RDF|hi");
+  ASSERT_STREQ(buf, "0RDF|hi");
+
+  /* Test handshake failure */
+  ASSERT((nc = mg_connect(&mgr, local_addr, cb4)) != NULL);
+  mg_set_protocol_http_websocket(nc);
+  buf[0] = '\0';
+  buf2[0] = '1';
+  buf2[1] = '\0';
+  nc->user_data = buf;
+  mg_send_websocket_handshake(nc, "/boo", NULL);
+  poll_until(&mgr, 1, c_str_ne, buf, (void *) "");
+  ASSERT_STREQ(buf, "code 403");
+
+  mg_mgr_free(&mgr);
 
   return NULL;
 }