From 529530501137df4c74e829586daab8f4b569ee83 Mon Sep 17 00:00:00 2001 From: Stefan Bratanov Date: Mon, 23 Oct 2023 15:23:07 +0100 Subject: [PATCH] Simplify loading/freeing trusted setup --- .../forkchoice/ForkChoiceTestExecutor.java | 15 ------ .../reference/phase0/kzg/KzgTestExecutor.java | 4 +- .../versions/deneb/helpers/KzgResolver.java | 5 +- .../main/java/tech/pegasys/teku/kzg/KZG.java | 11 ++-- .../pegasys/teku/kzg/ckzg4844/CKZG4844.java | 36 +++++++++---- .../java/tech/pegasys/teku/kzg/KZGTest.java | 51 +++++++------------ 6 files changed, 52 insertions(+), 70 deletions(-) diff --git a/eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/forkchoice/ForkChoiceTestExecutor.java b/eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/forkchoice/ForkChoiceTestExecutor.java index 1adc17c8440..c580a5e7657 100644 --- a/eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/forkchoice/ForkChoiceTestExecutor.java +++ b/eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/forkchoice/ForkChoiceTestExecutor.java @@ -42,13 +42,10 @@ import tech.pegasys.teku.infrastructure.async.SafeFuture; import tech.pegasys.teku.infrastructure.async.eventthread.InlineEventThread; import tech.pegasys.teku.infrastructure.unsigned.UInt64; -import tech.pegasys.teku.kzg.KZG; import tech.pegasys.teku.kzg.KZGProof; import tech.pegasys.teku.reference.TestDataUtils; import tech.pegasys.teku.reference.TestExecutor; import tech.pegasys.teku.spec.Spec; -import tech.pegasys.teku.spec.SpecMilestone; -import tech.pegasys.teku.spec.SpecVersion; import tech.pegasys.teku.spec.datastructures.attestation.ValidatableAttestation; import tech.pegasys.teku.spec.datastructures.blobs.versions.deneb.Blob; import tech.pegasys.teku.spec.datastructures.blocks.BeaconBlock; @@ -66,9 +63,7 @@ import tech.pegasys.teku.spec.executionlayer.ExecutionLayerChannelStub; import tech.pegasys.teku.spec.executionlayer.ExecutionPayloadStatus; import tech.pegasys.teku.spec.executionlayer.PayloadStatus; -import tech.pegasys.teku.spec.logic.common.helpers.MiscHelpers; import tech.pegasys.teku.spec.logic.common.statetransition.results.BlockImportResult; -import tech.pegasys.teku.spec.logic.versions.deneb.helpers.MiscHelpersDeneb; import tech.pegasys.teku.statetransition.forkchoice.ForkChoice; import tech.pegasys.teku.statetransition.forkchoice.ForkChoiceStateProvider; import tech.pegasys.teku.statetransition.forkchoice.MergeTransitionBlockValidator; @@ -161,8 +156,6 @@ spec, new SignedBlockAndState(anchorBlock, anchorState)), + "\nProtoarray data:\n" + protoArrayData, e); - } finally { - freeTrustedSetupIfRequired(spec); } } @@ -489,14 +482,6 @@ private void assertCheckpoint( .isEqualTo(new Checkpoint(expectedEpoch, expectedRoot)); } - private void freeTrustedSetupIfRequired(final Spec spec) { - Optional.ofNullable(spec.forMilestone(SpecMilestone.DENEB)) - .map(SpecVersion::miscHelpers) - .flatMap(MiscHelpers::toVersionDeneb) - .map(MiscHelpersDeneb::getKzg) - .ifPresent(KZG::freeTrustedSetup); - } - @SuppressWarnings({"unchecked", "TypeParameterUnusedInFormals"}) private static T get(final Map yamlData, final String key) { return (T) yamlData.get(key); diff --git a/eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/kzg/KzgTestExecutor.java b/eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/kzg/KzgTestExecutor.java index f89cb6c8a78..d895aec4a38 100644 --- a/eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/kzg/KzgTestExecutor.java +++ b/eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/kzg/KzgTestExecutor.java @@ -40,11 +40,9 @@ public final void runTest(final TestDefinition testDefinition) throws Throwable final SpecConfigDeneb specConfigDeneb = SpecConfigDeneb.required(networkConfig.getSpec().getGenesisSpecConfig()); - try { + try (kzg) { kzg.loadTrustedSetup(specConfigDeneb.getTrustedSetupPath().orElseThrow()); runTestImpl(testDefinition); - } finally { - kzg.freeTrustedSetup(); } } diff --git a/ethereum/spec/src/property-test/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/KzgResolver.java b/ethereum/spec/src/property-test/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/KzgResolver.java index 670b1f2bacf..12c827608ab 100644 --- a/ethereum/spec/src/property-test/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/KzgResolver.java +++ b/ethereum/spec/src/property-test/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/KzgResolver.java @@ -56,6 +56,7 @@ private KZG getKzgWithTrustedSetup() { private static class KzgAutoLoadFree implements Store.CloseOnReset { private static final String TRUSTED_SETUP = Resources.getResource(TrustedSetups.class, "trusted_setup.txt").toExternalForm(); + private final KZG kzg = CKZG4844.createInstance(); private KzgAutoLoadFree() { @@ -63,8 +64,8 @@ private KzgAutoLoadFree() { } @Override - public void close() { - kzg.freeTrustedSetup(); + public void close() throws Exception { + kzg.close(); } } } diff --git a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZG.java b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZG.java index a610405622d..8096024fde6 100644 --- a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZG.java +++ b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZG.java @@ -20,15 +20,13 @@ * This interface specifies all the KZG functions needed for the Deneb specification and is the * entry-point for all KZG operations in Teku. */ -public interface KZG { +public interface KZG extends AutoCloseable { KZG NOOP = new KZG() { - @Override - public void loadTrustedSetup(final String trustedSetupFile) throws KZGException {} @Override - public void freeTrustedSetup() throws KZGException {} + public void loadTrustedSetup(final String trustedSetupFile) throws KZGException {} @Override public boolean verifyBlobKzgProofBatch( @@ -49,12 +47,13 @@ public KZGProof computeBlobKzgProof(final Bytes blob, final KZGCommitment kzgCom throws KZGException { return KZGProof.INFINITY; } + + @Override + public void close() {} }; void loadTrustedSetup(String trustedSetupFile) throws KZGException; - void freeTrustedSetup() throws KZGException; - boolean verifyBlobKzgProofBatch( List blobs, List kzgCommitments, List kzgProofs) throws KZGException; 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 1d961011a9d..93a8906010d 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 @@ -54,6 +54,7 @@ private CKZG4844() { } } + /** Only one trusted setup at a time can be loaded. */ @Override public synchronized void loadTrustedSetup(final String trustedSetupFile) throws KZGException { if (loadedTrustedSetupFile.isPresent() @@ -62,6 +63,14 @@ public synchronized void loadTrustedSetup(final String trustedSetupFile) throws return; } try { + loadedTrustedSetupFile.ifPresent( + currentTrustedSetupFile -> { + LOG.debug( + "Freeing current trusted setup {} in order to load trusted setup from {}", + currentTrustedSetupFile, + trustedSetupFile); + freeTrustedSetup(); + }); final TrustedSetup trustedSetup = CKZG4844Utils.parseTrustedSetupFile(trustedSetupFile); final List g1Points = trustedSetup.g1Points(); final List g2Points = trustedSetup.g2Points(); @@ -77,17 +86,6 @@ public synchronized void loadTrustedSetup(final String trustedSetupFile) throws } } - @Override - public synchronized void freeTrustedSetup() throws KZGException { - try { - CKZG4844JNI.freeTrustedSetup(); - loadedTrustedSetupFile = Optional.empty(); - LOG.debug("Trusted setup was freed"); - } catch (final Exception ex) { - throw new KZGException("Failed to free trusted setup", ex); - } - } - @Override public boolean verifyBlobKzgProofBatch( final List blobs, @@ -128,4 +126,20 @@ public KZGProof computeBlobKzgProof(final Bytes blob, final KZGCommitment kzgCom "Failed to compute KZG proof for blob with commitment " + kzgCommitment, ex); } } + + /** Frees the current trusted setup if any is loaded */ + @Override + public void close() { + loadedTrustedSetupFile.ifPresent(__ -> freeTrustedSetup()); + } + + private void freeTrustedSetup() throws KZGException { + try { + CKZG4844JNI.freeTrustedSetup(); + loadedTrustedSetupFile = Optional.empty(); + LOG.debug("Trusted setup was freed"); + } catch (final Exception ex) { + throw new KZGException("Failed to free trusted setup", ex); + } + } } diff --git a/infrastructure/kzg/src/test/java/tech/pegasys/teku/kzg/KZGTest.java b/infrastructure/kzg/src/test/java/tech/pegasys/teku/kzg/KZGTest.java index 35cbe7b3c4e..91a7daeb675 100644 --- a/infrastructure/kzg/src/test/java/tech/pegasys/teku/kzg/KZGTest.java +++ b/infrastructure/kzg/src/test/java/tech/pegasys/teku/kzg/KZGTest.java @@ -34,7 +34,7 @@ import java.util.stream.Stream; import org.apache.tuweni.bytes.Bytes; import org.apache.tuweni.units.bigints.UInt256; -import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -53,38 +53,28 @@ public final class KZGTest { @BeforeAll public static void setUp() { kzg = CKZG4844.createInstance(); + loadTrustedSetup(); } - @AfterEach - public void cleanUpIfNeeded() { - try { - kzg.freeTrustedSetup(); - } catch (final KZGException ex) { - // NOOP - } - } - - @Test - public void testKzgLoadSameTrustedSetupTwice_shouldNotThrowException() { - loadTrustedSetup(); - loadTrustedSetup(); + @AfterAll + public static void cleanUp() throws Exception { + kzg.close(); } - @Test - public void testKzLoadDifferentTrustedSetupTwice_shouldThrowException() { - loadTrustedSetup(); - assertThrows(KZGException.class, () -> kzg.loadTrustedSetup("trusted_setup-not-existing.txt")); + private static void loadTrustedSetup() { + final String trustedSetup = + Resources.getResource(TrustedSetups.class, TRUSTED_SETUP_PATH).toExternalForm(); + kzg.loadTrustedSetup(trustedSetup); } @Test - public void testKzgFreeTrustedSetupTwice_shouldThrowException() { + public void testKzgLoadSameTrustedSetupTwice_shouldNotThrowException() { loadTrustedSetup(); - kzg.freeTrustedSetup(); - assertThrows(KZGException.class, kzg::freeTrustedSetup); } @Test - public void testUsageWithoutLoadedTrustedSetup_shouldThrowException() { + public void testUsageWithoutLoadedTrustedSetup_shouldThrowException() throws Exception { + kzg.close(); final List exceptions = List.of( assertThrows( @@ -102,11 +92,12 @@ public void testUsageWithoutLoadedTrustedSetup_shouldThrowException() { assertThat(exceptions) .allSatisfy( exception -> assertThat(exception).cause().hasMessage("Trusted Setup is not loaded.")); + // load trusted setup again for other tests + loadTrustedSetup(); } @Test public void testComputingAndVerifyingBatchProofs() { - loadTrustedSetup(); final int numberOfBlobs = 4; final List blobs = getSampleBlobs(numberOfBlobs); final List kzgCommitments = @@ -133,13 +124,11 @@ public void testComputingAndVerifyingBatchProofs() { @Test public void testVerifyingEmptyBatch() { - loadTrustedSetup(); assertThat(kzg.verifyBlobKzgProofBatch(List.of(), List.of(), List.of())).isTrue(); } @Test public void testComputingAndVerifyingBatchSingleProof() { - loadTrustedSetup(); final int numberOfBlobs = 1; final List blobs = getSampleBlobs(numberOfBlobs); final List kzgCommitments = @@ -167,7 +156,6 @@ public void testComputingAndVerifyingBatchSingleProof() { @Test public void testVerifyingBatchProofsThrowsIfSizesDoesntMatch() { - loadTrustedSetup(); final int numberOfBlobs = 4; final List blobs = getSampleBlobs(numberOfBlobs); final List kzgCommitments = @@ -209,7 +197,6 @@ public void testVerifyingBatchProofsThrowsIfSizesDoesntMatch() { "0x925668a49d06f4" }) public void testComputingProofWithIncorrectLengthBlobDoesNotCauseSegfault(final String blobHex) { - loadTrustedSetup(); final Bytes blob = Bytes.fromHexString(blobHex); final KZGException kzgException = @@ -241,6 +228,8 @@ public void incorrectTrustedSetupFilesShouldThrow(final String path) { final Throwable cause = assertThrows(KZGException.class, () -> kzg.loadTrustedSetup(trustedSetup)).getCause(); assertThat(cause.getMessage()).contains("Failed to parse trusted setup file"); + // reload real trusted setup for other tests + loadTrustedSetup(); } @Test @@ -252,6 +241,8 @@ public void monomialTrustedSetupFilesShouldThrow() { assertThat(kzgException.getMessage()).contains("Failed to load trusted setup"); assertThat(kzgException.getCause().getMessage()) .contains("There was an error while loading the Trusted Setup. (C_KZG_BADARGS)"); + // reload real trusted setup for other tests + loadTrustedSetup(); } @Test @@ -261,12 +252,6 @@ public void testInvalidLengthG2PointInNewTrustedSetup() { .hasMessage("Expected G2 point to be 96 bytes"); } - private void loadTrustedSetup() { - final String trustedSetup = - Resources.getResource(TrustedSetups.class, TRUSTED_SETUP_PATH).toExternalForm(); - kzg.loadTrustedSetup(trustedSetup); - } - private List getSampleBlobs(final int count) { return IntStream.range(0, count).mapToObj(__ -> getSampleBlob()).collect(Collectors.toList()); }