From 2a12b3fa4b8cdef53288aa2667df0dad2d7f2fc2 Mon Sep 17 00:00:00 2001
From: Mark Pulford <mark@kyne.com.au>
Date: Mon, 25 Apr 2011 19:15:56 +0930
Subject: Grow decode stack, prealloc strings during encode

- Check stack usage during decode to prevent crashing in excessively
  nested data structures.
- Preallocate the required memory for json_append_string().
---
 lua_json.c | 34 +++++++++++++++++++++-------------
 strbuf.c   |  4 +---
 strbuf.h   | 14 ++++++++++++--
 3 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/lua_json.c b/lua_json.c
index 1753e49..1d334ea 100644
--- a/lua_json.c
+++ b/lua_json.c
@@ -7,24 +7,18 @@
  *   json.null.
  * - Parsing comments is not support. According to json.org, this isn't
  *   part of the spec.
+ *
+ * Note: lua_json_decode() probably spends significant time rehashing
+ *       tables since it is difficult to know their size ahead of time.
+ *       Earlier JSON libaries didn't have this problem but the intermediate
+ *       storage (and their implementations) were much slower anyway..
  */
 
 /* FIXME:
  * - Ensure JSON data is UTF-8. Fail otherwise.
  *   - Alternatively, dynamically support Unicode in JSON string. Return current locale.
- * - Use lua_checkstack() to ensure there is enough stack space left to
- *   fulfill an operation. What happens if we don't, is that acceptible too?
- *   Does lua_checkstack grow the stack, or merely check if it is possible?
- */
-
-/* FIXME:
+ * - Consider implementing other Unicode standards.
  * - Option to encode non-printable characters? Only \" \\ are required
- * - Unicode?
- */
-
-/* FIXME:
- * - Review memory allocation handling and error returns.
- *   Ensure all memory is free. Including after exceptions.
  */
 
 #include <assert.h>
@@ -94,13 +88,19 @@ static void json_append_string(lua_State *l, strbuf_t *json, int lindex)
 
     str = lua_tolstring(l, lindex, &len);
 
+    /* Worst case is len * 6 (all unicode escapes).
+     * This buffer is reused constantly for small strings
+     * If there are any excess pages, they won't be hit anyway.
+     * This gains ~5% speedup. */
+    strbuf_ensure_empty_length(json, len * 6);
+
     strbuf_append_char(json, '\"');
     for (i = 0; i < len; i++) {
         p = json_escape_char(str[i]);
         if (p)
             strbuf_append_string(json, p);
         else
-            strbuf_append_char(json, str[i]);
+            strbuf_append_char_unsafe(json, str[i]);
     }
     strbuf_append_char(json, '\"');
 }
@@ -536,6 +536,10 @@ static void json_parse_object_context(lua_State *l, json_parse_t *json)
 {
     json_token_t token;
 
+    /* 3 slots required:
+     * .., table, key, value */
+    luaL_checkstack(l, 3, "too many nested data structures");
+
     lua_newtable(l);
 
     json_next_token(json, &token);
@@ -576,6 +580,10 @@ static void json_parse_array_context(lua_State *l, json_parse_t *json)
     json_token_t token;
     int i;
 
+    /* 2 slots required:
+     * .., table, value */
+    luaL_checkstack(l, 2, "too many nested data structures");
+
     lua_newtable(l);
 
     json_next_token(json, &token);
diff --git a/strbuf.c b/strbuf.c
index 227cb47..44145dd 100644
--- a/strbuf.c
+++ b/strbuf.c
@@ -152,9 +152,7 @@ void strbuf_resize(strbuf_t *s, int len)
 
 void strbuf_append_mem(strbuf_t *s, const char *c, int len)
 {
-    if (len > strbuf_empty_length(s))
-        strbuf_resize(s, s->length + len);
-
+    strbuf_ensure_empty_length(s, len);
     memcpy(s->buf + s->length, c, len);
     s->length += len;
 }
diff --git a/strbuf.h b/strbuf.h
index 914274c..d6e2d90 100644
--- a/strbuf.h
+++ b/strbuf.h
@@ -38,6 +38,7 @@ extern void strbuf_resize(strbuf_t *s, int len);
 static int strbuf_empty_length(strbuf_t *s);
 static int strbuf_length(strbuf_t *s);
 static char *strbuf_string(strbuf_t *s, int *len);
+static  void strbuf_ensure_empty_length(strbuf_t *s, int len);
 
 /* Update */
 extern void strbuf_append_fmt(strbuf_t *s, const char *format, ...);
@@ -53,6 +54,12 @@ static inline int strbuf_empty_length(strbuf_t *s)
     return s->size - s->length - 1;
 }
 
+static inline void strbuf_ensure_empty_length(strbuf_t *s, int len)
+{
+    if (len > strbuf_empty_length(s))
+        strbuf_resize(s, s->length + len);
+}
+
 static inline int strbuf_length(strbuf_t *s)
 {
     return s->length;
@@ -60,9 +67,12 @@ static inline int strbuf_length(strbuf_t *s)
 
 static inline void strbuf_append_char(strbuf_t *s, const char c)
 {
-    if (strbuf_empty_length(s) < 1)
-        strbuf_resize(s, s->length + 1);
+    strbuf_ensure_empty_length(s, 1);
+    s->buf[s->length++] = c;
+}
 
+static inline void strbuf_append_char_unsafe(strbuf_t *s, const char c)
+{
     s->buf[s->length++] = c;
 }
 
-- 
cgit v1.2.3-55-g6feb