Fix consensus payloads indexes (#1837)

neo-project · Aug 13, 2020 · fcb7ad1 · fcb7ad1
1 parent ca13358
commit fcb7ad1
Show file tree

Hide file tree

Showing 9 changed files with 40 additions and 26 deletions.
diff --git a/src/neo/Consensus/ConsensusContext.cs b/src/neo/Consensus/ConsensusContext.cs
@@ -148,10 +148,10 @@ public Block EnsureHeader()
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public uint GetPrimaryIndex(byte viewNumber)
+        public byte GetPrimaryIndex(byte viewNumber)
         {
             int p = ((int)Block.Index - viewNumber) % Validators.Length;
-            return p >= 0 ? (uint)p : (uint)(p + Validators.Length);
+            return p >= 0 ? (byte)p : (byte)(p + Validators.Length);
         }
 
         public bool Load()
@@ -198,7 +198,7 @@ private ConsensusPayload MakeSignedPayload(ConsensusMessage message)
                 Version = Block.Version,
                 PrevHash = Block.PrevHash,
                 BlockIndex = Block.Index,
-                ValidatorIndex = (ushort)MyIndex,
+                ValidatorIndex = (byte)MyIndex,
                 ConsensusMessage = message
             };
             SignPayload(payload);

diff --git a/src/neo/Consensus/RecoveryMessage.ChangeViewPayloadCompact.cs b/src/neo/Consensus/RecoveryMessage.ChangeViewPayloadCompact.cs
@@ -1,5 +1,6 @@
 using Neo.IO;
 using Neo.Network.P2P.Payloads;
+using System;
 using System.IO;
 
 namespace Neo.Consensus
