dev-npgsql/npgsql: Npgsql is the .NET data provider for PostgreSQL.

src/Npgsql/NpgsqlWriteBuffer.cs
592 строки20 KB
Shay Rojansky
Merge pull request from GHSA-x9vc-6hfv-hg8c
09 май 2024, 14:20
09 май 2024, 14:2067acbe0
using System;
using System.Buffers.Binary;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Text;
using System.Threading.Tasks;

#pragma warning disable CS1591 // Missing XML comment for publicly visible type or member
namespace Npgsql
{
    /// <summary>
    /// A buffer used by Npgsql to write data to the socket efficiently.
    /// Provides methods which encode different values types and tracks the current position.
    /// </summary>
    public sealed partial class NpgsqlWriteBuffer
    {
        #region Fields and Properties

        internal readonly NpgsqlConnector Connector;

        internal Stream Underlying { private get; set; }
        internal bool MessageLengthValidation { get; set; } = true;

        /// <summary>
        /// The total byte length of the buffer.
        /// </summary>
        internal int Size { get; private set; }

        bool _copyMode;
        internal Encoding TextEncoding { get; }

        public int WriteSpaceLeft => Size - WritePosition;

        internal readonly byte[] Buffer;
        readonly Encoder _textEncoder;

        internal int WritePosition;

        int _messageBytesFlushed;
        int? _messageLength;

        ParameterStream? _parameterStream;

        /// <summary>
        /// The minimum buffer size possible.
        /// </summary>
        internal const int MinimumSize = 4096;
        internal const int DefaultSize = 8192;

        #endregion

        #region Constructors

        internal NpgsqlWriteBuffer(NpgsqlConnector connector, Stream stream, int size, Encoding textEncoding)
        {
            if (size < MinimumSize)
                throw new ArgumentOutOfRangeException(nameof(size), size, "Buffer size must be at least " + MinimumSize);

            Connector = connector;
            Underlying = stream;
            Size = size;
            Buffer = new byte[Size];
            TextEncoding = textEncoding;
            _textEncoder = TextEncoding.GetEncoder();
        }

        #endregion

        #region I/O

        public async Task Flush(bool async)
        {
            if (_copyMode)
            {
                // In copy mode, we write CopyData messages. The message code has already been
                // written to the beginning of the buffer, but we need to go back and write the
                // length.
                if (WritePosition == 1)
                    return;
                var pos = WritePosition;
                WritePosition = 1;
                WriteInt32(pos - 1);
                WritePosition = pos;
            } else if (WritePosition == 0)
                return;
            else
                AdvanceMessageBytesFlushed(WritePosition);

            try
            {
                if (async)
                    await Underlying.WriteAsync(Buffer, 0, WritePosition);
                else
                    Underlying.Write(Buffer, 0, WritePosition);
            }
            catch (Exception e)
            {
                Connector.Break();
                throw new NpgsqlException("Exception while writing to stream", e);
            }

            try
            {
                if (async)
                    await Underlying.FlushAsync();
                else
                    Underlying.Flush();
            }
            catch (Exception e)
            {
                Connector.Break();
                throw new NpgsqlException("Exception while flushing stream", e);
            }

            NpgsqlEventSource.Log.BytesWritten(WritePosition);
            //NpgsqlEventSource.Log.RequestFailed();

            WritePosition = 0;
            if (_copyMode)
                WriteCopyDataHeader();
        }

        internal void Flush() => Flush(false).GetAwaiter().GetResult();

        #endregion

        #region Direct write

