using System.Buffers; using System.IO; namespace CompatBot.Utils; internal class BufferCopyStream : Stream, IDisposable { private readonly Stream baseStream; private bool usedForReads; private bool usedForWrites; private long position; private readonly int bufSize; private readonly byte[] writeBuf; private readonly byte[] readBuf; private int bufStart, bufLength; private readonly Lock sync = new(); private bool disposed; public BufferCopyStream(Stream? baseStream, int bufferSize = 4096) { if (baseStream == null) throw new ArgumentNullException(nameof(baseStream)); if (bufferSize < 1) throw new ArgumentException("Buffer size cannot be non-positive", nameof(bufferSize)); this.baseStream = baseStream; bufSize = bufferSize; writeBuf = ArrayPool.Shared.Rent(bufSize); readBuf = ArrayPool.Shared.Rent(16); } public override void Flush() => baseStream.Flush(); public override int Read(byte[] buffer, int offset, int count) { usedForReads = true; if (usedForWrites) throw new InvalidOperationException("Stream was used for writes before"); var useTempBuf = count < 16; const int mask = ~0b1111; try { int result; if (useTempBuf) { result = baseStream.Read(readBuf, 0, count); Buffer.BlockCopy(readBuf, 0, buffer, offset, result); } else result = baseStream.Read(buffer, offset, count & mask); // make count divisible by 16 to workaround mega client issues position += result; CopyToBuf(buffer, offset, result); return result; } catch (Exception e) { Config.Log.Error(e, $"Failed to read from the base stream: {nameof(buffer)}.Length={buffer.Length}, {nameof(offset)}={offset}, {nameof(count)}={count}"); throw; } } public override long Seek(long offset, SeekOrigin origin) => throw new NotImplementedException(); public override void SetLength(long value) => throw new NotImplementedException(); public override void Write(byte[] buffer, int offset, int count) { usedForWrites = true; if (usedForReads) throw new InvalidOperationException("Stream was used for reads before"); baseStream.Write(buffer, offset, count); position += count; CopyToBuf(buffer, offset, count); } protected override void Dispose(bool disposing) { if (!disposed) lock (sync) { if (!disposed) { ArrayPool.Shared.Return(writeBuf); ArrayPool.Shared.Return(readBuf); disposed = true; } } baseStream.Dispose(); base.Dispose(disposing); } void IDisposable.Dispose() { if (!disposed) lock (sync) { if (!disposed) { ArrayPool.Shared.Return(writeBuf); ArrayPool.Shared.Return(readBuf); disposed = true; } } baseStream.Dispose(); base.Dispose(); } private void CopyToBuf(byte[] buffer, int offset, int count) { if (count >= bufSize) { bufStart = 0; bufLength = bufSize; Buffer.BlockCopy(buffer, offset + count - bufSize, writeBuf, bufStart, bufLength); } else { // copy as much data as we can to the end of the buffer bufStart = (bufStart + bufLength) % bufSize; bufLength = Math.Min(bufSize - bufStart, count); Buffer.BlockCopy(buffer, offset, writeBuf, bufStart, bufLength); // if there's still more data, loop it around to the beginning if (bufLength < count) { Buffer.BlockCopy(buffer, offset + bufLength, writeBuf, 0, count-bufLength); bufLength = count; } } } public byte[] GetBufferedBytes() { var result = new byte[bufLength]; var partLength = Math.Min(bufSize - bufStart, bufLength); Buffer.BlockCopy(writeBuf, bufStart, result, 0, partLength); if (partLength < bufLength) Buffer.BlockCopy(writeBuf, 0, result, partLength, bufLength - partLength); return result; } public override bool CanRead => baseStream.CanRead; public override bool CanSeek => false; public override bool CanWrite => baseStream.CanWrite; public override long Length => baseStream.Length; public override long Position { get => position; set => throw new InvalidOperationException(); } }