Async websockets work

build/src/core.c

@@ -700,9 +700,7 @@ static int is_valid_http_method(const char *method) {
 static int parse_http_message(char *buf, int len, struct mg_connection *ri) {
   int is_request, n;
 
-  // Reset attributes. DO NOT TOUCH remote_ip, remote_port
-  ri->request_method = ri->uri = ri->http_version = NULL;
-  ri->num_headers = 0;
+  memset(ri, 0, sizeof(*ri));
   buf[len - 1] = '\0';
 
   // RFC says that all initial whitespaces should be ingored
@@ -891,6 +889,373 @@ int mg_write(struct mg_connection *c, const void *buf, int len) {
   return ret;
 }
 
+#if defined(USE_WEBSOCKET)
+
+static int is_big_endian(void) {
+  static const int n = 1;
+  return ((char *) &n)[0] == 0;
+}
+
+// START OF SHA-1 code
+// Copyright(c) By Steve Reid <steve@edmweb.com>
+#define SHA1HANDSOFF
+#if defined(__sun)
+#include "solarisfixes.h"
+#endif
+
+union char64long16 { unsigned char c[64]; uint32_t l[16]; };
+
+#define rol(value, bits) (((value) << (bits)) | ((value) >> (32 - (bits))))
+
+static uint32_t blk0(union char64long16 *block, int i) {
+  // Forrest: SHA expect BIG_ENDIAN, swap if LITTLE_ENDIAN
+  if (!is_big_endian()) {
+    block->l[i] = (rol(block->l[i], 24) & 0xFF00FF00) |
+      (rol(block->l[i], 8) & 0x00FF00FF);
+  }
+  return block->l[i];
+}
+
+#define blk(i) (block->l[i&15] = rol(block->l[(i+13)&15]^block->l[(i+8)&15] \
+    ^block->l[(i+2)&15]^block->l[i&15],1))
+#define R0(v,w,x,y,z,i) z+=((w&(x^y))^y)+blk0(block, i)+0x5A827999+rol(v,5);w=rol(w,30);
+#define R1(v,w,x,y,z,i) z+=((w&(x^y))^y)+blk(i)+0x5A827999+rol(v,5);w=rol(w,30);
+#define R2(v,w,x,y,z,i) z+=(w^x^y)+blk(i)+0x6ED9EBA1+rol(v,5);w=rol(w,30);
+#define R3(v,w,x,y,z,i) z+=(((w|x)&y)|(w&x))+blk(i)+0x8F1BBCDC+rol(v,5);w=rol(w,30);
+#define R4(v,w,x,y,z,i) z+=(w^x^y)+blk(i)+0xCA62C1D6+rol(v,5);w=rol(w,30);
+
+typedef struct {
+    uint32_t state[5];
+    uint32_t count[2];
+    unsigned char buffer[64];
+} SHA1_CTX;
+
+static void SHA1Transform(uint32_t state[5], const unsigned char buffer[64]) {
+  uint32_t a, b, c, d, e;
+  union char64long16 block[1];
+
+  memcpy(block, buffer, 64);
+  a = state[0];
+  b = state[1];
+  c = state[2];
+  d = state[3];
+  e = state[4];
+  R0(a,b,c,d,e, 0); R0(e,a,b,c,d, 1); R0(d,e,a,b,c, 2); R0(c,d,e,a,b, 3);
+  R0(b,c,d,e,a, 4); R0(a,b,c,d,e, 5); R0(e,a,b,c,d, 6); R0(d,e,a,b,c, 7);
+  R0(c,d,e,a,b, 8); R0(b,c,d,e,a, 9); R0(a,b,c,d,e,10); R0(e,a,b,c,d,11);
+  R0(d,e,a,b,c,12); R0(c,d,e,a,b,13); R0(b,c,d,e,a,14); R0(a,b,c,d,e,15);
+  R1(e,a,b,c,d,16); R1(d,e,a,b,c,17); R1(c,d,e,a,b,18); R1(b,c,d,e,a,19);
+  R2(a,b,c,d,e,20); R2(e,a,b,c,d,21); R2(d,e,a,b,c,22); R2(c,d,e,a,b,23);
+  R2(b,c,d,e,a,24); R2(a,b,c,d,e,25); R2(e,a,b,c,d,26); R2(d,e,a,b,c,27);
+  R2(c,d,e,a,b,28); R2(b,c,d,e,a,29); R2(a,b,c,d,e,30); R2(e,a,b,c,d,31);
+  R2(d,e,a,b,c,32); R2(c,d,e,a,b,33); R2(b,c,d,e,a,34); R2(a,b,c,d,e,35);
+  R2(e,a,b,c,d,36); R2(d,e,a,b,c,37); R2(c,d,e,a,b,38); R2(b,c,d,e,a,39);
+  R3(a,b,c,d,e,40); R3(e,a,b,c,d,41); R3(d,e,a,b,c,42); R3(c,d,e,a,b,43);
+  R3(b,c,d,e,a,44); R3(a,b,c,d,e,45); R3(e,a,b,c,d,46); R3(d,e,a,b,c,47);
+  R3(c,d,e,a,b,48); R3(b,c,d,e,a,49); R3(a,b,c,d,e,50); R3(e,a,b,c,d,51);
+  R3(d,e,a,b,c,52); R3(c,d,e,a,b,53); R3(b,c,d,e,a,54); R3(a,b,c,d,e,55);
+  R3(e,a,b,c,d,56); R3(d,e,a,b,c,57); R3(c,d,e,a,b,58); R3(b,c,d,e,a,59);
+  R4(a,b,c,d,e,60); R4(e,a,b,c,d,61); R4(d,e,a,b,c,62); R4(c,d,e,a,b,63);
+  R4(b,c,d,e,a,64); R4(a,b,c,d,e,65); R4(e,a,b,c,d,66); R4(d,e,a,b,c,67);
+  R4(c,d,e,a,b,68); R4(b,c,d,e,a,69); R4(a,b,c,d,e,70); R4(e,a,b,c,d,71);
+  R4(d,e,a,b,c,72); R4(c,d,e,a,b,73); R4(b,c,d,e,a,74); R4(a,b,c,d,e,75);
+  R4(e,a,b,c,d,76); R4(d,e,a,b,c,77); R4(c,d,e,a,b,78); R4(b,c,d,e,a,79);
+  state[0] += a;
+  state[1] += b;
+  state[2] += c;
+  state[3] += d;
+  state[4] += e;
+  a = b = c = d = e = 0;
+  memset(block, '\0', sizeof(block));
+}
+
+static void SHA1Init(SHA1_CTX* context) {
+  context->state[0] = 0x67452301;
+  context->state[1] = 0xEFCDAB89;
+  context->state[2] = 0x98BADCFE;
+  context->state[3] = 0x10325476;
+  context->state[4] = 0xC3D2E1F0;
+  context->count[0] = context->count[1] = 0;
+}
+
+static void SHA1Update(SHA1_CTX* context, const unsigned char* data,
+                       uint32_t len) {
+  uint32_t i, j;
+
+  j = context->count[0];
+  if ((context->count[0] += len << 3) < j)
+    context->count[1]++;
+  context->count[1] += (len>>29);
+  j = (j >> 3) & 63;
+  if ((j + len) > 63) {
+    memcpy(&context->buffer[j], data, (i = 64-j));
+    SHA1Transform(context->state, context->buffer);
+    for ( ; i + 63 < len; i += 64) {
+      SHA1Transform(context->state, &data[i]);
+    }
+    j = 0;
+  }
+  else i = 0;
+  memcpy(&context->buffer[j], &data[i], len - i);
+}
+
+static void SHA1Final(unsigned char digest[20], SHA1_CTX* context) {
+  unsigned i;
+  unsigned char finalcount[8], c;
+
+  for (i = 0; i < 8; i++) {
+    finalcount[i] = (unsigned char)((context->count[(i >= 4 ? 0 : 1)]
+                                     >> ((3-(i & 3)) * 8) ) & 255);
+  }
+  c = 0200;
+  SHA1Update(context, &c, 1);
+  while ((context->count[0] & 504) != 448) {
+    c = 0000;
+    SHA1Update(context, &c, 1);
+  }
+  SHA1Update(context, finalcount, 8);
+  for (i = 0; i < 20; i++) {
+    digest[i] = (unsigned char)
+      ((context->state[i>>2] >> ((3-(i & 3)) * 8) ) & 255);
+  }
+  memset(context, '\0', sizeof(*context));
+  memset(&finalcount, '\0', sizeof(finalcount));
+}
+// END OF SHA1 CODE
+
+static void base64_encode(const unsigned char *src, int src_len, char *dst) {
+  static const char *b64 =
+    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
+  int i, j, a, b, c;
+
+  for (i = j = 0; i < src_len; i += 3) {
+    a = src[i];
+    b = i + 1 >= src_len ? 0 : src[i + 1];
+    c = i + 2 >= src_len ? 0 : src[i + 2];
+
+    dst[j++] = b64[a >> 2];
+    dst[j++] = b64[((a & 3) << 4) | (b >> 4)];
+    if (i + 1 < src_len) {
+      dst[j++] = b64[(b & 15) << 2 | (c >> 6)];
+    }
+    if (i + 2 < src_len) {
+      dst[j++] = b64[c & 63];
+    }
+  }
+  while (j % 4 != 0) {
+    dst[j++] = '=';
+  }
+  dst[j++] = '\0';
+}
+
+static void send_websocket_handshake(struct mg_connection *conn,
+                                     const char *key) {
+  static const char *magic = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
+  char buf[500], sha[20], b64_sha[sizeof(sha) * 2];
+  SHA1_CTX sha_ctx;
+
+  snprintf(buf, sizeof(buf), "%s%s", key, magic);
+  SHA1Init(&sha_ctx);
+  SHA1Update(&sha_ctx, (unsigned char *) buf, strlen(buf));
+  SHA1Final((unsigned char *) sha, &sha_ctx);
+  base64_encode((unsigned char *) sha, sizeof(sha), b64_sha);
+  snprintf(buf, sizeof(buf), "%s%s%s",
+           "HTTP/1.1 101 Switching Protocols\r\n"
+           "Upgrade: websocket\r\n"
+           "Connection: Upgrade\r\n"
+           "Sec-WebSocket-Accept: ", b64_sha, "\r\n\r\n");
+
+  mg_write(conn, buf, strlen(buf));
+}
+
+static int deliver_websocket_frame(struct connection *conn) {
+  char *buf = conn->local_iobuf.buf;
+  int i, len, buf_len = conn->local_iobuf.len, frame_len = 0,
+      mask_len = 0, header_len = 0, data_len = 0, buffered = 0;
+
+  if (buf_len >= 2) {
+    len = buf[1] & 127;
+    mask_len = buf[1] & 128 ? 4 : 0;
+    if (len < 126 && buf_len >= mask_len) {
+      data_len = len;
+      header_len = 2 + mask_len;
+    } else if (len == 126 && buf_len >= 4 + mask_len) {
+      header_len = 4 + mask_len;
+      data_len = ((((int) buf[2]) << 8) + buf[3]);
+    } else if (buf_len >= 10 + mask_len) {
+      header_len = 10 + mask_len;
+      data_len = (((uint64_t) htonl(* (uint32_t *) &buf[2])) << 32) +
+        htonl(* (uint32_t *) &buf[6]);
+    }
+  }
+
+  frame_len = header_len + data_len;
+  buffered = frame_len > 0 && frame_len <= buf_len;
+
+  if (buffered) {
+    conn->mg_conn.content_len = data_len;
+    conn->mg_conn.content = buf + header_len;
+
+    // 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];
+      }
+    }
+
+    // Call the handler and remove frame from the iobuf
+    if (conn->endpoint.uh->handler(&conn->mg_conn)) {
+      conn->flags |= CONN_SPOOL_DONE;
+    }
+    memmove(buf, buf + frame_len, buf_len - frame_len);
+    conn->local_iobuf.len -= frame_len;
+  }
+
+  return buffered;
+}
+
+#if 0
+int mg_websocket_read(struct mg_connection *conn, int *bits, char **data) {
+  // Pointer to the beginning of the portion of the incoming websocket message
+  // queue. The original websocket upgrade request is never removed,
+  // so the queue begins after it.
+  unsigned char *buf = (unsigned char *) conn->buf + conn->request_len;
+  int n, stop = 0;
+  size_t i, len, mask_len, data_len, header_len, body_len;
+  char mask[4];
+
+  assert(conn->content_len == 0);
+
+  // Loop continuously, reading messages from the socket, invoking the callback,
+  // and waiting repeatedly until an error occurs.
+  while (!stop) {
+    header_len = 0;
+    // body_len is the length of the entire queue in bytes
+    // len is the length of the current message
+    // data_len is the length of the current message's data payload
+    // header_len is the length of the current message's header
+    if ((body_len = conn->data_len - conn->request_len) >= 2) {
+      len = buf[1] & 127;
+      mask_len = buf[1] & 128 ? 4 : 0;
+      if (len < 126 && body_len >= mask_len) {
+        data_len = len;
+        header_len = 2 + mask_len;
+      } else if (len == 126 && body_len >= 4 + mask_len) {
+        header_len = 4 + mask_len;
+        data_len = ((((int) buf[2]) << 8) + buf[3]);
+      } else if (body_len >= 10 + mask_len) {
+        header_len = 10 + mask_len;
+        data_len = (((uint64_t) htonl(* (uint32_t *) &buf[2])) << 32) +
+          htonl(* (uint32_t *) &buf[6]);
+      }
+    }
+
+    // Data layout is as follows:
+    //  conn->buf               buf
+    //     v                     v              frame1           | frame2
+    //     |---------------------|----------------|--------------|-------
+    //     |                     |<--header_len-->|<--data_len-->|
+    //     |<-conn->request_len->|<-----body_len----------->|
+    //     |<-------------------conn->data_len------------->|
+
+    if (header_len > 0) {
+      // Allocate space to hold websocket payload
+      if ((*data = malloc(data_len)) == NULL) {
+        // Allocation failed, exit the loop and then close the connection
+        // TODO: notify user about the failure
+        data_len = 0;
+        break;
+      }
+
+      // Save mask and bits, otherwise it may be clobbered by memmove below
+      *bits = buf[0];
+      memcpy(mask, buf + header_len - mask_len, mask_len);
+
+      // Read frame payload into the allocated buffer.
+      assert(body_len >= header_len);
+      if (data_len + header_len > body_len) {
+        len = body_len - header_len;
+        memcpy(*data, buf + header_len, len);
+        // TODO: handle pull error
+        pull_all(NULL, conn, *data + len, data_len - len);
+        conn->data_len = conn->request_len;
+      } else {
+        len = data_len + header_len;
+        memcpy(*data, buf + header_len, data_len);
+        memmove(buf, buf + len, body_len - len);
+        conn->data_len -= len;
+      }
+
+      // Apply mask if necessary
+      if (mask_len > 0) {
+        for (i = 0; i < data_len; i++) {
+          (*data)[i] ^= mask[i % 4];
+        }
+      }
+
+      return data_len;
+    } else {
+      // Buffering websocket request
+      if ((n = pull(NULL, conn, conn->buf + conn->data_len,
+                    conn->buf_size - conn->data_len)) <= 0) {
+        break;
+      }
+      conn->data_len += n;
+    }
+  }
+
+  return 0;
+}
+#endif
+
+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;
+    }
+
+    if (copy_len > 0) {
+      retval = mg_write(conn, copy, copy_len);
+    }
+    free(copy);
+
+    return retval;
+}
+#else
+static void send_websocket_handshake(struct mg_connection *conn,
+                                     const char *key) {
+  (void) key;
+  send_http_error(conn, "%s", "HTTP/1.1 501 Not Implemented\r\n\r\n");
+}
+#endif // !USE_WEBSOCKET
+
 static int is_error(int n) {
   return n == 0 || (n < 0 && errno != EINTR && errno != EAGAIN);
 }
