Add StrictMode (#392)

neo-project · Jan 28, 2021 · df2cc97 · df2cc97
1 parent eaf8d30
commit df2cc97
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 5 deletions.
diff --git a/src/neo-vm/BadScriptException.cs b/src/neo-vm/BadScriptException.cs
@@ -0,0 +1,10 @@
+using System;
+
+namespace Neo.VM
+{
+    public class BadScriptException : Exception
+    {
+        public BadScriptException() { }
+        public BadScriptException(string message) : base(message) { }
+    }
+}
diff --git a/src/neo-vm/Instruction.cs b/src/neo-vm/Instruction.cs
@@ -134,12 +134,12 @@ static Instruction()
 
         private Instruction(OpCode opcode)
         {
-            OpCode = opcode;
+            this.OpCode = opcode;
+            if (!Enum.IsDefined(opcode)) throw new BadScriptException();
         }
 
-        internal Instruction(byte[] script, int ip)
+        internal Instruction(byte[] script, int ip) : this((OpCode)script[ip++])
         {
-            OpCode = (OpCode)script[ip++];
             int operandSizePrefix = OperandSizePrefixTable[(int)OpCode];
             int operandSize = 0;
             switch (operandSizePrefix)

diff --git a/src/neo-vm/Script.cs b/src/neo-vm/Script.cs
@@ -1,3 +1,5 @@
+using Neo.VM.Types;
+using System;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Runtime.CompilerServices;
@@ -8,6 +10,7 @@ namespace Neo.VM
     public class Script
     {
         private readonly byte[] _value;
+        private readonly bool strictMode;
         private readonly Dictionary<int, Instruction> _instructions = new Dictionary<int, Instruction>();
 
         /// <summary>
@@ -40,16 +43,82 @@ public OpCode this[int index]
         /// Constructor
         /// </summary>
         /// <param name="script">Script</param>
-        public Script(byte[] script)
+        public Script(byte[] script) : this(script, false)
         {
-            _value = script;
+        }
+
+        public Script(byte[] script, bool strictMode)
+        {
+            this._value = script;
+            if (strictMode)
+            {
+                for (int ip = 0; ip < script.Length; ip += GetInstruction(ip).Size) { }
+                foreach (var (ip, instruction) in _instructions)
+                {
+                    switch (instruction.OpCode)
+                    {
+                        case OpCode.JMP:
+                        case OpCode.JMPIF:
+                        case OpCode.JMPIFNOT:
+                        case OpCode.JMPEQ:
+                        case OpCode.JMPNE:
+                        case OpCode.JMPGT:
+                        case OpCode.JMPGE:
+                        case OpCode.JMPLT:
+                        case OpCode.JMPLE:
+                        case OpCode.CALL:
+                        case OpCode.ENDTRY:
+                            if (!_instructions.ContainsKey(checked(ip + instruction.TokenI8)))
+                                throw new BadScriptException($"ip: {ip}, opcode: {instruction.OpCode}");
+                            break;
+                        case OpCode.PUSHA:
+                        case OpCode.JMP_L:
+                        case OpCode.JMPIF_L:
+                        case OpCode.JMPIFNOT_L:
+                        case OpCode.JMPEQ_L:
+                        case OpCode.JMPNE_L:
+                        case OpCode.JMPGT_L:
+                        case OpCode.JMPGE_L:
+                        case OpCode.JMPLT_L:
+                        case OpCode.JMPLE_L:
+                        case OpCode.CALL_L:
+                        case OpCode.ENDTRY_L:
+                            if (!_instructions.ContainsKey(checked(ip + instruction.TokenI32)))
+                                throw new BadScriptException($"ip: {ip}, opcode: {instruction.OpCode}");
+                            break;
+                        case OpCode.TRY:
+                            if (!_instructions.ContainsKey(checked(ip + instruction.TokenI8)))
+                                throw new BadScriptException($"ip: {ip}, opcode: {instruction.OpCode}");
+                            if (!_instructions.ContainsKey(checked(ip + instruction.TokenI8_1)))
+                                throw new BadScriptException($"ip: {ip}, opcode: {instruction.OpCode}");
+                            break;
+                        case OpCode.TRY_L:
+                            if (!_instructions.ContainsKey(checked(ip + instruction.TokenI32)))
+                                throw new BadScriptException($"ip: {ip}, opcode: {instruction.OpCode}");
+                            if (!_instructions.ContainsKey(checked(ip + instruction.TokenI32_1)))
+                                throw new BadScriptException($"ip: {ip}, opcode: {instruction.OpCode}");
+                            break;
+                        case OpCode.NEWARRAY_T:
+                        case OpCode.ISTYPE:
+                        case OpCode.CONVERT:
+                            StackItemType type = (StackItemType)instruction.TokenU8;
+                            if (!Enum.IsDefined(typeof(StackItemType), type))
+                                throw new BadScriptException();
+                            if (instruction.OpCode != OpCode.NEWARRAY_T && type == StackItemType.Any)
+                                throw new BadScriptException($"ip: {ip}, opcode: {instruction.OpCode}");
+                            break;
+                    }
+                }
+            }
+            this.strictMode = strictMode;
         }
 
         public Instruction GetInstruction(int ip)
         {
             if (ip >= Length) return Instruction.RET;
             if (!_instructions.TryGetValue(ip, out Instruction instruction))
             {
+                if (strictMode) throw new ArgumentException($"ip not found with strict mode", nameof(ip));
                 instruction = new Instruction(_value, ip);
                 _instructions.Add(ip, instruction);
             }

diff --git a/tests/neo-vm.Tests/UtScript.cs b/tests/neo-vm.Tests/UtScript.cs
@@ -27,6 +27,16 @@ public void Conversion()
             CollectionAssert.AreEqual(rawScript, scriptConversion);
         }
 
+        [TestMethod]
+        public void StrictMode()
+        {
+            var rawScript = new byte[] { (byte)OpCode.PUSH0, 0xFF };
+            Assert.ThrowsException<BadScriptException>(() => new Script(rawScript, true));
+
+            var script = new Script(rawScript, false);
+            Assert.AreEqual(2, script.Length);
+        }
+
         [TestMethod]
         public void Parse()
         {