Fix memory leak computing ECDH secrets

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

@@ -2114,15 +2114,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) {
@@ -2131,28 +2133,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");
     }