From adf89c8360f4a46c382b22daa063566513a71d28 Mon Sep 17 00:00:00 2001
From: dapplion <35266934+dapplion@users.noreply.github.com>
Date: Mon, 5 Apr 2021 22:38:12 +0200
Subject: [PATCH] Add validate key option to PublicKey.fromBytes()

---
 src/blst/publicKey.ts                   |  9 +++------
 src/interface.ts                        |  2 +-
 test/spec/fast_aggregate_verify.test.ts | 10 +++++++++-
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/blst/publicKey.ts b/src/blst/publicKey.ts
index 3e89b61..a020e1f 100644
--- a/src/blst/publicKey.ts
+++ b/src/blst/publicKey.ts
@@ -1,5 +1,5 @@
 import * as blst from "@chainsafe/blst";
-import {EmptyAggregateError, ZeroPublicKeyError} from "../errors";
+import {EmptyAggregateError} from "../errors";
 import {bytesToHex, hexToBytes} from "../helpers";
 import {PointFormat, PublicKey as IPublicKey} from "../interface";
 
@@ -9,12 +9,9 @@ export class PublicKey extends blst.PublicKey implements IPublicKey {
   }
 
   /** @param type Defaults to `CoordType.jacobian` */
-  static fromBytes(bytes: Uint8Array, type?: blst.CoordType): PublicKey {
+  static fromBytes(bytes: Uint8Array, type?: blst.CoordType, validate?: boolean): PublicKey {
     const pk = blst.PublicKey.fromBytes(bytes, type);
-    if (pk.value.is_inf()) {
-      throw new ZeroPublicKeyError();
-    }
-
+    if (validate) pk.keyValidate();
     return new PublicKey(pk.value);
   }
 
diff --git a/src/interface.ts b/src/interface.ts
index 05c752d..1a47144 100644
--- a/src/interface.ts
+++ b/src/interface.ts
@@ -29,7 +29,7 @@ export declare class SecretKey {
 
 export declare class PublicKey {
   /** @param type Only for impl `blst-native`. Defaults to `CoordType.jacobian` */
-  static fromBytes(bytes: Uint8Array, type?: CoordType): PublicKey;
+  static fromBytes(bytes: Uint8Array, type?: CoordType, validate?: boolean): PublicKey;
   static fromHex(hex: string): PublicKey;
   static aggregate(publicKeys: PublicKey[]): PublicKey;
   /** @param format Defaults to `PointFormat.compressed` */
diff --git a/test/spec/fast_aggregate_verify.test.ts b/test/spec/fast_aggregate_verify.test.ts
index 222eecd..ad5b9e7 100644
--- a/test/spec/fast_aggregate_verify.test.ts
+++ b/test/spec/fast_aggregate_verify.test.ts
@@ -3,6 +3,7 @@ import {describeDirectorySpecTest, InputType} from "@chainsafe/lodestar-spec-tes
 import {hexToBytes} from "../../src/helpers";
 import {SPEC_TESTS_DIR} from "../params";
 import {describeForAllImplementations} from "../switch";
+import {CoordType} from "@chainsafe/blst";
 
 interface IAggregateSigsVerifyTestCase {
   data: {
@@ -21,7 +22,14 @@ describeForAllImplementations((bls) => {
     path.join(SPEC_TESTS_DIR, "tests/general/phase0/bls/fast_aggregate_verify/small"),
     (testCase) => {
       const {pubkeys, message, signature} = testCase.data.input;
-      return bls.verifyAggregate(pubkeys.map(hexToBytes), hexToBytes(message), hexToBytes(signature));
+      try {
+        return bls.Signature.fromBytes(hexToBytes(signature)).verifyAggregate(
+          pubkeys.map((hex) => bls.PublicKey.fromBytes(hexToBytes(hex), CoordType.jacobian, true)),
+          hexToBytes(message)
+        );
+      } catch (e) {
+        return false;
+      }
     },
     {
       inputTypes: {data: InputType.YAML},