Check for counter wraparound when encrypting using MODE_CTR

- Add check_wraparound_func pointer to PCT_CounterObject
- Call check_wraparound_func from block_template.c

1 files changed, 19 insertions, 8 deletions

diff --git a/src/_counter.c b/src/_counter.c
index 6c4f845..9e14da4 100644
--- a/src/_counter.c
+++ b/src/_counter.c
@@ -142,11 +142,8 @@ _CounterObject_next_value(PCT_CounterObject *self, int little_endian)
     PyObject *y = NULL;
     PyObject *x = NULL;
 
-    if (self->carry && !self->allow_wraparound) {
-        PyErr_SetString(PyExc_OverflowError,
-                        "counter wrapped without allow_wraparound");
+    if (!self->check_wraparound_func(self))
         goto err_out;
-    }
 
     eight = PyInt_FromLong(8);
     if (!eight)
@@ -250,11 +247,8 @@ CounterObject_call(PCT_CounterObject *self, PyObject *args, PyObject *kwargs)
 {
     PyObject *retval;
 
-    if (self->carry && !self->allow_wraparound) {
-        PyErr_SetString(PyExc_OverflowError,
-                        "counter wrapped without allow_wraparound");
+    if (!self->check_wraparound_func(self))
         return NULL;
-    }
 
     retval = (PyObject *)PyString_FromStringAndSize((const char *)self->val, self->buf_size);
 
@@ -263,6 +257,17 @@ CounterObject_call(PCT_CounterObject *self, PyObject *args, PyObject *kwargs)
     return retval;
 }
 
+static int
+CounterObject_check_wraparound(PCT_CounterObject *self)
+{
+    if (self->carry && !self->allow_wraparound) {
+        PyErr_SetString(PyExc_OverflowError,
+                        "counter wrapped without allow_wraparound");
+        return 0;
+    }
+    return 1;
+}
+
 static PyMethodDef CounterLEObject_methods[] = {
     {"next_value", (PyCFunction)CounterLEObject_next_value, METH_VARARGS,
         "Get the numerical value of next value of the counter."},
@@ -386,6 +391,9 @@ CounterLE_new(PyObject *self, PyObject *args, PyObject *kwargs)
     /* Set the inc_func pointer */
     obj->inc_func = (void (*)(void *))CounterLEObject_increment;
 
+    /* Set the check_wraparound_func pointer */
+    obj->check_wraparound_func = (int (*)(void *))CounterObject_check_wraparound;
+
     /* Return the object */
     return (PyObject *)obj;
 }
@@ -412,6 +420,9 @@ CounterBE_new(PyObject *self, PyObject *args, PyObject *kwargs)
     /* Set the inc_func pointer */
     obj->inc_func = (void (*)(void *))CounterBEObject_increment;
 
+    /* Set the check_wraparound_func pointer */
+    obj->check_wraparound_func = (int (*)(void *))CounterObject_check_wraparound;
+
     /* Return the object */
     return (PyObject *)obj;
 }