diff --git a/jni/c/headers/java/fr_acinq_secp256k1_Secp256k1CFunctions.h b/jni/c/headers/java/fr_acinq_secp256k1_Secp256k1CFunctions.h
index 6268c13..93c5d76 100644
--- a/jni/c/headers/java/fr_acinq_secp256k1_Secp256k1CFunctions.h
+++ b/jni/c/headers/java/fr_acinq_secp256k1_Secp256k1CFunctions.h
@@ -270,10 +270,10 @@ JNIEXPORT jbyteArray JNICALL Java_fr_acinq_secp256k1_Secp256k1CFunctions_secp256
 /*
  * Class:     fr_acinq_secp256k1_Secp256k1CFunctions
  * Method:    secp256k1_frost_shares_gen
- * Signature: (J[BII[[B)[[[B
+ * Signature: (J[B[BII[[B)[[[B
  */
 JNIEXPORT jobjectArray JNICALL Java_fr_acinq_secp256k1_Secp256k1CFunctions_secp256k1_1frost_1shares_1gen
-  (JNIEnv *, jclass, jlong, jbyteArray, jint, jint, jobjectArray);
+  (JNIEnv *, jclass, jlong, jbyteArray, jbyteArray, jint, jint, jobjectArray);
 
 /*
  * Class:     fr_acinq_secp256k1_Secp256k1CFunctions
diff --git a/jni/c/src/fr_acinq_secp256k1_Secp256k1CFunctions.c b/jni/c/src/fr_acinq_secp256k1_Secp256k1CFunctions.c
index 98e6a6e..d4c4823 100644
--- a/jni/c/src/fr_acinq_secp256k1_Secp256k1CFunctions.c
+++ b/jni/c/src/fr_acinq_secp256k1_Secp256k1CFunctions.c
@@ -1354,13 +1354,14 @@ JNIEXPORT jbyteArray JNICALL Java_fr_acinq_secp256k1_Secp256k1CFunctions_secp256
  * Method:    secp256k1_frost_shares_gen
  * Signature: (J[B[BII[[B)[[[B
  */
