Limit size of memory buffer when reading request (#304)

Author: mikeharder

src/Microsoft.AspNetCore.Server.Kestrel/Http/Connection.cs

@@ -46,7 +46,26 @@ public Connection(ListenerContext context, UvStreamHandle socket) : base(context
 
             ConnectionId = GenerateConnectionId(Interlocked.Increment(ref _lastConnectionId));
 
-            _rawSocketInput = new SocketInput(Memory, ThreadPool);
+            var maxInputBufferLength = context.ServerInformation.MaxInputBufferLength;
+            if (maxInputBufferLength < -1 || maxInputBufferLength == 0)
+            {
+                var paramName = String.Join(".",
+                    nameof(context),
+                    nameof(context.ServerInformation),
+                    nameof(context.ServerInformation.MaxInputBufferLength));
+
+                throw new ArgumentOutOfRangeException(paramName, maxInputBufferLength, $"{paramName} must be positive or -1.");
+            }
+
+            if (maxInputBufferLength == -1)
+            {
+                _rawSocketInput = new SocketInput(Memory, ThreadPool);
+            }
+            else
+            {
+                _rawSocketInput = new SocketInput(Memory, ThreadPool, maxInputBufferLength, this, Thread);
+            }
+
             _rawSocketOutput = new SocketOutput(Thread, _socket, Memory, this, ConnectionId, Log, ThreadPool, WriteReqPool);
         }
 

src/Microsoft.AspNetCore.Server.Kestrel/Http/SocketInput.cs

@@ -29,13 +29,21 @@ public class SocketInput : ICriticalNotifyCompletion, IDisposable
 
         private int _consumingState;
 
+        private readonly BufferLengthConnectionController _bufferLengthConnectionController;
+
         public SocketInput(MemoryPool memory, IThreadPool threadPool)
         {
             _memory = memory;
             _threadPool = threadPool;
             _awaitableState = _awaitableIsNotCompleted;
         }
 
+        public SocketInput(MemoryPool memory, IThreadPool threadPool, long maxBufferLength, IConnectionControl connectionControl,
+            KestrelThread connectionThread) : this(memory, threadPool)
+        {
+            _bufferLengthConnectionController = new BufferLengthConnectionController(maxBufferLength, connectionControl, connectionThread);
+        }
+
         public bool RemoteIntakeFin { get; set; }
 
         public bool IsCompleted => ReferenceEquals(_awaitableState, _awaitableIsCompleted);
@@ -58,6 +66,9 @@ public MemoryPoolBlock IncomingStart()
 
         public void IncomingData(byte[] buffer, int offset, int count)
         {
+            // Must call Add() before bytes are available to consumer, to ensure that Length is >= 0
+            _bufferLengthConnectionController?.Add(count);
+
             if (count > 0)
             {
                 if (_tail == null)
@@ -85,6 +96,9 @@ public void IncomingData(byte[] buffer, int offset, int count)
 
         public void IncomingComplete(int count, Exception error)
         {
+            // Must call Add() before bytes are available to consumer, to ensure that Length is >= 0
+            _bufferLengthConnectionController?.Add(count);
+
             if (_pinned != null)
             {
                 _pinned.End += count;
@@ -167,10 +181,16 @@ public void ConsumingComplete(
 
             if (!consumed.IsDefault)
             {
+                var lengthConsumed = new MemoryPoolIterator(_head).GetLength(consumed);
+
                 returnStart = _head;
                 returnEnd = consumed.Block;
                 _head = consumed.Block;
                 _head.Start = consumed.Index;
+
+                // Must call Subtract() after bytes have been freed, to avoid producer starting too early and growing
+                // buffer beyond max length.
+                _bufferLengthConnectionController?.Subtract(lengthConsumed);
             }
 
             if (!examined.IsDefault &&
@@ -285,5 +305,72 @@ public void Dispose()
             _head = null;
             _tail = null;
         }
+
+        private class BufferLengthConnectionController
+        {
+            private readonly long _maxLength;
+            private readonly IConnectionControl _connectionControl;
+            private readonly KestrelThread _connectionThread;
+
+            private readonly object _lock = new object();
+
+            private long _length;
+            private bool _connectionPaused;
+
+            public BufferLengthConnectionController(long maxLength, IConnectionControl connectionControl, KestrelThread connectionThread) 
+            {
+                _maxLength = maxLength;
+                _connectionControl = connectionControl;
+                _connectionThread = connectionThread;
+            }
+
+            public long Length
+            {
+                get
+                {
+                    return _length;
+                }
+                set
+                {
+                    // Caller should ensure that bytes are never consumed before the producer has called Add()
+                    Debug.Assert(value >= 0);
+
+                    _length = value;
+                }
+            }
+
+            public void Add(int count)
+            {
+                // Add() should never be called while connection is paused, since ConnectionControl.Pause() runs on a libuv thread
+                // and should take effect immediately.
+                Debug.Assert(!_connectionPaused);
+
+                lock (_lock)
+                {
+                    Length += count;
+                    if (Length >= _maxLength)
+                    {
+                        _connectionPaused = true;
+                        _connectionControl.Pause();
+                    }
+                }
+            }
+
+            public void Subtract(int count)
+            {
+                lock (_lock)
+                {
+                    Length -= count;
+
+                    if (_connectionPaused && Length < _maxLength)
+                    {
+                        _connectionPaused = false;
+                        _connectionThread.Post(
+                            (connectionControl) => ((IConnectionControl)connectionControl).Resume(),
+                            _connectionControl);
+                    }
+                }
+            }
+        }
     }
 }

src/Microsoft.AspNetCore.Server.Kestrel/IKestrelServerInformation.cs

@@ -10,6 +10,14 @@ public interface IKestrelServerInformation
     {
         int ThreadCount { get; set; }
 
+        /// <summary>
+        /// The maximum number of bytes that will be buffered in memory when reading a request.
+        /// Default value is 1048576 (1 MB).
+        /// A custom value can be configured using the "kestrel.maxInputBufferLength" key in 
+        /// <seealso cref="Microsoft.Extensions.Configuration.IConfiguration"/>.
+        /// </summary>
+        long MaxInputBufferLength { get; set; }
+
         /// <summary>
         /// The amount of time after the server begins shutting down before connections will be forcefully closed.
         /// By default, Kestrel will wait 5 seconds for any ongoing requests to complete before terminating

src/Microsoft.AspNetCore.Server.Kestrel/KestrelServerInformation.cs

@@ -21,6 +21,7 @@ public KestrelServerInformation(IConfiguration configuration)
 
             Addresses = GetAddresses(configuration);
             ThreadCount = GetThreadCount(configuration);
+            MaxInputBufferLength = GetMaxInputBufferLength(configuration);
             ShutdownTimeout = GetShutdownTimeout(configuration);
             NoDelay = GetNoDelay(configuration);
             PoolingParameters = new KestrelServerPoolingParameters(configuration);
@@ -30,6 +31,8 @@ public KestrelServerInformation(IConfiguration configuration)
 
         public int ThreadCount { get; set; }
 
+        public long MaxInputBufferLength { get; set; }
+
         public TimeSpan ShutdownTimeout { get; set; }
 
         public bool NoDelay { get; set; }
@@ -96,6 +99,26 @@ private static int GetThreadCount(IConfiguration configuration)
             return ProcessorThreadCount;
         }
 
+        private static long GetMaxInputBufferLength(IConfiguration configuration)
+        {
+            const long defaultMaxInputBufferLength = 1024 * 1024;
+
+            var maxInputBufferLengthString = configuration["kestrel.maxInputBufferLength"];
+
+            if (string.IsNullOrEmpty(maxInputBufferLengthString))
+            {
+                return defaultMaxInputBufferLength;
+            }
+
+            long maxInputBufferLength;
+            if (long.TryParse(maxInputBufferLengthString, NumberStyles.Integer, CultureInfo.InvariantCulture, out maxInputBufferLength))
+            {
+                return maxInputBufferLength;
+            }
+
+            return defaultMaxInputBufferLength;
+        }
+
         private TimeSpan GetShutdownTimeout(IConfiguration configuration)
         {
             var shutdownTimeoutString = configuration["kestrel.shutdownTimout"];