        internal async Task DirectWrite(byte[] buffer, int offset, int count, bool async)
        {
            await Flush(async);

            if (_copyMode)
            {
                // Flush has already written the CopyData header for us, but write the CopyData
                // header to the socket with the write length before we can start writing the data directly.
                Debug.Assert(WritePosition == 5);

                WritePosition = 1;
                WriteInt32(checked(count + 4));
                WritePosition = 5;
                _copyMode = false;
                StartMessage(5);
                Flush();
                _copyMode = true;
                WriteCopyDataHeader();  // And ready the buffer after the direct write completes
            }
            else
            {
                Debug.Assert(WritePosition == 0);
                AdvanceMessageBytesFlushed(count);
            }

            try
            {
                if (async)
                    await Underlying.WriteAsync(buffer, offset, count);
                else
                    Underlying.Write(buffer, offset, count);
            }
            catch (Exception e)
            {
                Connector.Break();
                throw new NpgsqlException("Exception while writing to stream", e);
            }
        }

#if !NETSTANDARD2_0 && !NET461
        internal async Task DirectWrite(ReadOnlyMemory<byte> memory, bool async)
        {
            await Flush(async);

            if (_copyMode)
            {
                // Flush has already written the CopyData header for us, but write the CopyData
                // header to the socket with the write length before we can start writing the data directly.
                Debug.Assert(WritePosition == 5);

                WritePosition = 1;
                WriteInt32(checked(memory.Length + 4));
                WritePosition = 5;
                _copyMode = false;
                StartMessage(5);
                await Flush(async);
                _copyMode = true;
                WriteCopyDataHeader();  // And ready the buffer after the direct write completes
            }
            else
            {
                Debug.Assert(WritePosition == 0);
                AdvanceMessageBytesFlushed(memory.Length);
            }


            try
            {
                if (async)
                    await Underlying.WriteAsync(memory);
                else
                    Underlying.Write(memory.Span);
            }
            catch (Exception e)
            {
                Connector.Break();
                throw new NpgsqlException("Exception while writing to stream", e);
            }
        }
#endif

        #endregion Direct write

