feat: add methods for claiming VRF/Plain secondary slot

LimeChain · Oct 30, 2024 · d18f75b · d18f75b
1 parent 39be6d4
commit d18f75b
Showing 1 changed file with 55 additions and 17 deletions.
diff --git a/src/main/java/com/limechain/babe/Authorship.java b/src/main/java/com/limechain/babe/Authorship.java
@@ -18,6 +18,7 @@
 import java.math.BigInteger;
 import java.util.List;
 
+//TODO: Add logs for successfully claiming primary/secondary slot
 @NoArgsConstructor(access = AccessLevel.PRIVATE)
 public class Authorship {
 
@@ -53,27 +54,64 @@ public static BabePreDigest claimPrimarySlot(final byte[] randomness,
         return null;
     }
 
-    //TODO: Should return Plain/VRF Secondary PreDigest or null
-    //TODO: You can take epoch.config.allowed_slots.is_secondary_vrf_slots_allowed from BabeApiConfigurations
-    public static BabePreDigest claimSecondarySlotVrf(final byte[] randomness,
-                                                      final BigInteger slotNumber,
-                                                      final BigInteger epochNumber,
-                                                      final List<Authority> authorities,
-                                                      final Schnorrkel.KeyPair keyPair,
-                                                      final int authorityIndex) {
-
-        var secondarySlotAuthor = getSecondarySlotAuthor(randomness, slotNumber, authorities);
-        if (secondarySlotAuthor == null) {
+    public static BabePreDigest claimSecondarySlot(final byte[] randomness,
+                                                   final BigInteger slotNumber,
+                                                   final BigInteger epochNumber,
+                                                   final List<Authority> authorities,
+                                                   final Schnorrkel.KeyPair keyPair,
+                                                   final int authorityIndex,
+                                                   final boolean authorSecondaryVrfSlot) {
+
+        var secondarySlotAuthorIndex = getSecondarySlotAuthor(randomness, slotNumber, authorities);
+        if (secondarySlotAuthorIndex == null) {
             return null;
         }
 
-        //TODO: Add implementation
-        return null;
+        //Not our turn to propose
+        if (secondarySlotAuthorIndex != authorityIndex) {
+            return null;
+        }
+
+        if (authorSecondaryVrfSlot) {
+            return buildSecondaryVrfPreDigest(
+                    randomness,
+                    slotNumber,
+                    epochNumber,
+                    keyPair,
+                    authorityIndex
+            );
+        } else {
+            return new BabePreDigest(
+                    PreDigestType.BABE_SECONDARY_PLAIN,
+                    authorityIndex,
+                    slotNumber,
+                    null,
+                    null
+            );
+        }
+    }
+
+    private static BabePreDigest buildSecondaryVrfPreDigest(final byte[] randomness,
+                                                            final BigInteger slotNumber,
+                                                            final BigInteger epochNumber,
+                                                            final Schnorrkel.KeyPair keyPair,
+                                                            final int authorityIndex) {
+
+        var transcript = makeTranscript(randomness, slotNumber, epochNumber);
+        VrfOutputAndProof vrfOutputAndProof = Schnorrkel.getInstance().vrfSign(keyPair, transcript);
+
+        return new BabePreDigest(
+                PreDigestType.BABE_SECONDARY_VRF,
+                authorityIndex,
+                slotNumber,
+                vrfOutputAndProof.getOutput(),
+                vrfOutputAndProof.getProof()
+        );
     }
 
-    private static Authority getSecondarySlotAuthor(final byte[] randomness,
-                                                    final BigInteger slotNumber,
-                                                    final List<Authority> authorities) {
+    private static Integer getSecondarySlotAuthor(final byte[] randomness,
+                                                  final BigInteger slotNumber,
+                                                  final List<Authority> authorities) {
         if (authorities.isEmpty()) return null;
 
         byte[] concat = ByteArrayUtils.concatenate(randomness, slotNumber.toByteArray());
@@ -86,7 +124,7 @@ private static Authority getSecondarySlotAuthor(final byte[] randomness,
         var authorityIndex = rand.mod(authoritiesCount);
 
         if (authorityIndex.compareTo(authoritiesCount) < 0) {
-            return authorities.get(authorityIndex.intValue());
+            return authorityIndex.intValue();
         }
 
         return null;