Skip to content

Commit

Permalink
changes as per feedback
Browse files Browse the repository at this point in the history
  • Loading branch information
StefanBratanov committed Oct 23, 2023
1 parent db24e04 commit 262665a
Show file tree
Hide file tree
Showing 5 changed files with 65 additions and 73 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,8 @@ public final void runTest(final TestDefinition testDefinition) throws Throwable
final SpecConfigDeneb specConfigDeneb =
SpecConfigDeneb.required(networkConfig.getSpec().getGenesisSpecConfig());

try (kzg) {
kzg.loadTrustedSetup(specConfigDeneb.getTrustedSetupPath().orElseThrow());
runTestImpl(testDefinition);
}
kzg.loadTrustedSetup(specConfigDeneb.getTrustedSetupPath().orElseThrow());
runTestImpl(testDefinition);
}

private String extractNetwork(final String testName) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,9 +28,7 @@

/**
* This class provides a KZG instance with a loaded trusted setup that will automatically free the
* trusted setup when the property test is finished. It will re-use the same KZG instance for all
* iterations of the test, but it will create a new instance for each method. For a class with three
* property test methods, you can expect it to load/free three times.
* trusted setup when all the property tests are finished.
*/
class KzgResolver implements ResolveParameterHook {
public static final Tuple.Tuple2<Class<KzgResolver.KzgAutoLoadFree>, String> STORE_IDENTIFIER =
Expand All @@ -49,7 +47,7 @@ public PropagationMode propagateTo() {

private KZG getKzgWithTrustedSetup() {
final Store<KzgResolver.KzgAutoLoadFree> kzgStore =
Store.getOrCreate(STORE_IDENTIFIER, Lifespan.PROPERTY, KzgResolver.KzgAutoLoadFree::new);
Store.getOrCreate(STORE_IDENTIFIER, Lifespan.RUN, KzgResolver.KzgAutoLoadFree::new);
return kzgStore.get().kzg;
}

Expand All @@ -64,8 +62,8 @@ private KzgAutoLoadFree() {
}

@Override
public void close() throws Exception {
kzg.close();
public void close() {
kzg.freeTrustedSetup();
}
}
}
10 changes: 6 additions & 4 deletions infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZG.java
Original file line number Diff line number Diff line change
Expand Up @@ -20,14 +20,17 @@
* 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 extends AutoCloseable {
public interface KZG {

KZG NOOP =
new KZG() {

@Override
public void loadTrustedSetup(final String trustedSetupFile) throws KZGException {}

@Override
public void freeTrustedSetup() throws KZGException {}

@Override
public boolean verifyBlobKzgProofBatch(
final List<Bytes> blobs,
Expand All @@ -47,13 +50,12 @@ 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<Bytes> blobs, List<KZGCommitment> kzgCommitments, List<KZGProof> kzgProofs)
throws KZGException;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,17 @@ public synchronized void loadTrustedSetup(final String trustedSetupFile) throws
}
}

@Override
public 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<Bytes> blobs,
Expand Down Expand Up @@ -126,20 +137,4 @@ 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);
}
}
}
85 changes: 42 additions & 43 deletions infrastructure/kzg/src/test/java/tech/pegasys/teku/kzg/KZGTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@
import org.apache.tuweni.bytes.Bytes;
import org.apache.tuweni.units.bigints.UInt256;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
Expand All @@ -48,23 +48,22 @@ public final class KZGTest {
private static final Random RND = new Random(RANDOM_SEED);
private static final String TRUSTED_SETUP_PATH = "trusted_setup.txt";

private static KZG kzg;
private static final KZG KZG = CKZG4844.createInstance();

@BeforeAll
public static void setUp() {
kzg = CKZG4844.createInstance();
@BeforeEach
public void setUp() {
loadTrustedSetup();
}

@AfterAll
public static void cleanUp() throws Exception {
kzg.close();
public static void cleanUp() {
KZG.freeTrustedSetup();
}

private static void loadTrustedSetup() {
final String trustedSetup =
Resources.getResource(TrustedSetups.class, TRUSTED_SETUP_PATH).toExternalForm();
kzg.loadTrustedSetup(trustedSetup);
KZG.loadTrustedSetup(trustedSetup);
}

@Test
Expand All @@ -73,111 +72,115 @@ public void testKzgLoadSameTrustedSetupTwice_shouldNotThrowException() {
}

@Test
public void testUsageWithoutLoadedTrustedSetup_shouldThrowException() throws Exception {
kzg.close();
public void testKzgFreeTrustedSetupTwice_shouldThrowException() {
KZG.freeTrustedSetup();
assertThrows(KZGException.class, KZG::freeTrustedSetup);
}

@Test
public void testUsageWithoutLoadedTrustedSetup_shouldThrowException() {
KZG.freeTrustedSetup();
final List<KZGException> exceptions =
List.of(
assertThrows(
KZGException.class,
() ->
kzg.verifyBlobKzgProofBatch(
KZG.verifyBlobKzgProofBatch(
List.of(Bytes.fromHexString("0x", BYTES_PER_BLOB)),
List.of(KZGCommitment.infinity()),
List.of(KZGProof.INFINITY))),
assertThrows(KZGException.class, () -> kzg.blobToKzgCommitment(Bytes.EMPTY)),
assertThrows(KZGException.class, () -> KZG.blobToKzgCommitment(Bytes.EMPTY)),
assertThrows(
KZGException.class,
() -> kzg.computeBlobKzgProof(Bytes.EMPTY, KZGCommitment.infinity())));
() -> KZG.computeBlobKzgProof(Bytes.EMPTY, KZGCommitment.infinity())));

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() {
final int numberOfBlobs = 4;
final List<Bytes> blobs = getSampleBlobs(numberOfBlobs);
final List<KZGCommitment> kzgCommitments =
blobs.stream().map(kzg::blobToKzgCommitment).collect(Collectors.toList());
blobs.stream().map(KZG::blobToKzgCommitment).collect(Collectors.toList());
final List<KZGProof> kzgProofs =
Streams.zip(
kzgCommitments.stream(),
blobs.stream(),
(kzgCommitment, blob) -> kzg.computeBlobKzgProof(blob, kzgCommitment))
(kzgCommitment, blob) -> KZG.computeBlobKzgProof(blob, kzgCommitment))
.collect(Collectors.toList());
assertThat(kzg.verifyBlobKzgProofBatch(blobs, kzgCommitments, kzgProofs)).isTrue();
assertThat(KZG.verifyBlobKzgProofBatch(blobs, kzgCommitments, kzgProofs)).isTrue();

assertThat(
kzg.verifyBlobKzgProofBatch(getSampleBlobs(numberOfBlobs), kzgCommitments, kzgProofs))
KZG.verifyBlobKzgProofBatch(getSampleBlobs(numberOfBlobs), kzgCommitments, kzgProofs))
.isFalse();
assertThat(kzg.verifyBlobKzgProofBatch(blobs, getSampleCommitments(numberOfBlobs), kzgProofs))
assertThat(KZG.verifyBlobKzgProofBatch(blobs, getSampleCommitments(numberOfBlobs), kzgProofs))
.isFalse();
final List<KZGProof> invalidProofs =
getSampleBlobs(numberOfBlobs).stream()
.map((Bytes blob) -> kzg.computeBlobKzgProof(blob, kzg.blobToKzgCommitment(blob)))
.map((Bytes blob) -> KZG.computeBlobKzgProof(blob, KZG.blobToKzgCommitment(blob)))
.collect(Collectors.toList());
assertThat(kzg.verifyBlobKzgProofBatch(blobs, kzgCommitments, invalidProofs)).isFalse();
assertThat(KZG.verifyBlobKzgProofBatch(blobs, kzgCommitments, invalidProofs)).isFalse();
}

@Test
public void testVerifyingEmptyBatch() {
assertThat(kzg.verifyBlobKzgProofBatch(List.of(), List.of(), List.of())).isTrue();
assertThat(KZG.verifyBlobKzgProofBatch(List.of(), List.of(), List.of())).isTrue();
}

@Test
public void testComputingAndVerifyingBatchSingleProof() {
final int numberOfBlobs = 1;
final List<Bytes> blobs = getSampleBlobs(numberOfBlobs);
final List<KZGCommitment> kzgCommitments =
blobs.stream().map(kzg::blobToKzgCommitment).collect(Collectors.toList());
blobs.stream().map(KZG::blobToKzgCommitment).collect(Collectors.toList());
final List<KZGProof> kzgProofs =
Streams.zip(
kzgCommitments.stream(),
blobs.stream(),
(kzgCommitment, blob) -> kzg.computeBlobKzgProof(blob, kzgCommitment))
(kzgCommitment, blob) -> KZG.computeBlobKzgProof(blob, kzgCommitment))
.collect(Collectors.toList());
assertThat(kzgProofs.size()).isEqualTo(1);
assertThat(kzg.verifyBlobKzgProofBatch(blobs, kzgCommitments, kzgProofs)).isTrue();
assertThat(KZG.verifyBlobKzgProofBatch(blobs, kzgCommitments, kzgProofs)).isTrue();

assertThat(
kzg.verifyBlobKzgProofBatch(getSampleBlobs(numberOfBlobs), kzgCommitments, kzgProofs))
KZG.verifyBlobKzgProofBatch(getSampleBlobs(numberOfBlobs), kzgCommitments, kzgProofs))
.isFalse();
assertThat(kzg.verifyBlobKzgProofBatch(blobs, getSampleCommitments(numberOfBlobs), kzgProofs))
assertThat(KZG.verifyBlobKzgProofBatch(blobs, getSampleCommitments(numberOfBlobs), kzgProofs))
.isFalse();
final List<KZGProof> invalidProofs =
getSampleBlobs(numberOfBlobs).stream()
.map((Bytes blob) -> kzg.computeBlobKzgProof(blob, kzg.blobToKzgCommitment(blob)))
.map((Bytes blob) -> KZG.computeBlobKzgProof(blob, KZG.blobToKzgCommitment(blob)))
.collect(Collectors.toList());
assertThat(kzg.verifyBlobKzgProofBatch(blobs, kzgCommitments, invalidProofs)).isFalse();
assertThat(KZG.verifyBlobKzgProofBatch(blobs, kzgCommitments, invalidProofs)).isFalse();
}

