Fix memory leak on encoding errors when the buffer was resized

.github/workflows/deploy.yml

@@ -47,7 +47,7 @@ jobs:
           # Don't build with prerelease Python versions.
           CIBW_PROJECT_REQUIRES_PYTHON: ">=3.7,<3.11"
           # Run the test suite after each build.
-          CIBW_TEST_REQUIRES: "pytest"
+          CIBW_TEST_REQUIRES: "pytest psutil"
           CIBW_TEST_COMMAND: pytest {package}/tests
 
       - name: Upload as build artifacts
@@ -120,7 +120,7 @@ jobs:
           # Likewise, select only one Python version per job to speed this up.
           CIBW_BUILD: "${{ matrix.python-version }}-*"
           # Run the test suite after each build.
-          CIBW_TEST_REQUIRES: "pytest"
+          CIBW_TEST_REQUIRES: "pytest psutil"
           CIBW_TEST_COMMAND: pytest {package}/tests
 
       - name: Upload as build artifacts

.github/workflows/test.yml

@@ -31,6 +31,7 @@ jobs:
         run: |
           python -m pip install -U pip
           python -m pip install -U pytest
+          python -m pip install -U psutil
           python -m pip install .
         env:
           CFLAGS: '-DDEBUG'

.travis.yml

@@ -19,6 +19,7 @@ jobs:
 install:
   - pip install -U pip
   - pip install -U pytest
+  - pip install -U psutil
   - pip install .
 
 script: pytest

lib/ultrajsonenc.c

@@ -986,6 +986,11 @@ char *JSON_EncodeObject(JSOBJ obj, JSONObjectEncoder *enc, char *_buffer, size_t
 
   if (enc->errorMsg)
   {
+    if (enc->heap == 1)
+    {
+      // Buffer was realloc'd at some point, or no initial buffer was provided.
+      enc->free(enc->start);
+    }
     return NULL;
   }
 

python/objToJSON.c

@@ -876,19 +876,19 @@ PyObject* objToJSON(PyObject* self, PyObject *args, PyObject *kwargs)
 
   dconv_d2s_free(&encoder.d2s);
 
-  if (PyErr_Occurred())
+  if (encoder.errorMsg && !PyErr_Occurred())
   {
-    return NULL;
+    // If there is an error message and we don't already have a Python exception, set one.
+    PyErr_Format (PyExc_OverflowError, "%s", encoder.errorMsg);
   }
 
-  if (encoder.errorMsg)
+  if (PyErr_Occurred())
   {
     if (ret != buffer)
     {
       encoder.free (ret);
     }
 
-    PyErr_Format (PyExc_OverflowError, "%s", encoder.errorMsg);
     return NULL;
   }
 

tests/memory.py

@@ -0,0 +1,38 @@
+import sys
+
+import psutil
+
+
+def run(func, n):
+    # Run func once so any static variables etc. are allocated
+    try:
+        func()
+    except Exception:
+        pass
+
+    # Take a measurement
+    before = psutil.Process().memory_full_info().uss
+
+    # Run n times
+    for i in range(n):
+        try:
+            func()
+        except Exception:
+            pass
+
+    # Measure again and return success (less than 1% increase)
+    after = psutil.Process().memory_full_info().uss
+
+    print(f"USS before: {before}, after: {after}", file=sys.stderr)
+    return (after - before) / before < 0.01
+
+
+if __name__ == "__main__":
+    # Usage: python3 memory.py CODE [N]
+    #   CODE must define a function 'func'; it will be exec()'d!
+    #   N, if present, overrides the number of runs (default: 1000)
+    exec_globals = {}
+    exec(sys.argv[1], exec_globals)
+    n = int(sys.argv[2]) if sys.argv[2:] else 1000
+    if not run(exec_globals["func"], n):
+        sys.exit(1)

tests/test_ujson.py

@@ -3,7 +3,9 @@
 import io
 import json
 import math
+import os.path
 import re
+import subprocess
 import sys
 import uuid
 from collections import OrderedDict
@@ -1026,6 +1028,19 @@ def __str__(self):
         ujson.dumps({Obj(): 1}, sort_keys=sort_keys)
 
 
+def no_memory_leak(func_code, n=None):
+    code = f"import functools, ujson; func = {func_code}"
+    path = os.path.join(os.path.dirname(__file__), "memory.py")
+    n = [str(n)] if n is not None else []
+    p = subprocess.run(["python3", path, code] + n)
+    assert p.returncode == 0
+
+
+@pytest.mark.parametrize("input", ['["a" * 11000, b""]'])
+def test_no_memory_leak_encoding_errors(input):
+    no_memory_leak(f"functools.partial(ujson.dumps, {input})")
+
+
 """
 def test_decode_numeric_int_frc_overflow():
 input = "X.Y"