1 files changed, 41 insertions, 15 deletions

diff --git a/lib/crypto/c_src/rc4.c b/lib/crypto/c_src/rc4.c
index 483c87b04b..e423661097 100644
--- a/lib/crypto/c_src/rc4.c
+++ b/lib/crypto/c_src/rc4.c
@@ -25,15 +25,27 @@ ERL_NIF_TERM rc4_set_key(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
 #ifndef OPENSSL_NO_RC4
     ErlNifBinary key;
     ERL_NIF_TERM ret;
+    RC4_KEY *rc4_key;
 
     CHECK_NO_FIPS_MODE();
 
-    if (!enif_inspect_iolist_as_binary(env,argv[0], &key)) {
-	return enif_make_badarg(env);
-    }
-    RC4_set_key((RC4_KEY*)enif_make_new_binary(env, sizeof(RC4_KEY), &ret),
-		key.size, key.data);
+    ASSERT(argc == 1);
+
+    if (!enif_inspect_iolist_as_binary(env, argv[0], &key))
+        goto bad_arg;
+    if (key.size > INT_MAX)
+        goto bad_arg;
+
+    if ((rc4_key = (RC4_KEY*)enif_make_new_binary(env, sizeof(RC4_KEY), &ret)) == NULL)
+        goto err;
+
+    RC4_set_key(rc4_key, (int)key.size, key.data);
     return ret;
+
+ bad_arg:
+ err:
+    return enif_make_badarg(env);
+
 #else
     return enif_raise_exception(env, atom_notsup);
 #endif
@@ -45,20 +57,34 @@ ERL_NIF_TERM rc4_encrypt_with_state(ErlNifEnv* env, int argc, const ERL_NIF_TERM
     ErlNifBinary state, data;
     RC4_KEY* rc4_key;
     ERL_NIF_TERM new_state, new_data;
+    unsigned char *outp;
 
     CHECK_NO_FIPS_MODE();
 
-    if (!enif_inspect_iolist_as_binary(env,argv[0], &state)
-	|| state.size != sizeof(RC4_KEY)
-	|| !enif_inspect_iolist_as_binary(env,argv[1], &data)) {
-	return enif_make_badarg(env);
-    }
-    rc4_key = (RC4_KEY*)enif_make_new_binary(env, sizeof(RC4_KEY), &new_state);
+    ASSERT(argc == 2);
+
+    if (!enif_inspect_iolist_as_binary(env, argv[0], &state))
+        goto bad_arg;
+    if (state.size != sizeof(RC4_KEY))
+        goto bad_arg;
+    if (!enif_inspect_iolist_as_binary(env, argv[1], &data))
+        goto bad_arg;
+
+    if ((rc4_key = (RC4_KEY*)enif_make_new_binary(env, sizeof(RC4_KEY), &new_state)) == NULL)
+        goto err;
+    if ((outp = enif_make_new_binary(env, data.size, &new_data)) == NULL)
+        goto err;
+
     memcpy(rc4_key, state.data, sizeof(RC4_KEY));
-    RC4(rc4_key, data.size, data.data,
-	enif_make_new_binary(env, data.size, &new_data));
-    CONSUME_REDS(env,data);
-    return enif_make_tuple2(env,new_state,new_data);
+    RC4(rc4_key, data.size, data.data, outp);
+
+    CONSUME_REDS(env, data);
+    return enif_make_tuple2(env, new_state, new_data);
+
+ bad_arg:
+ err:
+    return enif_make_badarg(env);
+
 #else
     return enif_raise_exception(env, atom_notsup);
 #endif