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

Author: mikeharder

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

@@ -49,7 +49,15 @@ public Connection(ListenerContext context, UvStreamHandle socket) : base(context
 
             ConnectionId = GenerateConnectionId(Interlocked.Increment(ref _lastConnectionId));
 
-            _rawSocketInput = new SocketInput(Memory, ThreadPool);
+            if (ServerOptions.MaxInputBufferLength == -1)
+            {
+                _rawSocketInput = new SocketInput(Memory, ThreadPool);
+            }
+            else
+            {
+                _rawSocketInput = new SocketInput(Memory, ThreadPool, ServerOptions.MaxInputBufferLength, this, Thread);
+            }
+
             _rawSocketOutput = new SocketOutput(Thread, _socket, Memory, this, ConnectionId, Log, ThreadPool, WriteReqPool);
         }
 

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

@@ -32,13 +32,21 @@ public class SocketInput : ICriticalNotifyCompletion, IDisposable
         private bool _consuming;
         private bool _disposed;
 
+        private readonly BufferLengthConnectionController _bufferLengthConnectionController;
+
         public SocketInput(MemoryPool memory, IThreadPool threadPool)
         {
             _memory = memory;
             _threadPool = threadPool;
             _awaitableState = _awaitableIsNotCompleted;
         }
 
+        public SocketInput(MemoryPool memory, IThreadPool threadPool, int 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);
@@ -63,6 +71,9 @@ public void IncomingData(byte[] buffer, int offset, int count)
         {
             lock (_sync)
             {
+                // Must call Add() before bytes are available to consumer, to ensure that Length is >= 0
+                _bufferLengthConnectionController?.Add(count);
+
                 if (count > 0)
                 {
                     if (_tail == null)
@@ -93,6 +104,9 @@ public void IncomingComplete(int count, Exception error)
         {
             lock (_sync)
             {
+                // Must call Add() before bytes are available to consumer, to ensure that Length is >= 0
+                _bufferLengthConnectionController?.Add(count);
+
                 if (_pinned != null)
                 {
                     _pinned.End += count;
@@ -189,10 +203,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 &&
@@ -321,5 +341,73 @@ private static void ReturnBlocks(MemoryPoolBlock block, MemoryPoolBlock end)
                 returnBlock.Pool.Return(returnBlock);
             }
         }
+
+        private class BufferLengthConnectionController
+        {
+            private readonly int _maxLength;
+            private readonly IConnectionControl _connectionControl;
+            private readonly KestrelThread _connectionThread;
+
+            private readonly object _lock = new object();
+
+            private int _length;
+            private bool _connectionPaused;
+
+            public BufferLengthConnectionController(int maxLength, IConnectionControl connectionControl, KestrelThread connectionThread)
+            {
+                _maxLength = maxLength;
+                _connectionControl = connectionControl;
+                _connectionThread = connectionThread;
+            }
+
+            public int 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)
+            {
+                lock (_lock)
+                {
+                    // 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);
+
+                    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/KestrelServerOptions.cs

@@ -8,16 +8,32 @@ namespace Microsoft.AspNetCore.Server.Kestrel
 {
     public class KestrelServerOptions
     {
+        /// <summary>
+        /// Gets or sets whether the <c>Server</c> header should be included in each response.
+        /// </summary>
+        public bool AddServerHeader { get; set; } = true;
+
         public IServiceProvider ApplicationServices { get; set; }
 
         public IConnectionFilter ConnectionFilter { get; set; }
 
-        public bool NoDelay { get; set; } = true;
+        private int _maxInputBufferLength = 1024 * 1024;
+        public int MaxInputBufferLength
+        {
+            get
+            {
+                return _maxInputBufferLength;
+            }
+            set
+            {
+                if (value < -1 || value == 0)
+                {
+                    throw new ArgumentOutOfRangeException("value", "Value must be positive or -1.");
+                }
+            }
+        }
 
-        /// <summary>
-        /// Gets or sets whether the <c>Server</c> header should be included in each response.
-        /// </summary>
-        public bool AddServerHeader { get; set; } = true;
+        public bool NoDelay { get; set; } = true;
 
         /// <summary>
         /// The amount of time after the server begins shutting down before connections will be forcefully closed.