Added Digest Helper and PreDigestReader

src/main/java/com/limechain/babe/Authorship.java

@@ -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);
@@ -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");

src/main/java/com/limechain/babe/predigest/BabePreDigest.java

@@ -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;

src/main/java/com/limechain/babe/predigest/PreDigestType.java

@@ -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);
+    }
 }

src/main/java/com/limechain/babe/predigest/scale/PreDigestReader.java

@@ -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;
+    }
+}

src/main/java/com/limechain/babe/predigest/scale/PreDigestWriter.java

@@ -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());

src/main/java/com/limechain/network/protocol/blockannounce/BlockAnnounceEngine.java

@@ -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;
@@ -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;
@@ -25,7 +22,6 @@
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.util.Arrays;
 import java.util.logging.Level;
 
 @Log
@@ -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) {
@@ -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) {

src/main/java/com/limechain/network/protocol/warp/DigestHelper.java

@@ -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()));
+    }
+}

src/test/java/com/limechain/babe/scale/PreDigestScaleTest.java

@@ -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);
+    }
+}

src/test/java/com/limechain/network/protocol/blockannounce/BlockAnnounceEngineTest.java

@@ -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;
@@ -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)
@@ -53,7 +45,7 @@ class BlockAnnounceEngineTest {
     private WarpSyncState warpSyncState;
 
     @Mock
-    private EpochState epochState;
+    private DigestHelper digestHelper;
 
     @Mock
     private BlockAnnounceHandshake handshake;
@@ -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);
-    }
 }