Improve ApplicationEngine (#2796)

1. Added `CallingContext` to `ExecutionContextState`. This fixes a potential issue in `CalledByEntryCondition`. 2. Improved the return logic of dynamic calls. 3. Disallowed notifications in entry script.
neo-project · Aug 2, 2022 · 718fc9f · 718fc9f
1 parent 8e68c3f
commit 718fc9f
Show file tree

Hide file tree

Showing 8 changed files with 41 additions and 14 deletions.
diff --git a/src/Neo/Network/P2P/Payloads/Conditions/CalledByEntryCondition.cs b/src/Neo/Network/P2P/Payloads/Conditions/CalledByEntryCondition.cs
@@ -24,7 +24,10 @@ protected override void DeserializeWithoutType(ref MemoryReader reader, int maxN
 
         public override bool Match(ApplicationEngine engine)
         {
-            return engine.CallingScriptHash is null || engine.CallingScriptHash == engine.EntryScriptHash;
+            var state = engine.CurrentContext.GetState<ExecutionContextState>();
+            if (state.CallingContext is null) return true;
+            state = state.CallingContext.GetState<ExecutionContextState>();
+            return state.CallingContext is null;
         }
 
         protected override void SerializeWithoutType(BinaryWriter writer)

diff --git a/src/Neo/SmartContract/ApplicationEngine.Contract.cs b/src/Neo/SmartContract/ApplicationEngine.Contract.cs
@@ -82,7 +82,7 @@ protected internal void CallContract(UInt160 contractHash, string method, CallFl
             bool hasReturnValue = md.ReturnType != ContractParameterType.Void;
 
             ExecutionContext context = CallContractInternal(contract, md, callFlags, hasReturnValue, args);
-            if (!hasReturnValue) context.GetState<ExecutionContextState>().PushNullWhenReturn = true;
+            context.GetState<ExecutionContextState>().IsDynamicCall = true;
         }
 
         /// <summary>

diff --git a/src/Neo/SmartContract/ApplicationEngine.Runtime.cs b/src/Neo/SmartContract/ApplicationEngine.Runtime.cs
@@ -306,6 +306,8 @@ protected internal void RuntimeLog(byte[] state)
         protected internal void RuntimeNotify(byte[] eventName, Array state)
         {
             if (eventName.Length > MaxEventName) throw new ArgumentException(null, nameof(eventName));
+            if (CurrentContext.GetState<ExecutionContextState>().Contract is null)
+                throw new InvalidOperationException("Notifications are not allowed in dynamic scripts.");
             using MemoryStream ms = new(MaxNotificationSize);
             using BinaryWriter writer = new(ms, Utility.StrictUTF8, true);
             BinarySerializer.Serialize(writer, state, MaxNotificationSize);

diff --git a/src/Neo/SmartContract/ApplicationEngine.cs b/src/Neo/SmartContract/ApplicationEngine.cs
@@ -123,7 +123,15 @@ public partial class ApplicationEngine : ExecutionEngine
         /// <summary>
         /// The script hash of the calling contract. This field could be <see langword="null"/> if the current context is the entry context.
         /// </summary>
-        public UInt160 CallingScriptHash => CurrentContext?.GetState<ExecutionContextState>().CallingScriptHash;
+        public UInt160 CallingScriptHash
+        {
+            get
+            {
+                if (CurrentContext is null) return null;
+                var state = CurrentContext.GetState<ExecutionContextState>();
+                return state.NativeCallingScriptHash ?? state.CallingContext?.GetState<ExecutionContextState>().ScriptHash;
+            }
+        }
 
         /// <summary>
         /// The script hash of the entry context. This field could be <see langword="null"/> if no context is loaded to the engine.
@@ -224,15 +232,15 @@ private ExecutionContext CallContractInternal(ContractState contract, ContractMe
                 invocationCounter[contract.Hash] = 1;
             }
 
-            ExecutionContextState state = CurrentContext.GetState<ExecutionContextState>();
-            UInt160 callingScriptHash = state.ScriptHash;
+            ExecutionContext currentContext = CurrentContext;
+            ExecutionContextState state = currentContext.GetState<ExecutionContextState>();
             CallFlags callingFlags = state.CallFlags;
 
             if (args.Count != method.Parameters.Length) throw new InvalidOperationException($"Method {method} Expects {method.Parameters.Length} Arguments But Receives {args.Count} Arguments");
             if (hasReturnValue ^ (method.ReturnType != ContractParameterType.Void)) throw new InvalidOperationException("The return value type does not match.");
             ExecutionContext context_new = LoadContract(contract, method, flags & callingFlags);
             state = context_new.GetState<ExecutionContextState>();
-            state.CallingScriptHash = callingScriptHash;
+            state.CallingContext = currentContext;
 
             for (int i = args.Count - 1; i >= 0; i--)
                 context_new.EvaluationStack.Push(args[i]);
@@ -244,7 +252,7 @@ internal ContractTask CallFromNativeContract(UInt160 callingScriptHash, UInt160
         {
             ExecutionContext context_new = CallContractInternal(hash, method, CallFlags.All, false, args);
             ExecutionContextState state = context_new.GetState<ExecutionContextState>();
-            state.CallingScriptHash = callingScriptHash;
+            state.NativeCallingScriptHash = callingScriptHash;
             ContractTask task = new();
             contractTasks.Add(context_new, task.GetAwaiter());
             return task;
@@ -254,7 +262,7 @@ internal ContractTask<T> CallFromNativeContract<T>(UInt160 callingScriptHash, UI
         {
             ExecutionContext context_new = CallContractInternal(hash, method, CallFlags.All, true, args);
             ExecutionContextState state = context_new.GetState<ExecutionContextState>();
-            state.CallingScriptHash = callingScriptHash;
+            state.NativeCallingScriptHash = callingScriptHash;
             ContractTask<T> task = new();
             contractTasks.Add(context_new, task.GetAwaiter());
             return task;
@@ -273,7 +281,13 @@ protected override void ContextUnloaded(ExecutionContext context)
                     {
                         ExecutionContextState contextState = CurrentContext.GetState<ExecutionContextState>();
                         contextState.NotificationCount += state.NotificationCount;
-                        if (state.PushNullWhenReturn) Push(StackItem.Null);
+                        if (state.IsDynamicCall)
+                        {
+                            if (context.EvaluationStack.Count == 0)
+                                Push(StackItem.Null);
+                            else if (context.EvaluationStack.Count > 1)
+                                throw new NotSupportedException("Multiple return values are not allowed in cross-contract calls.");
+                        }
                     }
                 }
                 else

diff --git a/src/Neo/SmartContract/ExecutionContextState.cs b/src/Neo/SmartContract/ExecutionContextState.cs
@@ -24,9 +24,14 @@ public class ExecutionContextState
         public UInt160 ScriptHash { get; set; }
 
         /// <summary>
-        /// The script hash of the calling contract.
+        /// The calling context.
         /// </summary>
-        public UInt160 CallingScriptHash { get; set; }
+        public ExecutionContext CallingContext { get; set; }
+
+        /// <summary>
+        /// The script hash of the calling native contract. Used in native contracts only.
+        /// </summary>
+        internal UInt160 NativeCallingScriptHash { get; set; }
 
         /// <summary>
         /// The <see cref="ContractState"/> of the current context.
@@ -42,6 +47,6 @@ public class ExecutionContextState
 
         public int NotificationCount { get; set; }
 
-        public bool PushNullWhenReturn { get; set; }
+        public bool IsDynamicCall { get; set; }
     }
 }
diff --git a/tests/Neo.UnitTests/Extensions/NativeContractExtensions.cs b/tests/Neo.UnitTests/Extensions/NativeContractExtensions.cs
@@ -42,7 +42,7 @@ public static void UpdateContract(this DataCache snapshot, UInt160 callingScript
             // Fake calling script hash
             if (callingScriptHash != null)
             {
-                engine.CurrentContext.GetState<ExecutionContextState>().CallingScriptHash = callingScriptHash;
+                engine.CurrentContext.GetState<ExecutionContextState>().NativeCallingScriptHash = callingScriptHash;
                 engine.CurrentContext.GetState<ExecutionContextState>().ScriptHash = callingScriptHash;
             }
 
@@ -65,7 +65,7 @@ public static void DestroyContract(this DataCache snapshot, UInt160 callingScrip
             // Fake calling script hash
             if (callingScriptHash != null)
             {
-                engine.CurrentContext.GetState<ExecutionContextState>().CallingScriptHash = callingScriptHash;
+                engine.CurrentContext.GetState<ExecutionContextState>().NativeCallingScriptHash = callingScriptHash;
                 engine.CurrentContext.GetState<ExecutionContextState>().ScriptHash = callingScriptHash;
             }
 

diff --git a/tests/Neo.UnitTests/SmartContract/UT_ApplicationEngine.Runtime.cs b/tests/Neo.UnitTests/SmartContract/UT_ApplicationEngine.Runtime.cs
@@ -24,6 +24,7 @@ public void TestNotSupportedNotification()
         {
             using var engine = ApplicationEngine.Create(TriggerType.Application, null, null, TestBlockchain.TheNeoSystem.GenesisBlock, settings: TestBlockchain.TheNeoSystem.Settings, gas: 1100_00000000);
             engine.LoadScript(Array.Empty<byte>());
+            engine.CurrentContext.GetState<ExecutionContextState>().Contract = new();
 
             // circular
 

diff --git a/tests/Neo.UnitTests/SmartContract/UT_InteropService.cs b/tests/Neo.UnitTests/SmartContract/UT_InteropService.cs
@@ -93,6 +93,7 @@ public void Runtime_GetNotifications_Test()
                 // Execute
 
                 engine.LoadScript(script.ToArray());
+                engine.CurrentContext.GetState<ExecutionContextState>().Contract = new();
                 var currentScriptHash = engine.EntryScriptHash;
 
                 Assert.AreEqual(VMState.HALT, engine.Execute());
@@ -145,6 +146,7 @@ public void Runtime_GetNotifications_Test()
                 // Execute
 
                 engine.LoadScript(script.ToArray());
+                engine.CurrentContext.GetState<ExecutionContextState>().Contract = new();
                 var currentScriptHash = engine.EntryScriptHash;
 
                 Assert.AreEqual(VMState.HALT, engine.Execute());