Redo ReadNullTerminatedString (#3937)

roji · web-flow · commit da0fc0a802a6 · 2021-08-22T22:07:46.000+02:00
For safety with non-ASCII chars. Fixes #3660
diff --git a/src/Npgsql/Internal/NpgsqlReadBuffer.cs b/src/Npgsql/Internal/NpgsqlReadBuffer.cs
@@ -572,39 +572,69 @@ public ValueTask<string> ReadNullTerminatedString(bool async, CancellationToken
         /// </summary>
         ValueTask<string> ReadNullTerminatedString(Encoding encoding, bool async, CancellationToken cancellationToken = default)
         {
-            return ReadFromBuffer(this, encoding, out var s)
-                ? new ValueTask<string>(s)
-                : ReadLong(this, async, encoding, s);
-
-            static bool ReadFromBuffer(NpgsqlReadBuffer buffer, Encoding encoding, out string s)
+            for (var i = ReadPosition; i < FilledBytes; i++)
             {
-                var start = buffer.ReadPosition;
-                while (buffer.ReadPosition < buffer.FilledBytes)
+                if (Buffer[i] == 0)
                 {
-                    if (buffer.Buffer[buffer.ReadPosition++] == 0)
-                    {
-                        s = encoding.GetString(buffer.Buffer, start, buffer.ReadPosition - start - 1);
-                        return true;
-                    }
+                    var byteLen = i - ReadPosition;
+                    var result = new ValueTask<string>(encoding.GetString(Buffer, ReadPosition, byteLen));
+                    ReadPosition += byteLen + 1;
+                    return result;
                 }
-
-                s = encoding.GetString(buffer.Buffer, start, buffer.ReadPosition - start);
-                return false;
             }
 
-            static async ValueTask<string> ReadLong(NpgsqlReadBuffer buffer, bool async, Encoding encoding, string s)
+            return ReadLong(async);
+
+            async ValueTask<string> ReadLong(bool async)
             {
-                var builder = new StringBuilder(s);
-                bool complete;
-                do
+                var chunkSize = FilledBytes - ReadPosition;
+                var tempBuf = ArrayPool<byte>.Shared.Rent(chunkSize + 1024);
+
+                try
                 {
-                    await buffer.ReadMore(async);
-                    complete = ReadFromBuffer(buffer, encoding, out s);
-                    builder.Append(s);
-                }
-                while (!complete);
+                    bool foundTerminator;
+                    var byteLen = chunkSize;
+                    Array.Copy(Buffer, ReadPosition, tempBuf, 0, chunkSize);
+                    ReadPosition += chunkSize;
+
+                    do
+                    {
+                        await ReadMore(async);
+                        Debug.Assert(ReadPosition == 0);
 
-                return builder.ToString();
+                        foundTerminator = false;
+                        int i;
+                        for (i = 0; i < FilledBytes; i++)
+                        {
+                            if (Buffer[i] == 0)
+                            {
+                                foundTerminator = true;
+                                break;
+                            }
+                        }
+
+                        if (byteLen + i > tempBuf.Length)
+                        {
+                            var newTempBuf = ArrayPool<byte>.Shared.Rent(
+                                foundTerminator ? byteLen + i : byteLen + i + 1024);
+
+                            Array.Copy(tempBuf, 0, newTempBuf, 0, byteLen);
+                            ArrayPool<byte>.Shared.Return(tempBuf);
+                            tempBuf = newTempBuf;
+                        }
+
+                        Array.Copy(Buffer, 0, tempBuf, byteLen, i);
+                        byteLen += i;
+                        ReadPosition = i;
+                    } while (!foundTerminator);
+
+                    ReadPosition++;
+                    return encoding.GetString(tempBuf, 0, byteLen);
+                }
+                finally
+                {
+                    ArrayPool<byte>.Shared.Return(tempBuf);
+                }
             }
         }
 
diff --git a/test/Npgsql.Tests/ReadBufferTests.cs b/test/Npgsql.Tests/ReadBufferTests.cs
@@ -1,6 +1,5 @@
 ﻿using System;
 using System.IO;
-using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using Npgsql.Internal;
@@ -61,6 +60,8 @@ public void ReadNullTerminatedString_buffered_only()
                 .Write(PGUtil.UTF8Encoding.GetBytes(new string("bar")))
                 .WriteByte(0);
 
+            ReadBuffer.Ensure(1);
+
             Assert.That(ReadBuffer.ReadNullTerminatedString(), Is.EqualTo("foo"));
             Assert.That(ReadBuffer.ReadNullTerminatedString(), Is.EqualTo("bar"));
         }
@@ -69,6 +70,7 @@ public void ReadNullTerminatedString_buffered_only()
         public async Task ReadNullTerminatedString_with_io()
         {
             Writer.Write(PGUtil.UTF8Encoding.GetBytes(new string("Chunked ")));
+            ReadBuffer.Ensure(1);
             var task = ReadBuffer.ReadNullTerminatedString(async: true);
             Assert.That(!task.IsCompleted);
 
