ext/yajl/yajl_buf.c in yajl-ruby-1.4.1 vs ext/yajl/yajl_buf.c in yajl-ruby-1.4.2

@@ -33,47 +33,118 @@
 #include "yajl_buf.h"
 
 #include <assert.h>
 #include <stdlib.h>
 #include <string.h>
+#include <stdio.h>
 
 #define YAJL_BUF_INIT_SIZE 2048
 
 struct yajl_buf_t {
+    yajl_buf_state state;
     unsigned int len;
     unsigned int used;
     unsigned char * data;
     yajl_alloc_funcs * alloc;
 };
 
+static void *noop_realloc(void *ctx, void *ptr, unsigned int sz) {
+    fprintf(stderr, "Attempt to allocate on invalid yajl_buf_t\n");
+    abort();
+}
+static void *noop_malloc(void *ctx, unsigned int sz) { return noop_realloc(ctx, NULL, sz); }
+static void noop_free(void *ctx, void *ptr) { }
+
+static yajl_alloc_funcs noop_allocs = {
+    .malloc = &noop_malloc,
+    .realloc = &noop_realloc,
+    .free = &noop_free,
+};
+
+// A buffer to be returned if the initial allocation fails
+static struct yajl_buf_t buf_alloc_error = {
+    .state = yajl_buf_alloc_failed,
+    .alloc = &noop_allocs
+};
+
+#include <stdio.h>
+
+yajl_buf_state yajl_buf_err(yajl_buf buf)
+{
+    assert(buf);
+    return buf->state;
+}
+
 static
-void yajl_buf_ensure_available(yajl_buf buf, unsigned int want)
+yajl_buf_state yajl_buf_set_error(yajl_buf buf, yajl_buf_state err)
 {
+    buf->state = err;
+
+    // free and clear all data from the buffer
+    YA_FREE(buf->alloc, buf->data);
+    buf->len = 0;
+    buf->data = 0;
+    buf->used = 0;
+
+    return err;
+}
+
+static
+yajl_buf_state yajl_buf_ensure_available(yajl_buf buf, unsigned int want)
+{
     unsigned int need;
     
     assert(buf != NULL);
 
+    if (buf->state != yajl_buf_ok) {
+        return buf->state;
+    }
+
     /* first call */
     if (buf->data == NULL) {
         buf->len = YAJL_BUF_INIT_SIZE;
         buf->data = (unsigned char *) YA_MALLOC(buf->alloc, buf->len);
+        if (buf->data == NULL)  {
+            return yajl_buf_set_error(buf, yajl_buf_overflow);
+        }
+
         buf->data[0] = 0;
     }
 
+    if (want == 0) {
+        return yajl_buf_ok;
+    }
+
     need = buf->len;
 
     while (want >= (need - buf->used)) need <<= 1;
 
+    // Check for overflow
+    if (need < buf->used) {
+        return yajl_buf_set_error(buf, yajl_buf_overflow);
+    }
+
     if (need != buf->len) {
         buf->data = (unsigned char *) YA_REALLOC(buf->alloc, buf->data, need);
+
+        if (buf->data == NULL)  {
+            return yajl_buf_set_error(buf, yajl_buf_overflow);
+        }
+
         buf->len = need;
     }
+
+    return yajl_buf_ok;
 }
 
 yajl_buf yajl_buf_alloc(yajl_alloc_funcs * alloc)
 {
     yajl_buf b = YA_MALLOC(alloc, sizeof(struct yajl_buf_t));
+    if (b == NULL) {
+        return &buf_alloc_error;
+    }
+
     memset((void *) b, 0, sizeof(struct yajl_buf_t));
     b->alloc = alloc;
     return b;
 }
 
@@ -84,36 +155,46 @@
     YA_FREE(buf->alloc, buf);
 }
 
 void yajl_buf_append(yajl_buf buf, const void * data, unsigned int len)
 {
-    yajl_buf_ensure_available(buf, len);
+    if (yajl_buf_ensure_available(buf, len)) {
+        return;
+    }
     if (len > 0) {
         assert(data != NULL);
         memcpy(buf->data + buf->used, data, len);
         buf->used += len;
         buf->data[buf->used] = 0;
     }
 }
 
 void yajl_buf_clear(yajl_buf buf)
 {
+    assert(buf);
+    assert(!yajl_buf_err(buf));
     buf->used = 0;
     if (buf->data) buf->data[buf->used] = 0;
 }
 
 const unsigned char * yajl_buf_data(yajl_buf buf)
 {
+    assert(buf);
+    assert(!yajl_buf_err(buf));
     return buf->data;
 }
 
 unsigned int yajl_buf_len(yajl_buf buf)
 {
+    assert(buf);
+    assert(!yajl_buf_err(buf));
     return buf->used;
 }
 
 void
 yajl_buf_truncate(yajl_buf buf, unsigned int len)
 {
+    assert(buf);
+    assert(!yajl_buf_err(buf));
     assert(len <= buf->used);
     buf->used = len;
 }