Cache public key group membership status

bls/src/main/java/tech/pegasys/teku/bls/BLSPublicKey.java

@@ -40,7 +40,7 @@ public static BLSPublicKey empty() {
 
   /**
    * Aggregates list of PublicKeys, returns the public key that corresponds to G1 point at infinity
-   * if list is empty
+   * if list is empty, or if any of the public keys is infinity or not a G1 group member.
    *
    * @param publicKeys The list of public keys to aggregate
    * @return PublicKey The public key

bls/src/main/java/tech/pegasys/teku/bls/impl/blst/BlstPublicKey.java

@@ -34,6 +34,8 @@ class BlstPublicKey implements PublicKey {
               + "00000000000000000000000000000000"
               + "00000000000000000000000000000000");
 
+  private static final BlstPublicKey infinitePublicKey = fromBytes(INFINITY_COMPRESSED_BYTES);
+
   public static BlstPublicKey fromBytes(Bytes48 compressed) {
     try {
       P1_Affine ecPoint = new P1_Affine(compressed.toArrayUnsafe());
@@ -54,47 +56,60 @@ static BlstPublicKey fromPublicKey(PublicKey publicKey) {
   public static BlstPublicKey aggregate(List<BlstPublicKey> publicKeys) {
     checkArgument(publicKeys.size() > 0);
 
-    Optional<BlstPublicKey> maybeInfinitePubkey =
-        publicKeys.stream().filter(BlstPublicKey::isInfinity).findAny();
-    if (maybeInfinitePubkey.isPresent()) {
-      // if the Infinity is not a valid public key then aggregating with any
-      // non-valid pubkey should result to a non-valid pubkey
-      return maybeInfinitePubkey.get();
+    Optional<BlstPublicKey> maybeInvalidPubkey =
+        publicKeys.stream().filter(pk -> !pk.isValid()).findAny();
+    if (maybeInvalidPubkey.isPresent()) {
+      // Points not in the group and the point at infinity are not valid public keys. Aggregating
+      // with other points should result in a non-valid pubkey, the point at infinity.
+      return infinitePublicKey;
     }
 
+    // At this point, we know that the public keys are in the G1 group, so we use `blst.P1.add()`
+    // rather than `blst.P1.aggregate()` as the latter always performs the (slow) group check.
     P1 sum = new P1();
     for (BlstPublicKey publicKey : publicKeys) {
-      sum.aggregate(publicKey.ecPoint);
+      sum.add(publicKey.ecPoint);
     }
 
     return new BlstPublicKey(sum.to_affine());
   }
 
   final P1_Affine ecPoint;
-  private final Supplier<Boolean> isInfinity =
-      Suppliers.memoize(() -> toBytesCompressed().equals(INFINITY_COMPRESSED_BYTES));
+  private final Supplier<Boolean> isInfinity = Suppliers.memoize(() -> checkForInfinity());
+  private final Supplier<Boolean> isInGroup = Suppliers.memoize(() -> checkGroupMembership());
 
   public BlstPublicKey(P1_Affine ecPoint) {
     this.ecPoint = ecPoint;
   }
 
   @Override
   public void forceValidation() throws IllegalArgumentException {
-
-    if (!ecPoint.in_group()) {
+    if (!isValid()) {
       throw new IllegalArgumentException("Invalid PublicKey: " + toBytesCompressed());
     }
   }
 
-  @Override
-  public Bytes48 toBytesCompressed() {
-    return Bytes48.wrap(ecPoint.compress());
+  private boolean checkGroupMembership() {
+    return ecPoint.in_group();
+  }
+
+  private boolean checkForInfinity() {
+    return ecPoint.is_inf();
+  }
+
+  boolean isValid() {
+    return !isInfinity.get() && isInGroup.get();
   }
 
   boolean isInfinity() {
     return isInfinity.get();
   }
 
+  @Override
+  public Bytes48 toBytesCompressed() {
+    return Bytes48.wrap(ecPoint.compress());
+  }
+
   @Override
   public int hashCode() {
     return toBytesCompressed().hashCode();

bls/src/test/java/tech/pegasys/teku/bls/BLSPublicKeyTest.java

@@ -28,10 +28,13 @@
 
 abstract class BLSPublicKeyTest {
 
-  private static final Bytes InfinityPublicKey =
+  private static final Bytes infinityPublicKeyBytes =
       Bytes.fromHexString(
           "0xc00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000");
 
+  private static final BLSPublicKey infinityPublicKey =
+      BLSPublicKey.fromSSZBytes(infinityPublicKeyBytes);
+
   @Test
   void fromBytesCompressedValidate_okWhenValidBytes() {
     assertThatCode(
@@ -59,9 +62,8 @@ void succeedsWhenEqualsReturnsTrueForTheSameEmptyPublicKey() {
 
   @Test
   void succeedsWhenTwoInfinityPublicKeysAreEqual() {
-    // Infinity keys are valid G1 points, so pass the equality test
-    BLSPublicKey publicKey1 = BLSPublicKey.fromSSZBytes(InfinityPublicKey);
-    BLSPublicKey publicKey2 = BLSPublicKey.fromSSZBytes(InfinityPublicKey);
+    BLSPublicKey publicKey1 = new BLSPublicKey(infinityPublicKey.getPublicKey());
+    BLSPublicKey publicKey2 = new BLSPublicKey(infinityPublicKey.getPublicKey());
     assertEquals(publicKey1, publicKey2);
   }
 
@@ -112,7 +114,6 @@ void succeedsIfSerializationOfEmptyPublicKeyIsCorrect() {
 
   @Test
   void succeedsIfDeserializationOfInfinityPublicKeyIsCorrect() {
-    BLSPublicKey infinityPublicKey = BLSPublicKey.fromSSZBytes(InfinityPublicKey);
     byte[] pointBytes = new byte[48];
     pointBytes[0] = (byte) 0xc0;
     Bytes infinityBytesSsz =
@@ -163,7 +164,7 @@ void succeedsWhenRoundtripSSZReturnsTheSamePublicKey() {
 
   @Test
   void succeedsWhenRoundtripSSZReturnsTheInfinityPublicKey() {
-    BLSPublicKey publicKey1 = BLSPublicKey.fromSSZBytes(InfinityPublicKey);
+    BLSPublicKey publicKey1 = infinityPublicKey;
     BLSPublicKey publicKey2 = BLSPublicKey.fromSSZBytes(publicKey1.toSSZBytes());
     assertEquals(publicKey1, publicKey2);
   }
@@ -184,6 +185,36 @@ void aggregateSamePubKeys() {
     assertThat(aggrPk).isEqualTo(aggrPkGolden);
   }
 
+  @Test
+  void aggregateWithInfinitePubKeyShouldFail() {
+    BLSPublicKey pk =
+        BLSPublicKey.fromBytesCompressedValidate(
+            Bytes48.fromHexString(
+                "0x89ece308f9d1f0131765212deca99697b112d61f9be9a5f1f3780a51335b3ff981747a0b2ca2179b96d2c0c9024e5224"));
+
+    BLSPublicKey aggrPk = BLSPublicKey.aggregate(List.of(pk, infinityPublicKey));
+
+    assertThat(aggrPk).isEqualTo(infinityPublicKey);
+  }
+
+  @Test
+  void aggregateWithNotInGroupPubKeyShouldFail() {
+    // This one is in the G1 group
+    BLSPublicKey pk1 =
+        BLSPublicKey.fromBytesCompressedValidate(
+            Bytes48.fromHexString(
+                "0x89ece308f9d1f0131765212deca99697b112d61f9be9a5f1f3780a51335b3ff981747a0b2ca2179b96d2c0c9024e5224"));
+    // This one is on the curve but not in the G1 group
+    BLSPublicKey pk2 =
+        BLSPublicKey.fromBytesCompressed(
+            Bytes48.fromHexString(
+                "0x800000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000004"));
+
+    BLSPublicKey aggrPk = BLSPublicKey.aggregate(List.of(pk1, pk2));
+
+    assertThat(aggrPk).isEqualTo(infinityPublicKey);
+  }
+
   @Test
   public void toAbbreviatedString_shouldShowFirstSevenBytesOfPublicKey() {
     Bytes keyBytes =

bls/src/test/java/tech/pegasys/teku/bls/impl/blst/BlstPublicKeyTest.java

@@ -27,19 +27,18 @@ static void setup() {
     assertThat(BlstLoader.INSTANCE).isNotEmpty();
   }
 
-  // Blst library doesn't handle infinity pubkeys at the moment.
-  // Could enable the test when the issue https://github.com/supranational/blst/issues/11 is
-  // addressed
+  // The infinite public key is now considered to be invalid
   @Test
   void infinityPublicKey() {
-    BlstPublicKey inf1 =
-        BlstPublicKey.fromBytes(
-            Bytes48.fromHexString(
-                "0xc00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"));
-    Bytes48 bytes = inf1.toBytesCompressed();
-
-    BlstPublicKey publicKey = BlstPublicKey.fromBytes(bytes);
-    publicKey.forceValidation();
+    Bytes48 infinitePublicKeyBytes =
+        Bytes48.fromHexString(
+            "0xc00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000");
+    assertThatThrownBy(
+            () -> {
+              BlstPublicKey publicKey = BlstPublicKey.fromBytes(infinitePublicKeyBytes);
+              publicKey.forceValidation();
+            })
+        .isInstanceOf(IllegalArgumentException.class);
   }
 
   @Test