@Test
public void testVerifyingBatchProofsThrowsIfSizesDoesntMatch() {
final int numberOfBlobs = 4;
final List<Bytes> blobs = getSampleBlobs(numberOfBlobs);
final List<KZGCommitment> kzgCommitments =
blobs.stream().map(kzg::blobToKzgCommitment).collect(Collectors.toList());
blobs.stream().map(KZG::blobToKzgCommitment).collect(Collectors.toList());
final List<KZGProof> kzgProofs =
Streams.zip(
kzgCommitments.stream(),
blobs.stream(),
(kzgCommitment, blob) -> kzg.computeBlobKzgProof(blob, kzgCommitment))
(kzgCommitment, blob) -> KZG.computeBlobKzgProof(blob, kzgCommitment))
.collect(Collectors.toList());
final KZGException kzgException1 =
assertThrows(
KZGException.class,
() -> kzg.verifyBlobKzgProofBatch(blobs, kzgCommitments, List.of(kzgProofs.get(0))));
() -> KZG.verifyBlobKzgProofBatch(blobs, kzgCommitments, List.of(kzgProofs.get(0))));
final KZGException kzgException2 =
assertThrows(
KZGException.class,
() -> kzg.verifyBlobKzgProofBatch(blobs, List.of(kzgCommitments.get(0)), kzgProofs));
() -> KZG.verifyBlobKzgProofBatch(blobs, List.of(kzgCommitments.get(0)), kzgProofs));
final KZGException kzgException3 =
assertThrows(
KZGException.class,
() -> kzg.verifyBlobKzgProofBatch(List.of(blobs.get(0)), kzgCommitments, kzgProofs));
() -> KZG.verifyBlobKzgProofBatch(List.of(blobs.get(0)), kzgCommitments, kzgProofs));

