Refactored data buffering. Updated websocket example to have 2 message roundtrips.

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;
   }

examples/websocket_html_root/index.html

@@ -2,6 +2,11 @@
 <meta charset="utf-8" />  
 <title>WebSocket Test</title>  
 <script language="javascript" type="text/javascript">
+
+  // This javascript code creates a websocket to the URI "/foo",
+  // sends a message to it, waits for the reply, and sends an "exit" message.
+  // Server must terminate the conversation after receiving "exit" message.
+
   var writeToScreen = function(message) {
     var div = document.createElement('div');
     div.innerHTML = message;
@@ -22,7 +27,7 @@
     websocket.onmessage = function(ev) {
       writeToScreen('<span style="color: blue;">RESPONSE: ' + ev.data +
                     ' </span>');
-      websocket.close();
+      websocket.send('exit');
     };
     websocket.onerror = function(ev) {
       writeToScreen('<span style="color: red; ">ERROR: </span> ' + ev.data);