Merge pull request #137 from Sryborg/memory_leak_fix

[Fix] Memory Leak in py-tlsh
trendmicro · Mar 5, 2024 · 96536e3 · 96536e3
2 parents 25bc5d2 + 0a62e5d
commit 96536e3
Showing 1 changed file with 24 additions and 3 deletions.
diff --git a/py_ext/tlshmodule.cpp b/py_ext/tlshmodule.cpp
@@ -368,10 +368,17 @@ Tlsh_fromTlshStr(tlsh_TlshObject *self, PyObject *args)
 
     arg = PyTuple_GetItem(args, 0);
 #if PY_MAJOR_VERSION >= 3
-    if (!PyUnicode_Check(arg) || (arg = PyUnicode_AsASCIIString(arg)) == NULL) {
-      PyErr_SetString(PyExc_ValueError, "argument is not a TLSH hex string");
-      return NULL;
+    if (!PyUnicode_Check(arg)) {
+        PyErr_SetString(PyExc_ValueError, "argument is not a TLSH hex string");
+        return NULL;
     }
+
+    PyObject *asciiStr = PyUnicode_AsASCIIString(arg);
+    if (asciiStr == NULL) {
+        PyErr_SetString(PyExc_ValueError, "failed to convert argument to ASCII string");
+        return NULL;
+    }
+    arg = asciiStr;
 #else
     if (!PyString_Check(arg)) {
       PyErr_SetString(PyExc_ValueError, "argument is not a TLSH hex string");
@@ -381,20 +388,34 @@ Tlsh_fromTlshStr(tlsh_TlshObject *self, PyObject *args)
 
     if (PyBytes_AsStringAndSize(arg, &str, &len) == -1) {
         PyErr_SetString(PyExc_ValueError, "argument is not a TLSH hex string");
+#if PY_MAJOR_VERSION >= 3
+        Py_XDECREF(asciiStr);
+#endif
         return NULL;
     }
 
     if ((len != TLSH_STRING_LEN_REQ) && (len != TLSH_STRING_LEN_REQ-2)) {
         PyErr_SetString(PyExc_ValueError, "argument length incorrect: not a TLSH hex string");
+#if PY_MAJOR_VERSION >= 3
+        Py_XDECREF(asciiStr);
+#endif
         return NULL;
     }
 
     if (self->tlsh.fromTlshStr(str) != 0) {
         PyErr_SetString(PyExc_ValueError, "argument value incorrect: not a TLSH hex string");
+#if PY_MAJOR_VERSION >= 3
+        Py_XDECREF(asciiStr);
+#endif
         return NULL;
     }
+
     self->finalized = true;
 
+#if PY_MAJOR_VERSION >= 3
+    Py_XDECREF(asciiStr);
+#endif
+
     Py_RETURN_NONE;
 }