Add try/catch/finally to explicitly delete native structures

Consensys · Jul 30, 2020 · 913e484 · 913e484
1 parent 5dca043
commit 913e484
Show file tree

Hide file tree

Showing 6 changed files with 108 additions and 84 deletions.
diff --git a/bls/src/main/java/tech/pegasys/teku/bls/impl/blst/BlstBLS12381.java b/bls/src/main/java/tech/pegasys/teku/bls/impl/blst/BlstBLS12381.java
@@ -66,14 +66,18 @@ private static Random getRND() {
   }
 
   public static BlstSignature sign(BlstSecretKey secretKey, Bytes message) {
-    p2 p2Signature = new p2();
     p2 hash = HashToCurve.hashToG2(message);
-    blst.sign_pk_in_g1(p2Signature, hash, secretKey.getScalarVal());
-    p2_affine p2SignatureAffine = new p2_affine();
-    blst.p2_to_affine(p2SignatureAffine, p2Signature);
-    p2Signature.delete();
-    hash.delete();
-    return new BlstSignature(p2SignatureAffine, true);
+    p2 p2Signature = new p2();
+    try {
+      blst.sign_pk_in_g1(p2Signature, hash, secretKey.getScalarVal());
+      p2_affine p2SignatureAffine = new p2_affine();
+      blst.p2_to_affine(p2SignatureAffine, p2Signature);
+
+      return new BlstSignature(p2SignatureAffine, true);
+    } finally {
+      hash.delete();
+      p2Signature.delete();
+    }
   }
 
   public static boolean verify(BlstPublicKey publicKey, Bytes message, BlstSignature signature) {
@@ -141,10 +145,9 @@ BatchSemiAggregate blstPrepareBatchVerify(
 
     p2 g2Hash = HashToCurve.hashToG2(message);
     p2_affine p2Affine = new p2_affine();
-    blst.p2_to_affine(p2Affine, g2Hash);
-
     pairing ctx = new pairing();
     try {
+      blst.p2_to_affine(p2Affine, g2Hash);
       blst.pairing_init(ctx);
       BLST_ERROR ret =
           blst.pairing_mul_n_aggregate_pk_in_g1(
@@ -155,16 +158,17 @@ BatchSemiAggregate blstPrepareBatchVerify(
               nextBatchRandomMultiplier(),
               BATCH_RANDOM_BYTES * 8);
       if (ret != BLST_ERROR.BLST_SUCCESS) throw new IllegalArgumentException("Error: " + ret);
+
+      blst.pairing_commit(ctx);
+
+      return new BlstFiniteSemiAggregate(ctx);
     } catch (Exception e) {
       ctx.delete();
       throw e;
     } finally {
       g2Hash.delete();
       p2Affine.delete(); // not sure if its copied inside pairing_mul_n_aggregate_pk_in_g1
     }
-    blst.pairing_commit(ctx);
-
-    return new BlstFiniteSemiAggregate(ctx);
   }
 
   @Override
@@ -185,36 +189,38 @@ public BatchSemiAggregate prepareBatchVerify2(
 
   @Override
   public boolean completeBatchVerify(List<? extends BatchSemiAggregate> preparedList) {
-    boolean anyInvalidInfinity =
-        preparedList.stream()
-            .filter(a -> a instanceof BlstInfiniteSemiAggregate)
-            .map(a -> (BlstInfiniteSemiAggregate) a)
-            .anyMatch(a -> !a.isValid());
-
-    if (anyInvalidInfinity) {
-      return false;
-    }
-
-    List<BlstFiniteSemiAggregate> blstList =
-        preparedList.stream()
-            .filter(a -> a instanceof BlstFiniteSemiAggregate)
-            .map(b -> (BlstFiniteSemiAggregate) b)
-            .collect(Collectors.toList());
+    try {
+      boolean anyInvalidInfinity =
+          preparedList.stream()
+              .filter(a -> a instanceof BlstInfiniteSemiAggregate)
+              .map(a -> (BlstInfiniteSemiAggregate) a)
+              .anyMatch(a -> !a.isValid());
+
+      List<BlstFiniteSemiAggregate> blstList =
+          preparedList.stream()
+              .filter(a -> a instanceof BlstFiniteSemiAggregate)
+              .map(b -> (BlstFiniteSemiAggregate) b)
+              .collect(Collectors.toList());
+
+      if (blstList.isEmpty()) {
+        return !anyInvalidInfinity;
+      }
+      pairing ctx0 = blstList.get(0).getCtx();
+      boolean mergeRes = true;
+      for (int i = 1; i < blstList.size(); i++) {
+        BLST_ERROR ret = blst.pairing_merge(ctx0, blstList.get(i).getCtx());
+        mergeRes &= ret == BLST_ERROR.BLST_SUCCESS;
+      }
+
+      int boolRes = blst.pairing_finalverify(ctx0, null);
+      return mergeRes && boolRes != 0;
 
-    if (blstList.isEmpty()) {
-      return true;
-    }
-    pairing ctx0 = blstList.get(0).getCtx();
-    boolean mergeRes = true;
-    for (int i = 1; i < blstList.size(); i++) {
-      BLST_ERROR ret = blst.pairing_merge(ctx0, blstList.get(i).getCtx());
-      mergeRes &= ret == BLST_ERROR.BLST_SUCCESS;
-      blstList.get(i).release();
+    } finally {
+      preparedList.stream()
+          .filter(a -> a instanceof BlstFiniteSemiAggregate)
+          .map(b -> (BlstFiniteSemiAggregate) b)
+          .forEach(BlstFiniteSemiAggregate::release);
     }
-
-    int boolRes = blst.pairing_finalverify(ctx0, null);
-    blstList.get(0).release();
-    return mergeRes && boolRes != 0;
   }
 
   static BigInteger nextBatchRandomMultiplier() {

diff --git a/bls/src/main/java/tech/pegasys/teku/bls/impl/blst/BlstFiniteSemiAggregate.java b/bls/src/main/java/tech/pegasys/teku/bls/impl/blst/BlstFiniteSemiAggregate.java
@@ -54,6 +54,7 @@ public static BatchSemiAggregate merge(BatchSemiAggregate agg1, BatchSemiAggrega
   }
 
   pairing getCtx() {
+    if (released) throw new IllegalStateException("Attempting to use disposed BatchSemiAggregate");
     return ctx;
   }
 

diff --git a/bls/src/main/java/tech/pegasys/teku/bls/impl/blst/BlstPublicKey.java b/bls/src/main/java/tech/pegasys/teku/bls/impl/blst/BlstPublicKey.java
@@ -87,15 +87,18 @@ public static BlstPublicKey aggregate(List<BlstPublicKey> publicKeys) {
     }
 
     p1 sum = new p1();
-    blst.p1_from_affine(sum, finitePublicKeys.get(0).ecPoint);
-    for (int i = 1; i < finitePublicKeys.size(); i++) {
-      blst.p1_add_or_double_affine(sum, sum, finitePublicKeys.get(i).ecPoint);
+    try {
+      blst.p1_from_affine(sum, finitePublicKeys.get(0).ecPoint);
+      for (int i = 1; i < finitePublicKeys.size(); i++) {
+        blst.p1_add_or_double_affine(sum, sum, finitePublicKeys.get(i).ecPoint);
+      }
+      p1_affine res = new p1_affine();
+      blst.p1_to_affine(res, sum);
+
+      return new BlstPublicKey(res);
+    } finally {
+      sum.delete();
     }
-    p1_affine res = new p1_affine();
-    blst.p1_to_affine(res, sum);
-    sum.delete();
-
-    return new BlstPublicKey(res);
   }
 
   final p1_affine ecPoint;

diff --git a/bls/src/main/java/tech/pegasys/teku/bls/impl/blst/BlstSecretKey.java b/bls/src/main/java/tech/pegasys/teku/bls/impl/blst/BlstSecretKey.java
@@ -83,11 +83,15 @@ public BlstPublicKey derivePublicKey() {
       return BlstPublicKey.INFINITY;
     }
     p1 pk = new p1();
-    blst.sk_to_pk_in_g1(pk, getScalarVal());
-    p1_affine pkAffine = new p1_affine();
-    blst.p1_to_affine(pkAffine, pk);
-    pk.delete();
-    return new BlstPublicKey(pkAffine);
+    try {
+      blst.sk_to_pk_in_g1(pk, getScalarVal());
+      p1_affine pkAffine = new p1_affine();
+      blst.p1_to_affine(pkAffine, pk);
+
+      return new BlstPublicKey(pkAffine);
+    } finally {
+      pk.delete();
+    }
   }
 
   scalar getScalarVal() {

diff --git a/bls/src/main/java/tech/pegasys/teku/bls/impl/blst/BlstSignature.java b/bls/src/main/java/tech/pegasys/teku/bls/impl/blst/BlstSignature.java
@@ -55,8 +55,13 @@ public static BlstSignature fromBytes(Bytes compressed) {
         "Expected " + COMPRESSED_SIG_SIZE + " bytes of input but got %s",
         compressed.size());
     p2_affine ec2Point = new p2_affine();
-    BLST_ERROR rc = blst.p2_uncompress(ec2Point, compressed.toArrayUnsafe());
-    return new BlstSignature(ec2Point, rc == BLST_ERROR.BLST_SUCCESS);
+    try {
+      BLST_ERROR rc = blst.p2_uncompress(ec2Point, compressed.toArrayUnsafe());
+      return new BlstSignature(ec2Point, rc == BLST_ERROR.BLST_SUCCESS);
+    } catch (Exception e) {
+      ec2Point.delete();
+      throw e;
+    }
   }
 
   public static BlstSignature aggregate(List<BlstSignature> signatures) {
@@ -71,29 +76,29 @@ public static BlstSignature aggregate(List<BlstSignature> signatures) {
     }
 
     p2 sum = new p2();
-    blst.p2_from_affine(sum, finiteSignatures.get(0).ec2Point);
-    for (int i = 1; i < finiteSignatures.size(); i++) {
-      blst.p2_add_affine(sum, sum, finiteSignatures.get(i).ec2Point);
+    try {
+      blst.p2_from_affine(sum, finiteSignatures.get(0).ec2Point);
+      for (int i = 1; i < finiteSignatures.size(); i++) {
+        blst.p2_add_affine(sum, sum, finiteSignatures.get(i).ec2Point);
+      }
+      p2_affine res = new p2_affine();
+      blst.p2_to_affine(res, sum);
+
+      return new BlstSignature(res, true);
+    } finally {
+      sum.delete();
     }
-    p2_affine res = new p2_affine();
-    blst.p2_to_affine(res, sum);
-    sum.delete();
-    return new BlstSignature(res, true);
   }
 
-  private static pairing blstPrepareVerifyAggregated(
+  private static void blstPrepareVerifyAggregated(
       BlstPublicKey pubKey, Bytes message, pairing ctx, BlstSignature blstSignature) {
 
     p2 g2Hash = HashToCurve.hashToG2(message);
     p2_affine p2Affine = new p2_affine();
-    blst.p2_to_affine(p2Affine, g2Hash);
-
-    if (ctx == null) {
-      ctx = new pairing();
-      blst.pairing_init(ctx);
-    }
 
     try {
+      blst.p2_to_affine(p2Affine, g2Hash);
+
       BLST_ERROR ret =
           blst.pairing_aggregate_pk_in_g1(
               ctx,
@@ -104,15 +109,10 @@ private static pairing blstPrepareVerifyAggregated(
               HashToCurve.ETH2_DST.toArrayUnsafe(),
               null);
       if (ret != BLST_ERROR.BLST_SUCCESS) throw new IllegalArgumentException("Error: " + ret);
-    } catch (Exception e) {
-      ctx.delete();
-      throw e;
     } finally {
       g2Hash.delete();
       p2Affine.delete();
     }
-
-    return ctx;
   }
 
   private static boolean blstCompleteVerifyAggregated(pairing ctx) {
@@ -147,15 +147,20 @@ public Bytes toBytesUncompressed() {
   @Override
   public boolean verify(List<PublicKeyMessagePair> keysToMessages) {
 
-    pairing ctx = null;
-    for (int i = 0; i < keysToMessages.size(); i++) {
-      BlstPublicKey publicKey = (BlstPublicKey) keysToMessages.get(i).getPublicKey();
-      Bytes message = keysToMessages.get(i).getMessage();
-      BlstSignature signature = i == 0 ? this : null;
-      ctx = blstPrepareVerifyAggregated(publicKey, message, ctx, signature);
-    }
+    pairing ctx = new pairing();
 
-    return blstCompleteVerifyAggregated(ctx);
+    try {
+      blst.pairing_init(ctx);
+      for (int i = 0; i < keysToMessages.size(); i++) {
+        BlstPublicKey publicKey = (BlstPublicKey) keysToMessages.get(i).getPublicKey();
+        Bytes message = keysToMessages.get(i).getMessage();
+        BlstSignature signature = i == 0 ? this : null;
+        blstPrepareVerifyAggregated(publicKey, message, ctx, signature);
+      }
+      return blstCompleteVerifyAggregated(ctx);
+    } finally {
+      ctx.delete();
+    }
   }
 
   @Override

diff --git a/bls/src/main/java/tech/pegasys/teku/bls/impl/blst/HashToCurve.java b/bls/src/main/java/tech/pegasys/teku/bls/impl/blst/HashToCurve.java
@@ -30,7 +30,12 @@ static p2 hashToG2(Bytes message) {
 
   static p2 hashToG2(Bytes message, Bytes dst) {
     p2 p2Hash = new p2();
-    blst.hash_to_g2(p2Hash, message.toArray(), dst.toArray(), new byte[0]);
-    return p2Hash;
+    try {
+      blst.hash_to_g2(p2Hash, message.toArray(), dst.toArray(), new byte[0]);
+      return p2Hash;
+    } catch (Exception e) {
+      p2Hash.delete();
+      throw e;
+    }
   }
 }