From bb5c23a4b4d9e74be11097381ba20a903e6ad1a2 Mon Sep 17 00:00:00 2001
From: Mark Wooding <mdw@distorted.org.uk>
Date: Fri, 9 Nov 2018 12:28:16 +0000
Subject: [PATCH] bytestring.c: Implement indexing, slicing, concatenation and
 repeating.

Now these return bytestring objects, rather than Python strings.
---
 bytestring.c         | 108 ++++++++++++++++++++++++++++++++++++++++++++++++++-
 catacomb-python.h    |   1 +
 catacomb/__init__.py |   2 +-
 3 files changed, 108 insertions(+), 3 deletions(-)

diff --git a/bytestring.c b/bytestring.c
index 124ae9a..a7fe5cb 100644
--- a/bytestring.c
+++ b/bytestring.c
@@ -141,6 +141,91 @@ static PyObject *bytestring_pyrichcompare(PyObject *me,
   else RETURN_FALSE;
 }
 
+static PyObject *bytestring_pyconcat(PyObject *x, PyObject *y)
+{
+  const void *xv; Py_ssize_t xsz;
+  const void *yv; Py_ssize_t ysz;
+  PyObject *z = 0; char *zp; size_t zsz;
+
+  if (PyObject_AsReadBuffer(x, &xv, &xsz) ||
+      PyObject_AsReadBuffer(y, &yv, &ysz))
+    goto end;
+  zsz = (size_t)xsz + (size_t)ysz;
+  if (xsz < 0 || ysz < 0 || zsz < xsz) VALERR("too long");
+  z = bytestring_pywrap(0, zsz); zp = PyString_AS_STRING(z);
+  memcpy(zp, xv, xsz); memcpy(zp + xsz, yv, ysz);
+end:
+  return (z);
+}
+
+static PyObject *bytestring_pyrepeat(PyObject *me, Py_ssize_t n)
+{
+  const unsigned char *xp; size_t xsz;
+  PyObject *z = 0; char *zp; size_t zsz;
+
+  xp = (const unsigned char *)PyString_AS_STRING(me);
+  xsz = PyString_GET_SIZE(me);
+  if (n < 0 || xsz >= (size_t)-1/n) VALERR("too long");
+  zsz = n*xsz; z = bytestring_pywrap(0, zsz); zp = PyString_AS_STRING(z);
+  if (xsz == 1) memset(zp, *xp, zsz);
+  else while (zsz) { memcpy(zp, xp, xsz); zp += xsz; zsz -= xsz; }
+end:
+  return (z);
+}
+
+static PyObject *bytestring_pyitem(PyObject *me, Py_ssize_t i)
+{
+  PyObject *rc = 0;
+
+  if (i < 0 || i >= PyString_GET_SIZE(me)) IXERR("out of range");
+  rc = bytestring_pywrap(PyString_AS_STRING(me) + i, 1);
+end:
+  return (rc);
+}
+
+static PyObject *bytestring_pyslice(PyObject *me, Py_ssize_t i, Py_ssize_t j)
+{
+  PyObject *rc = 0;
+  size_t n = PyString_GET_SIZE(me);
+
+  if (i < 0) i = 0;
+  if (j < 0) j = 0;
+  else if (j > n) j = n;
+  if (j < i) i = j = 0;
+  if (i == 0 && j == n && me->ob_type == bytestring_pytype)
+    { Py_INCREF(me); rc = me; goto end; }
+  rc = bytestring_pywrap(PyString_AS_STRING(me) + i, j - i);
+end:
+  return (rc);
+}
+
+static PyObject *bytestring_pysubscript(PyObject *me, PyObject *ix)
+{
+  Py_ssize_t i, j, k, n;
+  const unsigned char *p;
+  unsigned char *q;
+  PyObject *rc = 0;
+
+  if (PyIndex_Check(ix)) {
+    i = PyNumber_AsSsize_t(ix, PyExc_IndexError);
+    if (i == -1 && PyErr_Occurred()) return (0);
+    if (i < 0) i += PyString_GET_SIZE(me);
+    rc = bytestring_pyitem(me, i);
+  } else if (PySlice_Check(ix)) {
+    if (PySlice_GetIndicesEx((PySliceObject *)ix, PyString_GET_SIZE(me),
+			     &i, &j, &k, &n))
+      return (0);
+    if (k == 1) return bytestring_pyslice(me, i, j);
+    rc = bytestring_pywrap(0, n);
+    p = (unsigned char *)PyString_AS_STRING(me) + i;
+    q = (unsigned char *)PyString_AS_STRING(rc);
+    while (n--) { *q++ = *p; p += k; }
+  } else
+    TYERR("wanted integer or slice");
+end:
+  return (rc);
+}
+
 #define BINOP(name, op)							\
   static PyObject *bytestring_py##name(PyObject *x, PyObject *y) {	\
     const void *xv, *yv;						\
