Revamp dh_generate_key_nif()

1 files changed, 143 insertions, 101 deletions

diff --git a/lib/crypto/c_src/dh.c b/lib/crypto/c_src/dh.c
index 0c18ad7a3f..10fd122658 100644
--- a/lib/crypto/c_src/dh.c
+++ b/lib/crypto/c_src/dh.c
@@ -24,122 +24,164 @@
 ERL_NIF_TERM dh_generate_key_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
 {/* (PrivKey|undefined, DHParams=[P,G], Mpint, Len|0) */
     DH *dh_params = NULL;
-    int mpint; /* 0 or 4 */
+    unsigned int mpint; /* 0 or 4 */
+    ERL_NIF_TERM head, tail;
+    BIGNUM *dh_p = NULL;
+    BIGNUM *dh_g = NULL;
+    BIGNUM *priv_key_in = NULL;
+    unsigned long len = 0;
+    unsigned char *pub_ptr, *prv_ptr;
+    int pub_len, prv_len;
+    ERL_NIF_TERM ret_pub, ret_prv, ret;
+    const BIGNUM *pub_key_gen, *priv_key_gen;
+#ifdef HAS_EVP_PKEY_CTX
+    EVP_PKEY_CTX *ctx = NULL;
+    EVP_PKEY *dhkey = NULL, *params = NULL;
+#endif
 
-    {
-        ERL_NIF_TERM head, tail;
-        BIGNUM
-            *dh_p = NULL,
-            *dh_g = NULL,
-            *priv_key_in = NULL;
-        unsigned long
-            len = 0;
-
-        if (!(get_bn_from_bin(env, argv[0], &priv_key_in)
-              || argv[0] == atom_undefined)
-            || !enif_get_list_cell(env, argv[1], &head, &tail)
-            || !get_bn_from_bin(env, head, &dh_p)
-            || !enif_get_list_cell(env, tail, &head, &tail)
-            || !get_bn_from_bin(env, head, &dh_g)
-            || !enif_is_empty_list(env, tail)
-            || !enif_get_int(env, argv[2], &mpint) || (mpint & ~4)
-            || !enif_get_ulong(env, argv[3], &len)
-
-            /* Load dh_params with values to use by the generator.
-               Mem mgmnt transfered from dh_p etc to dh_params */
-            || !(dh_params = DH_new())
-            || (priv_key_in && !DH_set0_key(dh_params, NULL, priv_key_in))
-            || !DH_set0_pqg(dh_params, dh_p, NULL, dh_g)
-            ) {
-            if (priv_key_in) BN_free(priv_key_in);
-            if (dh_p) BN_free(dh_p);
-            if (dh_g) BN_free(dh_g);
-            if (dh_params) DH_free(dh_params);
-            return enif_make_badarg(env);
-        }
+    if (argc != 4)
+        goto bad_arg;
 
-        if (len) {
-            if (len < BN_num_bits(dh_p))
-                DH_set_length(dh_params, len);
-            else {
-                if (priv_key_in) BN_free(priv_key_in);
-                if (dh_p) BN_free(dh_p);
-                if (dh_g) BN_free(dh_g);
-                if (dh_params) DH_free(dh_params);
-                return enif_make_badarg(env);
-            }
-        }
+    if (argv[0] != atom_undefined) {
+        if (!get_bn_from_bin(env, argv[0], &priv_key_in))
+            goto bad_arg;
+    }
+    if (!enif_get_list_cell(env, argv[1], &head, &tail))
+        goto bad_arg;
+    if (!get_bn_from_bin(env, head, &dh_p))
+        goto bad_arg;
+
+    if (!enif_get_list_cell(env, tail, &head, &tail))
+        goto bad_arg;
+    if (!get_bn_from_bin(env, head, &dh_g))
+        goto bad_arg;
+
+    if (!enif_is_empty_list(env, tail))
+        goto bad_arg;
+
+    if (!enif_get_uint(env, argv[2], &mpint))
+        goto bad_arg;
+    if (mpint != 0 && mpint != 4)
+        goto bad_arg;
+
+    if (!enif_get_ulong(env, argv[3], &len))
+        goto bad_arg;
+    if (len > LONG_MAX)
+        goto bad_arg;
+
+    /* Load dh_params with values to use by the generator.
+       Mem mgmnt transfered from dh_p etc to dh_params */
+    if ((dh_params = DH_new()) == NULL)
+        goto err;
+    if (priv_key_in) {
+        if (!DH_set0_key(dh_params, NULL, priv_key_in))
+            goto err;
+        /* On success, dh_params owns priv_key_in */
+        priv_key_in = NULL;
+    }
+    if (!DH_set0_pqg(dh_params, dh_p, NULL, dh_g))
+        goto err;
+    /* On success, dh_params owns dh_p and dh_g */
+    dh_p = NULL;
+    dh_g = NULL;
+
+    if (len) {
+        int bn_len;
+
+        if ((bn_len = BN_num_bits(dh_p)) < 0)
+            goto bad_arg;
+        if (len >= (size_t)bn_len)
+            goto bad_arg;
+
+        if (!DH_set_length(dh_params, (long)len))
+            goto err;
     }
 
 #ifdef HAS_EVP_PKEY_CTX
-    {
-        EVP_PKEY_CTX *ctx;
-        EVP_PKEY *dhkey, *params;
-        int success;
+    if ((params = EVP_PKEY_new()) == NULL)
+        goto err;
 
-        params = EVP_PKEY_new();
-        success = EVP_PKEY_set1_DH(params, dh_params);   /* set the key referenced by params to dh_params... */
-        DH_free(dh_params);                              /* ...dh_params (and params) must be freed */
-        if (!success) return atom_error;
+   /* set the key referenced by params to dh_params... */
+    if (EVP_PKEY_set1_DH(params, dh_params) != 1)
+        goto err;
 
-        ctx = EVP_PKEY_CTX_new(params, NULL);
-        EVP_PKEY_free(params);
-        if (!ctx) {
-            return atom_error;
-        }
+    if ((ctx = EVP_PKEY_CTX_new(params, NULL)) == NULL)
+        goto err;
 
-        if (!EVP_PKEY_keygen_init(ctx)) {
-            /* EVP_PKEY_CTX_free(ctx); */
-            return atom_error;
-        }
+    if (EVP_PKEY_keygen_init(ctx) != 1)
+        goto err;
 
-        dhkey = EVP_PKEY_new();
-        if (!EVP_PKEY_keygen(ctx, &dhkey)) {         /* "performs a key generation operation, the ... */
-                                                     /*... generated key is written to ppkey." (=last arg) */
-             /* EVP_PKEY_CTX_free(ctx); */
-             /* EVP_PKEY_free(dhkey); */
-             return atom_error;
-        }
+    if ((dhkey = EVP_PKEY_new()) == NULL)
+        goto err;
+
+    /* key gen op, key written to ppkey (=last arg) */
+    if (EVP_PKEY_keygen(ctx, &dhkey) != 1)
+        goto err;
+
+    DH_free(dh_params);
+    if ((dh_params = EVP_PKEY_get1_DH(dhkey)) == NULL)
+        goto err;
 
-        dh_params = EVP_PKEY_get1_DH(dhkey); /* return the referenced key. dh_params and dhkey must be freed */
-        EVP_PKEY_free(dhkey);
-        if (!dh_params) {
-            /* EVP_PKEY_CTX_free(ctx); */
-            return atom_error;
-        }
-        EVP_PKEY_CTX_free(ctx);
-    }
 #else
-    if (!DH_generate_key(dh_params)) return atom_error;
+    if (!DH_generate_key(dh_params))
+        goto err;
 #endif
-    {
-        unsigned char *pub_ptr, *prv_ptr;
-        int pub_len, prv_len;
-        ERL_NIF_TERM ret_pub, ret_prv;
-        const BIGNUM *pub_key_gen, *priv_key_gen;
-
-        DH_get0_key(dh_params,
-                    &pub_key_gen, &priv_key_gen); /* Get pub_key_gen and priv_key_gen.
-                                                     "The values point to the internal representation of
-                                                     the public key and private key values. This memory
-                                                     should not be freed directly." says man */
-        pub_len = BN_num_bytes(pub_key_gen);
-        prv_len = BN_num_bytes(priv_key_gen);
-        pub_ptr = enif_make_new_binary(env, pub_len+mpint, &ret_pub);
-        prv_ptr = enif_make_new_binary(env, prv_len+mpint, &ret_prv);
-        if (mpint) {
-            put_int32(pub_ptr, pub_len); pub_ptr += 4;
-            put_int32(prv_ptr, prv_len); prv_ptr += 4;
-        }
-        BN_bn2bin(pub_key_gen, pub_ptr);
-        BN_bn2bin(priv_key_gen, prv_ptr);
-        ERL_VALGRIND_MAKE_MEM_DEFINED(pub_ptr, pub_len);
-        ERL_VALGRIND_MAKE_MEM_DEFINED(prv_ptr, prv_len);
 
-        DH_free(dh_params);
+    DH_get0_key(dh_params, &pub_key_gen, &priv_key_gen);
+
+    if ((pub_len = BN_num_bytes(pub_key_gen)) < 0)
+        goto err;
+    if ((prv_len = BN_num_bytes(priv_key_gen)) < 0)
+        goto err;
 
-        return enif_make_tuple2(env, ret_pub, ret_prv);
+    if ((pub_ptr = enif_make_new_binary(env, (size_t)pub_len+mpint, &ret_pub)) == NULL)
+        goto err;
+    if ((prv_ptr = enif_make_new_binary(env, (size_t)prv_len+mpint, &ret_prv)) == NULL)
+        goto err;
+
+    if (mpint) {
+        put_int32(pub_ptr, pub_len);
+        pub_ptr += 4;
+
+        put_int32(prv_ptr, prv_len);
+        prv_ptr += 4;
     }
+
+    if (BN_bn2bin(pub_key_gen, pub_ptr) < 0)
+        goto err;
+    if (BN_bn2bin(priv_key_gen, prv_ptr) < 0)
+        goto err;
+
+    ERL_VALGRIND_MAKE_MEM_DEFINED(pub_ptr, pub_len);
+    ERL_VALGRIND_MAKE_MEM_DEFINED(prv_ptr, prv_len);
+
+    ret = enif_make_tuple2(env, ret_pub, ret_prv);
+    goto done;
+
+ bad_arg:
+ err:
+    ret = enif_make_badarg(env);
+
+ done:
+    if (priv_key_in)
+        BN_free(priv_key_in);
+    if (dh_p)
+        BN_free(dh_p);
+    if (dh_g)
+        BN_free(dh_g);
+    if (dh_params)
+        DH_free(dh_params);
+
+#ifdef HAS_EVP_PKEY_CTX
+    if (ctx)
+        EVP_PKEY_CTX_free(ctx);
+    if (dhkey)
+        EVP_PKEY_free(dhkey);
+    if (params)
+        EVP_PKEY_free(params);
+#endif
+
+    return ret;
 }
 
 ERL_NIF_TERM dh_compute_key_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])