Counter: raise OverflowError by default when the counter wraps around.

The old behaviour can be obtained by explicitly setting allow_wraparound=True
when invoking Counter.new

3 files changed, 21 insertions, 6 deletions

diff --git a/lib/Crypto/Util/Counter.py b/lib/Crypto/Util/Counter.py
index c3001e4..42dab42 100644
--- a/lib/Crypto/Util/Counter.py
+++ b/lib/Crypto/Util/Counter.py
@@ -28,7 +28,7 @@ from Crypto.Util import _counter
 import struct
 
 # Factory function
-def new(nbits, prefix="", suffix="", initial_value=1, overflow=0, little_endian=False, disable_shortcut=False):
+def new(nbits, prefix="", suffix="", initial_value=1, overflow=0, little_endian=False, allow_wraparound=False, disable_shortcut=False):
     # TODO: Document this
 
     # Sanity-check the message size
@@ -43,9 +43,9 @@ def new(nbits, prefix="", suffix="", initial_value=1, overflow=0, little_endian=
 
     initval = _encode(initial_value, nbytes, little_endian)
     if little_endian:
-        return _counter._newLE(str(prefix), str(suffix), initval, disable_shortcut=disable_shortcut)
+        return _counter._newLE(str(prefix), str(suffix), initval, allow_wraparound=allow_wraparound, disable_shortcut=disable_shortcut)
     else:
-        return _counter._newBE(str(prefix), str(suffix), initval, disable_shortcut=disable_shortcut)
+        return _counter._newBE(str(prefix), str(suffix), initval, allow_wraparound=allow_wraparound, disable_shortcut=disable_shortcut)
 
 def _encode(n, nbytes, little_endian=False):
     retval = []
diff --git a/src/_counter.c b/src/_counter.c
index 6964947..6c4f845 100644
--- a/src/_counter.c
+++ b/src/_counter.c
@@ -35,11 +35,12 @@ static int
 CounterObject_init(PCT_CounterObject *self, PyObject *args, PyObject *kwargs)
 {
     PyStringObject *prefix=NULL, *suffix=NULL, *initval=NULL;
+    int allow_wraparound = 0;
     int disable_shortcut = 0;
     Py_ssize_t size;
 
-    static char *kwlist[] = {"prefix", "suffix", "initval", "disable_shortcut", NULL};
-    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "SSS|i", kwlist, &prefix, &suffix, &initval, &disable_shortcut))
+    static char *kwlist[] = {"prefix", "suffix", "initval", "allow_wraparound", "disable_shortcut", NULL};
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "SSS|ii", kwlist, &prefix, &suffix, &initval, &allow_wraparound, &disable_shortcut))
         return -1;
 
     /* Check string size and set nbytes */
@@ -101,8 +102,9 @@ CounterObject_init(PCT_CounterObject *self, PyObject *args, PyObject *kwargs)
     memcpy(self->p, PyString_AS_STRING(initval), self->nbytes);
     memcpy(self->p + self->nbytes, PyString_AS_STRING(suffix), PyString_GET_SIZE(suffix));
 
-    /* Set shortcut_disabled */
+    /* Set shortcut_disabled and allow_wraparound */
     self->shortcut_disabled = disable_shortcut;
+    self->allow_wraparound = allow_wraparound;
 
     /* Clear the carry flag */
     self->carry = 0;
@@ -140,6 +142,12 @@ _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");
+        goto err_out;
+    }
+
     eight = PyInt_FromLong(8);
     if (!eight)
         goto err_out;
@@ -242,6 +250,12 @@ 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");
+        return NULL;
+    }
+
     retval = (PyObject *)PyString_FromStringAndSize((const char *)self->val, self->buf_size);
 
     self->inc_func(self);
diff --git a/src/_counter.h b/src/_counter.h
index 9735aab..faaf63f 100644
--- a/src/_counter.h
+++ b/src/_counter.h
@@ -38,6 +38,7 @@ typedef struct {
     void (*inc_func)(void *);   /* Pointer to the counter increment function */
     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 */
 } PCT_CounterObject;
 
 #endif /* PCT__COUNTER_H */