@@ -206,6 +291,25 @@ static PyNumberMethods bytestring_pynumber = {
   0,					/* @nb_hex@ */
 };
 
+static PySequenceMethods bytestring_pysequence = {
+  0,					/* @sq_length@ */
+  bytestring_pyconcat,			/* @sq_concat@ */
+  bytestring_pyrepeat,			/* @sq_repeat@ */
+  bytestring_pyitem,			/* @sq_item@ */
+  bytestring_pyslice,			/* @sq_slice@ */
+  0,					/* @sq_ass_item@ */
+  0,					/* @sq_ass_slice@ */
+  0,					/* @sq_contains@ */
+  0,					/* @sq_inplace_concat@ */
+  0,					/* @sq_inplace_repeat@ */
+};
+
+static PyMappingMethods bytestring_pymapping = {
+  0,					/* @mp_length@ */
+  bytestring_pysubscript,		/* @mp_subscript@ */
+  0,					/* @mp_ass_subscript@ */
+};
+
 static PyBufferProcs bytestring_pybuffer;
 
 static PyTypeObject bytestring_pytype_skel = {
@@ -221,8 +325,8 @@ static PyTypeObject bytestring_pytype_skel = {
   0,					/* @tp_compare@ */
   0,					/* @tp_repr@ */
   &bytestring_pynumber,			/* @tp_as_number@ */
-  0,					/* @tp_as_sequence@ */
-  0,					/* @tp_as_mapping@ */
+  &bytestring_pysequence,		/* @tp_as_sequence@ */
+  &bytestring_pymapping,		/* @tp_as_mapping@ */
   0,					/* @tp_hash@ */
   0,					/* @tp_call@ */
   0,					/* @tp_str@ */
diff --git a/catacomb-python.h b/catacomb-python.h
index bbb5fac..364a3c9 100644
--- a/catacomb-python.h
+++ b/catacomb-python.h
@@ -134,6 +134,7 @@
 } while (0)
 #define VALERR(str) EXCERR(PyExc_ValueError, str)
 #define TYERR(str) EXCERR(PyExc_TypeError, str)
+#define IXERR(str) EXCERR(PyExc_IndexError, str)
 #define ZDIVERR(str) EXCERR(PyExc_ZeroDivisionError, str)
 #define SYSERR(str) EXCERR(PyExc_SystemError, str)
 #define NIERR(str) EXCERR(PyExc_NotImplementedError, str)
diff --git a/catacomb/__init__.py b/catacomb/__init__.py
index 8177b2d..02e19d3 100644
--- a/catacomb/__init__.py
+++ b/catacomb/__init__.py
@@ -292,7 +292,7 @@ def secret_box(k, n, m):
   s = E.enczero(poly1305.masksz)
   y = E.encrypt(m)
   t = poly1305(r)(s).hash(y).done()
-  return ByteString(t + y)
+  return t + y
 
 def secret_unbox(k, n, c):
   E = xsalsa20(k).setiv(n)
-- 
2.11.0