Added mg_websocket_write(), per morgan3d pull request

examples/websocket.c

@@ -6,10 +6,8 @@
 #include "mongoose.h"
 
 static void websocket_ready_handler(struct mg_connection *conn) {
-  unsigned char buf[40];
+  static const char *message = "server ready";
-  buf[0] = 0x81;
+  mg_websocket_write(conn, WEBSOCKET_OPCODE_TEXT, message, strlen(message));
-  buf[1] = snprintf((char *) buf + 2, sizeof(buf) - 2, "%s", "server ready");
-  mg_write(conn, buf, 2 + buf[1]);
 }
 
 // Arguments:
@@ -18,33 +16,12 @@ static void websocket_ready_handler(struct mg_connection *conn) {
 //   data, data_len: payload data. Mask, if any, is already applied.
 static int websocket_data_handler(struct mg_connection *conn, int flags,
                                   char *data, size_t data_len) {
-  unsigned char reply[200];
+  (void) flags; // Unused
-  size_t i;
+  mg_websocket_write(conn, WEBSOCKET_OPCODE_TEXT, data, data_len);
-
-  (void) flags;
-
-  printf("rcv: [%.*s]\n", (int) data_len, data);
-
-  // Truncate echoed message, to simplify output code.
-  if (data_len > 125) {
-    data_len = 125;
-  }
-
-  // Prepare frame
-  reply[0] = 0x81;  // text, FIN set
-  reply[1] = data_len;
-
-  // Copy message from request to reply, applying the mask if required.
-  for (i = 0; i < data_len; i++) {
-    reply[i + 2] = data[i];
-  }
-
-  // Echo the message back to the client
-  mg_write(conn, reply, 2 + data_len);
 
   // Returning zero means stoping websocket conversation.
   // Close the conversation if client has sent us "exit" string.
-  return memcmp(reply + 2, "exit", 4);
+  return memcmp(data, "exit", 4);
 }
 
 int main(void) {

mongoose.c

@@ -4090,6 +4090,48 @@ static void read_websocket(struct mg_connection *conn) {
   }
 }
 
+int mg_websocket_write(struct mg_connection* conn, int opcode,
+                       const char *data, size_t data_len) {
+    unsigned char *copy;
+    size_t copy_len = 0;
+    int retval = -1;
+
+    if ((copy = (unsigned char *) malloc(data_len + 10)) == NULL) {
+      return -1;
+    }
+
+    copy[0] = 0x80 + (opcode & 0x0f);
+
+    // Frame format: http://tools.ietf.org/html/rfc6455#section-5.2
+    if (data_len < 126) {
+      // Inline 7-bit length field
+      copy[1] = data_len;
+      memcpy(copy + 2, data, data_len);
+      copy_len = 2 + data_len;
+    } else if (data_len <= 0xFFFF) {
+      // 16-bit length field
+      copy[1] = 126;
+      * (uint16_t *) (copy + 2) = htons(data_len);
+      memcpy(copy + 4, data, data_len);
+      copy_len = 4 + data_len;
+    } else {
+      // 64-bit length field
+      copy[1] = 127;
+      * (uint32_t *) (copy + 2) = htonl((uint64_t) data_len >> 32);
+      * (uint32_t *) (copy + 6) = htonl(data_len & 0xffffffff);
+      memcpy(copy + 10, data, data_len);
+      copy_len = 10 + data_len;
+    }
+
+    // Not thread safe
+    if (copy_len > 0) {
+      retval = mg_write(conn, copy, copy_len);
+    }
+    free(copy);
+
+    return retval;
+}
+
 static void handle_websocket_request(struct mg_connection *conn) {
   const char *version = mg_get_header(conn, "Sec-WebSocket-Version");
   if (version == NULL || strcmp(version, "13") != 0) {

mongoose.h

@@ -208,6 +208,28 @@ struct mg_request_info *mg_get_request_info(struct mg_connection *);
 int mg_write(struct mg_connection *, const void *buf, size_t len);
 
 
+// Send data to a websocket client wrapped in a websocket frame.
+// It is unsafe to read/write to this connection from another thread.
+// This function is available when mongoose is compiled with -DUSE_WEBSOCKET
+//
+// Return:
+//  0   when the connection has been closed
+//  -1  on error
+//  >0  number of bytes written on success
+int mg_websocket_write(struct mg_connection* conn, int opcode,
+                       const char *data, size_t data_len);
+
+// Opcodes, from http://tools.ietf.org/html/rfc6455
+enum {
+  WEBSOCKET_OPCODE_CONTINUATION = 0x0,
+  WEBSOCKET_OPCODE_TEXT = 0x1,
+  WEBSOCKET_OPCODE_BINARY = 0x2,
+  WEBSOCKET_OPCODE_CONNECTION_CLOSE = 0x8,
+  WEBSOCKET_OPCODE_PING = 0x9,
+  WEBSOCKET_OPCODE_PONG = 0xa
+};
+
+
 // Macros for enabling compiler-specific checks for printf-like arguments.
 #undef PRINTF_FORMAT_STRING
 #if _MSC_VER >= 1400