From 4e9587b09fdc6bba026e77127c41f5522bc6b3d6 Mon Sep 17 00:00:00 2001 From: Stefan Bratanov Date: Thu, 10 Aug 2023 22:30:20 +0100 Subject: [PATCH] Don't use ByteBuffer when flattening bytes (#7415) --- CHANGELOG.md | 2 +- .../tech/pegasys/teku/kzg/KZGCommitment.java | 11 +-- .../java/tech/pegasys/teku/kzg/KZGProof.java | 13 ++-- .../tech/pegasys/teku/kzg/TrustedSetup.java | 53 +++----------- .../pegasys/teku/kzg/ckzg4844/CKZG4844.java | 19 +++-- .../teku/kzg/ckzg4844/CKZG4844Utils.java | 70 +++++++++++++------ .../teku/kzg/ckzg4844/CKZG4844UtilsTest.java | 55 +++++++++++++++ 7 files changed, 139 insertions(+), 84 deletions(-) create mode 100644 infrastructure/kzg/src/test/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844UtilsTest.java diff --git a/CHANGELOG.md b/CHANGELOG.md index feac966bbb3..dcae3289bb7 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZGCommitment.java b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZGCommitment.java index e1786b86bef..ad1095d33f9 100644 --- a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZGCommitment.java +++ b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZGCommitment.java @@ -15,6 +15,7 @@ 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; @@ -22,8 +23,6 @@ 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) { diff --git a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZGProof.java b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZGProof.java index c5701a91396..ca47486cb83 100644 --- a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZGProof.java +++ b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZGProof.java @@ -15,6 +15,7 @@ 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; @@ -22,7 +23,6 @@ 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() { diff --git a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/TrustedSetup.java b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/TrustedSetup.java index 9865966d8b2..c48235360a0 100644 --- a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/TrustedSetup.java +++ b/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 g1Points; - private final List g2Points; +public record TrustedSetup(List g1Points, List g2Points) { - private void validateG2Point(Bytes g2Point) { - checkArgument(g2Point.size() == 96, "Expected G2 point to be 96 bytes"); - } - - public TrustedSetup(final List g1Points, final List g2Points) { + public TrustedSetup { + g1Points.forEach(this::validateG1Point); g2Points.forEach(this::validateG2Point); - this.g1Points = g1Points; - this.g2Points = g2Points; - } - - public List getG1Points() { - return g1Points; - } - - public List 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); } } diff --git a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844.java b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844.java index e59cfccf04b..9a5e3cb8e21 100644 --- a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844.java +++ b/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 g1Points = trustedSetup.g1Points(); + final List 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); diff --git a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844Utils.java b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844Utils.java index 2fbdf6b75d9..4293a285d97 100644 --- a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844Utils.java +++ b/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 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 blobs) { - return flattenBytes(blobs.stream(), CKZG4844JNI.getBytesPerBlob() * blobs.size()); + return flattenBytes(blobs, CKZG4844JNI.getBytesPerBlob() * blobs.size()); } public static byte[] flattenCommitments(final List commitments) { - final Stream 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 kzgProofs) { - final Stream 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 g1Points) { + return flattenBytes(g1Points, CKZG4844.G1_POINT_SIZE * g1Points.size()); + } + + public static byte[] flattenG2Points(final List 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 g1Points = new ArrayList<>(); + final List 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 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 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 toFlatten, final int expectedSize) { + return flattenBytes(toFlatten, Bytes::toArrayUnsafe, expectedSize); + } + + private static byte[] flattenBytes( + final List toFlatten, final Function 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; } } diff --git a/infrastructure/kzg/src/test/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844UtilsTest.java b/infrastructure/kzg/src/test/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844UtilsTest.java new file mode 100644 index 00000000000..73a95505ef7 --- /dev/null +++ b/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 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 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"); + } +}