Revamp poly1305_nif()

* Fix a possible memory leak where enif_alloc_binary() wasn't always
  converted into a term or freed in the error path.

1 files changed, 40 insertions, 28 deletions

diff --git a/lib/crypto/c_src/poly1305.c b/lib/crypto/c_src/poly1305.c
index 3e2bcfa60e..da434e4eb9 100644
--- a/lib/crypto/c_src/poly1305.c
+++ b/lib/crypto/c_src/poly1305.c
@@ -25,54 +25,66 @@ ERL_NIF_TERM poly1305_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
 {/* (Key, Text) */
 #ifdef HAVE_POLY1305
     ErlNifBinary key_bin, text, ret_bin;
-    ERL_NIF_TERM ret = atom_error;
+    ERL_NIF_TERM ret;
     EVP_PKEY *key = NULL;
     EVP_MD_CTX *mctx = NULL;
     EVP_PKEY_CTX *pctx = NULL;
     const EVP_MD *md = NULL;
     size_t size;
-    int type;
+    int ret_bin_alloc = 0;
 
-    type = EVP_PKEY_POLY1305;
+    if (argc != 2)
+        goto bad_arg;
+    if (!enif_inspect_binary(env, argv[0], &key_bin))
+        goto bad_arg;
+    if (key_bin.size != 32)
+        goto bad_arg;
+    if (!enif_inspect_binary(env, argv[1], &text))
+        goto bad_arg;
 
-    if (!enif_inspect_binary(env, argv[0], &key_bin) ||
-        !(key_bin.size == 32) ) {
-        return enif_make_badarg(env);
-    }
-
-    if (!enif_inspect_binary(env, argv[1], &text) ) {
-        return enif_make_badarg(env);
-    }
-
-    key = EVP_PKEY_new_raw_private_key(type, /*engine*/ NULL, key_bin.data,  key_bin.size);
-
-    if (!key ||
-        !(mctx = EVP_MD_CTX_new()) ||
-        !EVP_DigestSignInit(mctx, &pctx, md, /*engine*/ NULL, key) ||
-        !EVP_DigestSignUpdate(mctx, text.data, text.size)) {
+    if ((key = EVP_PKEY_new_raw_private_key(EVP_PKEY_POLY1305, /*engine*/ NULL, key_bin.data,  key_bin.size)) == NULL)
         goto err;
-    }
 
-    if (!EVP_DigestSignFinal(mctx, NULL, &size) ||
-        !enif_alloc_binary(size, &ret_bin) ||
-        !EVP_DigestSignFinal(mctx, ret_bin.data, &size)) {
+    if ((mctx = EVP_MD_CTX_new()) == NULL)
+        goto err;
+    if (EVP_DigestSignInit(mctx, &pctx, md, /*engine*/ NULL, key) != 1)
+        goto err;
+    if (EVP_DigestSignUpdate(mctx, text.data, text.size) != 1)
         goto err;
-    }
 
-    if ((size != ret_bin.size) &&
-        !enif_realloc_binary(&ret_bin, size)) {
+    if (EVP_DigestSignFinal(mctx, NULL, &size) != 1)
+        goto err;
+    if (!enif_alloc_binary(size, &ret_bin))
         goto err;
+    ret_bin_alloc = 1;
+    if (EVP_DigestSignFinal(mctx, ret_bin.data, &size) != 1)
+        goto err;
+
+    if (size != ret_bin.size) {
+        if (!enif_realloc_binary(&ret_bin, size))
+            goto err;
     }
 
     ret = enif_make_binary(env, &ret_bin);
+    ret_bin_alloc = 0;
+    goto done;
+
+ bad_arg:
+    return enif_make_badarg(env);
 
  err:
-    EVP_MD_CTX_free(mctx);
-    EVP_PKEY_free(key);
+    if (ret_bin_alloc)
+        enif_release_binary(&ret_bin);
+    ret = atom_error;
+
+ done:
+    if (mctx)
+        EVP_MD_CTX_free(mctx);
+    if (key)
+        EVP_PKEY_free(key);
     return ret;
 
 #else
     return atom_notsup;
 #endif
 }
-