Skip to content
This repository was archived by the owner on Dec 18, 2018. It is now read-only.

Commit 9645d2e

Browse files
committed
Limit size of memory buffer when reading request (#304)
1 parent bd1c815 commit 9645d2e

File tree

4 files changed

+138
-1
lines changed

4 files changed

+138
-1
lines changed

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

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,26 @@ public Connection(ListenerContext context, UvStreamHandle socket) : base(context
4646

4747
ConnectionId = GenerateConnectionId(Interlocked.Increment(ref _lastConnectionId));
4848

49-
_rawSocketInput = new SocketInput(Memory, ThreadPool);
49+
var maxInputBufferLength = context.ServerInformation.MaxInputBufferLength;
50+
if (maxInputBufferLength < -1 || maxInputBufferLength == 0)
51+
{
52+
var paramName = String.Join(".",
53+
nameof(context),
54+
nameof(context.ServerInformation),
55+
nameof(context.ServerInformation.MaxInputBufferLength));
56+
57+
throw new ArgumentOutOfRangeException(paramName, maxInputBufferLength, $"{paramName} must be positive or -1.");
58+
}
59+
60+
if (maxInputBufferLength == -1)
61+
{
62+
_rawSocketInput = new SocketInput(Memory, ThreadPool);
63+
}
64+
else
65+
{
66+
_rawSocketInput = new SocketInput(Memory, ThreadPool, maxInputBufferLength, this, Thread);
67+
}
68+
5069
_rawSocketOutput = new SocketOutput(Thread, _socket, Memory, this, ConnectionId, Log, ThreadPool, WriteReqPool);
5170
}
5271

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

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,21 @@ public class SocketInput : ICriticalNotifyCompletion, IDisposable
2929

3030
private int _consumingState;
3131

32+
private readonly BufferLengthConnectionController _bufferLengthConnectionController;
33+
3234
public SocketInput(MemoryPool memory, IThreadPool threadPool)
3335
{
3436
_memory = memory;
3537
_threadPool = threadPool;
3638
_awaitableState = _awaitableIsNotCompleted;
3739
}
3840

41+
public SocketInput(MemoryPool memory, IThreadPool threadPool, long maxBufferLength, IConnectionControl connectionControl,
42+
KestrelThread connectionThread) : this(memory, threadPool)
43+
{
44+
_bufferLengthConnectionController = new BufferLengthConnectionController(maxBufferLength, connectionControl, connectionThread);
45+
}
46+
3947
public bool RemoteIntakeFin { get; set; }
4048

4149
public bool IsCompleted => ReferenceEquals(_awaitableState, _awaitableIsCompleted);
@@ -58,6 +66,9 @@ public MemoryPoolBlock IncomingStart()
5866

5967
public void IncomingData(byte[] buffer, int offset, int count)
6068
{
69+
// Must call Add() before bytes are available to consumer, to ensure that Length is >= 0
70+
_bufferLengthConnectionController?.Add(count);
71+
6172
if (count > 0)
6273
{
6374
if (_tail == null)
@@ -85,6 +96,9 @@ public void IncomingData(byte[] buffer, int offset, int count)
8596

8697
public void IncomingComplete(int count, Exception error)
8798
{
99+
// Must call Add() before bytes are available to consumer, to ensure that Length is >= 0
100+
_bufferLengthConnectionController?.Add(count);
101+
88102
if (_pinned != null)
89103
{
90104
_pinned.End += count;
@@ -167,10 +181,16 @@ public void ConsumingComplete(
167181

168182
if (!consumed.IsDefault)
169183
{
184+
var lengthConsumed = new MemoryPoolIterator(_head).GetLength(consumed);
185+
170186
returnStart = _head;
171187
returnEnd = consumed.Block;
172188
_head = consumed.Block;
173189
_head.Start = consumed.Index;
190+
191+
// Must call Subtract() after bytes have been freed, to avoid producer starting too early and growing
192+
// buffer beyond max length.
193+
_bufferLengthConnectionController?.Subtract(lengthConsumed);
174194
}
175195

176196
if (!examined.IsDefault &&
@@ -285,5 +305,72 @@ public void Dispose()
285305
_head = null;
286306
_tail = null;
287307
}
308+
309+
private class BufferLengthConnectionController
310+
{
311+
private readonly long _maxLength;
312+
private readonly IConnectionControl _connectionControl;
313+
private readonly KestrelThread _connectionThread;
314+
315+
private readonly object _lock = new object();
316+
317+
private long _length;
318+
private bool _connectionPaused;
319+
320+
public BufferLengthConnectionController(long maxLength, IConnectionControl connectionControl, KestrelThread connectionThread)
321+
{
322+
_maxLength = maxLength;
323+
_connectionControl = connectionControl;
324+
_connectionThread = connectionThread;
325+
}
326+
327+
public long Length
328+
{
329+
get
330+
{
331+
return _length;
332+
}
333+
set
334+
{
335+
// Caller should ensure that bytes are never consumed before the producer has called Add()
336+
Debug.Assert(value >= 0);
337+
338+
_length = value;
339+
}
340+
}
341+
342+
public void Add(int count)
343+
{
344+
// Add() should never be called while connection is paused, since ConnectionControl.Pause() runs on a libuv thread
345+
// and should take effect immediately.
346+
Debug.Assert(!_connectionPaused);
347+
348+
lock (_lock)
349+
{
350+
Length += count;
351+
if (Length >= _maxLength)
352+
{
353+
_connectionPaused = true;
354+
_connectionControl.Pause();
355+
}
356+
}
357+
}
358+
359+
public void Subtract(int count)
360+
{
361+
lock (_lock)
362+
{
363+
Length -= count;
364+
365+
if (_connectionPaused && Length < _maxLength)
366+
{
367+
_connectionPaused = false;
368+
_connectionThread.Post(
369+
(connectionControl) => ((IConnectionControl)connectionControl).Resume(),
370+
_connectionControl);
371+
}
372+
}
373+
}
374+
}
288375
}
289376
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,14 @@ public interface IKestrelServerInformation
1010
{
1111
int ThreadCount { get; set; }
1212

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

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ public KestrelServerInformation(IConfiguration configuration)
2121

2222
Addresses = GetAddresses(configuration);
2323
ThreadCount = GetThreadCount(configuration);
24+
MaxInputBufferLength = GetMaxInputBufferLength(configuration);
2425
ShutdownTimeout = GetShutdownTimeout(configuration);
2526
NoDelay = GetNoDelay(configuration);
2627
PoolingParameters = new KestrelServerPoolingParameters(configuration);
@@ -30,6 +31,8 @@ public KestrelServerInformation(IConfiguration configuration)
3031

3132
public int ThreadCount { get; set; }
3233

34+
public long MaxInputBufferLength { get; set; }
35+
3336
public TimeSpan ShutdownTimeout { get; set; }
3437

3538
public bool NoDelay { get; set; }
@@ -96,6 +99,26 @@ private static int GetThreadCount(IConfiguration configuration)
9699
return ProcessorThreadCount;
97100
}
98101

102+
private static long GetMaxInputBufferLength(IConfiguration configuration)
103+
{
104+
const long defaultMaxInputBufferLength = 1024 * 1024;
105+
106+
var maxInputBufferLengthString = configuration["kestrel.maxInputBufferLength"];
107+
108+
if (string.IsNullOrEmpty(maxInputBufferLengthString))
109+
{
110+
return defaultMaxInputBufferLength;
111+
}
112+
113+
long maxInputBufferLength;
114+
if (long.TryParse(maxInputBufferLengthString, NumberStyles.Integer, CultureInfo.InvariantCulture, out maxInputBufferLength))
115+
{
116+
return maxInputBufferLength;
117+
}
118+
119+
return defaultMaxInputBufferLength;
120+
}
121+
99122
private TimeSpan GetShutdownTimeout(IConfiguration configuration)
100123
{
101124
var shutdownTimeoutString = configuration["kestrel.shutdownTimout"];

0 commit comments

Comments
 (0)