Set max block size (#953)

* Draft

* Take into account p2p max payload

* Move max allowed to policy

* Check block size on prepResponse

* Change reason

* Prevent overflow

* Optimization

* Reduce the length of the array

* Organizing consensus code

* Revert "Organizing consensus code"

This reverts commit 94a29dc.

* Remove Policy UT

* Resolve Policy conflicts

* prepare unit test

* Small unit test

* More ut

* Add one check

* Clean using

* Organizing consensus and comments

* Split unit test

* UT check block size

* Clean

* Expected witness size

* optimize

* Remove fakeWitness

* Format comments

* rename var

* Add (..)

* Remove SetKey method

* Centralize expected block size

* Optimize

* Fix

* Add one test

* Optimize `EnsureMaxBlockSize()`

* Fix unit tests

* Rename

* Indent

* Vitor suggestion

* Merge with Scoped signatures

* Remove extra line

* Revert "Remove extra line"

This reverts commit e134881.

* Remove extra line

neo.UnitTests/Consensus/UT_ConsensusContext.cs

@@ -0,0 +1,163 @@
+using Akka.TestKit.Xunit2;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Moq;
+using Neo.Consensus;
+using Neo.IO;
+using Neo.Network.P2P.Payloads;
+using Neo.SmartContract;
+using Neo.SmartContract.Native;
+using Neo.Wallets;
+using System;
+using System.Linq;
+
+namespace Neo.UnitTests.Consensus
+{
+
+    [TestClass]
+    public class UT_ConsensusContext : TestKit
+    {
+        ConsensusContext _context;
+        KeyPair[] _validatorKeys;
+
+        [TestInitialize]
+        public void TestSetup()
+        {
+            TestBlockchain.InitializeMockNeoSystem();
+
+            var rand = new Random();
+            var mockWallet = new Mock<Wallet>();
+            mockWallet.Setup(p => p.GetAccount(It.IsAny<UInt160>())).Returns<UInt160>(p => new TestWalletAccount(p));
+
+            // Create dummy validators
+
+            _validatorKeys = new KeyPair[7];
+            for (int x = 0; x < _validatorKeys.Length; x++)
+            {
+                var pk = new byte[32];
+                rand.NextBytes(pk);
+
+                _validatorKeys[x] = new KeyPair(pk);
+            }
+
+            _context = new ConsensusContext(mockWallet.Object, TestBlockchain.GetStore())
+            {
+                Validators = _validatorKeys.Select(u => u.PublicKey).ToArray()
+            };
+            _context.Reset(0);
+        }
+
+        [TestCleanup]
+        public void Cleanup()
+        {
+            Shutdown();
+        }
+
+        [TestMethod]
+        public void TestMaxBlockSize_Good()
+        {
+            // Only one tx, is included
+
+            var tx1 = CreateTransactionWithSize(200);
+            _context.EnsureMaxBlockSize(new Transaction[] { tx1 });
+            EnsureContext(_context, tx1);
+
+            // All txs included
+
+            var max = (int)NativeContract.Policy.GetMaxTransactionsPerBlock(_context.Snapshot);
+            var txs = new Transaction[max];
+
+            for (int x = 0; x < max; x++) txs[x] = CreateTransactionWithSize(100);
+
+            _context.EnsureMaxBlockSize(txs);
+            EnsureContext(_context, txs);
+        }
+
+        [TestMethod]
+        public void TestMaxBlockSize_Exceed()
+        {
+            // Two tx, the last one exceed the size rule, only the first will be included
+
+            var tx1 = CreateTransactionWithSize(200);
+            var tx2 = CreateTransactionWithSize(256 * 1024);
+            _context.EnsureMaxBlockSize(new Transaction[] { tx1, tx2 });
+            EnsureContext(_context, tx1);
+
+            // Exceed txs number, just MaxTransactionsPerBlock included
+
+            var max = (int)NativeContract.Policy.GetMaxTransactionsPerBlock(_context.Snapshot);
+            var txs = new Transaction[max + 1];
+
+            for (int x = 0; x < max; x++) txs[x] = CreateTransactionWithSize(100);
+
+            _context.EnsureMaxBlockSize(txs);
+            EnsureContext(_context, txs.Take(max).ToArray());
+        }
+
+        private Transaction CreateTransactionWithSize(int v)
+        {
+            var r = new Random();
+            var tx = new Transaction()
+            {
+                Cosigners = new Cosigner[0],
+                Attributes = new TransactionAttribute[0],
+                NetworkFee = 0,
+                Nonce = (uint)Environment.TickCount,
+                Script = new byte[0],
+                Sender = UInt160.Zero,
+                SystemFee = 0,
+                ValidUntilBlock = (uint)r.Next(),
+                Version = 0,
+                Witnesses = new Witness[0],
+            };
+
+            // Could be higher (few bytes) if varSize grows
+            tx.Script = new byte[v - tx.Size];
+            return tx;
+        }
+
+        private Block SignBlock(ConsensusContext context)
+        {
+            context.Block.MerkleRoot = null;
+
+            for (int x = 0; x < _validatorKeys.Length; x++)
+            {
+                _context.MyIndex = x;
+
+                var com = _context.MakeCommit();
+                _context.CommitPayloads[_context.MyIndex] = com;
+            }
+
+            // Manual block sign
+
+            Contract contract = Contract.CreateMultiSigContract(context.M, context.Validators);
+            ContractParametersContext sc = new ContractParametersContext(context.Block);
+            for (int i = 0, j = 0; i < context.Validators.Length && j < context.M; i++)
+            {
+                if (context.CommitPayloads[i]?.ConsensusMessage.ViewNumber != context.ViewNumber) continue;
+                sc.AddSignature(contract, context.Validators[i], context.CommitPayloads[i].GetDeserializedMessage<Commit>().Signature);
+                j++;
+            }
+            context.Block.Witness = sc.GetWitnesses()[0];
+            context.Block.Transactions = context.TransactionHashes.Select(p => context.Transactions[p]).ToArray();
+            return context.Block;
+        }
+
+        private void EnsureContext(ConsensusContext context, params Transaction[] expected)
+        {
+            // Check all tx
+
+            Assert.AreEqual(expected.Length, context.Transactions.Count);
+            Assert.IsTrue(expected.All(tx => context.Transactions.ContainsKey(tx.Hash)));
+
+            Assert.AreEqual(expected.Length, context.TransactionHashes.Length);
+            Assert.IsTrue(expected.All(tx => context.TransactionHashes.Count(t => t == tx.Hash) == 1));
+
+            // Ensure length
+
+            var block = SignBlock(context);
+
+            Assert.AreEqual(context.GetExpectedBlockSize(), block.Size);
+            Assert.IsTrue(block.Size < NativeContract.Policy.GetMaxBlockSize(context.Snapshot));
+        }
+    }
+}

neo.UnitTests/SmartContract/Native/UT_PolicyContract.cs

@@ -32,12 +32,16 @@ public void Check_Initialize()
 
             NativeContract.Policy.Initialize(new ApplicationEngine(TriggerType.Application, null, snapshot, 0)).Should().BeTrue();
 
-            (keyCount + 3).Should().Be(snapshot.Storages.GetChangeSet().Count());
+            (keyCount + 4).Should().Be(snapshot.Storages.GetChangeSet().Count());
 
             var ret = NativeContract.Policy.Call(snapshot, "getMaxTransactionsPerBlock");
             ret.Should().BeOfType<VM.Types.Integer>();
             ret.GetBigInteger().Should().Be(512);
 
+            ret = NativeContract.Policy.Call(snapshot, "getMaxBlockSize");
+            ret.Should().BeOfType<VM.Types.Integer>();
+            ret.GetBigInteger().Should().Be(1024 * 256);
+
             ret = NativeContract.Policy.Call(snapshot, "getFeePerByte");
             ret.Should().BeOfType<VM.Types.Integer>();
             ret.GetBigInteger().Should().Be(1000);
@@ -47,6 +51,52 @@ public void Check_Initialize()
             ((VM.Types.Array)ret).Count.Should().Be(0);
         }
 
