pk: Simplify code

mbedtls/pk.pxd

@@ -383,10 +383,6 @@ cdef extern from "mbedtls/pk.h" nogil:
 
 cdef class CipherBase:
     cdef mbedtls_pk_context _ctx
-    cdef bytes _write(
-        self,
-        int (*fun)(mbedtls_pk_context*, unsigned char*, size_t),
-        size_t)
 
 
 cdef class RSA(CipherBase):

mbedtls/pk.pyx

@@ -18,6 +18,7 @@ __license__ = "MIT License"
 
 
 from libc.stdlib cimport malloc, free
+from libc.string cimport memset
 
 cimport mbedtls.mpi as _mpi
 cimport mbedtls.pk as _pk
@@ -380,29 +381,37 @@ cdef class CipherBase:
         """
         raise NotImplementedError
 
-    cdef bytes _write(self, int (*fun)(_pk.mbedtls_pk_context *,
-                                       unsigned char *, size_t),
-                      size_t olen):
-        cdef unsigned char[:] buf = bytearray(olen * b"\0")
-        cdef int ret = fun(&self._ctx, &buf[0], buf.size)
-        check_error(ret)
-        # DER format: `ret` is the size of the buffer, offset from the end.
-        # PEM format: `ret` is zero.
-        if not ret:
-            ret = olen
-        # cast unsigned char[:] -> bytearray -> bytes
-        return bytes(bytearray(buf[olen - ret:olen]))
-
     def _private_to_DER(self):
         if not self._has_private():
             return b""
-        return self._write(&_pk.mbedtls_pk_write_key_der, PRV_DER_MAX_BYTES)
+        cdef int olen
+        cdef size_t osize = PRV_DER_MAX_BYTES
+        cdef unsigned char *output = <unsigned char *>malloc(
+            osize * sizeof(unsigned char))
+        if not output:
+            raise MemoryError()
+        try:
+            olen = check_error(
+                _pk.mbedtls_pk_write_key_der(&self._ctx, output, osize))
+            return bytes(output[osize - olen:osize])
+        finally:
+            free(output)
 
     def _private_to_PEM(self):
         if not self._has_private():
             return ""
-        return self._write(&_pk.mbedtls_pk_write_key_pem,
-                           PRV_DER_MAX_BYTES * 4 // 3 + 100).decode("ascii")
+        cdef size_t osize = PRV_DER_MAX_BYTES * 4 // 3 + 100
+        cdef unsigned char *output = <unsigned char *>malloc(
+            osize * sizeof(unsigned char))
+        if not output:
+            raise MemoryError()
+        memset(output, 0, osize)
+        try:
+            check_error(
+                _pk.mbedtls_pk_write_key_pem(&self._ctx, output, osize))
+            return bytes(output[0:osize]).decode("ascii")
+        finally:
+            free(output)
 
     def export_key(self, format="DER"):
         """Return the private key.
@@ -422,13 +431,34 @@ cdef class CipherBase:
     def _public_to_DER(self):
         if not self._has_public():
             return b""
-        return self._write(&_pk.mbedtls_pk_write_pubkey_der, PUB_DER_MAX_BYTES)
+        cdef int olen
+        cdef size_t osize = PRV_DER_MAX_BYTES
+        cdef unsigned char *output = <unsigned char *>malloc(
+            osize * sizeof(unsigned char))
+        if not output:
+            raise MemoryError()
+        try:
+            olen = check_error(
+                _pk.mbedtls_pk_write_pubkey_der(&self._ctx, output, osize))
+            return bytes(output[osize - olen:osize])
+        finally:
+            free(output)
 
     def _public_to_PEM(self):
         if not self._has_public():
             return ""
-        return self._write(&_pk.mbedtls_pk_write_pubkey_pem,
-                           PUB_DER_MAX_BYTES * 4 // 3 + 100).decode("ascii")
+        cdef size_t osize = PRV_DER_MAX_BYTES * 4 // 3 + 100
+        cdef unsigned char *output = <unsigned char *>malloc(
+            osize * sizeof(unsigned char))
+        if not output:
+            raise MemoryError()
+        memset(output, 0, osize)
+        try:
+            check_error(
+                _pk.mbedtls_pk_write_pubkey_pem(&self._ctx, output, osize))
+            return bytes(output[0:osize]).decode("ascii")
+        finally:
+            free(output)
 
     def export_public_key(self, format="DER"):
         """Return the public key.