@@ -8,20 +9,22 @@ partial class RecoveryMessage
     {
         public class ChangeViewPayloadCompact : ISerializable
         {
-            public ushort ValidatorIndex;
+            public byte ValidatorIndex;
             public byte OriginalViewNumber;
             public ulong Timestamp;
             public byte[] InvocationScript;
 
             int ISerializable.Size =>
-                sizeof(ushort) +                //ValidatorIndex
+                sizeof(byte) +                  //ValidatorIndex
                 sizeof(byte) +                  //OriginalViewNumber
                 sizeof(ulong) +                 //Timestamp
                 InvocationScript.GetVarSize();  //InvocationScript
 
             void ISerializable.Deserialize(BinaryReader reader)
             {
-                ValidatorIndex = reader.ReadUInt16();
+                ValidatorIndex = reader.ReadByte();
+                if (ValidatorIndex >= ProtocolSettings.Default.ValidatorsCount)
+                    throw new FormatException();
                 OriginalViewNumber = reader.ReadByte();
                 Timestamp = reader.ReadUInt64();
                 InvocationScript = reader.ReadVarBytes(1024);

diff --git a/src/neo/Consensus/RecoveryMessage.CommitPayloadCompact.cs b/src/neo/Consensus/RecoveryMessage.CommitPayloadCompact.cs
@@ -1,5 +1,6 @@
 using Neo.IO;
 using Neo.Network.P2P.Payloads;
+using System;
 using System.IO;
 
 namespace Neo.Consensus
@@ -9,20 +10,22 @@ partial class RecoveryMessage
         public class CommitPayloadCompact : ISerializable
         {
             public byte ViewNumber;
-            public ushort ValidatorIndex;
+            public byte ValidatorIndex;
             public byte[] Signature;
             public byte[] InvocationScript;
 
             int ISerializable.Size =>
                 sizeof(byte) +                  //ViewNumber
-                sizeof(ushort) +                //ValidatorIndex
+                sizeof(byte) +                  //ValidatorIndex
                 Signature.Length +              //Signature
                 InvocationScript.GetVarSize();  //InvocationScript
 
             void ISerializable.Deserialize(BinaryReader reader)
             {
                 ViewNumber = reader.ReadByte();
-                ValidatorIndex = reader.ReadUInt16();
+                ValidatorIndex = reader.ReadByte();
+                if (ValidatorIndex >= ProtocolSettings.Default.ValidatorsCount)
+                    throw new FormatException();
                 Signature = reader.ReadFixedBytes(64);
                 InvocationScript = reader.ReadVarBytes(1024);
             }

diff --git a/src/neo/Consensus/RecoveryMessage.PreparationPayloadCompact.cs b/src/neo/Consensus/RecoveryMessage.PreparationPayloadCompact.cs
@@ -1,5 +1,6 @@
 using Neo.IO;
 using Neo.Network.P2P.Payloads;
+using System;
 using System.IO;
 
 namespace Neo.Consensus
@@ -8,16 +9,18 @@ partial class RecoveryMessage
     {
         public class PreparationPayloadCompact : ISerializable
         {
-            public ushort ValidatorIndex;
+            public byte ValidatorIndex;
             public byte[] InvocationScript;
 
             int ISerializable.Size =>
-                sizeof(ushort) +                //ValidatorIndex
+                sizeof(byte) +                  //ValidatorIndex
                 InvocationScript.GetVarSize();  //InvocationScript
 
             void ISerializable.Deserialize(BinaryReader reader)
             {
-                ValidatorIndex = reader.ReadUInt16();
+                ValidatorIndex = reader.ReadByte();
+                if (ValidatorIndex >= ProtocolSettings.Default.ValidatorsCount)
+                    throw new FormatException();
                 InvocationScript = reader.ReadVarBytes(1024);
             }
 

diff --git a/src/neo/Consensus/RecoveryMessage.cs b/src/neo/Consensus/RecoveryMessage.cs
@@ -90,14 +90,14 @@ internal ConsensusPayload[] GetCommitPayloadsFromRecoveryMessage(ConsensusContex
         internal ConsensusPayload GetPrepareRequestPayload(ConsensusContext context, ConsensusPayload payload)
         {
             if (PrepareRequestMessage == null) return null;
-            if (!PreparationMessages.TryGetValue((int)context.Block.ConsensusData.PrimaryIndex, out RecoveryMessage.PreparationPayloadCompact compact))
+            if (!PreparationMessages.TryGetValue(context.Block.ConsensusData.PrimaryIndex, out PreparationPayloadCompact compact))
                 return null;
             return new ConsensusPayload
             {
                 Version = payload.Version,
                 PrevHash = payload.PrevHash,
                 BlockIndex = payload.BlockIndex,
-                ValidatorIndex = (ushort)context.Block.ConsensusData.PrimaryIndex,
+                ValidatorIndex = context.Block.ConsensusData.PrimaryIndex,
                 ConsensusMessage = PrepareRequestMessage,
                 Witness = new Witness
                 {

diff --git a/src/neo/Network/P2P/Payloads/ConsensusData.cs b/src/neo/Network/P2P/Payloads/ConsensusData.cs
@@ -1,13 +1,14 @@
 using Neo.Cryptography;
 using Neo.IO;
 using Neo.IO.Json;
+using System;
 using System.IO;
 
 namespace Neo.Network.P2P.Payloads
 {
     public class ConsensusData : ISerializable
     {
-        public uint PrimaryIndex;
+        public byte PrimaryIndex;
         public ulong Nonce;
 
         private UInt256 _hash = null;
@@ -23,17 +24,19 @@ public UInt256 Hash
             }
         }
 
-        public int Size => IO.Helper.GetVarSize((int)PrimaryIndex) + sizeof(ulong);
+        public int Size => sizeof(byte) + sizeof(ulong);
 
         void ISerializable.Deserialize(BinaryReader reader)
         {
-            PrimaryIndex = (uint)reader.ReadVarInt((ulong)ProtocolSettings.Default.ValidatorsCount - 1);
+            PrimaryIndex = reader.ReadByte();
+            if (PrimaryIndex >= ProtocolSettings.Default.ValidatorsCount)
+                throw new FormatException();
             Nonce = reader.ReadUInt64();
         }
 
         void ISerializable.Serialize(BinaryWriter writer)
         {
-            writer.WriteVarInt(PrimaryIndex);
+            writer.Write(PrimaryIndex);
             writer.Write(Nonce);
         }
 

diff --git a/src/neo/Network/P2P/Payloads/ConsensusPayload.cs b/src/neo/Network/P2P/Payloads/ConsensusPayload.cs
@@ -15,7 +15,7 @@ public class ConsensusPayload : IInventory
         public uint Version;
         public UInt256 PrevHash;
         public uint BlockIndex;
-        public ushort ValidatorIndex;
+        public byte ValidatorIndex;
         public byte[] Data;
         public Witness Witness;
 
@@ -57,7 +57,7 @@ public UInt256 Hash
             sizeof(uint) +      //Version
             PrevHash.Size +     //PrevHash
             sizeof(uint) +      //BlockIndex
-            sizeof(ushort) +    //ValidatorIndex
+            sizeof(byte) +      //ValidatorIndex
             Data.GetVarSize() + //Data
             1 + Witness.Size;   //Witness
 
@@ -91,7 +91,9 @@ void IVerifiable.DeserializeUnsigned(BinaryReader reader)
             Version = reader.ReadUInt32();
             PrevHash = reader.ReadSerializable<UInt256>();
             BlockIndex = reader.ReadUInt32();
-            ValidatorIndex = reader.ReadUInt16();
+            ValidatorIndex = reader.ReadByte();
+            if (ValidatorIndex >= ProtocolSettings.Default.ValidatorsCount)
+                throw new FormatException();
             Data = reader.ReadVarBytes();
         }
 

diff --git a/tests/neo.UnitTests/Consensus/UT_ChangeViewPayloadCompact.cs b/tests/neo.UnitTests/Consensus/UT_ChangeViewPayloadCompact.cs
@@ -13,10 +13,10 @@ public class UT_ChangeViewPayloadCompact
         public void Size_Get()
         {
             var test = new RecoveryMessage.ChangeViewPayloadCompact() { Timestamp = 1, ValidatorIndex = 1, InvocationScript = new byte[0], OriginalViewNumber = 1 };
-            ((ISerializable)test).Size.Should().Be(12);
+            ((ISerializable)test).Size.Should().Be(11);
 
             test = new RecoveryMessage.ChangeViewPayloadCompact() { Timestamp = 1, ValidatorIndex = 1, InvocationScript = new byte[1024], OriginalViewNumber = 1 };
-            ((ISerializable)test).Size.Should().Be(1038);
+            ((ISerializable)test).Size.Should().Be(1037);
         }
 
         [TestMethod]

diff --git a/tests/neo.UnitTests/Consensus/UT_Consensus.cs b/tests/neo.UnitTests/Consensus/UT_Consensus.cs
@@ -444,7 +444,7 @@ public void ConsensusService_SingleNodeActors_OnStart_PrepReq_PrepResponses_Comm
         /// <param name="vI">new ValidatorIndex for the cpToCopy
         /// <param name="kp">KeyPair that will be used for signing the Commit message used for creating blocks
         /// <param name="blockHashToSign">HashCode of the Block that is being produced and current being signed
-        public ConsensusPayload GetCommitPayloadModifiedAndSignedCopy(ConsensusPayload cpToCopy, ushort vI, KeyPair kp, byte[] blockHashToSign)
+        public ConsensusPayload GetCommitPayloadModifiedAndSignedCopy(ConsensusPayload cpToCopy, byte vI, KeyPair kp, byte[] blockHashToSign)
         {
             var cpCommitTemp = cpToCopy.ToArray().AsSerializable<ConsensusPayload>();
             cpCommitTemp.ValidatorIndex = vI;
@@ -463,7 +463,7 @@ public ConsensusPayload GetCommitPayloadModifiedAndSignedCopy(ConsensusPayload c
         /// </summary>
         /// <param name="cpToCopy">ConsensusPayload that will be modified
         /// <param name="vI">new ValidatorIndex for the cpToCopy
-        public ConsensusPayload GetPayloadAndModifyValidator(ConsensusPayload cpToCopy, ushort vI)
+        public ConsensusPayload GetPayloadAndModifyValidator(ConsensusPayload cpToCopy, byte vI)
         {
             var cpTemp = cpToCopy.ToArray().AsSerializable<ConsensusPayload>();
             cpTemp.ValidatorIndex = vI;
@@ -910,7 +910,7 @@ public void TestSerializeAndDeserializeRecoveryMessageWithoutChangeViewsWithComm
             copiedMsg.CommitMessages.Should().BeEquivalentTo(msg.CommitMessages);
         }
 
-        private static ConsensusPayload MakeSignedPayload(ConsensusContext context, ConsensusMessage message, ushort validatorIndex, byte[] witnessInvocationScript)
+        private static ConsensusPayload MakeSignedPayload(ConsensusContext context, ConsensusMessage message, byte validatorIndex, byte[] witnessInvocationScript)
         {
             return new ConsensusPayload
             {