Skip to content

Commit

Permalink
Added Digest Helper and PreDigestReader (#590)
Browse files Browse the repository at this point in the history
Added Digest Helper class for processing different types of header
digests and PreDigestReader
  • Loading branch information
osrib authored Nov 1, 2024
1 parent d1c1d08 commit c6271a5
Show file tree
Hide file tree
Showing 11 changed files with 188 additions and 127 deletions.
4 changes: 2 additions & 2 deletions src/main/java/com/limechain/babe/Authorship.java
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ public static BabePreDigest claimPrimarySlot(final byte[] randomness,
final BigInteger slotNumber,
final BigInteger epochNumber,
final Schnorrkel.KeyPair keyPair,
final int authorityIndex,
final long authorityIndex,
final BigInteger threshold) {

var transcript = makeTranscript(randomness, slotNumber, epochNumber);
Expand Down Expand Up @@ -84,7 +84,7 @@ public static BigInteger calculatePrimaryThreshold(

private static double getBabeConstant(@NotNull Pair<BigInteger, BigInteger> constant,
@NotNull List<Authority> authorities,
int authorityIndex) {
long authorityIndex) {

if (BigInteger.ZERO.equals(constant.getValue1())) {
throw new IllegalArgumentException("Invalid authority index provided");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,16 @@

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;

import java.math.BigInteger;

@Data
@AllArgsConstructor
@NoArgsConstructor
public class BabePreDigest {
private PreDigestType type;
private int authorityIndex;
private long authorityIndex;
private BigInteger slotNumber;
private byte[] vrfOutput;
private byte[] vrfProof;
Expand Down
9 changes: 9 additions & 0 deletions src/main/java/com/limechain/babe/predigest/PreDigestType.java
Original file line number Diff line number Diff line change
Expand Up @@ -13,4 +13,13 @@ public enum PreDigestType {
PreDigestType(int value) {
this.value = (byte) value;
}

public static PreDigestType getByValue(byte value) {
for (PreDigestType type : PreDigestType.values()) {
if (type.value == value) {
return type;
}
}
throw new IllegalArgumentException("Unknown PreDigestType value: " + value);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
package com.limechain.babe.predigest.scale;

import com.limechain.babe.predigest.BabePreDigest;
import com.limechain.babe.predigest.PreDigestType;
import io.emeraldpay.polkaj.scale.ScaleCodecReader;
import io.emeraldpay.polkaj.scale.ScaleReader;
import io.emeraldpay.polkaj.scale.reader.UInt64Reader;

import static io.emeraldpay.polkaj.schnorrkel.VrfOutputAndProof.OUTPUT_BYTE_LEN;
import static io.emeraldpay.polkaj.schnorrkel.VrfOutputAndProof.PROOF_BYTE_LEN;

public class PreDigestReader implements ScaleReader<BabePreDigest> {
@Override
public BabePreDigest read(ScaleCodecReader reader) {
BabePreDigest preDigest = new BabePreDigest();
PreDigestType type = PreDigestType.getByValue(reader.readByte());
preDigest.setType(type);
preDigest.setAuthorityIndex(reader.readUint32());
preDigest.setSlotNumber(new UInt64Reader().read(reader));
if (!PreDigestType.BABE_SECONDARY_PLAIN.equals(type)) {
preDigest.setVrfOutput(reader.readByteArray(OUTPUT_BYTE_LEN));
preDigest.setVrfProof(reader.readByteArray(PROOF_BYTE_LEN));
}
return preDigest;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -4,18 +4,24 @@
import com.limechain.babe.predigest.BabePreDigest;
import io.emeraldpay.polkaj.scale.ScaleCodecWriter;
import io.emeraldpay.polkaj.scale.ScaleWriter;
import io.emeraldpay.polkaj.scale.writer.UInt64Writer;

import java.io.IOException;


public class PreDigestWriter implements ScaleWriter<BabePreDigest> {

private final UInt64Writer uint64Writer;

public PreDigestWriter() {
this.uint64Writer = new UInt64Writer();
}

@Override
public void write(ScaleCodecWriter writer, BabePreDigest preDigest) throws IOException {
PreDigestType type = preDigest.getType();
writer.writeByte(type.getValue());
writer.writeUint32(preDigest.getAuthorityIndex());
writer.writeUint128(preDigest.getSlotNumber());
uint64Writer.write(writer, preDigest.getSlotNumber());
if (!PreDigestType.BABE_SECONDARY_PLAIN.equals(type)) {
writer.writeByteArray(preDigest.getVrfOutput());
writer.writeByteArray(preDigest.getVrfProof());
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
package com.limechain.network.protocol.blockannounce;

import com.limechain.babe.state.EpochState;
import com.limechain.exception.scale.ScaleEncodingException;
import com.limechain.network.ConnectionManager;
import com.limechain.network.protocol.blockannounce.messages.BlockAnnounceHandshake;
Expand All @@ -9,9 +8,7 @@
import com.limechain.network.protocol.blockannounce.scale.BlockAnnounceHandshakeScaleReader;
import com.limechain.network.protocol.blockannounce.scale.BlockAnnounceHandshakeScaleWriter;
import com.limechain.network.protocol.blockannounce.scale.BlockAnnounceMessageScaleReader;
import com.limechain.network.protocol.warp.dto.ConsensusEngine;
import com.limechain.network.protocol.warp.dto.DigestType;
import com.limechain.network.protocol.warp.dto.HeaderDigest;
import com.limechain.network.protocol.warp.DigestHelper;
import com.limechain.rpc.server.AppBean;
import com.limechain.storage.block.BlockState;
import com.limechain.sync.warpsync.WarpSyncState;
Expand All @@ -25,7 +22,6 @@

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.logging.Level;

@Log
Expand All @@ -37,13 +33,13 @@ public class BlockAnnounceEngine {
protected ConnectionManager connectionManager;
protected WarpSyncState warpSyncState;
protected BlockAnnounceHandshakeBuilder handshakeBuilder;
protected EpochState epochState;
protected DigestHelper digestHelper;

public BlockAnnounceEngine() {
connectionManager = ConnectionManager.getInstance();
warpSyncState = AppBean.getBean(WarpSyncState.class);
handshakeBuilder = new BlockAnnounceHandshakeBuilder();
epochState = AppBean.getBean(EpochState.class);
digestHelper = new DigestHelper();
}

public void receiveRequest(byte[] msg, Stream stream) {
Expand Down Expand Up @@ -98,20 +94,7 @@ private void handleBlockAnnounce(byte[] msg, PeerId peerId) {
if (BlockState.getInstance().isInitialized()) {
BlockState.getInstance().addBlockToBlockTree(announce.getHeader());
}
updateEpochStateIfBabeConsensusMessageExists(announce);
}

//TODO: We would need DigestHelper class that would encapsulate logic needed for all types message handling HeaderDigest[]
public void updateEpochStateIfBabeConsensusMessageExists(BlockAnnounceMessage announce) {
HeaderDigest[] headerDigests = announce.getHeader().getDigest();
if (headerDigests != null && headerDigests.length > 0) {
Arrays.stream(headerDigests)
.filter(headerDigest -> DigestType.CONSENSUS_MESSAGE.equals(headerDigest.getType()) &&
ConsensusEngine.BABE.equals(headerDigest.getId()))
.findFirst()
.map(HeaderDigest::getMessage)
.ifPresent(message -> epochState.updateNextEpochBlockConfig(message));
}
digestHelper.handleHeaderDigests(announce.getHeader().getDigest());
}

public void writeHandshakeToStream(Stream stream, PeerId peerId) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
package com.limechain.network.protocol.warp;

import com.limechain.babe.predigest.BabePreDigest;
import com.limechain.babe.predigest.scale.PreDigestReader;
import com.limechain.babe.state.EpochState;
import com.limechain.network.protocol.warp.dto.ConsensusEngine;
import com.limechain.network.protocol.warp.dto.DigestType;
import com.limechain.network.protocol.warp.dto.HeaderDigest;
import com.limechain.rpc.server.AppBean;
import com.limechain.utils.scale.ScaleUtils;

import java.util.Arrays;
import java.util.Optional;

/**
* Helper class for processing different types of header digests
*/
public class DigestHelper {
private EpochState epochState;

public DigestHelper() {
this.epochState = AppBean.getBean(EpochState.class);
}

public void handleHeaderDigests(HeaderDigest[] headerDigests) {
if (headerDigests == null || headerDigests.length == 0) {
return;
}
updateEpochStateIfBabeConsensusMessageExists(headerDigests);
handleBabePreRuntimeDigest(headerDigests);
}

private void updateEpochStateIfBabeConsensusMessageExists(HeaderDigest[] headerDigests) {
Arrays.stream(headerDigests)
.filter(headerDigest -> DigestType.CONSENSUS_MESSAGE.equals(headerDigest.getType()) &&
ConsensusEngine.BABE.equals(headerDigest.getId()))
.findFirst()
.map(HeaderDigest::getMessage)
.ifPresent(message -> epochState.updateNextEpochBlockConfig(message));
}

private Optional<BabePreDigest> handleBabePreRuntimeDigest(HeaderDigest[] headerDigests) {
return Arrays.stream(headerDigests)
.filter(headerDigest -> DigestType.PRE_RUNTIME.equals(headerDigest.getType()) &&
ConsensusEngine.BABE.equals(headerDigest.getId()))
.findFirst()
.map(HeaderDigest::getMessage)
.map(message -> ScaleUtils.Decode.decode(message, new PreDigestReader()));
}
}
36 changes: 36 additions & 0 deletions src/test/java/com/limechain/babe/scale/PreDigestScaleTest.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
package com.limechain.babe.scale;

import com.limechain.babe.predigest.BabePreDigest;
import com.limechain.babe.predigest.PreDigestType;
import com.limechain.babe.predigest.scale.PreDigestReader;
import com.limechain.babe.predigest.scale.PreDigestWriter;
import com.limechain.utils.StringUtils;
import com.limechain.utils.scale.ScaleUtils;
import org.junit.jupiter.api.Test;

import java.math.BigInteger;

import static io.emeraldpay.polkaj.schnorrkel.VrfOutputAndProof.OUTPUT_BYTE_LEN;
import static io.emeraldpay.polkaj.schnorrkel.VrfOutputAndProof.PROOF_BYTE_LEN;
import static org.junit.jupiter.api.Assertions.*;

public class PreDigestScaleTest {

@Test
public void testEncodeAndDecode_BABE_SECONDARY_VRF() {
byte[] bytes = StringUtils.hexToBytes("0x03b7010000286f301100000000424dc7bc71c9b0f3a15b6aba389471fff57e154dbf3dc046f47513a38375ce4cfe421df446b2b7a0f5b1b2e69c810a6ba36787c0301e6636df3e45688116e005363dfe9922b3b7cc7b80e2fdab8b0b98dcfb4b96cc470fb35f753debda40aa0e");
BabePreDigest preDigest = ScaleUtils.Decode.decode(bytes, new PreDigestReader());
assertEquals(PreDigestType.BABE_SECONDARY_VRF, preDigest.getType());
assertEquals(439, preDigest.getAuthorityIndex());
assertEquals(BigInteger.valueOf(288386856), preDigest.getSlotNumber());
byte[] vrfOutput = preDigest.getVrfOutput();
byte[] vrfProof = preDigest.getVrfProof();
assertNotNull(vrfOutput);
assertNotNull(vrfProof);
assertEquals(OUTPUT_BYTE_LEN, vrfOutput.length);
assertEquals(PROOF_BYTE_LEN, vrfProof.length);

byte[] encode = ScaleUtils.Encode.encode(new PreDigestWriter(), preDigest);
assertArrayEquals(bytes, encode);
}
}
73 changes: 0 additions & 73 deletions src/test/java/com/limechain/babe/scale/PreDigestWriterTest.java

This file was deleted.

Original file line number Diff line number Diff line change
@@ -1,16 +1,13 @@
package com.limechain.network.protocol.blockannounce;

import com.limechain.babe.state.EpochState;
import com.limechain.network.ConnectionManager;
import com.limechain.network.protocol.blockannounce.messages.BlockAnnounceHandshake;
import com.limechain.network.protocol.blockannounce.messages.BlockAnnounceHandshakeBuilder;
import com.limechain.network.protocol.blockannounce.messages.BlockAnnounceMessage;
import com.limechain.network.protocol.blockannounce.scale.BlockAnnounceHandshakeScaleWriter;
import com.limechain.network.protocol.blockannounce.scale.BlockAnnounceMessageScaleReader;
import com.limechain.network.protocol.warp.DigestHelper;
import com.limechain.network.protocol.warp.dto.BlockHeader;
import com.limechain.network.protocol.warp.dto.ConsensusEngine;
import com.limechain.network.protocol.warp.dto.DigestType;
import com.limechain.network.protocol.warp.dto.HeaderDigest;
import com.limechain.sync.warpsync.WarpSyncState;
import io.emeraldpay.polkaj.scale.ScaleCodecReader;
import io.emeraldpay.polkaj.scale.ScaleCodecWriter;
Expand All @@ -28,12 +25,7 @@

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockConstruction;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;
import static org.mockito.Mockito.*;

@SuppressWarnings("unused")
@ExtendWith(MockitoExtension.class)
Expand All @@ -53,7 +45,7 @@ class BlockAnnounceEngineTest {
private WarpSyncState warpSyncState;

@Mock
private EpochState epochState;
private DigestHelper digestHelper;

@Mock
private BlockAnnounceHandshake handshake;
Expand Down Expand Up @@ -157,21 +149,4 @@ void receiveBlockAnnounceWhenConnectedShouldSyncMessage() {
verify(warpSyncState).syncBlockAnnounce(blockAnnounceMessage);
}
}

@Test
void updateEpochStateIfBabeConsensusMessageExistsShouldUpdateEpochState() {
byte[] message = new byte[]{1, 2, 3};
BlockAnnounceMessage blockAnnounceMessage = mock(BlockAnnounceMessage.class);
BlockHeader blockHeader = mock(BlockHeader.class);
when(blockAnnounceMessage.getHeader()).thenReturn(blockHeader);
HeaderDigest babeDigest = mock(HeaderDigest.class);
when(babeDigest.getType()).thenReturn(DigestType.CONSENSUS_MESSAGE);
when(babeDigest.getId()).thenReturn(ConsensusEngine.BABE);
when(babeDigest.getMessage()).thenReturn(message);
when(blockHeader.getDigest()).thenReturn(new HeaderDigest[]{babeDigest});

blockAnnounceEngine.updateEpochStateIfBabeConsensusMessageExists(blockAnnounceMessage);

verify(epochState).updateNextEpochBlockConfig(message);
}
}
Loading

0 comments on commit c6271a5

Please sign in to comment.