Fix MODE_CTR memory leak under Python 3

The leak arose from the string creation in this line:

    PyObject_HasAttr(counter, PyUnicode_FromString("__PCT_CTR_SHORTCUT__"))

This commit replaces the __PCT_CTR_SHORTCUT__ hack with code that
imports the _counter module and checks the appropriate types.

2 files changed, 34 insertions, 18 deletions

diff --git a/src/_counter.c b/src/_counter.c
index 2d6aa48..c1f1e5f 100644
--- a/src/_counter.c
+++ b/src/_counter.c
@@ -374,7 +374,7 @@ CounterBEObject_getattr(PyObject *s, char *name)
 }
 
 static PyTypeObject
-my_CounterLEType = {
+PCT_CounterLEType = {
 #ifdef IS_PY3K
 	PyVarObject_HEAD_INIT(NULL, 0)  /* deferred type init for compilation on Windows, type will be filled in at runtime */
 #else
@@ -422,7 +422,7 @@ my_CounterLEType = {
 };
 
 static PyTypeObject
-my_CounterBEType = {
+PCT_CounterBEType = {
 #ifdef IS_PY3K
 	PyVarObject_HEAD_INIT(NULL, 0)  /* deferred type init for compilation on Windows, type will be filled in at runtime */
 #else
@@ -478,7 +478,7 @@ CounterLE_new(PyObject *self, PyObject *args, PyObject *kwargs)
     PCT_CounterObject *obj = NULL;
 
     /* Create the new object */
-    obj = PyObject_New(PCT_CounterObject, &my_CounterLEType);
+    obj = PyObject_New(PCT_CounterObject, &PCT_CounterLEType);
     if (obj == NULL) {
         return NULL;
     }
@@ -504,7 +504,7 @@ CounterBE_new(PyObject *self, PyObject *args, PyObject *kwargs)
     PCT_CounterObject *obj = NULL;
 
     /* Create the new object */
-    obj = PyObject_New(PCT_CounterObject, &my_CounterBEType);
+    obj = PyObject_New(PCT_CounterObject, &PCT_CounterBEType);
     if (obj == NULL) {
         return NULL;
     }
@@ -559,25 +559,33 @@ init_counter(void)
 
     /* TODO - Is the error handling here correct? */
 #ifdef IS_PY3K
-	/* PyType_Ready automatically fills in ob_type with &PyType_Type if it's not already set */
-	if (PyType_Ready(&my_CounterLEType) < 0)
-		return NULL;
-	if (PyType_Ready(&my_CounterBEType) < 0)
-		return NULL;
+    /* PyType_Ready automatically fills in ob_type with &PyType_Type if it's not already set */
+    if (PyType_Ready(&PCT_CounterLEType) < 0)
+        return NULL;
+    if (PyType_Ready(&PCT_CounterBEType) < 0)
+        return NULL;
 
     /* Initialize the module */
     m = PyModule_Create(&moduledef);
     if (m == NULL)
         return NULL;
 
-	return m;
 #else
     m = Py_InitModule("_counter", module_methods);
     if (m == NULL)
         return;
-		
-	my_CounterLEType.ob_type = &PyType_Type;
-    my_CounterBEType.ob_type = &PyType_Type;
+
+    PCT_CounterLEType.ob_type = &PyType_Type;
+    PCT_CounterBEType.ob_type = &PyType_Type;
+#endif
+
+    /* Add the counter types to the module so that epydoc can see them, and so
+     * that we can access them in the block cipher modules. */
+    PyObject_SetAttrString(m, "CounterBE", (PyObject *)&PCT_CounterBEType);
+    PyObject_SetAttrString(m, "CounterLE", (PyObject *)&PCT_CounterLEType);
+
+#ifdef IS_PY3K
+    return m;
 #endif
 }
 
diff --git a/src/block_template.c b/src/block_template.c
index 8ef09ac..3e19eca 100644
--- a/src/block_template.c
+++ b/src/block_template.c
@@ -49,6 +49,11 @@
 #endif
 #define _MODULE_STRING _XSTR(MODULE_NAME)
 
+/* Object references for the counter_shortcut */
+static PyObject *_counter_module = NULL;
+static PyTypeObject *PCT_CounterBEType = NULL;
+static PyTypeObject *PCT_CounterLEType = NULL;
+
 typedef struct 
 {
 	PyObject_HEAD 
@@ -182,11 +187,7 @@ ALGnew(PyObject *self, PyObject *args, PyObject *kwdict)
 			PyErr_SetString(PyExc_TypeError,
 					"'counter' keyword parameter is required with CTR mode");
 			return NULL;
-#ifdef IS_PY3K
-		} else if (PyObject_HasAttr(counter, PyUnicode_FromString("__PCT_CTR_SHORTCUT__"))) {
-#else
-		} else if (PyObject_HasAttrString(counter, "__PCT_CTR_SHORTCUT__")) {
-#endif
+		} else if (counter->ob_type == PCT_CounterBEType || counter->ob_type == PCT_CounterLEType) {
 			counter_shortcut = 1;
 		} else if (!PyCallable_Check(counter)) {
 			PyErr_SetString(PyExc_ValueError, 
@@ -800,6 +801,13 @@ _MODULE_NAME (void)
 	PyModule_AddIntConstant(m, "block_size", BLOCK_SIZE);
 	PyModule_AddIntConstant(m, "key_size", KEY_SIZE);
 
+	Py_CLEAR(_counter_module);
+	_counter_module = PyImport_ImportModule("Crypto.Util._counter");
+	if (_counter_module) {
+		PCT_CounterBEType = (PyTypeObject *)PyObject_GetAttrString(_counter_module, "CounterBE");
+		PCT_CounterLEType = (PyTypeObject *)PyObject_GetAttrString(_counter_module, "CounterLE");
+	}
+
 	/* Check for errors */
 	if (PyErr_Occurred())
 		Py_FatalError("can't initialize module " _MODULE_STRING);