Browse Source

Add memory usage limit option to Engine (#482)

Fixes #480
Nihal Talur 7 years ago
parent
commit
f90d6c1947

+ 8 - 0
Jint.Tests/Runtime/EngineTests.cs

@@ -647,6 +647,14 @@ namespace Jint.Tests.Runtime
             );
         }
 
+        [Fact]
+        public void ShouldThrowMemoryLimitExceeded()
+        {
+            Assert.Throws<MemoryLimitExceededException>(
+                () => new Engine(cfg => cfg.LimitMemory(2048)).Execute("a=[]; while(true){ a.push(0); }")
+            );
+        }
+
         [Fact]
         public void ShouldThrowTimeout()
         {

+ 45 - 2
Jint/Engine.cs

@@ -42,6 +42,7 @@ namespace Jint
         private readonly ExecutionContextStack _executionContexts;
         private JsValue _completionValue = JsValue.Undefined;
         private int _statementsCount;
+        private long _initialMemoryUsage;
         private long _timeoutTicks;
         private INode _lastSyntaxNode;
 
@@ -79,6 +80,16 @@ namespace Jint
 
         internal JintCallStack CallStack = new JintCallStack();
 
+        static Engine()
+        {
+            var methodInfo = typeof(GC).GetMethod("GetAllocatedBytesForCurrentThread");
+
+            if (methodInfo != null)
+            {
+                GetAllocatedBytesForCurrentThread =  (Func<long>)Delegate.CreateDelegate(typeof(Func<long>), null, methodInfo);
+            }
+        }
+
         public Engine() : this(null)
         {
         }
@@ -163,7 +174,7 @@ namespace Jint
             // gather some options as fields for faster checks
             _isDebugMode = Options.IsDebugMode;
             _isStrict = Options.IsStrict;
-            _maxStatements = Options.MaxStatementCount;
+            _maxStatements = Options._MaxStatements;
             _referenceResolver = Options.ReferenceResolver;
 
             ReferencePool = new ReferencePool();
@@ -251,6 +262,8 @@ namespace Jint
         }
         #endregion
 
+        private static readonly Func<long> GetAllocatedBytesForCurrentThread;
+
         public void EnterExecutionContext(
             LexicalEnvironment lexicalEnvironment,
             LexicalEnvironment variableEnvironment,
@@ -314,6 +327,14 @@ namespace Jint
             _statementsCount = 0;
         }
 
+        public void ResetMemoryUsage()
+        {
+            if (GetAllocatedBytesForCurrentThread != null)
+            {
+                _initialMemoryUsage = GetAllocatedBytesForCurrentThread();
+            }
+        }
+
         public void ResetTimeoutTicks()
         {
             var timeoutIntervalTicks = Options._TimeoutInterval.Ticks;
@@ -342,6 +363,12 @@ namespace Jint
         public Engine Execute(Program program)
         {
             ResetStatementsCount();
+            
+            if (Options._MemoryLimit > 0)
+            {
+                ResetMemoryUsage();
+            }
+            
             ResetTimeoutTicks();
             ResetLastStatement();
             ResetCallStack();
@@ -388,6 +415,22 @@ namespace Jint
                 ThrowTimeoutException();
             }
 
+            if (Options._MemoryLimit > 0)
+            {
+                if (GetAllocatedBytesForCurrentThread != null)
+                {
+                    var memoryUsage = GetAllocatedBytesForCurrentThread() - _initialMemoryUsage;
+                    if (memoryUsage > Options._MemoryLimit)
+                    {
+                        throw new MemoryLimitExceededException($"Script has allocated {memoryUsage} but is limited to {Options._MemoryLimit}");
+                    }
+                }
+                else
+                {
+                    throw new PlatformNotSupportedException("The current platform doesn't support MemoryLimit.");
+                }
+            }
+
             _lastSyntaxNode = statement;
 
             if (_isDebugMode)
@@ -980,4 +1023,4 @@ namespace Jint
             throw new JavaScriptException(TypeError);
         }
     }
-}
+}

+ 1 - 1
Jint/Jint.csproj

@@ -9,4 +9,4 @@
   <ItemGroup>
     <PackageReference Include="Esprima" Version="1.0.0-beta-1026" />
   </ItemGroup>
-</Project>
+</Project>

+ 14 - 1
Jint/Options.cs

@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.Globalization;
 using System.Linq;
 using System.Reflection;
+using System.Runtime.CompilerServices;
 using Jint.Native;
 using Jint.Runtime.Interop;
 
@@ -16,6 +17,7 @@ namespace Jint
         private bool _allowClr;
         private readonly List<IObjectConverter> _objectConverters = new List<IObjectConverter>();
         private int _maxStatements;
+        private long _memoryLimit;
         private int _maxRecursionDepth = -1;
         private TimeSpan _timeoutInterval;
         private CultureInfo _culture = CultureInfo.CurrentCulture;
@@ -112,6 +114,11 @@ namespace Jint
             _maxStatements = maxStatements;
             return this;
         }
+        public Options LimitMemory(long memoryLimit)
+        {
+            _memoryLimit = memoryLimit;
+            return this;
+        }
 
         public Options TimeoutInterval(TimeSpan timeoutInterval)
         {
@@ -168,7 +175,13 @@ namespace Jint
 
         internal List<IObjectConverter> _ObjectConverters => _objectConverters;
 
-        internal int MaxStatementCount => _maxStatements;
+        internal long _MemoryLimit => _memoryLimit;
+
+        internal int _MaxStatements
+        {
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            get { return _maxStatements; }
+        }
 
         internal int MaxRecursionDepth => _maxRecursionDepth;
 

+ 15 - 0
Jint/Runtime/MemoryLimitExceededException.cs

@@ -0,0 +1,15 @@
+using System;
+
+namespace Jint.Runtime
+{
+    public class MemoryLimitExceededException : Exception
+    {
+        public MemoryLimitExceededException() : base()
+        {
+        }
+
+        public MemoryLimitExceededException(string message) : base(message)
+        {
+        }
+    }
+}