diff --git a/examples/websocket.c b/examples/websocket.c index c0a3c5bf..509f8bf8 100644 --- a/examples/websocket.c +++ b/examples/websocket.c @@ -7,37 +7,53 @@ static void *callback(enum mg_event event, struct mg_connection *conn) { if (event == MG_WEBSOCKET_READY) { - static const char *hello = "hello from mongoose! waiting for message ..."; - char frame[2]; - - // Prepare websocket frame. - frame[0] = 0x81; // text frame - frame[1] = strlen(hello); // length is < 126 - - // Write frame and a text message - mg_write(conn, frame, sizeof(frame)); - mg_write(conn, hello, strlen(hello)); - return ""; + unsigned char buf[40]; + buf[0] = 0x81; + buf[1] = snprintf((char *) buf + 2, sizeof(buf) - 2, "%s", "server ready"); + mg_write(conn, buf, 2 + buf[1]); + return ""; // MG_WEBSOCKET_READY return value is ignored } else if (event == MG_WEBSOCKET_MESSAGE) { - unsigned char buf[500], reply[500]; - int len, msg_len, i, mask_len, xor; + unsigned char buf[200], reply[200]; + int n, i, mask_len, xor, msg_len, len; - // Read message from the client and echo it back - if ((len = mg_read(conn, buf, sizeof(buf))) > 8) { - msg_len = buf[1] & 127; - mask_len = (buf[1] & 128) ? 4 : 0; - if (msg_len < 126) { - reply[0] = 0x81; // text, FIN set - reply[1] = msg_len; - for (i = 0; i < msg_len; i++) { - xor = mask_len == 0 ? 0 : buf[2 + (i % 4)]; - reply[i + 2] = buf[i + 2 + mask_len] ^ xor; + // Read message from the client. + // Accept only small (<126 bytes) messages. + len = 0; + msg_len = mask_len = 0; + for (;;) { + if ((n = mg_read(conn, buf + len, sizeof(buf) - len)) <= 0) { + return ""; // Read error, close websocket + } + len += n; + if (len >= 2) { + msg_len = buf[1] & 127; + mask_len = (buf[1] & 128) ? 4 : 0; + if (msg_len > 125) { + return ""; // Message is too long, close websocket + } + // If we've buffered the whole message, exit the loop + if (len >= 2 + mask_len + msg_len) { + break; } - mg_write(conn, reply, 2 + msg_len); } } - return ""; // Return non-NULL: stop websocket conversation + // Prepare frame + reply[0] = 0x81; // text, FIN set + reply[1] = msg_len; + + // Copy message from request to reply, applying the mask if required. + for (i = 0; i < msg_len; i++) { + xor = mask_len == 0 ? 0 : buf[2 + (i % 4)]; + reply[i + 2] = buf[i + 2 + mask_len] ^ xor; + } + + // Echo the message back to the client + mg_write(conn, reply, 2 + msg_len); + + // Return non-NULL means stoping websocket conversation. + // Close the conversation if client has sent us "exit" string. + return memcmp(reply + 2, "exit", 4) == 0 ? "" : NULL; } else { return NULL; } diff --git a/examples/websocket_html_root/index.html b/examples/websocket_html_root/index.html index 2457b02c..7b46d4d6 100644 --- a/examples/websocket_html_root/index.html +++ b/examples/websocket_html_root/index.html @@ -2,6 +2,11 @@ WebSocket Test