Fixes GetEntryScriptHash and GetCallingScriptHash (2.x) (#138)

neo-project · May 2, 2019 · 38c5dc0 · 38c5dc0
1 parent 0d13635
commit 38c5dc0
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 15 deletions.
diff --git a/src/neo-vm/ExecutionContext.cs b/src/neo-vm/ExecutionContext.cs
@@ -62,15 +62,19 @@ public byte[] ScriptHash
             }
         }
 
+        public byte[] CallingScriptHash { get; }
+
         /// <summary>
         /// Constructor
         /// </summary>
         /// <param name="script">Script</param>
+        /// <param name="callingScriptHash">Script hash of the calling script</param>
         /// <param name="rvcount">Number of items to be returned</param>
-        internal ExecutionContext(Script script, int rvcount)
+        internal ExecutionContext(Script script, byte[] callingScriptHash, int rvcount)
         {
             this.RVCount = rvcount;
             this.Script = script;
+            this.CallingScriptHash = callingScriptHash;
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]

diff --git a/src/neo-vm/ExecutionEngine.cs b/src/neo-vm/ExecutionEngine.cs
@@ -64,8 +64,7 @@ public class ExecutionEngine : IDisposable
         public RandomAccessStack<ExecutionContext> InvocationStack { get; } = new RandomAccessStack<ExecutionContext>();
         public RandomAccessStack<StackItem> ResultStack { get; } = new RandomAccessStack<StackItem>();
         public ExecutionContext CurrentContext => InvocationStack.Peek();
-        public ExecutionContext CallingContext => InvocationStack.Count > 1 ? InvocationStack.Peek(1) : null;
-        public ExecutionContext EntryContext => InvocationStack.Peek(InvocationStack.Count - 1);
+        public byte[] EntryScriptHash { get; private set; }
         public VMState State { get; internal protected set; } = VMState.BREAK;
 
         public ExecutionEngine(IScriptContainer container, ICrypto crypto, IScriptTable table = null, IInteropService service = null)
@@ -285,7 +284,7 @@ private bool ExecuteInstruction()
                     case OpCode.CALL:
                         {
                             if (!CheckMaxInvocationStack()) return false;
-                            ExecutionContext context_call = LoadScript(context.Script);
+                            ExecutionContext context_call = LoadScript(context.Script, context.ScriptHash);
                             context_call.InstructionPointer = context.InstructionPointer + instruction.TokenI16;
                             if (context_call.InstructionPointer < 0 || context_call.InstructionPointer > context_call.Script.Length) return false;
                             context.EvaluationStack.CopyTo(context_call.EvaluationStack);
@@ -328,7 +327,7 @@ private bool ExecuteInstruction()
                             {
                                 script_hash = context.EvaluationStack.Pop().GetByteArray();
                             }
-                            ExecutionContext context_new = LoadScriptByHash(script_hash);
+                            ExecutionContext context_new = LoadScriptByHash(script_hash, context.ScriptHash);
                             if (context_new == null) return false;
                             context.EvaluationStack.CopyTo(context_new.EvaluationStack);
                             if (instruction.OpCode == OpCode.TAILCALL)
@@ -1171,7 +1170,7 @@ private bool ExecuteInstruction()
                             int rvcount = instruction.Operand[0];
                             int pcount = instruction.Operand[1];
                             if (context.EvaluationStack.Count < pcount) return false;
-                            ExecutionContext context_call = LoadScript(context.Script, rvcount);
+                            ExecutionContext context_call = LoadScript(context.Script, context.ScriptHash, rvcount);
                             context_call.InstructionPointer = context.InstructionPointer + instruction.TokenI16_1 + 2;
                             if (context_call.InstructionPointer < 0 || context_call.InstructionPointer > context_call.Script.Length) return false;
                             context.EvaluationStack.CopyTo(context_call.EvaluationStack, pcount);
@@ -1208,7 +1207,7 @@ private bool ExecuteInstruction()
                                 script_hash = instruction.ReadBytes(2, 20);
                             }
 
-                            ExecutionContext context_new = LoadScriptByHash(script_hash, rvcount);
+                            ExecutionContext context_new = LoadScriptByHash(script_hash, context.ScriptHash, rvcount);
                             if (context_new == null) return false;
                             context.EvaluationStack.CopyTo(context_new.EvaluationStack, pcount);
                             if (instruction.OpCode == OpCode.CALL_ET || instruction.OpCode == OpCode.CALL_EDT)
@@ -1238,24 +1237,26 @@ private bool ExecuteInstruction()
             return true;
         }
 
-        public ExecutionContext LoadScript(byte[] script, int rvcount = -1)
+        public ExecutionContext LoadScript(byte[] script, byte[] callingScriptHash = null, int rvcount = -1)
         {
-            return LoadScript(new Script(Crypto, script), rvcount);
+            return LoadScript(new Script(Crypto, script), callingScriptHash, rvcount);
         }
 
-        protected virtual ExecutionContext LoadScript(Script script, int rvcount = -1)
+        protected virtual ExecutionContext LoadScript(Script script, byte[] callingScriptHash = null, int rvcount = -1)
         {
-            ExecutionContext context = new ExecutionContext(script, rvcount);
+            ExecutionContext context = new ExecutionContext(script, callingScriptHash, rvcount);
+            if (EntryScriptHash is null)
+                EntryScriptHash = context.ScriptHash;
             InvocationStack.Push(context);
             return context;
         }
 
-        private ExecutionContext LoadScriptByHash(byte[] hash, int rvcount = -1)
+        private ExecutionContext LoadScriptByHash(byte[] hash, byte[] callingScriptHash, int rvcount = -1)
         {
             if (table == null) return null;
             byte[] script = table.GetScript(hash);
             if (script == null) return null;
-            return LoadScript(new Script(hash, script), rvcount);
+            return LoadScript(new Script(hash, script), callingScriptHash, rvcount);
         }
 
         protected virtual bool PostExecuteInstruction(Instruction instruction)

diff --git a/tests/neo-vm.Tests/Types/InteropService.cs b/tests/neo-vm.Tests/Types/InteropService.cs
@@ -11,12 +11,12 @@ public bool Invoke(byte[] method, ExecutionEngine engine)
             {
                 case "System.ExecutionEngine.GetEntryScriptHash":
                     {
-                        engine.CurrentContext.EvaluationStack.Push(engine.EntryContext.ScriptHash);
+                        engine.CurrentContext.EvaluationStack.Push(engine.EntryScriptHash);
                         return true;
                     }
                 case "System.ExecutionEngine.GetCallingScriptHash":
                     {
-                        engine.CurrentContext.EvaluationStack.Push(engine.CallingContext.ScriptHash);
+                        engine.CurrentContext.EvaluationStack.Push(engine.CurrentContext.CallingScriptHash ?? new byte[0]);
                         return true;
                     }
                 case "System.ExecutionEngine.GetExecutingScriptHash":