add more sanity-checking of length prefix

tnetstring/__init__.py

@@ -225,14 +225,19 @@ def load(file):
     function promises not to read more data than necessary.
     """
     #  Read the length prefix one char at a time.
+    #  Note that the netstring spec explicitly forbids padding zeros.
     c = file.read(1)
     if not c.isdigit():
         raise ValueError("not a tnetstring: missing or invalid length prefix")
     datalen = ord(c) - ord("0")
     c = file.read(1)
-    while c.isdigit():
-        datalen = (10 * datalen) + (ord(c) - ord("0"))
-        c = file.read(1)
+    if datalen != 0:
+        while c.isdigit():
+            datalen = (10 * datalen) + (ord(c) - ord("0"))
+            if datalen > 999999999:
+                errmsg = "not a tnetstring: absurdly large length prefix"
+                raise ValueError(errmsg)
+            c = file.read(1)
     if c != ":":
         raise ValueError("not a tnetstring: missing or invalid length prefix")
     #  Now we can read and parse the payload.

tnetstring/_tnetstring.c

@@ -12,6 +12,7 @@
 #include <Python.h>
 
 
+#define TNS_MAX_LENGTH 999999999
 #include "tns_core.c"
 
 
@@ -81,26 +82,44 @@ _tnetstring_load(PyObject* self, PyObject *args)
   }
   c = PyString_AS_STRING(res)[0];
   Py_DECREF(res); res = NULL;
+  //  Note that the netsring spec explicitly forbids padding zeroes.
+  //  If the first char is zero, it must be the only char.
   if(c < '0' || c > '9') {
       PyErr_SetString(PyExc_ValueError,
                       "Not a tnetstring: invlaid or missing length prefix");
       goto error;
-  }
-  do {
-      datalen = (10 * datalen) + (c - '0');
+  } else if (c == '0') {
       res = PyObject_CallMethodObjArgs(file, methnm, metharg, NULL);
       if(res == NULL) {
           goto error;
       }
       Py_INCREF(res);
       if(!PyString_Check(res) || !PyString_GET_SIZE(res)) {
           PyErr_SetString(PyExc_ValueError,
-                          "Not a tnetstring: invlaid or missing length prefix");
+                      "Not a tnetstring: invlaid or missing length prefix");
           goto error;
       }
       c = PyString_AS_STRING(res)[0];
       Py_DECREF(res); res = NULL;
-  } while(c >= '0' && c <= '9');
+  } else {
+      do {
+          datalen = (10 * datalen) + (c - '0');
+          check(datalen < TNS_MAX_LENGTH,
+                "Not a tnetstring: absurdly large length prefix"); 
+          res = PyObject_CallMethodObjArgs(file, methnm, metharg, NULL);
+          if(res == NULL) {
+              goto error;
+          }
+          Py_INCREF(res);
+          if(!PyString_Check(res) || !PyString_GET_SIZE(res)) {
+              PyErr_SetString(PyExc_ValueError,
+                        "Not a tnetstring: invlaid or missing length prefix");
+              goto error;
+          }
+          c = PyString_AS_STRING(res)[0];
+          Py_DECREF(res); res = NULL;
+      } while(c >= '0' && c <= '9');
+  }
 
   //  Validate end-of-length-prefix marker.
   if(c != ':') {

tnetstring/tests/test_format.py

@@ -92,18 +92,31 @@ def test_roundtrip_file_examples(self):
         for data, expect in FORMAT_EXAMPLES.items():
             s = StringIO.StringIO()
             s.write(data)
+            s.write("OK")
             s.seek(0)
             self.assertEqual(expect,tnetstring.load(s))
+            self.assertEqual("OK",s.read())
             s = StringIO.StringIO()
             tnetstring.dump(expect,s)
+            s.write("OK")
             s.seek(0)
             self.assertEqual(expect,tnetstring.load(s))
+            self.assertEqual("OK",s.read())
 
     def test_roundtrip_file_random(self):
         for _ in xrange(500):
             v = get_random_object()
             s = StringIO.StringIO()
             tnetstring.dump(v,s)
+            s.write("OK")
             s.seek(0)
             self.assertEqual(v,tnetstring.load(s))
+            self.assertEqual("OK",s.read())
+
+    def test_error_on_absurd_lengths(self):
+        s = StringIO.StringIO()
+        s.write("1000000000:pwned!,")
+        s.seek(0)
+        self.assertRaises(ValueError,tnetstring.load,s)
+        self.assertEquals(s.read(1),":")
 

tnetstring/tns_core.c

@@ -18,6 +18,10 @@
 
 #include "tns_core.h"
 
+#ifndef TNS_MAX_LENGTH
+#define TNS_MAX_LENGTH 999999999
+#endif
+
 //  These are our internal-use functions.
 
 static int tns_parse_dict(void *dict, const char *data, size_t len);
@@ -277,6 +281,14 @@ tns_strtosz(const char *data, size_t len, size_t *sz, char **end)
   value = c - '0';
   pos++;
 
+  //  The netstring spec explicitly forbits padding zeros.
+  //  If it's a zero, we must be at end.
+  if(value == 0) {
+      *sz = value;
+      *end = (char*) pos;
+      return 0;
+  }
+
   //  Consume all other digits.
   while(pos < eod) {
       c = *pos;

tools/shootout.py

@@ -29,33 +29,33 @@ def add_test(v):
 
 def thrash_tnetstring():
     for obj, tns, json, msh in TESTS:
-#        tnetstring.dumps(obj)
+        tnetstring.dumps(obj)
         assert tnetstring.loads(tns) == obj
-#        assert tnetstring.loads(tnetstring.dumps(obj)) == obj
+        assert tnetstring.loads(tnetstring.dumps(obj)) == obj
 
 def thrash_cjson():
     for obj, tns, json, msh in TESTS:
-#        cjson.encode(obj)
+        cjson.encode(obj)
         assert cjson.decode(json) == obj
-#        assert cjson.decode(cjson.encode(obj)) == obj
+        assert cjson.decode(cjson.encode(obj)) == obj
 
 def thrash_yajl():
     for obj, tns, json, msh in TESTS:
-#        yajl.dumps(obj)
+        yajl.dumps(obj)
         assert yajl.loads(json) == obj
-#        assert yajl.loads(yajl.dumps(obj)) == obj
+        assert yajl.loads(yajl.dumps(obj)) == obj
 
 def thrash_ujson():
     for obj, tns, json, msh in TESTS:
-#        ujson.dumps(obj)
+        ujson.dumps(obj)
         assert ujson.loads(json) == obj
-#        assert ujson.loads(ujson.dumps(obj)) == obj
+        assert ujson.loads(ujson.dumps(obj)) == obj
 
 def thrash_marshal():
     for obj, tns, json, msh in TESTS:
-#        marshal.dumps(obj)
+        marshal.dumps(obj)
         assert marshal.loads(msh) == obj
-#        assert marshal.loads(marshal.dumps(obj)) == obj
+        assert marshal.loads(marshal.dumps(obj)) == obj
 
 
 if __name__ == "__main__":