        #region Write Simple

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteSByte(sbyte value) => Write(value);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteByte(byte value) => Write(value);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        internal void WriteInt16(int value)
            => WriteInt16((short)value, false);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteInt16(short value)
            => WriteInt16(value, false);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteInt16(short value, bool littleEndian)
            => Write(littleEndian == BitConverter.IsLittleEndian ? value : BinaryPrimitives.ReverseEndianness(value));

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteUInt16(ushort value)
            => WriteUInt16(value, false);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteUInt16(ushort value, bool littleEndian)
            => Write(littleEndian == BitConverter.IsLittleEndian ? value : BinaryPrimitives.ReverseEndianness(value));

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteInt32(int value)
            => WriteInt32(value, false);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteInt32(int value, bool littleEndian)
            => Write(littleEndian == BitConverter.IsLittleEndian ? value : BinaryPrimitives.ReverseEndianness(value));

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteUInt32(uint value)
            => WriteUInt32(value, false);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteUInt32(uint value, bool littleEndian)
            => Write(littleEndian == BitConverter.IsLittleEndian ? value : BinaryPrimitives.ReverseEndianness(value));

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteInt64(long value)
            => WriteInt64(value, false);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteInt64(long value, bool littleEndian)
            => Write(littleEndian == BitConverter.IsLittleEndian ? value : BinaryPrimitives.ReverseEndianness(value));

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteUInt64(ulong value)
            => WriteUInt64(value, false);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteUInt64(ulong value, bool littleEndian)
            => Write(littleEndian == BitConverter.IsLittleEndian ? value : BinaryPrimitives.ReverseEndianness(value));

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteSingle(float value)
            => WriteSingle(value, false);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteSingle(float value, bool littleEndian)
            => WriteInt32(Unsafe.As<float, int>(ref value), littleEndian);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteDouble(double value)
            => WriteDouble(value, false);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void WriteDouble(double value, bool littleEndian)
            => WriteInt64(Unsafe.As<double, long>(ref value), littleEndian);

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        void Write<T>(T value)
        {
            if (Unsafe.SizeOf<T>() > WriteSpaceLeft)
                ThrowNotSpaceLeft();

            Unsafe.WriteUnaligned(ref Buffer[WritePosition], value);
            WritePosition += Unsafe.SizeOf<T>();
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        static void ThrowNotSpaceLeft()
            => throw new InvalidOperationException("There is not enough space left in the buffer.");

        public Task WriteString(string s, int byteLen, bool async)
            => WriteString(s, s.Length, byteLen, async);

        public Task WriteString(string s, int charLen, int byteLen, bool async)
        {
            if (byteLen <= WriteSpaceLeft)
            {
                WriteString(s, charLen);
                return Task.CompletedTask;
            }
            return WriteStringLong();

            async Task WriteStringLong()
            {
                Debug.Assert(byteLen > WriteSpaceLeft);
                if (byteLen <= Size)
                {
                    // String can fit entirely in an empty buffer. Flush and retry rather than
                    // going into the partial writing flow below (which requires ToCharArray())
                    await Flush(async);
                    WriteString(s, charLen);
                }
                else
                {
                    var charPos = 0;
                    while (true)
                    {
                        WriteStringChunked(s, charPos, charLen - charPos, true, out var charsUsed, out var completed);
                        if (completed)
                            break;
                        await Flush(async);
                        charPos += charsUsed;
                    }
                }
            }
        }

        internal Task WriteChars(char[] chars, int offset, int charLen, int byteLen, bool async)
        {
            if (byteLen <= WriteSpaceLeft)
            {
                WriteChars(chars, offset, charLen);
                return Task.CompletedTask;
            }
            return WriteCharsLong();

            async Task WriteCharsLong()
            {
                Debug.Assert(byteLen > WriteSpaceLeft);
                if (byteLen <= Size)
                {
                    // String can fit entirely in an empty buffer. Flush and retry rather than
                    // going into the partial writing flow below (which requires ToCharArray())
                    await Flush(async);
                    WriteChars(chars, offset, charLen);
                }
                else
                {
                    var charPos = 0;

                    while (true)
                    {
                        WriteStringChunked(chars, charPos + offset, charLen - charPos, true, out var charsUsed, out var completed);
                        if (completed)
                            break;
                        await Flush(async);
                        charPos += charsUsed;
                    }
                }
            }
        }

        public void WriteString(string s, int len = 0)
        {
            Debug.Assert(TextEncoding.GetByteCount(s) <= WriteSpaceLeft);
            WritePosition += TextEncoding.GetBytes(s, 0, len == 0 ? s.Length : len, Buffer, WritePosition);
        }

        internal void WriteChars(char[] chars, int offset, int len)
        {
            var charCount = len == 0 ? chars.Length : len;
            Debug.Assert(TextEncoding.GetByteCount(chars, 0, charCount) <= WriteSpaceLeft);
            WritePosition += TextEncoding.GetBytes(chars, offset, charCount, Buffer, WritePosition);
        }

        public void WriteBytes(ReadOnlySpan<byte> buf)
        {
            Debug.Assert(buf.Length <= WriteSpaceLeft);
            buf.CopyTo(new Span<byte>(Buffer, WritePosition, Buffer.Length - WritePosition));
            WritePosition += buf.Length;
        }

        public void WriteBytes(byte[] buf, int offset, int count)
            => WriteBytes(new ReadOnlySpan<byte>(buf, offset, count));

        public Task WriteBytesRaw(byte[] bytes, bool async)
        {
            if (bytes.Length <= WriteSpaceLeft)
            {
                WriteBytes(bytes);
                return Task.CompletedTask;
            }
            return WriteBytesLong();

            async Task WriteBytesLong()
            {
                if (bytes.Length <= Size)
                {
                    // value can fit entirely in an empty buffer. Flush and retry rather than
                    // going into the partial writing flow below
                    await Flush(async);
                    WriteBytes(bytes);
                }
                else
                {
                    var remaining = bytes.Length;
                    do
                    {
                        if (WriteSpaceLeft == 0)
                            await Flush(async);
                        var writeLen = Math.Min(remaining, WriteSpaceLeft);
                        var offset = bytes.Length - remaining;
                        WriteBytes(bytes, offset, writeLen);
                        remaining -= writeLen;
                    }
                    while (remaining > 0);
                }
            }
        }

        public void WriteNullTerminatedString(string s)
        {
            Debug.Assert(s.All(c => c < 128), "Method only supports ASCII strings");
            Debug.Assert(WriteSpaceLeft >= s.Length + 1);
            WritePosition += Encoding.ASCII.GetBytes(s, 0, s.Length, Buffer, WritePosition);
            WriteByte(0);
        }

        #endregion

        #region Write Complex

        public Stream GetStream()
        {
            if (_parameterStream == null)
                _parameterStream = new ParameterStream(this);

            _parameterStream.Init();
            return _parameterStream;
        }

        internal void WriteStringChunked(char[] chars, int charIndex, int charCount,
                                         bool flush, out int charsUsed, out bool completed)
        {
            if (WriteSpaceLeft < _textEncoder.GetByteCount(chars, charIndex, char.IsHighSurrogate(chars[charIndex]) ? 2 : 1, flush: false))
            {
                charsUsed = 0;
                completed = false;
                return;
            }

            _textEncoder.Convert(chars, charIndex, charCount, Buffer, WritePosition, WriteSpaceLeft,
                                 flush, out charsUsed, out var bytesUsed, out completed);
            WritePosition += bytesUsed;
        }

        internal unsafe void WriteStringChunked(string s, int charIndex, int charCount,
                                                bool flush, out int charsUsed, out bool completed)
        {
            int bytesUsed;

            fixed (char* sPtr = s)
            fixed (byte* bufPtr = Buffer)
            {
                if (WriteSpaceLeft < _textEncoder.GetByteCount(sPtr + charIndex, char.IsHighSurrogate(*(sPtr + charIndex)) ? 2 : 1, flush: false))
                {
                    charsUsed = 0;
                    completed = false;
                    return;
                }

                _textEncoder.Convert(sPtr + charIndex, charCount, bufPtr + WritePosition, WriteSpaceLeft,
                                     flush, out charsUsed, out bytesUsed, out completed);
            }

            WritePosition += bytesUsed;
        }

        #endregion

        #region Copy

        internal void StartCopyMode()
        {
            _copyMode = true;
            Size -= 5;
            WriteCopyDataHeader();
        }

        internal void EndCopyMode()
        {
            // EndCopyMode is usually called after a Flush which ended the last CopyData message.
            // That Flush also wrote the header for another CopyData which we clear here.
            _copyMode = false;
            Size += 5;
            Clear();
        }

        void WriteCopyDataHeader()
        {
            Debug.Assert(_copyMode);
            Debug.Assert(WritePosition == 0);
            WriteByte((byte)BackendMessageCode.CopyData);
            // Leave space for the message length
            WriteInt32(0);
        }

        #endregion

        #region Misc

        internal void StartMessage(int messageLength)
        {
            if (!MessageLengthValidation)
                return;

            if (_messageLength != null && _messageBytesFlushed != _messageLength && WritePosition != -_messageBytesFlushed + _messageLength)
                Throw();

            // Add negative WritePosition to compensate for previous message(s) written without flushing.
            _messageBytesFlushed = -WritePosition;
            _messageLength = messageLength;

            void Throw()
            {
                Connector.Break();
                throw new OverflowException("Did not write the amount of bytes the message length specified");
            }
        }

        void AdvanceMessageBytesFlushed(int count)
        {
            if (!MessageLengthValidation)
                return;

            if (count < 0 || _messageLength is null || (long)_messageBytesFlushed + count > _messageLength)
                Throw();

            _messageBytesFlushed += count;

            void Throw()
            {
                if (count < 0)
                    throw new ArgumentOutOfRangeException(nameof(count), "Can't advance by a negative count");

                if (_messageLength is null)
                {
                    Connector.Break();
                    new InvalidOperationException("No message was started");
                }

                if ((long)_messageBytesFlushed + count > _messageLength)
                {
                    Connector.Break();
                    throw new OverflowException("Tried to write more bytes than the message length specified");
                }
            }
        }

        internal void Clear()
        {
            WritePosition = 0;
            _messageLength = null;
        }

        /// <summary>
        /// Returns all contents currently written to the buffer (but not flushed).
        /// Useful for pre-generating messages.
        /// </summary>
        internal byte[] GetContents()
        {
            var buf = new byte[WritePosition];
            Array.Copy(Buffer, buf, WritePosition);
            return buf;
        }

        #endregion
    }
}