1 files changed, 40 insertions, 32 deletions

diff --git a/src/crypto/ecdh/ecdh.c b/src/crypto/ecdh/ecdh.c
index 4a1964a..14856db 100644
--- a/src/crypto/ecdh/ecdh.c
+++ b/src/crypto/ecdh/ecdh.c
@@ -75,27 +75,33 @@
 
 
 int ECDH_compute_key(void *out, size_t outlen, const EC_POINT *pub_key,
-                     EC_KEY *priv_key,
-                     void *(*kdf)(const void *in, size_t inlen, void *out,
-                                  size_t *outlen)) {
-  const BIGNUM *const priv = EC_KEY_get0_private_key(priv_key);
-  if (priv == NULL) {
-    OPENSSL_PUT_ERROR(ECDH, ECDH_R_NO_PRIVATE_VALUE);
-    return -1;
-  }
+                     EC_KEY *priv_key, void *(*KDF)(const void *in, size_t inlen,
+                                                void *out, size_t *outlen)) {
+  BN_CTX *ctx;
+  EC_POINT *tmp = NULL;
+  BIGNUM *x = NULL, *y = NULL;
+  const BIGNUM *priv;
+  const EC_GROUP *group;
+  int ret = -1;
+  size_t buflen;
+  uint8_t *buf = NULL;
 
-  BN_CTX *ctx = BN_CTX_new();
-  if (ctx == NULL) {
-    return -1;
+  if ((ctx = BN_CTX_new()) == NULL) {
+    goto err;
   }
   BN_CTX_start(ctx);
+  x = BN_CTX_get(ctx);
+  y = BN_CTX_get(ctx);
 
-  int ret = -1;
-  size_t buflen = 0;
-  uint8_t *buf = NULL;
+  priv = EC_KEY_get0_private_key(priv_key);
+  if (priv == NULL) {
+    OPENSSL_PUT_ERROR(ECDH, ECDH_R_NO_PRIVATE_VALUE);
+    goto err;
+  }
+
+  group = EC_KEY_get0_group(priv_key);
 
-  const EC_GROUP *const group = EC_KEY_get0_group(priv_key);
-  EC_POINT *tmp = EC_POINT_new(group);
+  tmp = EC_POINT_new(group);
   if (tmp == NULL) {
     OPENSSL_PUT_ERROR(ECDH, ERR_R_MALLOC_FAILURE);
     goto err;
@@ -106,13 +112,7 @@ int ECDH_compute_key(void *out, size_t outlen, const EC_POINT *pub_key,
     goto err;
   }
 
-  BIGNUM *x = BN_CTX_get(ctx);
-  if (!x) {
-    OPENSSL_PUT_ERROR(ECDH, ERR_R_MALLOC_FAILURE);
-    goto err;
-  }
-
-  if (!EC_POINT_get_affine_coordinates_GFp(group, tmp, x, NULL, ctx)) {
+  if (!EC_POINT_get_affine_coordinates_GFp(group, tmp, x, y, ctx)) {
     OPENSSL_PUT_ERROR(ECDH, ECDH_R_POINT_ARITHMETIC_FAILURE);
     goto err;
   }
@@ -129,25 +129,33 @@ int ECDH_compute_key(void *out, size_t outlen, const EC_POINT *pub_key,
     goto err;
   }
 
-  if (kdf != NULL) {
-    if (kdf(buf, buflen, out, &outlen) == NULL) {
+  if (KDF != 0) {
+    if (KDF(buf, buflen, out, &outlen) == NULL) {
       OPENSSL_PUT_ERROR(ECDH, ECDH_R_KDF_FAILED);
       goto err;
     }
+    ret = outlen;
   } else {
     /* no KDF, just copy as much as we can */
-    if (buflen < outlen) {
+    if (outlen > buflen) {
       outlen = buflen;
     }
     memcpy(out, buf, outlen);
+    ret = outlen;
   }
 
-  ret = outlen;
-
 err:
-  OPENSSL_free(buf);
-  EC_POINT_free(tmp);
-  BN_CTX_end(ctx);
-  BN_CTX_free(ctx);
+  if (tmp) {
+    EC_POINT_free(tmp);
+  }
+  if (ctx) {
+    BN_CTX_end(ctx);
+  }
+  if (ctx) {
+    BN_CTX_free(ctx);
+  }
+  if (buf) {
+    OPENSSL_free(buf);
+  }
   return ret;
 }