Don't use ByteBuffer when flattening bytes (Consensys#7415)

CHANGELOG.md

@@ -14,6 +14,6 @@ For information on changes in released versions of Teku, see the [releases page]
 
 ### Additions and Improvements
 
-- Update attestation subnet subscriptions strategy according to [the spec changes](https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/p2p-interface.md#attestation-subnet-subscription). All nodes (including non-validating ones) will subscribe to 2 subnets regardless of the number of validators
+- Update attestation subnet subscriptions strategy according to [the spec changes](https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/p2p-interface.md#attestation-subnet-subscription). All nodes (including non-validating ones) will subscribe to 2 subnets regardless of the number of validators.
 
 ### Bug Fixes

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZGCommitment.java

@@ -15,15 +15,14 @@
 
 import static com.google.common.base.Preconditions.checkArgument;
 
+import ethereum.ckzg4844.CKZG4844JNI;
 import java.util.Objects;
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes48;
 import org.apache.tuweni.ssz.SSZ;
 
 public final class KZGCommitment {
 
-  public static final int KZG_COMMITMENT_SIZE = 48;
-
   /**
    * Creates 0x00..00 for point-at-infinity
    *
@@ -39,12 +38,14 @@ public static KZGCommitment fromHexString(final String hexString) {
 
   public static KZGCommitment fromSSZBytes(final Bytes bytes) {
     checkArgument(
-        bytes.size() == KZG_COMMITMENT_SIZE,
-        "Expected " + KZG_COMMITMENT_SIZE + " bytes but received %s.",
+        bytes.size() == CKZG4844JNI.BYTES_PER_COMMITMENT,
+        "Expected " + CKZG4844JNI.BYTES_PER_COMMITMENT + " bytes but received %s.",
         bytes.size());
     return SSZ.decode(
         bytes,
-        reader -> new KZGCommitment(Bytes48.wrap(reader.readFixedBytes(KZG_COMMITMENT_SIZE))));
+        reader ->
+            new KZGCommitment(
+                Bytes48.wrap(reader.readFixedBytes(CKZG4844JNI.BYTES_PER_COMMITMENT))));
   }
 
   public static KZGCommitment fromBytesCompressed(final Bytes48 bytes) {

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZGProof.java

@@ -15,14 +15,14 @@
 
 import static com.google.common.base.Preconditions.checkArgument;
 
+import ethereum.ckzg4844.CKZG4844JNI;
 import java.util.Objects;
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes48;
 import org.apache.tuweni.ssz.SSZ;
 
 public final class KZGProof {
 
-  public static final int KZG_PROOF_SIZE = 48;
   /** Creates 0x00..00 for point-at-infinity */
   public static final KZGProof INFINITY = KZGProof.fromBytesCompressed(Bytes48.ZERO);
 
@@ -32,11 +32,12 @@ public static KZGProof fromHexString(final String hexString) {
 
   public static KZGProof fromSSZBytes(final Bytes bytes) {
     checkArgument(
-        bytes.size() == KZG_PROOF_SIZE,
-        "Expected " + KZG_PROOF_SIZE + " bytes but received %s.",
+        bytes.size() == CKZG4844JNI.BYTES_PER_PROOF,
+        "Expected " + CKZG4844JNI.BYTES_PER_PROOF + " bytes but received %s.",
         bytes.size());
     return SSZ.decode(
-        bytes, reader -> new KZGProof(Bytes48.wrap(reader.readFixedBytes(KZG_PROOF_SIZE))));
+        bytes,
+        reader -> new KZGProof(Bytes48.wrap(reader.readFixedBytes(CKZG4844JNI.BYTES_PER_PROOF))));
   }
 
   public static KZGProof fromBytesCompressed(final Bytes48 bytes) throws IllegalArgumentException {
@@ -66,8 +67,8 @@ public Bytes48 getBytesCompressed() {
     return bytesCompressed;
   }
 
-  public byte[] toArray() {
-    return bytesCompressed.toArray();
+  public byte[] toArrayUnsafe() {
+    return bytesCompressed.toArrayUnsafe();
   }
 
   public String toAbbreviatedString() {

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/TrustedSetup.java

@@ -14,57 +14,26 @@
 package tech.pegasys.teku.kzg;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static tech.pegasys.teku.kzg.ckzg4844.CKZG4844.G1_POINT_SIZE;
+import static tech.pegasys.teku.kzg.ckzg4844.CKZG4844.G2_POINT_SIZE;
 
-import com.google.common.base.MoreObjects;
 import java.util.List;
-import java.util.Objects;
 import org.apache.tuweni.bytes.Bytes;
-import org.apache.tuweni.bytes.Bytes48;
 
-public class TrustedSetup {
-  private final List<Bytes48> g1Points;
-  private final List<Bytes> g2Points;
+public record TrustedSetup(List<Bytes> g1Points, List<Bytes> g2Points) {
 
-  private void validateG2Point(Bytes g2Point) {
-    checkArgument(g2Point.size() == 96, "Expected G2 point to be 96 bytes");
-  }
-
-  public TrustedSetup(final List<Bytes48> g1Points, final List<Bytes> g2Points) {
+  public TrustedSetup {
+    g1Points.forEach(this::validateG1Point);
     g2Points.forEach(this::validateG2Point);
-    this.g1Points = g1Points;
-    this.g2Points = g2Points;
-  }
-
-  public List<Bytes48> getG1Points() {
-    return g1Points;
-  }
-
-  public List<Bytes> getG2Points() {
-    return g2Points;
-  }
-
-  @Override
-  public boolean equals(final Object o) {
-    if (this == o) {
-      return true;
-    }
-    if (o == null || getClass() != o.getClass()) {
-      return false;
-    }
-    final TrustedSetup that = (TrustedSetup) o;
-    return Objects.equals(g1Points, that.g1Points) && Objects.equals(g2Points, that.g2Points);
   }
 
-  @Override
-  public int hashCode() {
-    return Objects.hash(g1Points, g2Points);
+  private void validateG1Point(final Bytes g1Point) {
+    checkArgument(
+        g1Point.size() == G1_POINT_SIZE, "Expected G1 point to be %s bytes", G1_POINT_SIZE);
   }
 
-  @Override
-  public String toString() {
-    return MoreObjects.toStringHelper(this)
-        .add("g1Points", g1Points)
-        .add("g2Points", g2Points)
-        .toString();
+  private void validateG2Point(final Bytes g2Point) {
+    checkArgument(
+        g2Point.size() == G2_POINT_SIZE, "Expected G2 point to be %s bytes", G2_POINT_SIZE);
   }
 }

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844.java

@@ -35,6 +35,9 @@
  */
 public final class CKZG4844 implements KZG {
 
+  public static final int G1_POINT_SIZE = 48;
+  public static final int G2_POINT_SIZE = 96;
+
   private static final Logger LOG = LogManager.getLogger();
 
   private static CKZG4844 instance;
@@ -90,7 +93,7 @@ public synchronized void loadTrustedSetup(final String trustedSetupFilePath) thr
     try {
       final TrustedSetup trustedSetup = CKZG4844Utils.parseTrustedSetupFile(trustedSetupFilePath);
       loadTrustedSetup(trustedSetup);
-    } catch (IOException ex) {
+    } catch (final IOException ex) {
       throw new KZGException("Failed to load trusted setup from file: " + trustedSetupFilePath, ex);
     }
   }
@@ -103,11 +106,13 @@ public void loadTrustedSetup(final TrustedSetup trustedSetup) throws KZGExceptio
       return;
     }
     try {
+      final List<Bytes> g1Points = trustedSetup.g1Points();
+      final List<Bytes> g2Points = trustedSetup.g2Points();
       CKZG4844JNI.loadTrustedSetup(
-          CKZG4844Utils.flattenBytes(trustedSetup.getG1Points()),
-          trustedSetup.getG1Points().size(),
-          CKZG4844Utils.flattenBytes(trustedSetup.getG2Points()),
-          trustedSetup.getG2Points().size());
+          CKZG4844Utils.flattenG1Points(g1Points),
+          g1Points.size(),
+          CKZG4844Utils.flattenG2Points(g2Points),
+          g2Points.size());
       loadedTrustedSetupHash = Optional.of(trustedSetup.hashCode());
       LOG.debug("Loaded trusted setup: {}", trustedSetup);
     } catch (final Exception ex) {
@@ -135,9 +140,9 @@ public boolean verifyBlobKzgProofBatch(
     try {
       final byte[] blobsBytes = CKZG4844Utils.flattenBlobs(blobs);
       final byte[] commitmentsBytes = CKZG4844Utils.flattenCommitments(kzgCommitments);
-      final byte[] proofBytes = CKZG4844Utils.flattenProofs(kzgProofs);
+      final byte[] proofsBytes = CKZG4844Utils.flattenProofs(kzgProofs);
       return CKZG4844JNI.verifyBlobKzgProofBatch(
-          blobsBytes, commitmentsBytes, proofBytes, blobs.size());
+          blobsBytes, commitmentsBytes, proofsBytes, blobs.size());
     } catch (final Exception ex) {
       throw new KZGException(
           "Failed to verify blobs and commitments against KZG proofs " + kzgProofs, ex);

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844Utils.java

@@ -13,19 +13,18 @@
 
 package tech.pegasys.teku.kzg.ckzg4844;
 
+import com.google.common.base.Preconditions;
 import ethereum.ckzg4844.CKZG4844JNI;
 import java.io.BufferedReader;
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
-import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.stream.Stream;
+import java.util.function.Function;
 import org.apache.tuweni.bytes.Bytes;
-import org.apache.tuweni.bytes.Bytes48;
 import tech.pegasys.teku.infrastructure.http.UrlSanitizer;
 import tech.pegasys.teku.infrastructure.io.resource.ResourceLoader;
 import tech.pegasys.teku.kzg.KZGCommitment;
@@ -34,26 +33,30 @@
 
 public class CKZG4844Utils {
 
-  private static final int G2_POINT_SIZE = 96;
-
-  public static byte[] flattenBytes(final List<? extends Bytes> bytes) {
-    return flattenBytes(
-        bytes.stream(), bytes.stream().map(Bytes::size).reduce(Integer::sum).orElse(0));
-  }
+  private static final int MAX_BYTES_TO_FLATTEN = 100_663_296; // ~100.66 MB or 768 blobs
 
   public static byte[] flattenBlobs(final List<Bytes> blobs) {
-    return flattenBytes(blobs.stream(), CKZG4844JNI.getBytesPerBlob() * blobs.size());
+    return flattenBytes(blobs, CKZG4844JNI.getBytesPerBlob() * blobs.size());
   }
 
   public static byte[] flattenCommitments(final List<KZGCommitment> commitments) {
-    final Stream<Bytes> commitmentsBytes =
-        commitments.stream().map(KZGCommitment::getBytesCompressed);
-    return flattenBytes(commitmentsBytes, KZGCommitment.KZG_COMMITMENT_SIZE * commitments.size());
+    return flattenBytes(
+        commitments,
+        KZGCommitment::toArrayUnsafe,
+        CKZG4844JNI.BYTES_PER_COMMITMENT * commitments.size());
   }
 
   public static byte[] flattenProofs(final List<KZGProof> kzgProofs) {
-    final Stream<Bytes> kzgProofBytes = kzgProofs.stream().map(KZGProof::getBytesCompressed);
-    return flattenBytes(kzgProofBytes, KZGProof.KZG_PROOF_SIZE * kzgProofs.size());
+    return flattenBytes(
+        kzgProofs, KZGProof::toArrayUnsafe, CKZG4844JNI.BYTES_PER_PROOF * kzgProofs.size());
+  }
+
+  public static byte[] flattenG1Points(final List<Bytes> g1Points) {
+    return flattenBytes(g1Points, CKZG4844.G1_POINT_SIZE * g1Points.size());
+  }
+
+  public static byte[] flattenG2Points(final List<Bytes> g2Points) {
+    return flattenBytes(g2Points, CKZG4844.G2_POINT_SIZE * g2Points.size());
   }
 
   public static TrustedSetup parseTrustedSetupFile(final String filePath) throws IOException {
@@ -70,27 +73,48 @@ public static TrustedSetup parseTrustedSetupFile(final String filePath) throws I
       // Number of G2 points
       final int g2Size = Integer.parseInt(reader.readLine());
       // List of G1 points, one on each new line
-      final List<Bytes48> g1Points = new ArrayList<>();
+      final List<Bytes> g1Points = new ArrayList<>();
       for (int i = 0; i < g1Size; i++) {
-        final Bytes48 g1Point = Bytes48.fromHexString(reader.readLine());
+        final Bytes g1Point = Bytes.fromHexString(reader.readLine(), CKZG4844.G1_POINT_SIZE);
         g1Points.add(g1Point);
       }
       // List of G2 points, one on each new line
       final List<Bytes> g2Points = new ArrayList<>();
       for (int i = 0; i < g2Size; i++) {
-        final Bytes g2Point = Bytes.fromHexString(reader.readLine(), G2_POINT_SIZE);
+        final Bytes g2Point = Bytes.fromHexString(reader.readLine(), CKZG4844.G2_POINT_SIZE);
         g2Points.add(g2Point);
       }
 
       return new TrustedSetup(g1Points, g2Points);
-    } catch (Exception ex) {
+    } catch (final Exception ex) {
       throw new IOException(String.format("Failed to parse trusted setup file\n: %s", filePath));
     }
   }
 
-  private static byte[] flattenBytes(final Stream<? extends Bytes> bytes, final int capacity) {
-    final ByteBuffer buffer = ByteBuffer.allocate(capacity);
-    bytes.map(Bytes::toArray).forEach(buffer::put);
-    return buffer.array();
+  private static byte[] flattenBytes(final List<Bytes> toFlatten, final int expectedSize) {
+    return flattenBytes(toFlatten, Bytes::toArrayUnsafe, expectedSize);
+  }
+
+  private static <T> byte[] flattenBytes(
+      final List<T> toFlatten, final Function<T, byte[]> bytesConverter, final int expectedSize) {
+    Preconditions.checkArgument(
+        expectedSize <= MAX_BYTES_TO_FLATTEN,
+        "Maximum of %s bytes can be flattened, but %s were requested",
+        MAX_BYTES_TO_FLATTEN,
+        expectedSize);
+    final byte[] flattened = new byte[expectedSize];
+    int destPos = 0;
+    for (final T data : toFlatten) {
+      final byte[] bytes = bytesConverter.apply(data);
+      System.arraycopy(bytes, 0, flattened, destPos, bytes.length);
+      destPos += bytes.length;
+    }
+    if (destPos != expectedSize) {
+      throw new IllegalArgumentException(
+          String.format(
+              "The actual bytes to flatten (%d) was not the same as the expected size specified (%d)",
+              destPos, expectedSize));
+    }
+    return flattened;
   }
 }

infrastructure/kzg/src/test/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844UtilsTest.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright Consensys Software Inc., 2022
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+
+package tech.pegasys.teku.kzg.ckzg4844;
+
+import static ethereum.ckzg4844.CKZG4844JNI.getBytesPerBlob;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import ethereum.ckzg4844.CKZG4844JNI;
+import ethereum.ckzg4844.CKZG4844JNI.Preset;
+import java.util.List;
+import java.util.stream.IntStream;
+import org.apache.tuweni.bytes.Bytes;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+class CKZG4844UtilsTest {
+
+  @BeforeAll
+  public static void setUp() {
+    CKZG4844JNI.loadNativeLibrary(Preset.MAINNET);
+  }
+
+  @Test
+  public void testFlattenBlobsWithUnexpectedSizeThrows() {
+    final int blobCount = 42;
+    final List<Bytes> blobs = IntStream.range(0, blobCount).mapToObj(__ -> Bytes.of()).toList();
+    final IllegalArgumentException exception =
+        assertThrows(IllegalArgumentException.class, () -> CKZG4844Utils.flattenBlobs(blobs));
+    assertThat(exception)
+        .hasMessage(
+            "The actual bytes to flatten (0) was not the same as the expected size specified (5505024)");
+  }
+
+  @Test
+  public void testFlattenBlobsWithBigBoySizeThrows() {
+    final int blobCount = Integer.MAX_VALUE / getBytesPerBlob();
+    final List<Bytes> blobs = IntStream.range(0, blobCount).mapToObj(__ -> Bytes.of()).toList();
+    final IllegalArgumentException exception =
+        assertThrows(IllegalArgumentException.class, () -> CKZG4844Utils.flattenBlobs(blobs));
+    assertThat(exception)
+        .hasMessage("Maximum of 100663296 bytes can be flattened, but 2147352576 were requested");
+  }
+}