+        [TestMethod]
+        public void Check_SetMaxBlockSize()
+        {
+            var snapshot = Store.GetSnapshot().Clone();
+
+            // Fake blockchain
+
+            snapshot.PersistingBlock = new Block() { Index = 1000, PrevHash = UInt256.Zero };
+            snapshot.Blocks.Add(UInt256.Zero, new Neo.Ledger.TrimmedBlock() { NextConsensus = UInt160.Zero });
+
+            NativeContract.Policy.Initialize(new ApplicationEngine(TriggerType.Application, null, snapshot, 0)).Should().BeTrue();
+
+            // Without signature
+
+            var ret = NativeContract.Policy.Call(snapshot, new Nep5NativeContractExtensions.ManualWitness(null),
+                "setMaxBlockSize", new ContractParameter(ContractParameterType.Integer) { Value = 1024 });
+            ret.Should().BeOfType<VM.Types.Boolean>();
+            ret.GetBoolean().Should().BeFalse();
+
+            ret = NativeContract.Policy.Call(snapshot, "getMaxBlockSize");
+            ret.Should().BeOfType<VM.Types.Integer>();
+            ret.GetBigInteger().Should().Be(1024 * 256);
+
+            // More than expected
+
+            ret = NativeContract.Policy.Call(snapshot, new Nep5NativeContractExtensions.ManualWitness(UInt160.Zero),
+                 "setMaxBlockSize", new ContractParameter(ContractParameterType.Integer) { Value = Neo.Network.P2P.Message.PayloadMaxSize });
+            ret.Should().BeOfType<VM.Types.Boolean>();
+            ret.GetBoolean().Should().BeFalse();
+
+            ret = NativeContract.Policy.Call(snapshot, "getMaxBlockSize");
+            ret.Should().BeOfType<VM.Types.Integer>();
+            ret.GetBigInteger().Should().Be(1024 * 256);
+
+            // With signature
+
+            ret = NativeContract.Policy.Call(snapshot, new Nep5NativeContractExtensions.ManualWitness(UInt160.Zero),
+                "setMaxBlockSize", new ContractParameter(ContractParameterType.Integer) { Value = 1024 });
+            ret.Should().BeOfType<VM.Types.Boolean>();
+            ret.GetBoolean().Should().BeTrue();
+
+            ret = NativeContract.Policy.Call(snapshot, "getMaxBlockSize");
+            ret.Should().BeOfType<VM.Types.Integer>();
+            ret.GetBigInteger().Should().Be(1024);
+        }
+
         [TestMethod]
         public void Check_SetMaxTransactionsPerBlock()
         {

neo/Consensus/ChangeViewReason.cs

@@ -7,5 +7,6 @@ public enum ChangeViewReason : byte
         TxNotFound = 0x2,
         TxRejectedByPolicy = 0x3,
         TxInvalid = 0x4,
+        BlockRejectedByPolicy = 0x5
     }
 }