-JNIEXPORT jobjectArray JNICALL Java_fr_acinq_secp256k1_Secp256k1CFunctions_secp256k1_1frost_1shares_1gen(JNIEnv *penv, jclass clazz, jlong jctx, jbyteArray jseed32, jint jthreshold, jint jn_participants, jobjectArray jids33)
+JNIEXPORT jobjectArray JNICALL Java_fr_acinq_secp256k1_Secp256k1CFunctions_secp256k1_1frost_1shares_1gen
+  (JNIEnv *penv, jclass clazz, jlong jctx, jbyteArray jpok64, jbyteArray jseed32, jint jthreshold, jint jn_participants, jobjectArray jids33)
 {
   secp256k1_context *ctx = (secp256k1_context *)jctx;
 
   secp256k1_frost_share *shares;
   secp256k1_pubkey* vss_commitment;
-  unsigned char pok64[64];
+  jbyte* pok64;
 
   size_t size;
 
@@ -1388,36 +1389,40 @@ JNIEXPORT jobjectArray JNICALL Java_fr_acinq_secp256k1_Secp256k1CFunctions_secp2
 
   size = (*penv)->GetArrayLength(penv, jids33);
   CHECKRESULT(size != jn_participants, "invalid ids33 size");
+
+  CHECKRESULT((*penv)->GetArrayLength(penv, jpok64) != 64, "pok64 length must be 64 bytes");
+
+  shares = calloc(jn_participants, sizeof(secp256k1_frost_share*));
+  vss_commitment = calloc(jthreshold, sizeof(secp256k1_pubkey*));
+
   for (i = 0; i < jn_participants; i++)
   {
     jbyteArray id33 = (jbyteArray)(*penv)->GetObjectArrayElement(penv, jids33, i);
     // TODO: Check id33 size is 33...
-    ids33[i] = (*penv)->GetByteArrayElements(penv, id33, 0);
+    ids33[i] = (*penv)->GetByteArrayElements(penv, id33, 0); // TODO: use setElement
   }
 
   int result = 0;
 
+  pok64 = (*penv)->GetByteArrayElements(penv, jpok64, 0);
   result = secp256k1_frost_shares_gen(
     ctx,
     shares,
     vss_commitment,
-    pok64,
+    (unsigned char *) pok64,
     seed32,
     jthreshold,
     jn_participants,
     ids33
   );
-
-  (*penv)->ReleaseByteArrayElements(penv, jseed32, seed32, 0);
+  (*penv)->ReleaseByteArrayElements(penv, jpok64, pok64, 0);
   CHECKRESULT(!result, "secp256k1_frost_shares_gen failed");
 
-  jbyte* jpok64;
-
-  jobjectArray output = (*penv)->NewObjectArray(penv, 3, jobjectArray, NULL);
+  jobjectArray output = (*penv)->NewObjectArray(penv, 2, jobjectArray, NULL);
 
   output[0] = (*penv)->NewObjectArray(penv, jn_participants, jbyteArray, NULL);
+
   output[1] = (*penv)->NewObjectArray(penv, jthreshold, jbyteArray, NULL);
-  output[2] = (*penv)->NewObjectArray(penv, 1, jbyteArray, NULL);
 
   // TODO: Copy over the required data...
   return output;
@@ -1576,22 +1581,20 @@ JNIEXPORT jint JNICALL Java_fr_acinq_secp256k1_Secp256k1CFunctions_secp256k1_1fr
 
   CHECKRESULT(jthreshold <= 0, "threshold can't be 0");
 
-  share = calloc(1, sizeof(secp256k1_frost_share));
   size = (*penv)->GetArrayLength(penv, jshare);
   CHECKRESULT1(size != sizeof(secp256k1_frost_share), "invalid share size", free(&share));
   in32 = (*penv)->GetByteArrayElements(penv, jshare, 0);
-  result = secp256k1_frost_share_parse(ctx, share, (unsigned char *)in32);
+  result = secp256k1_frost_share_parse(ctx, &share, (unsigned char *)in32);
   (*penv)->ReleaseByteArrayElements(penv, jshare, in32, 0);
   CHECKRESULT1(!result, "secp256k1_frost_share_parse failed", free(&share));
 
   id33 = (*penv)->GetByteArrayElements(penv, jid33, 0);
 
-  vss_commitment = calloc(1, sizeof(secp256k1_pubkey));
   jpub = (jbyteArray)(*penv)->GetObjectArrayElement(penv, jvss_commitment, i);
   size = (*penv)->GetArrayLength(penv, jpub);
   CHECKRESULT1((size != 33) && (size != 65), "invalid public key size", free_pubkeys(pubkeys, count));
   jpub = (*penv)->GetByteArrayElements(penv, jpub, 0);
-  result = secp256k1_ec_pubkey_parse(ctx, vss_commitment, (unsigned char *)jpub, size);
+  result = secp256k1_ec_pubkey_parse(ctx, &vss_commitment, (unsigned char *)jpub, size);
   (*penv)->ReleaseByteArrayElements(penv, jpub, jpub, 0);
   CHECKRESULT1(!result, "secp256k1_ec_pubkey_parse failed", free(vss_commitment));
 
@@ -1876,11 +1879,10 @@ JNIEXPORT jbyteArray JNICALL Java_fr_acinq_secp256k1_Secp256k1CFunctions_secp256
   copy_bytes_from_java(penv, jsession_id32, size, session_id32);
 
   if (jshare != NULL) {
-    share = calloc(1, sizeof(secp256k1_frost_share));
     size = (*penv)->GetArrayLength(penv, jshare);
     //      TODO:  CHECKRESULT1(size != fr_acinq_secp256k1_Secp256k1CFunctions_SECP256K1_MUSIG_PUBLIC_NONCE_SIZE, "invalid public nonce size", free_nonces(pubnonces, count));
     in32 = (*penv)->GetByteArrayElements(penv, jshare, 0);
-    result = secp256k1_frost_share_parse(ctx, share, (unsigned char *)in32);
+    result = secp256k1_frost_share_parse(ctx, &share, (unsigned char *)in32);
     (*penv)->ReleaseByteArrayElements(penv, jshare, in32, 0);
     //  TODO: CHECKRESULT1(!result, "secp256k1_frost_share_parse failed", free_shares(shares, count));
 
@@ -2096,13 +2098,12 @@ JNIEXPORT jbyteArray JNICALL Java_fr_acinq_secp256k1_Secp256k1CFunctions_secp256
 
   copy_bytes_from_java(penv, jsecnonce, fr_acinq_secp256k1_Secp256k1CFunctions_SECP256K1_MUSIG_SECRET_NONCE_SIZE, secnonce.data);
 
-  agg_share = calloc(1, sizeof(secp256k1_frost_share));
   size = (*penv)->GetArrayLength(penv, jagg_share);
-  CHECKRESULT1(size != sizeof(secp256k1_frost_share), "invalid agg_share size", free_shares(shares, count));
+  CHECKRESULT(size != sizeof(secp256k1_frost_share), "invalid agg_share size");
   in32 = (*penv)->GetByteArrayElements(penv, jagg_share, 0);
-  result = secp256k1_frost_share_parse(ctx, agg_share, (unsigned char *)in32);
+  result = secp256k1_frost_share_parse(ctx, &agg_share, (unsigned char *)in32);
   (*penv)->ReleaseByteArrayElements(penv, jagg_share, in32, 0);
-  CHECKRESULT1(!result, "secp256k1_frost_share_parse failed", free_shares(shares, count));
+  CHECKRESULT(!result, "secp256k1_frost_share_parse failed");
 
   CHECKRESULT((*penv)->GetArrayLength(penv, jsession) != fr_acinq_secp256k1_Secp256k1CFunctions_SECP256K1_MUSIG_SESSION_SIZE, "invalid session size");
   copy_bytes_from_java(penv, jsession, fr_acinq_secp256k1_Secp256k1CFunctions_SECP256K1_MUSIG_SESSION_SIZE, session.data);
diff --git a/jni/src/main/java/fr/acinq/secp256k1/Secp256k1CFunctions.java b/jni/src/main/java/fr/acinq/secp256k1/Secp256k1CFunctions.java
index 6bb910d..bf04d87 100644
--- a/jni/src/main/java/fr/acinq/secp256k1/Secp256k1CFunctions.java
+++ b/jni/src/main/java/fr/acinq/secp256k1/Secp256k1CFunctions.java
@@ -134,7 +134,7 @@ public class Secp256k1CFunctions {
      *      [1] vss_commitment: pointer to the VSS commitment
      *      [2] pok64: pointer to the proof of knowledge
      */
-    public static native byte[][][] secp256k1_frost_shares_gen(long ctx, byte[] seed32, int threshold, int total_signers, byte[][] ids33);
+    public static native byte[][][] secp256k1_frost_shares_gen(long ctx, byte[] pok64, byte[] seed32, int threshold, int total_signers, byte[][] ids33);
 
     /**
      *  Aggregates shares
diff --git a/jni/src/main/kotlin/fr/acinq/secp256k1/NativeSecp256k1.kt b/jni/src/main/kotlin/fr/acinq/secp256k1/NativeSecp256k1.kt
index 8bba854..298d318 100644
--- a/jni/src/main/kotlin/fr/acinq/secp256k1/NativeSecp256k1.kt
+++ b/jni/src/main/kotlin/fr/acinq/secp256k1/NativeSecp256k1.kt
@@ -135,8 +135,11 @@ public object NativeSecp256k1 : Secp256k1 {
         totalSigners: Int,
         ids33: Array<ByteArray>
     ): Triple<Array<ByteArray>, Array<ByteArray>, ByteArray> {
+        val pok64 = ByteArray(64)
+
         val result = Secp256k1CFunctions.secp256k1_frost_shares_gen(
             Secp256k1Context.getContext(),
+            pok64,
             seed32,
             threshold,
             totalSigners,
@@ -146,7 +149,7 @@ public object NativeSecp256k1 : Secp256k1 {
         return Triple(
             result[0],
             result[1],
-            result[2].first() // This is bad code...
+            pok64
         )
     }