浏览代码

Check against long.MaxValue instead of int in memory limit configuration (#2068)

Marko Lahma 7 月之前
父节点
当前提交
a4a51e3183
共有 2 个文件被更改,包括 28 次插入20 次删除
  1. 1 1
      Jint/Constraints/ConstraintsOptionsExtensions.cs
  2. 27 19
      Jint/Constraints/MemoryLimitConstraint.cs

+ 1 - 1
Jint/Constraints/ConstraintsOptionsExtensions.cs

@@ -27,7 +27,7 @@ public static class ConstraintsOptionsExtensions
     {
         options.WithoutConstraint(x => x is MemoryLimitConstraint);
 
-        if (memoryLimit > 0 && memoryLimit < int.MaxValue)
+        if (memoryLimit > 0 && memoryLimit < long.MaxValue)
         {
             options.Constraint(new MemoryLimitConstraint(memoryLimit));
         }

+ 27 - 19
Jint/Constraints/MemoryLimitConstraint.cs

@@ -4,19 +4,22 @@ namespace Jint.Constraints;
 
 public sealed class MemoryLimitConstraint : Constraint
 {
-    private static readonly Func<long>? GetAllocatedBytesForCurrentThread;
     private readonly long _memoryLimit;
     private long _initialMemoryUsage;
 
+#if !NET8_0_OR_GREATER
+    private static readonly Func<long>? _getAllocatedBytesForCurrentThread;
+
     static MemoryLimitConstraint()
     {
         var methodInfo = typeof(GC).GetMethod("GetAllocatedBytesForCurrentThread");
 
         if (methodInfo != null)
         {
-            GetAllocatedBytesForCurrentThread = (Func<long>)Delegate.CreateDelegate(typeof(Func<long>), null, methodInfo);
+            _getAllocatedBytesForCurrentThread = (Func<long>) Delegate.CreateDelegate(typeof(Func<long>), null, methodInfo);
         }
     }
+#endif
 
     internal MemoryLimitConstraint(long memoryLimit)
     {
@@ -25,28 +28,33 @@ public sealed class MemoryLimitConstraint : Constraint
 
     public override void Check()
     {
-        if (_memoryLimit > 0)
+        if (_memoryLimit <= 0)
         {
-            if (GetAllocatedBytesForCurrentThread != null)
-            {
-                var memoryUsage = GetAllocatedBytesForCurrentThread() - _initialMemoryUsage;
-                if (memoryUsage > _memoryLimit)
-                {
-                    ExceptionHelper.ThrowMemoryLimitExceededException($"Script has allocated {memoryUsage} but is limited to {_memoryLimit}");
-                }
-            }
-            else
-            {
-                ExceptionHelper.ThrowPlatformNotSupportedException("The current platform doesn't support MemoryLimit.");
-            }
+            return;
+        }
+
+#if NET8_0_OR_GREATER
+        var usage = GC.GetAllocatedBytesForCurrentThread();
+#else
+        if (_getAllocatedBytesForCurrentThread == null)
+        {
+            ExceptionHelper.ThrowPlatformNotSupportedException("The current platform doesn't support MemoryLimit.");
+        }
+
+        var usage = _getAllocatedBytesForCurrentThread();
+#endif
+        if (usage - _initialMemoryUsage > _memoryLimit)
+        {
+            ExceptionHelper.ThrowMemoryLimitExceededException($"Script has allocated {usage - _initialMemoryUsage} but is limited to {_memoryLimit}");
         }
     }
 
     public override void Reset()
     {
-        if (GetAllocatedBytesForCurrentThread != null)
-        {
-            _initialMemoryUsage = GetAllocatedBytesForCurrentThread();
-        }
+#if NET8_0_OR_GREATER
+        _initialMemoryUsage = GC.GetAllocatedBytesForCurrentThread();
+#else
+        _initialMemoryUsage = _getAllocatedBytesForCurrentThread?.Invoke() ?? 0;
+#endif
     }
 }