neo/Consensus/ConsensusContext.cs

@@ -6,6 +6,7 @@
 using Neo.Persistence;
 using Neo.SmartContract;
 using Neo.SmartContract.Native;
+using Neo.VM;
 using Neo.Wallets;
 using System;
 using System.Collections.Generic;
@@ -38,10 +39,11 @@ internal class ConsensusContext : IDisposable, ISerializable
 
         public Snapshot Snapshot { get; private set; }
         private KeyPair keyPair;
+        private int _witnessSize;
         private readonly Wallet wallet;
         private readonly Store store;
         private readonly Random random = new Random();
-
+
         public int F => (Validators.Length - 1) / 3;
         public int M => Validators.Length - F;
         public bool IsPrimary => MyIndex == Block.ConsensusData.PrimaryIndex;
@@ -203,19 +205,84 @@ private void SignPayload(ConsensusPayload payload)
                 return;
             }
             payload.Witness = sc.GetWitnesses()[0];
+        }
+
+        /// <summary>
+        /// Return the expected block size
+        /// </summary>
+        internal int GetExpectedBlockSize()
+        {
+            return GetExpectedBlockSizeWithoutTransactions(Transactions.Count) + // Base size
+                Transactions.Values.Sum(u => u.Size);   // Sum Txs
+        }
+
+        /// <summary>
+        /// Return the expected block size without txs
+        /// </summary>
+        /// <param name="expectedTransactions">Expected transactions</param>
+        internal int GetExpectedBlockSizeWithoutTransactions(int expectedTransactions)
+        {
+            var blockSize =
+                // BlockBase
+                sizeof(uint) +       //Version
+                UInt256.Length +     //PrevHash
+                UInt256.Length +     //MerkleRoot
+                sizeof(ulong) +      //Timestamp
+                sizeof(uint) +       //Index
+                UInt160.Length +     //NextConsensus
+                1 +                  //
+                _witnessSize;        //Witness
+
+            blockSize +=
+                // Block
+                Block.ConsensusData.Size +                      //ConsensusData
+                IO.Helper.GetVarSize(expectedTransactions + 1); //Transactions count
+
+            return blockSize;
+        }
+
+        /// <summary>
+        /// Prevent that block exceed the max size
+        /// </summary>
+        /// <param name="txs">Ordered transactions</param>
+        internal void EnsureMaxBlockSize(IEnumerable<Transaction> txs)
+        {
+            uint maxBlockSize = NativeContract.Policy.GetMaxBlockSize(Snapshot);
+            uint maxTransactionsPerBlock = NativeContract.Policy.GetMaxTransactionsPerBlock(Snapshot);
+
+            // Limit Speaker proposal to the limit `MaxTransactionsPerBlock` or all available transactions of the mempool
+            txs = txs.Take((int)maxTransactionsPerBlock);
+            List<UInt256> hashes = new List<UInt256>();
+            Transactions = new Dictionary<UInt256, Transaction>();
+            Block.Transactions = new Transaction[0];
+
+            // We need to know the expected block size
+
+            var blockSize = GetExpectedBlockSizeWithoutTransactions(txs.Count());
+
+            // Iterate transaction until reach the size
+
+            foreach (Transaction tx in txs)
+            {
+                // Check if maximum block size has been already exceeded with the current selected set
+                blockSize += tx.Size;
+                if (blockSize > maxBlockSize) break;
+
+                hashes.Add(tx.Hash);
+                Transactions.Add(tx.Hash, tx);
+            }
+
+            TransactionHashes = hashes.ToArray();
         }
 
         public ConsensusPayload MakePrepareRequest()
         {
             byte[] buffer = new byte[sizeof(ulong)];
             random.NextBytes(buffer);
-            List<Transaction> transactions = Blockchain.Singleton.MemPool.GetSortedVerifiedTransactions()
-                .Take((int)NativeContract.Policy.GetMaxTransactionsPerBlock(Snapshot))
-                .ToList();
-            TransactionHashes = transactions.Select(p => p.Hash).ToArray();
-            Transactions = transactions.ToDictionary(p => p.Hash);
-            Block.Timestamp = Math.Max(TimeProvider.Current.UtcNow.ToTimestampMS(), PrevHeader.Timestamp + 1);
             Block.ConsensusData.Nonce = BitConverter.ToUInt64(buffer, 0);
+            EnsureMaxBlockSize(Blockchain.Singleton.MemPool.GetSortedVerifiedTransactions());
+            Block.Timestamp = Math.Max(TimeProvider.Current.UtcNow.ToTimestampMS(), PrevHeader.Timestamp + 1);
+
             return PreparationPayloads[MyIndex] = MakeSignedPayload(new PrepareRequest
             {
                 Timestamp = Block.Timestamp,
@@ -279,7 +346,24 @@ public void Reset(byte viewNumber)
                     NextConsensus = Blockchain.GetConsensusAddress(NativeContract.NEO.GetValidators(Snapshot).ToArray()),
                     ConsensusData = new ConsensusData()
                 };
+                var pv = Validators;
                 Validators = NativeContract.NEO.GetNextBlockValidators(Snapshot);
+                if (_witnessSize == 0 || (pv != null && pv.Length != Validators.Length))
+                {
+                    // Compute the expected size of the witness
+                    using (ScriptBuilder sb = new ScriptBuilder())
+                    {
+                        for (int x = 0; x < M; x++)
+                        {
+                            sb.EmitPush(new byte[64]);
+                        }
+                        _witnessSize = new Witness
+                        {
+                            InvocationScript = sb.ToArray(),
+                            VerificationScript = Contract.CreateMultiSigRedeemScript(M, Validators)
+                        }.Size;
+                    }
+                }
                 MyIndex = -1;
                 ChangeViewPayloads = new ConsensusPayload[Validators.Length];
                 LastChangeViewPayloads = new ConsensusPayload[Validators.Length];

neo/Consensus/ConsensusService.cs

@@ -80,6 +80,14 @@ private bool AddTransaction(Transaction tx, bool verify)
                 // previously sent prepare request, then we don't want to send a prepare response.
                 if (context.IsPrimary || context.WatchOnly) return true;
 
+                // Check maximum block size via Native Contract policy
+                if (context.GetExpectedBlockSize() > NativeContract.Policy.GetMaxBlockSize(context.Snapshot))
+                {
+                    Log($"rejected block: {context.Block.Index}{Environment.NewLine} The size exceed the policy", LogLevel.Warning);
+                    RequestChangeView(ChangeViewReason.BlockRejectedByPolicy);
+                    return false;
+                }
+
                 // Timeout extension due to prepare response sent
                 // around 2*15/M=30.0/5 ~ 40% block time (for M=5)
                 ExtendTimerByFactor(2);