Stream.of(kzgException1, kzgException2, kzgException3)
.forEach(
Expand All @@ -201,7 +204,7 @@ public void testComputingProofWithIncorrectLengthBlobDoesNotCauseSegfault(final

final KZGException kzgException =
assertThrows(
KZGException.class, () -> kzg.computeBlobKzgProof(blob, kzg.blobToKzgCommitment(blob)));
KZGException.class, () -> KZG.computeBlobKzgProof(blob, KZG.blobToKzgCommitment(blob)));

assertThat(kzgException)
.cause()
Expand All @@ -226,23 +229,19 @@ public void testComputingProofWithIncorrectLengthBlobDoesNotCauseSegfault(final
public void incorrectTrustedSetupFilesShouldThrow(final String path) {
final String trustedSetup = Resources.getResource(TrustedSetups.class, path).toExternalForm();
final Throwable cause =
assertThrows(KZGException.class, () -> kzg.loadTrustedSetup(trustedSetup)).getCause();
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
public void monomialTrustedSetupFilesShouldThrow() {
final String trustedSetup =
Resources.getResource(TrustedSetups.class, "trusted_setup_monomial.txt").toExternalForm();
final KZGException kzgException =
assertThrows(KZGException.class, () -> kzg.loadTrustedSetup(trustedSetup));
assertThrows(KZGException.class, () -> KZG.loadTrustedSetup(trustedSetup));
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
Expand Down Expand Up @@ -271,7 +270,7 @@ private List<KZGCommitment> getSampleCommitments(final int count) {
}

private KZGCommitment getSampleCommitment() {
return kzg.blobToKzgCommitment(getSampleBlob());
return KZG.blobToKzgCommitment(getSampleBlob());
}

private UInt256 randomBLSFieldElement() {
Expand Down

0 comments on commit 262665a

Please sign in to comment.