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

src/Microsoft.AspNetCore.Server.Kestrel/Filter/Internal/FilteredStreamAdapter.cs

@@ -24,9 +24,10 @@ public FilteredStreamAdapter(
             Stream filteredStream,
             MemoryPool memory,
             IKestrelTrace logger,
-            IThreadPool threadPool)
+            IThreadPool threadPool,
+            IBufferSizeControl bufferSizeControl)
         {
-            SocketInput = new SocketInput(memory, threadPool);
+            SocketInput = new SocketInput(memory, threadPool, bufferSizeControl);
             SocketOutput = new StreamSocketOutput(connectionId, filteredStream, memory, logger);
 
             _connectionId = connectionId;

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

@@ -0,0 +1,83 @@
+using System.Diagnostics;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
+{
+    public class BufferSizeControl : IBufferSizeControl
+    {
+        private readonly long _maxSize;
+        private readonly IConnectionControl _connectionControl;
+        private readonly KestrelThread _connectionThread;
+
+        private readonly object _lock = new object();
+
+        private long _size;
+        private bool _connectionPaused;
+
+        public BufferSizeControl(long maxSize, IConnectionControl connectionControl, KestrelThread connectionThread)
+        {
+            _maxSize = maxSize;
+            _connectionControl = connectionControl;
+            _connectionThread = connectionThread;
+        }
+
+        private long Size
+        {
+            get
+            {
+                return _size;
+            }
+            set
+            {
+                // Caller should ensure that bytes are never consumed before the producer has called Add()
+                Debug.Assert(value >= 0);
+                _size = value;
+            }
+        }
+
+        public void Add(int count)
+        {
+            Debug.Assert(count >= 0);
+
+            if (count == 0)
+            {
+                // No-op and avoid taking lock to reduce contention
+                return;
+            }
+
+            lock (_lock)
+            {
+                Size += count;
+                if (!_connectionPaused && Size >= _maxSize)
+                {
+                    _connectionPaused = true;
+                    _connectionThread.Post(
+                        (connectionControl) => ((IConnectionControl)connectionControl).Pause(),
+                        _connectionControl);
+                }
+            }
+        }
+
+        public void Subtract(int count)
+        {
+            Debug.Assert(count >= 0);
+
+            if (count == 0)
+            {
+                // No-op and avoid taking lock to reduce contention
+                return;
+            }
+
+            lock (_lock)
+            {
+                Size -= count;
+                if (_connectionPaused && Size < _maxSize)
+                {
+                    _connectionPaused = false;
+                    _connectionThread.Post(
+                        (connectionControl) => ((IConnectionControl)connectionControl).Resume(),
+                        _connectionControl);
+                }
+            }
+        }
+    }
+}

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

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Server.Kestrel.Filter;
@@ -41,6 +42,8 @@ public class Connection : ConnectionContext, IConnectionControl
         private ConnectionState _connectionState;
         private TaskCompletionSource<object> _socketClosedTcs;
 
+        private BufferSizeControl _bufferSizeControl;
+
         public Connection(ListenerContext context, UvStreamHandle socket) : base(context)
         {
             _socket = socket;
@@ -49,7 +52,12 @@ public Connection(ListenerContext context, UvStreamHandle socket) : base(context
 
             ConnectionId = GenerateConnectionId(Interlocked.Increment(ref _lastConnectionId));
 
-            _rawSocketInput = new SocketInput(Memory, ThreadPool);
+            if (ServerOptions.MaxRequestBufferSize.HasValue)
+            {
+                _bufferSizeControl = new BufferSizeControl(ServerOptions.MaxRequestBufferSize.Value, this, Thread);
+            }
+
+            _rawSocketInput = new SocketInput(Memory, ThreadPool, _bufferSizeControl);
             _rawSocketOutput = new SocketOutput(Thread, _socket, Memory, this, ConnectionId, Log, ThreadPool, WriteReqPool);
         }
 
@@ -217,7 +225,7 @@ private void ApplyConnectionFilter()
 
                     if (_filterContext.Connection != _libuvStream)
                     {
-                        _filteredStreamAdapter = new FilteredStreamAdapter(ConnectionId, _filterContext.Connection, Memory, Log, ThreadPool);
+                        _filteredStreamAdapter = new FilteredStreamAdapter(ConnectionId, _filterContext.Connection, Memory, Log, ThreadPool, _bufferSizeControl);
 
                         SocketInput = _filteredStreamAdapter.SocketInput;
                         SocketOutput = _filteredStreamAdapter.SocketOutput;
@@ -316,7 +324,17 @@ void IConnectionControl.Pause()
         void IConnectionControl.Resume()
         {
             Log.ConnectionResume(ConnectionId);
-            _socket.ReadStart(_allocCallback, _readCallback, this);
+            try
+            {
+                _socket.ReadStart(_allocCallback, _readCallback, this);
+            }
+            catch (UvException)
+            {
+                // ReadStart() can throw a UvException in some cases (e.g. socket is no longer connected).
+                // This should be treated the same as OnRead() seeing a "normalDone" condition.
+                Log.ConnectionReadFin(ConnectionId);
+                _rawSocketInput.IncomingComplete(0, null);
+            }
         }
 
         void IConnectionControl.End(ProduceEndType endType)

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

@@ -0,0 +1,13 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
+{
+    public interface IBufferSizeControl
+    {
+        void Add(int count);
+        void Subtract(int count);
+    }
+}

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

@@ -18,6 +18,7 @@ public class SocketInput : ICriticalNotifyCompletion, IDisposable
 
         private readonly MemoryPool _memory;
         private readonly IThreadPool _threadPool;
+        private readonly IBufferSizeControl _bufferSizeControl;
         private readonly ManualResetEventSlim _manualResetEvent = new ManualResetEventSlim(false, 0);
 
         private Action _awaitableState;
@@ -32,10 +33,11 @@ public class SocketInput : ICriticalNotifyCompletion, IDisposable
         private bool _consuming;
         private bool _disposed;
 
-        public SocketInput(MemoryPool memory, IThreadPool threadPool)
+        public SocketInput(MemoryPool memory, IThreadPool threadPool, IBufferSizeControl bufferSizeControl = null)
         {
             _memory = memory;
             _threadPool = threadPool;
+            _bufferSizeControl = bufferSizeControl;
             _awaitableState = _awaitableIsNotCompleted;
         }
 
@@ -63,6 +65,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
+                _bufferSizeControl?.Add(count);
+
                 if (count > 0)
                 {
                     if (_tail == null)
@@ -93,6 +98,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
+                _bufferSizeControl?.Add(count);
+
                 if (_pinned != null)
                 {
                     _pinned.End += count;
@@ -189,10 +197,21 @@ public void ConsumingComplete(
                 {
                     if (!consumed.IsDefault)
                     {
+                        // Compute lengthConsumed before modifying _head or consumed
+                        var lengthConsumed = 0;
+                        if (_bufferSizeControl != null)
+                        {
+                            lengthConsumed = new MemoryPoolIterator(_head).GetLength(consumed);
+                        }
+
                         returnStart = _head;
                         returnEnd = consumed.Block;
                         _head = consumed.Block;
                         _head.Start = consumed.Index;
+
+                        // Must call Subtract() after _head has been advanced, to avoid producer starting too early and growing
+                        // buffer beyond max length.
+                        _bufferSizeControl?.Subtract(lengthConsumed);
                     }
 
                     if (!examined.IsDefault &&

src/Microsoft.AspNetCore.Server.Kestrel/Internal/Networking/UvStreamHandle.cs

@@ -105,16 +105,16 @@ public void ReadStart(
             }
         }
 
+        // UvStreamHandle.ReadStop() should be idempotent to match uv_read_stop()
         public void ReadStop()
         {
-            if (!_readVitality.IsAllocated)
+            if (_readVitality.IsAllocated)
             {
-                throw new InvalidOperationException("TODO: ReadStart must be called before ReadStop may be called");
+                _readVitality.Free();
             }
             _allocCallback = null;
             _readCallback = null;
             _readState = null;
-            _readVitality.Free();
             _uv.read_stop(this);
         }
 

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

@@ -8,16 +8,40 @@ namespace Microsoft.AspNetCore.Server.Kestrel
 {
     public class KestrelServerOptions
     {
+        // Matches the default client_max_body_size in nginx.  Also large enough that most requests
+        // should be under the limit.
+        private long? _maxRequestBufferSize = 1024 * 1024;
+
+        /// <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;
-
         /// <summary>
-        /// Gets or sets whether the <c>Server</c> header should be included in each response.
+        /// Maximum size of the request buffer.  Default is 1,048,576 bytes (1 MB).
+        /// If value is null, the size of the request buffer is unlimited.
         /// </summary>
-        public bool AddServerHeader { get; set; } = true;
+        public long? MaxRequestBufferSize
+        {
+            get
+            {
+                return _maxRequestBufferSize;
+            }
+            set
+            {
+                if (value.HasValue && value.Value <= 0)
+                {
+                    throw new ArgumentOutOfRangeException("value", "Value must be null or a positive integer.");
+                }
+                _maxRequestBufferSize = value;
+            }
+        }
+
+        public bool NoDelay { get; set; } = true;
 
         /// <summary>
         /// The amount of time after the server begins shutting down before connections will be forcefully closed.

test/Microsoft.AspNetCore.Server.Kestrel.FunctionalTests/IWebHostPortExtensions.cs

@@ -10,16 +10,35 @@ namespace Microsoft.AspNetCore.Hosting
 {
     public static class IWebHostPortExtensions
     {
+        public static string GetHost(this IWebHost host)
+        {
+            return host.GetUris().First().Host;
+        }
+
         public static int GetPort(this IWebHost host)
         {
             return host.GetPorts().First();
         }
 
+        public static int GetPort(this IWebHost host, string scheme)
+        {
+            return host.GetUris()
+                .Where(u => u.Scheme.Equals(scheme, StringComparison.OrdinalIgnoreCase))
+                .Select(u => u.Port)
+                .First();
+        }
+
         public static IEnumerable<int> GetPorts(this IWebHost host)
+        {
+            return host.GetUris()
+                .Select(u => u.Port);
+        }
+
+        public static IEnumerable<Uri> GetUris(this IWebHost host)
         {
             return host.ServerFeatures.Get<IServerAddressesFeature>().Addresses
                 .Select(a => a.Replace("://+", "://localhost"))
-                .Select(a => (new Uri(a)).Port);
+                .Select(a => new Uri(a));
         }
     }
 }