From 3998aa86eb7d339ae9994815e4f13314b4f3a3ee Mon Sep 17 00:00:00 2001
From: Kieran O'Mahony <kieranom@gmail.com>
Date: Mon, 11 Jun 2012 23:38:53 +0100
Subject: [PATCH] check if malloc succeeds to avoid segmentation faults when
 ujson runs out of memory

---
 lib/ultrajsondec.c |  8 ++++++++
 lib/ultrajsonenc.c | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/lib/ultrajsondec.c b/lib/ultrajsondec.c
index 275b1a1..0b2f304 100644
--- a/lib/ultrajsondec.c
+++ b/lib/ultrajsondec.c
@@ -420,12 +420,20 @@ FASTCALL_ATTR JSOBJ FASTCALL_MSVC decode_string ( struct DecoderState *ds)
         if (ds->escHeap)
         {
             ds->escStart = (wchar_t *) ds->dec->realloc (ds->escStart, newSize * sizeof(wchar_t));
+            if (!ds->escStart)
+            {
+                return SetError(ds, -1, "Could not reserve memory block");
+            }
         }
         else
         {
             wchar_t *oldStart = ds->escStart;
             ds->escHeap = 1;
             ds->escStart = (wchar_t *) ds->dec->malloc (newSize * sizeof(wchar_t));
+            if (!ds->escStart)
+            {
+                return SetError(ds, -1, "Could not reserve memory block");
+            }
             memcpy (ds->escStart, oldStart, escLen * sizeof(wchar_t));
         }
 
diff --git a/lib/ultrajsonenc.c b/lib/ultrajsonenc.c
index b497ace..94491c4 100644
--- a/lib/ultrajsonenc.c
+++ b/lib/ultrajsonenc.c
@@ -105,12 +105,24 @@ void Buffer_Realloc (JSONObjectEncoder *enc, size_t cbNeeded)
     if (enc->heap)
     {
         enc->start = (char *) enc->realloc (enc->start, newSize);
+        if (!enc->start)
+        {
+            SetError (NULL, enc, "Could not reserve memory block");
+            return;
+        }
+        
     }
     else
     {
         char *oldStart = enc->start;
         enc->heap = 1;
         enc->start = (char *) enc->malloc (newSize);
+        if (!enc->start)
+        {
+            SetError (NULL, enc, "Could not reserve memory block");
+            return;
+        }
+        
         memcpy (enc->start, oldStart, offset);
     }
     enc->offset = enc->start + offset;
@@ -629,6 +641,10 @@ void encode(JSOBJ obj, JSONObjectEncoder *enc, const char *name, size_t cbName)
     */
 
     Buffer_Reserve(enc, 256 + (((cbName / 4) + 1) * 12));
+    if (enc->errorMsg)
+    {
+        return;
+    }
 
     if (name)
     {
@@ -780,6 +796,12 @@ void encode(JSOBJ obj, JSONObjectEncoder *enc, const char *name, size_t cbName)
         {
             value = enc->getStringValue(obj, &tc, &szlen);
             Buffer_Reserve(enc, ((szlen / 4) + 1) * 12);
+            if (enc->errorMsg)
+            {
+                enc->endTypeContext(obj, &tc);
+                return;
+            }
+            
             Buffer_AppendCharUnchecked (enc, '\"');
 
 
@@ -836,6 +858,12 @@ char *JSON_EncodeObject(JSOBJ obj, JSONObjectEncoder *enc, char *_buffer, size_t
     {
         _cbBuffer = 32768;
         enc->start = (char *) enc->malloc (_cbBuffer);
+        if (!enc->start)
+        {
+            SetError(obj, enc, "Could not reserve memory block");
+            return NULL;
+        }
+        
         enc->heap = 1;
     }
     else
@@ -851,6 +879,11 @@ char *JSON_EncodeObject(JSOBJ obj, JSONObjectEncoder *enc, char *_buffer, size_t
     encode (obj, enc, NULL, 0);
     
     Buffer_Reserve(enc, 1);
+    if (enc->errorMsg)
+    {
+        return NULL;
+    }
+    
     Buffer_AppendCharUnchecked(enc, '\0');
 
     return enc->start;