C wsproxy: seq numbers and decode multiple frames.

websocket.c

@@ -15,6 +15,7 @@
 #include <arpa/inet.h>
 #include <openssl/err.h>
 #include <openssl/ssl.h>
+#include <resolv.h>      /* base64 encode/decode */
 #include "websocket.h"
 
 const char server_handshake[] = "HTTP/1.1 101 Web Socket Protocol Handshake\r\n\
@@ -27,6 +28,16 @@ WebSocket-Protocol: sample\r\n\
 
 const char policy_response[] = "<cross-domain-policy><allow-access-from domain=\"*\" to-ports=\"*\" /></cross-domain-policy>\n";
 
+/*
+ * Global state
+ *
+ *   Warning: not thread safe
+ */
+int ssl_initialized = 0;
+char *tbuf, *cbuf, *tbuf_tmp, *cbuf_tmp;
+unsigned int bufsize, dbufsize;
+client_settings_t client_settings;
+
 void traffic(char * token) {
     fprintf(stdout, "%s", token);
     fflush(stdout);
@@ -47,9 +58,6 @@ void fatal(char *msg)
  * SSL Wrapper Code
  */
 
-/*   Warning: not thread safe */
-int ssl_initialized = 0;
-
 ssize_t ws_recv(ws_ctx_t *ctx, void *buf, size_t len) {
     if (ctx->ssl) {
         //printf("SSL recv\n");
@@ -147,7 +155,56 @@ int ws_socket_free(ws_ctx_t *ctx) {
 /* ------------------------------------------------------- */
 
 
-ws_ctx_t *do_handshake(int sock, client_settings_t *client_settings) {
+int encode(u_char const *src, size_t srclength, char *target, size_t targsize) {
+    int sz = 0, len = 0;
+    target[sz++] = '\x00';
+    if (client_settings.do_seq_num) {
+        sz += sprintf(target+sz, "%d:", client_settings.seq_num);
+        client_settings.seq_num++;
+    }
+    if (client_settings.do_b64encode) {
+        len = __b64_ntop(src, srclength, target+sz, targsize-sz);
+    } else {
+        fatal("UTF-8 not yet implemented");
+    }
+    if (len < 0) {
+        return len;
+    }
+    sz += len;
+    target[sz++] = '\xff';
+    return sz;
+}
+
+int decode(char *src, size_t srclength, u_char *target, size_t targsize) {
+    char *start, *end;
+    int len, retlen = 0;
+    if ((src[0] != '\x00') || (src[srclength-1] != '\xff')) {
+        fprintf(stderr, "WebSocket framing error\n");
+        return -1;
+    }
+    start = src+1; // Skip '\x00' start
+    do {
+        /* We may have more than one frame */
+        end = strchr(start, '\xff');
+        if (end < (src+srclength-1)) {
+            printf("More than one frame to decode\n");
+        }
+        *end = '\x00';
+        if (client_settings.do_b64encode) {
+            len = __b64_pton(start, target+retlen, targsize-retlen);
+        } else {
+            fatal("UTF-8 not yet implemented");
+        }
+        if (len < 0) {
+            return len;
+        }
+        retlen += len;
+        start = end + 2; // Skip '\xff' end and '\x00' start 
+    } while (end < (src+srclength-1));
+    return retlen;
+}
+
+ws_ctx_t *do_handshake(int sock) {
     char handshake[4096], response[4096];
     char *scheme, *line, *path, *host, *origin;
     char *args_start, *args_end, *arg_idx;
@@ -155,8 +212,9 @@ ws_ctx_t *do_handshake(int sock, client_settings_t *client_settings) {
     ws_ctx_t * ws_ctx;
 
     // Reset settings
-    client_settings->b64encode = 0;
-    client_settings->seq_num = 0;
+    client_settings.do_b64encode = 0;
+    client_settings.do_seq_num = 0;
+    client_settings.seq_num = 0;
 
     len = recv(sock, handshake, 1024, MSG_PEEK);
     handshake[len] = 0;
@@ -211,12 +269,12 @@ ws_ctx_t *do_handshake(int sock, client_settings_t *client_settings) {
         arg_idx = strstr(args_start, "b64encode");
         if (arg_idx && arg_idx < args_end) {
             //printf("setting b64encode\n");
-            client_settings->b64encode = 1;
+            client_settings.do_b64encode = 1;
         }
         arg_idx = strstr(args_start, "seq_num");
         if (arg_idx && arg_idx < args_end) {
             //printf("setting seq_num\n");
-            client_settings->seq_num = 1;
+            client_settings.do_seq_num = 1;
         }
     }
 
@@ -228,12 +286,22 @@ ws_ctx_t *do_handshake(int sock, client_settings_t *client_settings) {
 }
 
 void start_server(int listen_port,
-                  void (*handler)(ws_ctx_t*),
-                  client_settings_t *client_settings) {
+                  void (*handler)(ws_ctx_t*)) {
     int lsock, csock, clilen, sopt = 1;
     struct sockaddr_in serv_addr, cli_addr;
     ws_ctx_t *ws_ctx;
 
+    /* Initialize buffers */
+    bufsize = 65536;
+    if (! (tbuf = malloc(bufsize)) )
+            { fatal("malloc()"); }
+    if (! (cbuf = malloc(bufsize)) )
+            { fatal("malloc()"); }
+    if (! (tbuf_tmp = malloc(bufsize)) )
+            { fatal("malloc()"); }
+    if (! (cbuf_tmp = malloc(bufsize)) )
+            { fatal("malloc()"); }
+
     lsock = socket(AF_INET, SOCK_STREAM, 0);
     if (lsock < 0) { error("ERROR creating listener socket"); }
     bzero((char *) &serv_addr, sizeof(serv_addr));
@@ -256,8 +324,23 @@ void start_server(int listen_port,
             error("ERROR on accept");
         }
         printf("Got client connection from %s\n", inet_ntoa(cli_addr.sin_addr));
-        ws_ctx = do_handshake(csock, client_settings);
-        if (ws_ctx == NULL) { continue; }
+        ws_ctx = do_handshake(csock);
+        if (ws_ctx == NULL) {
+            close(csock);
+            continue;
+        }
+
+        /* Calculate dbufsize based on client_settings */
+        if (client_settings.do_b64encode) {
+            /* base64 is 4 bytes for every 3
+             *    20 for WS '\x00' / '\xff', seq_num and good measure  */
+            dbufsize = (bufsize * 3)/4 - 20;
+        } else {
+            fatal("UTF-8 not yet implemented");
+            /* UTF-8 encoding is up to 2X larger */
+            dbufsize = (bufsize/2) - 15;
+        }
+
         handler(ws_ctx);
         close(csock);
     }