@@ -899,7 +1264,7 @@ static void write_to_client(struct connection *conn) {
   struct iobuf *io = &conn->remote_iobuf;
   int n = send(conn->client_sock, io->buf, io->len, 0);
 
-  //DBG(("Written %d of %d(%d): [%.*s]", n, io->len, io->size, 0, io->buf));
+  DBG(("Written %d of %d(%d): [%.*s]", n, io->len, io->size, 0, io->buf));
 
   if (is_error(n)) {
     conn->flags |= CONN_CLOSE;
@@ -1152,8 +1517,14 @@ static void send_http_error(struct connection *conn, const char *fmt, ...) {
 }
 
 static void call_uri_handler_if_data_is_buffered(struct connection *conn) {
-  if (conn->local_iobuf.len >= conn->mg_conn.content_len) {
-    conn->endpoint.uh->handler(&conn->mg_conn);
+  struct iobuf *loc = &conn->local_iobuf;
+  struct mg_connection *c = &conn->mg_conn;
+
+  c->content = loc->buf;
+  if (conn->mg_conn.is_websocket) {
+    do { } while (deliver_websocket_frame(conn));
+  } else if (loc->len >= c->content_len) {
+    conn->endpoint.uh->handler(c);
     close_local_endpoint(conn);
   }
 }
@@ -1203,6 +1574,15 @@ static void send_continue_if_expected(struct connection *conn) {
   }
 }
 
+static void send_websocket_handshake_if_requested(struct mg_connection *conn) {
+  const char *ver = mg_get_header(conn, "Sec-WebSocket-Version"),
+        *key = mg_get_header(conn, "Sec-WebSocket-Key");
+  if (ver != NULL && key != NULL) {
+    conn->is_websocket = 1;
+    send_websocket_handshake(conn, key);
+  }
+}
+
 static void process_request(struct connection *conn) {
   struct iobuf *io = &conn->local_iobuf;
 
@@ -1226,6 +1606,7 @@ static void process_request(struct connection *conn) {
     // Invalid request, or request is too big: close the connection
     conn->flags |= CONN_CLOSE;
   } else if (conn->request_len > 0 && conn->endpoint_type == EP_NONE) {
+    send_websocket_handshake_if_requested(&conn->mg_conn);
     send_continue_if_expected(conn);
     open_local_endpoint(conn);
   } else if (conn->endpoint_type == EP_USER) {
@@ -1475,18 +1856,20 @@ struct mg_server *mg_create_server(void *server_data) {
 // End of library, start of the application code
 
 static void iterate_callback(struct mg_connection *c, void *param) {
-  if (c->connection_param != NULL) {
+  if (c->is_websocket) {
     char buf[20];
     int len = snprintf(buf, sizeof(buf), "%d", * (int *) param);
-    mg_write(c, buf, len);
+    mg_websocket_write(c, 1, buf, len);
   }
 }
 
+// This thread sends heartbeats to all websocket connections with 1s interval.
+// The heartbeat message is simply an iteration counter.
 static void *timer_thread(void *param) {
   struct mg_server *server = (struct mg_server *) param;
   int i;
 
-  for (i = 0; i < 1000; i++) {
+  for (i = 0; i < 9999999; i++) {
     sleep(1);
     mg_iterate_over_connections(server, iterate_callback, &i);
   }
@@ -1494,28 +1877,28 @@ static void *timer_thread(void *param) {
   return NULL;
 }
 
-static int websocket_handler(struct mg_connection *conn) {
-  char headers[500], content[500];
-  int headers_len, content_len;
+// This handler is called for each incoming websocket frame, one or more
+// times for connection lifetime.
+static int handler(struct mg_connection *conn) {
+  static const char oops[] = "HTTP/1.0 200 OK\r\n\r\nwebsocket data expected\n";
 
-  content_len = snprintf(content, sizeof(content), "%s %s, POST len %d\n",
-                         conn->request_method, conn->uri, conn->content_len);
-  headers_len = snprintf(headers, sizeof(headers), "HTTP/1.0 200 OK\r\n"
-                         "Content-Length: %d\r\n\r\n", content_len);
+  if (!conn->is_websocket) {
+    mg_write(conn, oops, sizeof(oops) - 1);
+    return 1;
+  }
 
-  mg_write(conn, headers, headers_len);
-  mg_write(conn, content, content_len);
+  mg_websocket_write(conn, 1, conn->content, conn->content_len);
 
-  return 1;
+  DBG(("WS msg len: %d", conn->content_len));
+  return conn->content_len == 4 && !memcmp(conn->content, "exit", 4);
 }
 
-
 int main(void) {
   struct mg_server *server = mg_create_server(NULL);
 
   mg_set_option(server, "listening_port", "8080");
   mg_set_option(server, "document_root", ".");
-  mg_add_uri_handler(server, "/ws", websocket_handler);
+  mg_add_uri_handler(server, "/ws", handler);
   mg_start_thread(timer_thread, server);
 
   printf("Started on port %s\n", mg_get_option(server, "listening_port"));

build/src/core.h

@@ -46,6 +46,7 @@ struct mg_connection {
   char *content;              // POST (or websocket message) data, or NULL
   int content_len;            // content length
 
+  int is_websocket;           // Connection is a websocket connection
   void *server_param;         // Parameter passed to mg_add_uri_handler()
   void *connection_param;     // Placeholder for connection-specific data
 };