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

3 files changed, 24 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;
 }
diff --git a/src/_counter.h b/src/_counter.h
index faaf63f..af401cd 100644
--- a/src/_counter.h
+++ b/src/_counter.h
@@ -39,6 +39,7 @@ typedef struct {
     int shortcut_disabled;  /* This gets set to a non-zero value when the shortcut mechanism is disabled */
     int carry;         /* This gets set by Counter*Object_increment when the counter wraps around */
     int allow_wraparound;   /* When this is false, we raise OverflowError on next_value() or __call__() when the counter wraps around */
+    int (*check_wraparound_func)(void *);    /* Pointer to the functon that checks for overflow. Returns zero and sets OverflowError if the check fails. */
 } PCT_CounterObject;
 
 #endif /* PCT__COUNTER_H */
diff --git a/src/block_template.c b/src/block_template.c
index 7548b21..6b9a338 100644
--- a/src/block_template.c
+++ b/src/block_template.c
@@ -409,6 +409,10 @@ ALG_Encrypt(ALGobject *self, PyObject *args)
 				 * and manipulate the counter directly. */
 
 				PCT_CounterObject *ctr = (PCT_CounterObject *)(self->counter);
+				if (!ctr->check_wraparound_func(ctr)) {
+					free(buffer);
+					return NULL;
+				}
 				if (ctr->buf_size != BLOCK_SIZE) {
 					PyErr_Format(PyExc_TypeError,
 						     "CTR counter function returned "