From f87eee57d3ad99d09260c0bc2d32ec23d2010a45 Mon Sep 17 00:00:00 2001
From: Dmitrii Shmatko <zilm13@users.noreply.github.com>
Date: Mon, 20 Feb 2023 16:02:34 +0400
Subject: [PATCH] Fix getting SEGFAULT in Java with incorrect parameters (#148)

* Java bindings: Fix SEGFAULT on incorrect parameters/wrong file

* Prettify allocate KZGSettings
---
 bindings/java/c_kzg_4844_jni.c                | 34 ++++++++++++-------
 .../ethereum/ckzg4844/CKZG4844JNITest.java    | 28 ++++++++++++++-
 2 files changed, 48 insertions(+), 14 deletions(-)

diff --git a/bindings/java/c_kzg_4844_jni.c b/bindings/java/c_kzg_4844_jni.c
index 0b27f7f..370c4d1 100644
--- a/bindings/java/c_kzg_4844_jni.c
+++ b/bindings/java/c_kzg_4844_jni.c
@@ -11,7 +11,10 @@ void reset_trusted_setup()
 {
   if (settings)
   {
-    free_trusted_setup(settings);
+    if (settings->fs)
+    {
+      free_trusted_setup(settings);
+    }
     free(settings);
     settings = NULL;
   }
@@ -39,6 +42,21 @@ void throw_invalid_size_exception(JNIEnv *env, const char *prefix, size_t size,
   throw_c_kzg_exception(env, C_KZG_BADARGS, message);
 }
 
+KZGSettings *allocate_settings(JNIEnv *env) {
+  KZGSettings *s = malloc(sizeof(KZGSettings));
+  if (s == NULL)
+  {
+    throw_exception(env, "Failed to allocate memory for the Trusted Setup.");
+  }
+  else
+  {
+    s->fs = NULL;
+    s->g1_values = NULL;
+    s->g2_values = NULL;
+  }
+  return s;
+}
+
 JNIEXPORT jint JNICALL Java_ethereum_ckzg4844_CKZG4844JNI_getFieldElementsPerBlob(JNIEnv *env, jclass thisCls)
 {
   return (jint)FIELD_ELEMENTS_PER_BLOB;
@@ -52,12 +70,7 @@ JNIEXPORT void JNICALL Java_ethereum_ckzg4844_CKZG4844JNI_loadTrustedSetup__Ljav
     return;
   }
 
-  settings = malloc(sizeof(KZGSettings));
-  if (settings == NULL)
-  {
-    throw_exception(env, "Failed to allocate memory for the Trusted Setup.");
-    return;
-  }
+  settings = allocate_settings(env);
 
   const char *file_native = (*env)->GetStringUTFChars(env, file, 0);
 
@@ -92,12 +105,7 @@ JNIEXPORT void JNICALL Java_ethereum_ckzg4844_CKZG4844JNI_loadTrustedSetup___3BJ
     return;
   }
 
-  settings = malloc(sizeof(KZGSettings));
-  if (settings == NULL)
-  {
-    throw_exception(env, "Failed to allocate memory for the Trusted Setup.");
-    return;
-  }
+  settings = allocate_settings(env);
 
   jbyte *g1_native = (*env)->GetByteArrayElements(env, g1, NULL);
   jbyte *g2_native = (*env)->GetByteArrayElements(env, g2, NULL);
diff --git a/bindings/java/src/test/java/ethereum/ckzg4844/CKZG4844JNITest.java b/bindings/java/src/test/java/ethereum/ckzg4844/CKZG4844JNITest.java
index 1012853..4d7b452 100644
--- a/bindings/java/src/test/java/ethereum/ckzg4844/CKZG4844JNITest.java
+++ b/bindings/java/src/test/java/ethereum/ckzg4844/CKZG4844JNITest.java
@@ -20,7 +20,7 @@ public class CKZG4844JNITest {
   private enum TrustedSetupSource {
     FILE,
     PARAMETERS,
-    RESOURCE;
+    RESOURCE
   }
 
   private static final Preset PRESET;
@@ -239,6 +239,32 @@ public class CKZG4844JNITest {
     assertExceptionIsTrustedSetupIsNotLoaded(exception);
   }
 
+  @Test
+  public void shouldThrowExceptionOnIncorrectTrustedSetupParameters() {
+    final LoadTrustedSetupParameters parameters =
+        TestUtils.createLoadTrustedSetupParameters(TRUSTED_SETUP_FILE_BY_PRESET.get(PRESET));
+    final CKZGException ckzgException =
+        assertThrows(
+            CKZGException.class,
+            () ->
+                CKZG4844JNI.loadTrustedSetup(
+                    parameters.getG1(),
+                    parameters.getG1Count() + 1,
+                    parameters.getG2(),
+                    parameters.getG2Count()));
+    assertTrue(ckzgException.getMessage().contains("C_KZG_BADARGS"));
+  }
+
+  @Test
+  public void shouldThrowExceptionOnIncorrectTrustedSetupFromFile() {
+    final Preset incorrectPreset = PRESET == Preset.MAINNET ? Preset.MINIMAL : Preset.MAINNET;
+    final CKZGException ckzgException =
+        assertThrows(
+            CKZGException.class,
+            () -> CKZG4844JNI.loadTrustedSetup(TRUSTED_SETUP_FILE_BY_PRESET.get(incorrectPreset)));
+    assertTrue(ckzgException.getMessage().contains("C_KZG_BADARGS"));
+  }
+
   private void assertExceptionIsTrustedSetupIsNotLoaded(final RuntimeException exception) {
     assertEquals("Trusted Setup is not loaded.", exception.getMessage());
   }