Fix memory leak computing ECDH secrets (#393)

The context allocated in the method `XECKEY.computeECDHSecret` was never
freed when a key was successfully generated. This update frees memory
associated with the context prior to return of the secret key bytes.

Whitespace and formatting was also done to make use of brackets for if
statements.

Fixes #387

Signed-off-by: Jason Katonica <katonica@us.ibm.com>

Author: jasonkatonica

src/main/java/com/ibm/crypto/plus/provider/XDHKeyAgreement.java

@@ -110,7 +110,7 @@ protected Key engineDoPhase(Key key, boolean lastPhase)
             this.secret = XECKey.computeECDHSecret(provider.getOCKContext(), genCtx,
                     ockXecKeyPub.getPKeyId(), ockXecKeyPriv.getPKeyId(), secrectBufferSize);
         } catch (OCKException e) {
-            throw new IllegalStateException(e.getMessage());
+            throw new IllegalStateException("Failed to generate secret", e);
         } catch (Exception e) {
             throw new InvalidKeyException("Failed to generate secret", e);
         }

src/main/native/ECKey.c

@@ -2118,15 +2118,17 @@ JNIEXPORT jbyteArray JNICALL Java_com_ibm_crypto_plus_provider_ock_NativeInterfa
   jbyteArray           secretBytes = NULL;
   unsigned char *      secretBytesNative = NULL;
   jboolean             isCopy = 0;
-  jbyteArray           retSecretBytes = NULL;
   size_t               secret_key_len = 0;
+  int                  rc = 0;
 
-  
-  if( debug ) gslogFunctionEntry(functionName);
+  if (debug) {
+    gslogFunctionEntry(functionName);
+  }
 
   gen_ctx = ICC_EVP_PKEY_CTX_new(ockCtx,(ICC_EVP_PKEY *) ockPrivXecKey,NULL); /* Set private key */
-  if(gen_ctx == NULL) throwOCKException(env, 0, "NULL from ICC_EVP_PKEY_CTX_new"); 
-  else {
+  if (NULL == gen_ctx) {
+     throwOCKException(env, 0, "NULL from ICC_EVP_PKEY_CTX_new");
+  } else {
     ICC_EVP_PKEY_derive_init(ockCtx, gen_ctx);
     ICC_EVP_PKEY_derive_set_peer(ockCtx, gen_ctx, ockPubXecKey); /* Set public key */
     if (secretBufferSize > 0) {
@@ -2135,28 +2137,43 @@ JNIEXPORT jbyteArray JNICALL Java_com_ibm_crypto_plus_provider_ock_NativeInterfa
         ICC_EVP_PKEY_derive(ockCtx, gen_ctx, NULL, &secret_key_len); /* Get secret key size */
     }
     secretBytes = (*env)->NewByteArray(env, secret_key_len); /* Create Java secret bytes array with size */
-    if( secretBytes == NULL ) throwOCKException(env, 0, "NewByteArray failed"); 
-    else {
+    if (NULL == secretBytes) {
+      throwOCKException(env, 0, "NewByteArray failed"); 
+    } else {
       secretBytesNative = (unsigned char*)((*env)->GetPrimitiveArrayCritical(env, secretBytes, &isCopy));
-      if( secretBytesNative == NULL ) throwOCKException(env, 0, "NULL from GetPrimitiveArrayCritical"); 
-      else {
-        ICC_EVP_PKEY_derive(ockCtx, gen_ctx, secretBytesNative, &secret_key_len);
-        retSecretBytes = secretBytes;
-        if( secretBytesNative != NULL ) (*env)->ReleasePrimitiveArrayCritical(env, secretBytes, secretBytesNative, 0);
-        if((secretBytes != NULL) && (retSecretBytes == NULL)) (*env)->DeleteLocalRef(env, secretBytes);
-        if( debug ) gslogFunctionExit(functionName);
-        return retSecretBytes;
+      if (NULL == secretBytesNative) {
+        throwOCKException(env, 0, "NULL from GetPrimitiveArrayCritical");
+      } else {
+        rc = ICC_EVP_PKEY_derive(ockCtx, gen_ctx, secretBytesNative, &secret_key_len);
+        if (ICC_OSSL_SUCCESS != rc ) {
+          throwOCKException(env, 0, "ICC_EVP_PKEY_derive failed to derive a key");
+        }
+        ICC_EVP_PKEY_CTX_free(ockCtx, gen_ctx);
+        (*env)->ReleasePrimitiveArrayCritical(env, secretBytes, secretBytesNative, 0);
+        if (debug) {
+          gslogFunctionExit(functionName);
+        }
+        return secretBytes;
       }
     }
-    if (gen_ctx != NULL) {
-      ICC_EVP_PKEY_CTX_free(ockCtx,gen_ctx);
-      gen_ctx = NULL;
-    }
   }
 
-  if( secretBytesNative != NULL ) (*env)->ReleasePrimitiveArrayCritical(env, secretBytes, secretBytesNative, 0);
-  if((secretBytes != NULL) && (retSecretBytes == NULL)) (*env)->DeleteLocalRef(env, secretBytes);
-  if( debug ) gslogFunctionExit(functionName);
+  if (NULL != gen_ctx) {
+    ICC_EVP_PKEY_CTX_free(ockCtx, gen_ctx);
+  }
+
+  if (NULL != secretBytesNative) {
+    (*env)->ReleasePrimitiveArrayCritical(env, secretBytes, secretBytesNative, 0);
+  }
+
+  if (NULL != secretBytes) {
+    (*env)->DeleteLocalRef(env, secretBytes);
+  }
+
+  if (debug) {
+    gslogFunctionExit(functionName);
+  }
+
   return NULL;
 }
 

src/test/java/ibm/jceplus/junit/base/BaseTestXDH.java

@@ -31,7 +31,7 @@
 import java.util.Arrays;
 import javax.crypto.KeyAgreement;
 import org.junit.jupiter.api.Test;
-import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class BaseTestXDH extends BaseTestJunit5 {
 
@@ -337,18 +337,10 @@ private void testSmallOrder(String name, String a_pri, String b_pub, String resu
             throws Exception {
 
         try {
-            //System.out.println("Pub - "+b_pub);
             runDiffieHellmanTest(name, a_pri, b_pub, result);
-        } catch (InvalidKeyException ex) {
-            assertTrue(true);
-            return;
-        } catch (InvalidKeySpecException ex) {
-            assertTrue(true);
+        } catch (IllegalStateException ex) {
             return;
-        } catch (Exception e1) {
-            System.out.println(e1.getMessage());
         }
-
         throw new RuntimeException("No exception on small-order point");
     }