Add ABI check when importing _counter from block_template

3 files changed, 26 insertions, 0 deletions

diff --git a/src/_counter.c b/src/_counter.c
index 729fc9e..f82c378 100644
--- a/src/_counter.c
+++ b/src/_counter.c
@@ -32,6 +32,13 @@
 #define PyLong_FromLong PyInt_FromLong
 #endif
 
+/* Deal with old API in Python 2.1 */
+#if PYTHON_API_VERSION < 1011
+#define PyModule_AddIntConstant(m,n,v) {PyObject *o=PyInt_FromLong(v); \
+           if (o!=NULL) \
+             {PyDict_SetItemString(PyModule_GetDict(m),n,o); Py_DECREF(o);}}
+#endif
+
 /* NB: This can be called multiple times for a given object, via the __init__ method.  Be careful. */
 static int
 CounterObject_init(PCT_CounterObject *self, PyObject *args, PyObject *kwargs)
@@ -566,6 +573,9 @@ init_counter(void)
     PyObject_SetAttrString(m, "CounterBE", (PyObject *)&PCT_CounterBEType);
     PyObject_SetAttrString(m, "CounterLE", (PyObject *)&PCT_CounterLEType);
 
+    /* Allow block_template.c to do an ABI check */
+    PyModule_AddIntConstant(m, "_PCT_CTR_ABI_VERSION", PCT_CTR_ABI_VERSION);
+
 #ifdef IS_PY3K
     return m;
 #endif
diff --git a/src/_counter.h b/src/_counter.h
index e08757a..df2091f 100644
--- a/src/_counter.h
+++ b/src/_counter.h
@@ -26,6 +26,8 @@
 
 #include "pycrypto_common.h"
 
+#define PCT_CTR_ABI_VERSION   1   /* Bump this whenever the structure below changes */
+
 typedef struct {
     PyObject_HEAD
     PyBytesObject *prefix;     /* Prefix (useful for a nonce) */
diff --git a/src/block_template.c b/src/block_template.c
index c655509..e1f3f2f 100644
--- a/src/block_template.c
+++ b/src/block_template.c
@@ -70,6 +70,8 @@ typedef struct
 #ifdef IS_PY3K
 static PyTypeObject ALGtype;
 #define is_ALGobject(v)		(Py_TYPE(v) == &ALGtype)
+#define PyInt_CheckExact	PyLong_CheckExact
+#define PyInt_AS_LONG		PyLong_AS_LONG
 #else
 staticforward PyTypeObject ALGtype;
 #define is_ALGobject(v)		((v)->ob_type == &ALGtype)
@@ -765,6 +767,7 @@ static struct PyModuleDef moduledef = {
 #define PyModule_AddIntConstant(m,n,v) {PyObject *o=PyInt_FromLong(v); \
            if (o!=NULL) \
              {PyDict_SetItemString(PyModule_GetDict(m),n,o); Py_DECREF(o);}}
+#define PyInt_CheckExact	PyInt_Check
 #endif
 
 
@@ -807,6 +810,17 @@ _MODULE_NAME (void)
 	if (_counter_module) {
 		PCT_CounterBEType = (PyTypeObject *)PyObject_GetAttrString(_counter_module, "CounterBE");
 		PCT_CounterLEType = (PyTypeObject *)PyObject_GetAttrString(_counter_module, "CounterLE");
+
+		/* Simple ABI version check in case the user doesn't re-compile all of
+		 * the modules during an upgrade. */
+		PyObject *abiver = PyObject_GetAttrString(_counter_module, "_PCT_CTR_ABI_VERSION");
+		if (PCT_CounterBEType == NULL || PyType_Check((PyObject *)PCT_CounterBEType) < 0 ||
+			 PCT_CounterLEType == NULL || PyType_Check((PyObject *)PCT_CounterLEType) < 0 ||
+			 abiver == NULL || PyInt_CheckExact(abiver) < 0 || PyInt_AS_LONG(abiver) != PCT_CTR_ABI_VERSION)
+		{
+			PyErr_SetString(PyExc_ImportError, "Crypto.Util._counter ABI mismatch.  Was PyCrypto incorrectly compiled?");
+		}
+		Py_CLEAR(abiver);
 	}
 
 	/* Check for errors */