From 935fe0cec4ad7d62c89cba8b2b667a60ab8336f3 Mon Sep 17 00:00:00 2001
From: JustAnotherArchivist <JustAnotherArchivist@users.noreply.github.com>
Date: Wed, 13 Apr 2022 00:04:24 +0000
Subject: [PATCH] Fix segmentation fault when an exception is raised while
 converting a dict key to a string

Fixes #522
---
 python/objToJSON.c  |  9 +++++++++
 tests/test_ujson.py | 10 ++++++++++
 2 files changed, 19 insertions(+)

diff --git a/python/objToJSON.c b/python/objToJSON.c
index 8f4be736..d7f6cb92 100644
--- a/python/objToJSON.c
+++ b/python/objToJSON.c
@@ -257,6 +257,11 @@ static int Dict_iterNext(JSOBJ obj, JSONTypeContext *tc)
     itemNameTmp = GET_TC(tc)->itemName;
     GET_TC(tc)->itemName = PyObject_Str(GET_TC(tc)->itemName);
     Py_DECREF(itemNameTmp);
+    if (PyErr_Occurred())
+    {
+      PRINTMARK();
+      return -1;
+    }
     itemNameTmp = GET_TC(tc)->itemName;
     GET_TC(tc)->itemName = PyUnicode_AsUTF8String (GET_TC(tc)->itemName);
     Py_DECREF(itemNameTmp);
@@ -332,6 +337,10 @@ static int SortedDict_iterNext(JSOBJ obj, JSONTypeContext *tc)
       else if (!PyBytes_Check(key))
       {
         key = PyObject_Str(key);
+        if (PyErr_Occurred())
+        {
+          goto error;
+        }
         keyTmp = key;
         key = PyUnicode_AsUTF8String(key);
         Py_DECREF(keyTmp);
diff --git a/tests/test_ujson.py b/tests/test_ujson.py
index ed8a2b2d..19596571 100644
--- a/tests/test_ujson.py
+++ b/tests/test_ujson.py
@@ -996,6 +996,16 @@ def convert(self):
     )
 
 
+@pytest.mark.parametrize("sort_keys", [False, True])
+def test_obj_str_exception(sort_keys):
+    class Obj:
+        def __str__(self):
+            raise NotImplementedError
+
+    with pytest.raises(NotImplementedError):
+        ujson.dumps({Obj(): 1}, sort_keys=sort_keys)
+
+
 """
 def test_decode_numeric_int_frc_overflow():
 input = "X.Y"