Fix disposables in ApplicationEngine (neo-project#2191)

cloud8little · Jan 24, 2021 · 8859ab1 · 8859ab1
1 parent 736a4e1
commit 8859ab1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/src/neo/SmartContract/ApplicationEngine.Storage.cs b/src/neo/SmartContract/ApplicationEngine.Storage.cs
@@ -65,17 +65,13 @@ protected internal byte[] Get(StorageContext context, byte[] key)
         protected internal IIterator Find(StorageContext context, byte[] prefix)
         {
             byte[] prefix_key = StorageKey.CreateSearchPrefix(context.Id, prefix);
-            StorageIterator iterator = new StorageIterator(Snapshot.Storages.Find(prefix_key).GetEnumerator());
-            Disposables.Add(iterator);
-            return iterator;
+            return new StorageIterator(Snapshot.Storages.Find(prefix_key).GetEnumerator());
         }
 
         protected internal IEnumerator FindKeys(StorageContext context, byte[] prefix, byte removePrefix)
         {
             byte[] prefix_key = StorageKey.CreateSearchPrefix(context.Id, prefix);
-            StorageKeyEnumerator enumerator = new StorageKeyEnumerator(Snapshot.Storages.Find(prefix_key).Select(p => p.Key).GetEnumerator(), removePrefix);
-            Disposables.Add(enumerator);
-            return enumerator;
+            return new StorageKeyEnumerator(Snapshot.Storages.Find(prefix_key).Select(p => p.Key).GetEnumerator(), removePrefix);
         }
 
         protected internal void Put(StorageContext context, byte[] key, byte[] value)

diff --git a/src/neo/SmartContract/ApplicationEngine.cs b/src/neo/SmartContract/ApplicationEngine.cs
@@ -144,6 +144,7 @@ public ExecutionContext LoadScript(Script script, ushort pcount = 0, int rvcount
 
         protected internal StackItem Convert(object value)
         {
+            if (value is IDisposable disposable) Disposables.Add(disposable);
             return value switch
             {
                 null => StackItem.Null,