From 59007ee6973b73126b9de42f7a3c9b71cf076cc6 Mon Sep 17 00:00:00 2001
From: MMGen <mmgen@tuta.io>
Date: Wed, 31 Oct 2018 18:14:32 +0000
Subject: [PATCH] py3port: rewrite secp256k1 extension module

---
 extmod/secp256k1mod.c | 69 ++++++++++++++++++++++++++++++++++++-------
 1 file changed, 58 insertions(+), 11 deletions(-)

diff --git a/extmod/secp256k1mod.c b/extmod/secp256k1mod.c
index 3ea1bb93..2b4d4404 100755
--- a/extmod/secp256k1mod.c
+++ b/extmod/secp256k1mod.c
@@ -23,39 +23,86 @@ static PyObject * priv2pub(PyObject *self, PyObject *args) {
 	const unsigned char * privkey;
 	const int klen;
 	const int compressed;
-	if (!PyArg_ParseTuple(args, "t#I", &privkey, &klen, &compressed))
+	if (!PyArg_ParseTuple(args, "y#I", &privkey, &klen, &compressed)) {
+		PyErr_SetString(PyExc_ValueError, "Unable to parse extension mod arguments");
 		return NULL;
+	}
 	if (klen != 32) {
 		PyErr_SetString(PyExc_ValueError, "Private key length not 32 bytes");
 		return NULL;
 	}
 	secp256k1_pubkey pubkey;
-	size_t pubkeyclen = compressed == 1 ? 33: 65;
+	size_t pubkeyclen = compressed == 1 ? 33 : 65;
 	unsigned char pubkeyc[pubkeyclen];
 	static secp256k1_context *ctx = NULL;
 	if (ctx == NULL) {
 	/*	puts ("Initializing context"); */
 		ctx = secp256k1_context_create(SECP256K1_CONTEXT_SIGN | SECP256K1_CONTEXT_VERIFY);
 	}
+	if (ctx == NULL) {
+		PyErr_SetString(PyExc_RuntimeError, "Context initialization failed");
+		return NULL;
+	}
 	if (secp256k1_ec_pubkey_create(ctx, &pubkey, privkey) != 1) {
 		PyErr_SetString(PyExc_RuntimeError, "Public key creation failed");
 		return NULL;
 	}
 	if (secp256k1_ec_pubkey_serialize(ctx, pubkeyc, &pubkeyclen, &pubkey,
-			compressed == 1 ? SECP256K1_EC_COMPRESSED: SECP256K1_EC_UNCOMPRESSED) != 1) {
+			compressed == 1 ? SECP256K1_EC_COMPRESSED : SECP256K1_EC_UNCOMPRESSED) != 1) {
 		PyErr_SetString(PyExc_RuntimeError, "Public key serialization failed");
 		return NULL;
 	}
-	return Py_BuildValue("s#", pubkeyc,pubkeyclen);
+	return Py_BuildValue("y#", pubkeyc,pubkeyclen);
 }
 
-static PyMethodDef secp256k1Methods[] = {
-	{"priv2pub", priv2pub, METH_VARARGS, "Generate pubkey from privkey using libsecp256k1"},
-	{NULL, NULL, 0, NULL} /* Sentinel */
+/* https://docs.python.org/3/howto/cporting.html */
+
+struct module_state {
+    PyObject *error;
 };
 
-PyMODINIT_FUNC initsecp256k1(void) {
-	PyObject *m;
-	m = Py_InitModule("secp256k1", secp256k1Methods);
-	if (m == NULL) return;
+#define GETSTATE(m) ((struct module_state*)PyModule_GetState(m))
+
+static PyMethodDef secp256k1_methods[] = {
+	{"priv2pub", priv2pub, METH_VARARGS, "Generate pubkey from privkey using libsecp256k1"},
+    {NULL, NULL}
+};
+
+static int secp256k1_traverse(PyObject *m, visitproc visit, void *arg) {
+    Py_VISIT(GETSTATE(m)->error);
+    return 0;
+}
+
+static int secp256k1_clear(PyObject *m) {
+    Py_CLEAR(GETSTATE(m)->error);
+    return 0;
+}
+
+static struct PyModuleDef moduledef = {
+        PyModuleDef_HEAD_INIT,
+        "secp256k1",
+        NULL,
+        sizeof(struct module_state),
+        secp256k1_methods,
+        NULL,
+        secp256k1_traverse,
+        secp256k1_clear,
+        NULL
+};
+
+#define INITERROR return NULL
+
+PyMODINIT_FUNC PyInit_secp256k1(void) {
+    PyObject *module = PyModule_Create(&moduledef);
+
+    if (module == NULL)
+        INITERROR;
+    struct module_state *st = GETSTATE(module);
+
+    st->error = PyErr_NewException("secp256k1.Error", NULL, NULL);
+    if (st->error == NULL) {
+        Py_DECREF(module);
+        INITERROR;
+    }
+    return module;
 }