98 lines
3.7 KiB
C#
98 lines
3.7 KiB
C#
using System;
|
|
using System.IO;
|
|
using System.Threading;
|
|
using System.Threading.Tasks;
|
|
|
|
namespace Unity.Pipeline
|
|
{
|
|
/// <summary>
|
|
/// Read-only stream wrapper that enforces a maximum number of bytes read from an inner stream.
|
|
/// Once cumulative reads exceed <c>maxBytes</c>, the next read throws
|
|
/// <see cref="RequestTooLargeException"/>. Used to cap HTTP request bodies even when the
|
|
/// Content-Length header is absent or untruthful (e.g. chunked transfer-encoding), so an
|
|
/// oversized body cannot be buffered into memory.
|
|
/// </summary>
|
|
internal sealed class MaxLengthStream : Stream
|
|
{
|
|
private readonly Stream m_Inner;
|
|
private readonly long m_MaxBytes;
|
|
private readonly bool m_LeaveOpen;
|
|
private long m_TotalRead;
|
|
|
|
/// <param name="leaveOpen">
|
|
/// When true, disposing this wrapper does not dispose <paramref name="inner"/>. The server
|
|
/// uses this so it can keep reading (draining) the raw request stream after the wrapper is
|
|
/// disposed on an over-limit read.
|
|
/// </param>
|
|
public MaxLengthStream(Stream inner, long maxBytes, bool leaveOpen = false)
|
|
{
|
|
m_Inner = inner ?? throw new ArgumentNullException(nameof(inner));
|
|
m_MaxBytes = maxBytes;
|
|
m_LeaveOpen = leaveOpen;
|
|
}
|
|
|
|
public override int Read(byte[] buffer, int offset, int count)
|
|
{
|
|
var read = m_Inner.Read(buffer, offset, count);
|
|
Account(read);
|
|
return read;
|
|
}
|
|
|
|
public override async Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
|
|
{
|
|
var read = await m_Inner.ReadAsync(buffer, offset, count, cancellationToken);
|
|
Account(read);
|
|
return read;
|
|
}
|
|
|
|
// Throw as soon as the running total exceeds the cap. Exactly maxBytes is allowed.
|
|
private void Account(int read)
|
|
{
|
|
if (read <= 0)
|
|
return;
|
|
|
|
m_TotalRead += read;
|
|
if (m_TotalRead > m_MaxBytes)
|
|
throw new RequestTooLargeException(m_MaxBytes);
|
|
}
|
|
|
|
public override bool CanRead => true;
|
|
public override bool CanSeek => false;
|
|
public override bool CanWrite => false;
|
|
public override long Length => throw new NotSupportedException();
|
|
public override long Position
|
|
{
|
|
get => m_TotalRead;
|
|
set => throw new NotSupportedException();
|
|
}
|
|
|
|
public override void Flush() { }
|
|
public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
|
|
public override void SetLength(long value) => throw new NotSupportedException();
|
|
public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();
|
|
|
|
protected override void Dispose(bool disposing)
|
|
{
|
|
// Own the inner stream's lifetime (matching the previous `using (StreamReader(InputStream))`)
|
|
// unless the caller asked to leave it open so it can keep reading from it afterwards.
|
|
if (disposing && !m_LeaveOpen)
|
|
m_Inner.Dispose();
|
|
base.Dispose(disposing);
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Thrown by <see cref="MaxLengthStream"/> when a read would exceed the configured byte cap.
|
|
/// </summary>
|
|
internal sealed class RequestTooLargeException : Exception
|
|
{
|
|
public long MaxBytes { get; }
|
|
|
|
public RequestTooLargeException(long maxBytes)
|
|
: base($"Request body exceeds the maximum allowed size of {maxBytes} bytes.")
|
|
{
|
|
MaxBytes = maxBytes;
|
|
}
|
|
}
|
|
}
|