Skip to content

Commit

Permalink
Fix handling of WS handshake error response
Browse files Browse the repository at this point in the history
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
  • Loading branch information
Deomid Ryabkov authored and cesantabot committed Feb 14, 2019
1 parent 148d3f0 commit 1e9fabe
Show file tree
Hide file tree
Showing 8 changed files with 108 additions and 58 deletions.
4 changes: 3 additions & 1 deletion docs/c-api/mg_http.h/mg_set_protocol_http_websocket.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 *`

Expand Down
12 changes: 9 additions & 3 deletions examples/websocket_chat_client/websocket_chat_client.c
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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: {
Expand Down Expand Up @@ -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;
}
Expand Down
29 changes: 19 additions & 10 deletions mongoose.c
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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));
}
}
Expand Down Expand Up @@ -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);
}

Expand Down
6 changes: 4 additions & 2 deletions mongoose.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 *`
*
Expand Down
27 changes: 17 additions & 10 deletions src/mg_http.c
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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));
}
}
Expand Down
6 changes: 4 additions & 2 deletions src/mg_http.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 *`
*
Expand Down
2 changes: 2 additions & 0 deletions src/mg_http_websocket.c
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down
80 changes: 50 additions & 30 deletions test/unit_test.c
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -2780,13 +2754,48 @@ 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;

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");
Expand All @@ -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 = "-"; */
Expand All @@ -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;
}
Expand Down

0 comments on commit 1e9fabe

Please sign in to comment.