From 7519b2ef3a646a859751df2e8ce54a330d426af9 Mon Sep 17 00:00:00 2001
From: Dmitry Frank <mail@dmitryfrank.com>
Date: Thu, 1 Feb 2018 13:50:47 +0200
Subject: [PATCH] Improve websocket implementation
CL: Mongoose Web Server: Websocket: Respond to Ping with Pong
CL: Mongoose Web Server: Websocket: Properly close a connection with Close frame (in response to a client's close and when protocol failure is detected)
CL: Mongoose Web Server: Websocket: Fix support of fragmented messages
CL: Mongoose Web Server: Websocket: Add support for control frames interjected in the middle of a fragmented message
PUBLISHED_FROM=e2b3794aaacc64633540c493194cccc62afa2077
---
docs/c-api/http.h/mg_send_websocket_framev.md | 6 +-
mongoose.c | 190 +++++++++++++-----
mongoose.h | 6 +-
3 files changed, 144 insertions(+), 58 deletions(-)
diff --git a/docs/c-api/http.h/mg_send_websocket_framev.md b/docs/c-api/http.h/mg_send_websocket_framev.md
index 06e094601..ae7da3805 100644
--- a/docs/c-api/http.h/mg_send_websocket_framev.md
+++ b/docs/c-api/http.h/mg_send_websocket_framev.md
@@ -7,8 +7,6 @@ signature: |
const struct mg_str *strings, int num_strings);
---
-Sends multiple websocket frames.
-
-Like `mg_send_websocket_frame()`, but composes a frame from multiple
-*buffers.
+Like `mg_send_websocket_frame()`, but composes a single frame from multiple
+buffers.
diff --git a/mongoose.c b/mongoose.c
index a2b8fee6c..e2037ee50 100644
--- a/mongoose.c
+++ b/mongoose.c
@@ -5734,7 +5734,13 @@ struct mg_reverse_proxy_data {
};
struct mg_ws_proto_data {
- size_t reass_len; /* Defragmented size of the frame so far. */
+ /*
+ * Defragmented size of the frame so far.
+ *
+ * First byte of nc->recv_mbuf.buf is an op, the rest of the data is
+ * defragmented data.
+ */
+ size_t reass_len;
};
struct mg_http_proto_data {
@@ -9569,12 +9575,23 @@ MG_INTERNAL void mg_handle_put(struct mg_connection *nc, const char *path,
#define MG_WEBSOCKET_PING_INTERVAL_SECONDS 5
#endif
+#define FLAGS_MASK_FIN (1 << 7)
+#define FLAGS_MASK_OP 0x0f
+
static int mg_is_ws_fragment(unsigned char flags) {
- return (flags & 0x80) == 0 || (flags & 0x0f) == 0;
+ return (flags & FLAGS_MASK_FIN) == 0 ||
+ (flags & FLAGS_MASK_OP) == WEBSOCKET_OP_CONTINUE;
}
static int mg_is_ws_first_fragment(unsigned char flags) {
- return (flags & 0x80) == 0 && (flags & 0x0f) != 0;
+ return (flags & FLAGS_MASK_FIN) == 0 &&
+ (flags & FLAGS_MASK_OP) != WEBSOCKET_OP_CONTINUE;
+}
+
+static int mg_is_ws_control_frame(unsigned char flags) {
+ unsigned char op = (flags & FLAGS_MASK_OP);
+ return op == WEBSOCKET_OP_CLOSE || op == WEBSOCKET_OP_PING ||
+ op == WEBSOCKET_OP_PONG;
}
static void mg_handle_incoming_websocket_frame(struct mg_connection *nc,
@@ -9591,94 +9608,166 @@ static struct mg_ws_proto_data *mg_ws_get_proto_data(struct mg_connection *nc) {
return (htd != NULL ? &htd->ws_data : NULL);
}
+/*
+ * Sends a Close websocket frame with the given data, and closes the underlying
+ * connection. If `len` is ~0, strlen(data) is used.
+ */
+static void mg_ws_close(struct mg_connection *nc, const void *data,
+ size_t len) {
+ if ((int) len == ~0) {
+ len = strlen((const char *) data);
+ }
+ mg_send_websocket_frame(nc, WEBSOCKET_OP_CLOSE, data, len);
+ nc->flags |= MG_F_SEND_AND_CLOSE;
+}
+
static int mg_deliver_websocket_data(struct mg_connection *nc) {
/* Using unsigned char *, cause of integer arithmetic below */
- uint64_t i, data_len = 0, frame_len = 0, buf_len = nc->recv_mbuf.len, len,
- mask_len = 0, header_len = 0;
- unsigned char *p = (unsigned char *) nc->recv_mbuf.buf, *buf = p,
- *e = p + buf_len;
+ uint64_t i, data_len = 0, frame_len = 0, new_data_len = nc->recv_mbuf.len,
+ len, mask_len = 0, header_len = 0;
struct mg_ws_proto_data *wsd = mg_ws_get_proto_data(nc);
- int ok;
- int reass = buf_len > 0 && mg_is_ws_fragment(p[0]) &&
- !(nc->flags & MG_F_WEBSOCKET_NO_DEFRAG);
+ unsigned char *new_data = (unsigned char *) nc->recv_mbuf.buf,
+ *e = (unsigned char *) nc->recv_mbuf.buf + nc->recv_mbuf.len;
+ uint8_t flags;
+ int ok, reass;
+
+ if (wsd->reass_len > 0) {
+ /*
+ * We already have some previously received data which we need to
+ * reassemble and deliver to the client code when we get the final
+ * fragment.
+ *
+ * NOTE: it doesn't mean that the current message must be a continuation:
+ * it might be a control frame (Close, Ping or Pong), which should be
+ * handled without breaking the fragmented message.
+ */
+
+ size_t existing_len = wsd->reass_len;
+ assert(new_data_len >= existing_len);
- /* If that's a continuation frame that must be reassembled, handle it */
- if (reass && !mg_is_ws_first_fragment(p[0]) && buf_len >= 1 &&
- buf_len >= 1 + wsd->reass_len) {
- buf += 1 + wsd->reass_len;
- buf_len -= 1 + wsd->reass_len;
+ new_data += existing_len;
+ new_data_len -= existing_len;
+ }
+
+ flags = new_data[0];
+
+ reass = new_data_len > 0 && mg_is_ws_fragment(flags) &&
+ !(nc->flags & MG_F_WEBSOCKET_NO_DEFRAG);
+
+ if (reass && mg_is_ws_control_frame(flags)) {
+ /*
+ * Control frames can't be fragmented, so if we encounter fragmented
+ * control frame, close connection immediately.
+ */
+ mg_ws_close(nc, "fragmented control frames are illegal", ~0);
+ return 0;
+ } else if (new_data_len > 0 && !reass && !mg_is_ws_control_frame(flags) &&
+ wsd->reass_len > 0) {
+ /*
+ * When in the middle of a fragmented message, only the continuations
+ * and control frames are allowed.
+ */
+ mg_ws_close(nc, "non-continuation in the middle of a fragmented message",
+ ~0);
+ return 0;
}
- if (buf_len >= 2) {
- len = buf[1] & 0x7f;
- mask_len = buf[1] & 0x80 ? 4 : 0;
- if (len < 126 && buf_len >= mask_len) {
+ if (new_data_len >= 2) {
+ len = new_data[1] & 0x7f;
+ mask_len = new_data[1] & FLAGS_MASK_FIN ? 4 : 0;
+ if (len < 126 && new_data_len >= mask_len) {
data_len = len;
header_len = 2 + mask_len;
- } else if (len == 126 && buf_len >= 4 + mask_len) {
+ } else if (len == 126 && new_data_len >= 4 + mask_len) {
header_len = 4 + mask_len;
- data_len = ntohs(*(uint16_t *) &buf[2]);
- } else if (buf_len >= 10 + mask_len) {
+ data_len = ntohs(*(uint16_t *) &new_data[2]);
+ } else if (new_data_len >= 10 + mask_len) {
header_len = 10 + mask_len;
- data_len = (((uint64_t) ntohl(*(uint32_t *) &buf[2])) << 32) +
- ntohl(*(uint32_t *) &buf[6]);
+ data_len = (((uint64_t) ntohl(*(uint32_t *) &new_data[2])) << 32) +
+ ntohl(*(uint32_t *) &new_data[6]);
}
}
frame_len = header_len + data_len;
- ok = (frame_len > 0 && frame_len <= buf_len);
+ ok = (frame_len > 0 && frame_len <= new_data_len);
/* Check for overflow */
if (frame_len < header_len || frame_len < data_len) {
ok = 0;
- nc->flags |= MG_F_CLOSE_IMMEDIATELY;
+ mg_ws_close(nc, "overflowed message", ~0);
}
if (ok) {
+ size_t cleanup_len = 0;
struct websocket_message wsm;
wsm.size = (size_t) data_len;
- wsm.data = buf + header_len;
- wsm.flags = buf[0];
+ wsm.data = new_data + header_len;
+ wsm.flags = flags;
/* Apply mask if necessary */
if (mask_len > 0) {
for (i = 0; i < data_len; i++) {
- buf[i + header_len] ^= (buf + header_len - mask_len)[i % 4];
+ new_data[i + header_len] ^= (new_data + header_len - mask_len)[i % 4];
}
}
if (reass) {
- /* On first fragmented frame, nullify size */
- if (mg_is_ws_first_fragment(wsm.flags)) {
- p[0] &= ~0x0f; /* Next frames will be treated as continuation */
- buf = p + 1;
- wsd->reass_len = 0;
+ /* This is a message fragment */
+
+ if (mg_is_ws_first_fragment(flags)) {
+ /*
+ * On the first fragmented frame, skip the first byte (op) and also
+ * reset size to 1 (op), it'll be incremented with the data len below.
+ */
+ new_data += 1;
+ wsd->reass_len = 1 /* op */;
}
/* Append this frame to the reassembled buffer */
- memmove(buf, wsm.data, e - wsm.data);
+ memmove(new_data, wsm.data, e - wsm.data);
wsd->reass_len += wsm.size;
- nc->recv_mbuf.len -= wsm.data - buf;
+ nc->recv_mbuf.len -= wsm.data - new_data;
+
+ if (flags & FLAGS_MASK_FIN) {
+ /* On last fragmented frame - call user handler and remove data */
+ wsm.flags = FLAGS_MASK_FIN | nc->recv_mbuf.buf[0];
+ wsm.data = (unsigned char *) nc->recv_mbuf.buf + 1 /* op */;
+ wsm.size = wsd->reass_len - 1 /* op */;
+ cleanup_len = wsd->reass_len;
+ wsd->reass_len = 0;
- /* On last fragmented frame - call user handler and remove data */
- if (wsm.flags & 0x80) {
- wsm.data = p + 1;
- wsm.size = wsd->reass_len;
+ /* Pass reassembled message to the client code. */
mg_handle_incoming_websocket_frame(nc, &wsm);
- mbuf_remove(&nc->recv_mbuf, 1 + wsd->reass_len);
- wsd->reass_len = 0;
+ mbuf_remove(&nc->recv_mbuf, cleanup_len); /* Cleanup frame */
}
} else {
- /* TODO(lsm): properly handle OOB control frames during defragmentation */
+ /*
+ * This is a complete message, not a fragment. It might happen in between
+ * of a fragmented message (in this case, WebSocket protocol requires
+ * current message to be a control frame).
+ */
+ cleanup_len = (size_t) frame_len;
+
+ /* First of all, check if we need to react on a control frame. */
+ switch (flags & FLAGS_MASK_OP) {
+ case WEBSOCKET_OP_PING:
+ mg_send_websocket_frame(nc, WEBSOCKET_OP_PONG, wsm.data, wsm.size);
+ break;
+
+ case WEBSOCKET_OP_CLOSE:
+ mg_ws_close(nc, wsm.data, wsm.size);
+ break;
+ }
+
+ /* Pass received message to the client code. */
mg_handle_incoming_websocket_frame(nc, &wsm);
- mbuf_remove(&nc->recv_mbuf, (size_t) frame_len); /* Cleanup frame */
- wsd->reass_len = 0;
- }
- /* If the frame is not reassembled - client closes and close too */
- if (!reass && (buf[0] & 0x0f) == WEBSOCKET_OP_CLOSE) {
- nc->flags |= MG_F_SEND_AND_CLOSE;
+ /* Cleanup frame */
+ memmove(nc->recv_mbuf.buf + wsd->reass_len,
+ nc->recv_mbuf.buf + wsd->reass_len + cleanup_len,
+ nc->recv_mbuf.len - wsd->reass_len - cleanup_len);
+ nc->recv_mbuf.len -= cleanup_len;
}
}
@@ -9723,7 +9812,8 @@ static void mg_send_ws_header(struct mg_connection *nc, int op, size_t len,
int header_len;
unsigned char header[10];
- header[0] = (op & WEBSOCKET_DONT_FIN ? 0x0 : 0x80) + (op & 0x0f);
+ header[0] =
+ (op & WEBSOCKET_DONT_FIN ? 0x0 : FLAGS_MASK_FIN) | (op & FLAGS_MASK_OP);
if (len < 126) {
header[1] = (unsigned char) len;
header_len = 2;
diff --git a/mongoose.h b/mongoose.h
index bfdba362b..c54bcbf7a 100644
--- a/mongoose.h
+++ b/mongoose.h
@@ -4527,10 +4527,8 @@ void mg_send_websocket_frame(struct mg_connection *nc, int op_and_flags,
const void *data, size_t data_len);
/*
- * Sends multiple websocket frames.
- *
- * Like `mg_send_websocket_frame()`, but composes a frame from multiple
- *buffers.
+ * Like `mg_send_websocket_frame()`, but composes a single frame from multiple
+ * buffers.
*/
void mg_send_websocket_framev(struct mg_connection *nc, int op_and_flags,
const struct mg_str *strings, int num_strings);
--
GitLab