1 files changed, 106 insertions, 40 deletions

diff --git a/lib/crypto/c_src/crypto.c b/lib/crypto/c_src/crypto.c
index a42de9adb1..aa99f2236e 100644
--- a/lib/crypto/c_src/crypto.c
+++ b/lib/crypto/c_src/crypto.c
@@ -242,6 +242,7 @@ static ERL_NIF_TERM aes_cfb_8_crypt(ErlNifEnv* env, int argc, const ERL_NIF_TERM
 static ERL_NIF_TERM aes_cfb_128_crypt(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM aes_ctr_encrypt(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM aes_ctr_stream_encrypt(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
+static ERL_NIF_TERM aes_ecb_crypt(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM rand_bytes_1(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM strong_rand_bytes_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM rand_bytes_3(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
@@ -379,6 +380,7 @@ static ErlNifFunc nif_funcs[] = {
     {"aes_ctr_decrypt", 3, aes_ctr_encrypt},
     {"aes_ctr_stream_encrypt", 2, aes_ctr_stream_encrypt},
     {"aes_ctr_stream_decrypt", 2, aes_ctr_stream_encrypt},
+    {"aes_ecb_crypt", 3, aes_ecb_crypt},
     {"rand_bytes", 1, rand_bytes_1},
     {"strong_rand_bytes_nif", 1, strong_rand_bytes_nif},
     {"rand_bytes", 3, rand_bytes_3},
@@ -410,7 +412,7 @@ static ErlNifFunc nif_funcs[] = {
     {"bf_ecb_crypt", 3, bf_ecb_crypt},
     {"blowfish_ofb64_encrypt", 3, blowfish_ofb64_encrypt},
 
-    {"ec_key_generate", 1, ec_key_generate},
+    {"ec_key_generate", 2, ec_key_generate},
     {"ecdsa_sign_nif", 4, ecdsa_sign_nif},
     {"ecdsa_verify_nif", 5, ecdsa_verify_nif},
     {"ecdh_compute_key_nif", 3, ecdh_compute_key_nif},
@@ -1686,14 +1688,15 @@ static ERL_NIF_TERM aes_cfb_8_crypt(ErlNifEnv* env, int argc, const ERL_NIF_TERM
     int new_ivlen = 0;
     ERL_NIF_TERM ret;
 
-    if (!enif_inspect_iolist_as_binary(env, argv[0], &key) || key.size != 16
+    if (!enif_inspect_iolist_as_binary(env, argv[0], &key)
+	|| !(key.size == 16 || key.size == 24 || key.size == 32)
 	|| !enif_inspect_binary(env, argv[1], &ivec) || ivec.size != 16
 	|| !enif_inspect_iolist_as_binary(env, argv[2], &text)) {
 	return enif_make_badarg(env);
     }
 
     memcpy(ivec_clone, ivec.data, 16);
-    AES_set_encrypt_key(key.data, 128, &aes_key);
+    AES_set_encrypt_key(key.data, key.size * 8, &aes_key);
     AES_cfb8_encrypt((unsigned char *) text.data,
 		       enif_make_new_binary(env, text.size, &ret), 
 		       text.size, &aes_key, ivec_clone, &new_ivlen,
@@ -1712,14 +1715,15 @@ static ERL_NIF_TERM aes_cfb_128_crypt(ErlNifEnv* env, int argc, const ERL_NIF_TE
 
     CHECK_OSE_CRYPTO();
 
-    if (!enif_inspect_iolist_as_binary(env, argv[0], &key) || key.size != 16
+    if (!enif_inspect_iolist_as_binary(env, argv[0], &key)
+	|| !(key.size == 16 || key.size == 24 || key.size == 32)
 	|| !enif_inspect_binary(env, argv[1], &ivec) || ivec.size != 16
 	|| !enif_inspect_iolist_as_binary(env, argv[2], &text)) {
 	return enif_make_badarg(env);
     }
 
     memcpy(ivec_clone, ivec.data, 16);
-    AES_set_encrypt_key(key.data, 128, &aes_key);
+    AES_set_encrypt_key(key.data, key.size * 8, &aes_key);
     AES_cfb128_encrypt((unsigned char *) text.data,
 		       enif_make_new_binary(env, text.size, &ret), 
 		       text.size, &aes_key, ivec_clone, &new_ivlen,
@@ -2032,6 +2036,38 @@ static ERL_NIF_TERM chacha20_poly1305_decrypt(ErlNifEnv* env, int argc, const ER
 #endif
 }
 
+static ERL_NIF_TERM aes_ecb_crypt(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
+{/* (Key, Data, IsEncrypt) */    
+    ErlNifBinary key_bin, data_bin;
+    AES_KEY aes_key;
+    int i;
+    unsigned char* ret_ptr;
+    ERL_NIF_TERM ret;    
+
+    CHECK_OSE_CRYPTO();
+
+    if (!enif_inspect_iolist_as_binary(env, argv[0], &key_bin)
+    || (key_bin.size != 16 && key_bin.size != 32)
+    || !enif_inspect_iolist_as_binary(env, argv[1], &data_bin)
+    || data_bin.size % 16 != 0) {
+    return enif_make_badarg(env);
+    }
+
+    if (argv[2] == atom_true) {
+        i = AES_ENCRYPT;
+        AES_set_encrypt_key(key_bin.data, key_bin.size*8, &aes_key);
+    }
+    else {
+        i = AES_DECRYPT;
+        AES_set_decrypt_key(key_bin.data, key_bin.size*8, &aes_key);
+    }
+
+    ret_ptr = enif_make_new_binary(env, data_bin.size, &ret);
+    AES_ecb_encrypt(data_bin.data, ret_ptr, &aes_key, i);
+    CONSUME_REDS(env,data_bin);
+    return ret;
+}
+
 static ERL_NIF_TERM rand_bytes_1(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
 {/* (Bytes) */     
     unsigned bytes;
@@ -2464,11 +2500,12 @@ done:
 static ERL_NIF_TERM aes_cbc_crypt(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
 {/* (Key, IVec, Data, IsEncrypt) */
     ErlNifBinary key_bin, ivec_bin, data_bin;
-    AES_KEY aes_key;
     unsigned char ivec[16];
-    int i;
+    int enc, i = 0, outlen = 0;
+    EVP_CIPHER_CTX *ctx = NULL;
+    const EVP_CIPHER *cipher = NULL;
     unsigned char* ret_ptr;
-    ERL_NIF_TERM ret;    
+    ERL_NIF_TERM ret;
 
     CHECK_OSE_CRYPTO();
 
@@ -2482,20 +2519,44 @@ static ERL_NIF_TERM aes_cbc_crypt(ErlNifEnv* env, int argc, const ERL_NIF_TERM a
 	return enif_make_badarg(env);
     }
 
-    if (argv[3] == atom_true) {
-	i = AES_ENCRYPT;
-	AES_set_encrypt_key(key_bin.data, key_bin.size*8, &aes_key);
-    }
-    else {
-	i = AES_DECRYPT;
-	AES_set_decrypt_key(key_bin.data, key_bin.size*8, &aes_key);
-    }
+    if (argv[3] == atom_true)
+	enc = 1;
+    else
+	enc = 0;
+
+    if (!(ctx = EVP_CIPHER_CTX_new()))
+	return enif_make_badarg(env);
+
+    if (key_bin.size == 16)
+	cipher = EVP_aes_128_cbc();
+    else if (key_bin.size == 32)
+	cipher = EVP_aes_256_cbc();
+
+    memcpy(ivec, ivec_bin.data, 16); /* writeable copy */
+
+    /* openssl docs say we need to leave at least 3 blocks available
+       at the end of the buffer for EVP calls. let's be safe */
+    ret_ptr = enif_make_new_binary(env, data_bin.size + 16*3, &ret);
+
+    if (EVP_CipherInit_ex(ctx, cipher, NULL, key_bin.data, ivec, enc) != 1)
+	return enif_make_badarg(env);
+
+    /* disable padding, we only handle whole blocks */
+    EVP_CIPHER_CTX_set_padding(ctx, 0);
+
+    if (EVP_CipherUpdate(ctx, ret_ptr, &i, data_bin.data, data_bin.size) != 1)
+	return enif_make_badarg(env);
+    outlen += i;
+    if (EVP_CipherFinal_ex(ctx, ret_ptr + outlen, &i) != 1)
+	return enif_make_badarg(env);
+    outlen += i;
+
+    EVP_CIPHER_CTX_free(ctx);
 
-    ret_ptr = enif_make_new_binary(env, data_bin.size, &ret);
-    memcpy(ivec, ivec_bin.data, 16); /* writable copy */
-    AES_cbc_encrypt(data_bin.data, ret_ptr, data_bin.size, &aes_key, ivec, i);
     CONSUME_REDS(env,data_bin);
-    return ret;
+
+    /* the garbage collector is going to love this */
+    return enif_make_sub_binary(env, ret, 0, outlen);
 }
 
 static ERL_NIF_TERM aes_ige_crypt_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
@@ -3689,32 +3750,37 @@ out:
 static ERL_NIF_TERM ec_key_generate(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
 {
 #if defined(HAVE_EC)
-    EC_KEY *key = ec_key_new(env, argv[0]);
+    EC_KEY *key;
+    const EC_GROUP *group;
+    const EC_POINT *public_key;
+    ERL_NIF_TERM priv_key;
+    ERL_NIF_TERM pub_key = atom_undefined;
 
     CHECK_OSE_CRYPTO();
 
-    if (key && EC_KEY_generate_key(key)) {
-	const EC_GROUP *group;
-	const EC_POINT *public_key;
-	ERL_NIF_TERM priv_key;
-	ERL_NIF_TERM pub_key = atom_undefined;
-
-	group = EC_KEY_get0_group(key);
-	public_key = EC_KEY_get0_public_key(key);
+    if (!get_ec_key(env, argv[0], argv[1], atom_undefined, &key))
+	goto badarg;
 
-	if (group && public_key) {
-	    pub_key = point2term(env, group, public_key,
-				 EC_KEY_get_conv_form(key));
-	}
-	priv_key = bn2term(env, EC_KEY_get0_private_key(key));
-	EC_KEY_free(key);
-	return enif_make_tuple2(env, pub_key, priv_key);
+    if (argv[1] == atom_undefined) {
+	if (!EC_KEY_generate_key(key))
+	    goto badarg;
     }
-    else {
-	if (key)
-	    EC_KEY_free(key);
-	return enif_make_badarg(env);
+
+    group = EC_KEY_get0_group(key);
+    public_key = EC_KEY_get0_public_key(key);
+
+    if (group && public_key) {
+	pub_key = point2term(env, group, public_key,
+			     EC_KEY_get_conv_form(key));
     }
+    priv_key = bn2term(env, EC_KEY_get0_private_key(key));
+    EC_KEY_free(key);
+    return enif_make_tuple2(env, pub_key, priv_key);
+
+badarg:
+    if (key)
+	EC_KEY_free(key);
+    return enif_make_badarg(env);
 #else
     return atom_notsup;
 #endif