129 lines
4.0 KiB
C#
129 lines
4.0 KiB
C#
namespace MareSynchronos.Utils;
|
|
|
|
// Limits the number of bytes read/written to an underlying stream
|
|
public class LimitedStream : Stream
|
|
{
|
|
private readonly Stream _stream;
|
|
private long _estimatedPosition = 0;
|
|
public long MaxPosition { get; private init; }
|
|
public bool DisposeUnderlying { get; set; } = true;
|
|
|
|
public Stream UnderlyingStream { get => _stream; }
|
|
|
|
public LimitedStream(Stream underlyingStream, long byteLimit)
|
|
{
|
|
_stream = underlyingStream;
|
|
try
|
|
{
|
|
_estimatedPosition = _stream.Position;
|
|
}
|
|
catch { }
|
|
MaxPosition = _estimatedPosition + byteLimit;
|
|
}
|
|
|
|
protected override void Dispose(bool disposing)
|
|
{
|
|
base.Dispose(disposing);
|
|
if (!DisposeUnderlying)
|
|
return;
|
|
_stream.Dispose();
|
|
}
|
|
|
|
public override bool CanRead => _stream.CanRead;
|
|
public override bool CanSeek => _stream.CanSeek;
|
|
public override bool CanWrite => _stream.CanWrite;
|
|
public override long Length => _stream.Length;
|
|
|
|
public override long Position { get => _stream.Position; set => _stream.Position = _estimatedPosition = value; }
|
|
|
|
public override void Flush()
|
|
{
|
|
_stream.Flush();
|
|
}
|
|
|
|
public override int Read(byte[] buffer, int offset, int count)
|
|
{
|
|
int remainder = (int)long.Clamp(MaxPosition - _estimatedPosition, 0, int.MaxValue);
|
|
|
|
if (count > remainder)
|
|
count = remainder;
|
|
|
|
int n = _stream.Read(buffer, offset, count);
|
|
_estimatedPosition += n;
|
|
return n;
|
|
}
|
|
|
|
public async override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
|
|
{
|
|
int remainder = (int)long.Clamp(MaxPosition - _estimatedPosition, 0, int.MaxValue);
|
|
|
|
if (count > remainder)
|
|
count = remainder;
|
|
|
|
#pragma warning disable CA1835
|
|
int n = await _stream.ReadAsync(buffer, offset, count, cancellationToken).ConfigureAwait(false);
|
|
#pragma warning restore CA1835
|
|
_estimatedPosition += n;
|
|
return n;
|
|
}
|
|
|
|
public async override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
|
|
{
|
|
int remainder = (int)long.Clamp(MaxPosition - _estimatedPosition, 0, int.MaxValue);
|
|
|
|
if (buffer.Length > remainder)
|
|
buffer = buffer[..remainder];
|
|
|
|
int n = await _stream.ReadAsync(buffer, cancellationToken).ConfigureAwait(false);
|
|
_estimatedPosition += n;
|
|
return n;
|
|
}
|
|
|
|
public override long Seek(long offset, SeekOrigin origin)
|
|
{
|
|
long result = _stream.Seek(offset, origin);
|
|
_estimatedPosition = result;
|
|
return result;
|
|
}
|
|
|
|
public override void SetLength(long value)
|
|
{
|
|
_stream.SetLength(value);
|
|
}
|
|
|
|
public override void Write(byte[] buffer, int offset, int count)
|
|
{
|
|
int remainder = (int)long.Clamp(MaxPosition - _estimatedPosition, 0, int.MaxValue);
|
|
|
|
if (count > remainder)
|
|
count = remainder;
|
|
|
|
_stream.Write(buffer, offset, count);
|
|
_estimatedPosition += count;
|
|
}
|
|
|
|
public async override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
|
|
{
|
|
int remainder = (int)long.Clamp(MaxPosition - _estimatedPosition, 0, int.MaxValue);
|
|
|
|
if (count > remainder)
|
|
count = remainder;
|
|
|
|
#pragma warning disable CA1835
|
|
await _stream.WriteAsync(buffer, offset, count, cancellationToken).ConfigureAwait(false);
|
|
#pragma warning restore CA1835
|
|
_estimatedPosition += count;
|
|
}
|
|
|
|
public async override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
|
|
{
|
|
int remainder = (int)long.Clamp(MaxPosition - _estimatedPosition, 0, int.MaxValue);
|
|
|
|
if (buffer.Length > remainder)
|
|
buffer = buffer[..remainder];
|
|
|
|
await _stream.WriteAsync(buffer, cancellationToken).ConfigureAwait(false);
|
|
_estimatedPosition += buffer.Length;
|
|
}
|
|
}
|