From f37638b46b829edb3bee5b164ec708e6308681af Mon Sep 17 00:00:00 2001
From: Shargon <shargon@gmail.com>
Date: Sat, 3 Oct 2020 16:08:12 +0200
Subject: [PATCH] Fix Witness limits (#1958)

* Fix Witness max size

* Fix comment

* Fix comment

* Change it to use constants

* Changed to 1024

* Fix UT
---
 src/neo/Network/P2P/Payloads/Witness.cs       | 18 ++++++++---
 .../Network/P2P/Payloads/UT_Witness.cs        | 31 ++++++++++++-------
 2 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/src/neo/Network/P2P/Payloads/Witness.cs b/src/neo/Network/P2P/Payloads/Witness.cs
index d54158050d..7ae2d51d70 100644
--- a/src/neo/Network/P2P/Payloads/Witness.cs
+++ b/src/neo/Network/P2P/Payloads/Witness.cs
@@ -8,6 +8,17 @@ namespace Neo.Network.P2P.Payloads
 {
     public class Witness : ISerializable
     {
+        /// <summary>
+        /// This is designed to allow a MultiSig 21/11 (committee)
+        /// Invocation = 11 * (64 + 2) = 726
+        /// </summary>
+        private const int MaxInvocationScript = 1024;
+
+        /// <summary>
+        /// Verification = m + (PUSH_PubKey * 21) + length + null + syscall = 1 + ((2 + 33) * 21) + 2 + 1 + 5 = 744
+        /// </summary>
+        private const int MaxVerificationScript = 1024;
+
         public byte[] InvocationScript;
         public byte[] VerificationScript;
 
@@ -32,11 +43,8 @@ public virtual UInt160 ScriptHash
 
         void ISerializable.Deserialize(BinaryReader reader)
         {
-            // This is designed to allow a MultiSig 10/10 (around 1003 bytes) ~1024 bytes
-            // Invocation = 10 * 64 + 10 = 650 ~ 664  (exact is 653)
-            InvocationScript = reader.ReadVarBytes(663);
-            // Verification = 10 * 33 + 10 = 340 ~ 360   (exact is 351)
-            VerificationScript = reader.ReadVarBytes(361);
+            InvocationScript = reader.ReadVarBytes(MaxInvocationScript);
+            VerificationScript = reader.ReadVarBytes(MaxVerificationScript);
         }
 
         void ISerializable.Serialize(BinaryWriter writer)
diff --git a/tests/neo.UnitTests/Network/P2P/Payloads/UT_Witness.cs b/tests/neo.UnitTests/Network/P2P/Payloads/UT_Witness.cs
index ff47beef29..9681276a1c 100644
--- a/tests/neo.UnitTests/Network/P2P/Payloads/UT_Witness.cs
+++ b/tests/neo.UnitTests/Network/P2P/Payloads/UT_Witness.cs
@@ -28,13 +28,13 @@ public void InvocationScript_Get()
             uut.InvocationScript.Should().BeNull();
         }
 
-        private Witness PrepareDummyWitness(int maxAccounts)
+        private Witness PrepareDummyWitness(int pubKeys, int m)
         {
-            var address = new WalletAccount[maxAccounts];
-            var wallets = new NEP6Wallet[maxAccounts];
-            var walletsUnlocks = new IDisposable[maxAccounts];
+            var address = new WalletAccount[pubKeys];
+            var wallets = new NEP6Wallet[pubKeys];
+            var walletsUnlocks = new IDisposable[pubKeys];
 
-            for (int x = 0; x < maxAccounts; x++)
+            for (int x = 0; x < pubKeys; x++)
             {
                 wallets[x] = TestUtils.GenerateTestWallet();
                 walletsUnlocks[x] = wallets[x].Unlock("123");
@@ -43,9 +43,9 @@ private Witness PrepareDummyWitness(int maxAccounts)
 
             // Generate multisignature
 
-            var multiSignContract = Contract.CreateMultiSigContract(maxAccounts, address.Select(a => a.GetKey().PublicKey).ToArray());
+            var multiSignContract = Contract.CreateMultiSigContract(m, address.Select(a => a.GetKey().PublicKey).ToArray());
 
-            for (int x = 0; x < maxAccounts; x++)
+            for (int x = 0; x < pubKeys; x++)
             {
                 wallets[x].CreateAccount(multiSignContract, address[x].GetKey());
             }
@@ -69,7 +69,7 @@ private Witness PrepareDummyWitness(int maxAccounts)
                 Witnesses = new Witness[0]
             });
 
-            for (int x = 0; x < maxAccounts; x++)
+            for (int x = 0; x < m; x++)
             {
                 Assert.IsTrue(wallets[x].Sign(data));
             }
@@ -81,7 +81,7 @@ private Witness PrepareDummyWitness(int maxAccounts)
         [TestMethod]
         public void MaxSize_OK()
         {
-            var witness = PrepareDummyWitness(10);
+            var witness = PrepareDummyWitness(10, 10);
 
             // Check max size
 
@@ -100,11 +100,20 @@ public void MaxSize_OK()
         [TestMethod]
         public void MaxSize_Error()
         {
-            var witness = PrepareDummyWitness(11);
+            var witness = new Witness
+            {
+                InvocationScript = new byte[1025],
+                VerificationScript = new byte[10]
+            };
+
+            // Check max size
+
+            Assert.ThrowsException<FormatException>(() => witness.ToArray().AsSerializable<Witness>());
 
             // Check max size
 
-            Assert.IsTrue(witness.Size > 1024);
+            witness.InvocationScript = new byte[10];
+            witness.VerificationScript = new byte[1025];
             Assert.ThrowsException<FormatException>(() => witness.ToArray().AsSerializable<Witness>());
         }