Fix possible nre while writing to RawCopyStream after break (#5213)

vonzshik · web-flow · commit d8c70e561740 · 2023-08-25T18:12:30.000+03:00
Fixes #5209
diff --git a/src/Npgsql/NpgsqlRawCopyStream.cs b/src/Npgsql/NpgsqlRawCopyStream.cs
@@ -143,25 +143,17 @@ public override void Write(ReadOnlySpan<byte> buffer)
             return;
         }
 
-        try
-        {
-            // Value is too big, flush.
-            Flush();
-
-            if (buffer.Length <= _writeBuf.WriteSpaceLeft)
-            {
-                _writeBuf.WriteBytes(buffer);
-                return;
-            }
+        // Value is too big, flush.
+        Flush();
 
-            // Value is too big even after a flush - bypass the buffer and write directly.
-            _writeBuf.DirectWrite(buffer);
-        }
-        catch (Exception e)
+        if (buffer.Length <= _writeBuf.WriteSpaceLeft)
         {
-            _connector.Break(e);
-            throw;
+            _writeBuf.WriteBytes(buffer);
+            return;
         }
+
+        // Value is too big even after a flush - bypass the buffer and write directly.
+        _writeBuf.DirectWrite(buffer);
     }
 
 #if NETSTANDARD2_0
@@ -188,25 +180,17 @@ async ValueTask WriteAsyncInternal(ReadOnlyMemory<byte> buffer, CancellationToke
                 return;
             }
 
-            try
-            {
-                // Value is too big, flush.
-                await FlushAsync(true, cancellationToken);
-
-                if (buffer.Length <= _writeBuf.WriteSpaceLeft)
-                {
-                    _writeBuf.WriteBytes(buffer.Span);
-                    return;
-                }
+            // Value is too big, flush.
+            await FlushAsync(true, cancellationToken);
 
-                // Value is too big even after a flush - bypass the buffer and write directly.
-                await _writeBuf.DirectWrite(buffer, true, cancellationToken);
-            }
-            catch (Exception e)
+            if (buffer.Length <= _writeBuf.WriteSpaceLeft)
             {
-                _connector.Break(e);
-                throw;
+                _writeBuf.WriteBytes(buffer.Span);
+                return;
             }
+
+            // Value is too big even after a flush - bypass the buffer and write directly.
+            await _writeBuf.DirectWrite(buffer, true, cancellationToken);
         }
     }
 
@@ -580,4 +564,4 @@ public ValueTask DisposeAsync()
         Dispose();
         return default;
     }
-}
+}
diff --git a/test/Npgsql.Tests/CopyTests.cs b/test/Npgsql.Tests/CopyTests.cs
@@ -8,6 +8,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Npgsql.Internal;
+using Npgsql.Tests.Support;
 using NpgsqlTypes;
 using NUnit.Framework;
 using static Npgsql.Tests.TestUtil;
@@ -1127,6 +1128,29 @@ public async Task Copy_to_is_not_supported_in_regular_command_execution()
         Assert.That(() => conn.ExecuteNonQuery($@"COPY {table} (foo) TO stdin"), Throws.Exception.TypeOf<NotSupportedException>());
     }
 
+    [Test, IssueLink("https://github.com/npgsql/npgsql/issues/5209")]
+    [Platform(Exclude = "MacOsX", Reason = "Write might not throw an exception")]
+    public async Task RawBinaryCopy_write_nre([Values] bool async)
+    {
+        await using var postmasterMock = PgPostmasterMock.Start(ConnectionString);
+        await using var dataSource = CreateDataSource(postmasterMock.ConnectionString);
+        await using var conn = await dataSource.OpenConnectionAsync();
+
+        var server = await postmasterMock.WaitForServerConnection();
+        await server
+            .WriteCopyInResponse(isBinary: true)
+            .FlushAsync();
+
+        await using var stream = await conn.BeginRawBinaryCopyAsync("COPY SomeTable (field_text, field_int4) FROM STDIN");
+        server.Close();
+        var value = Encoding.UTF8.GetBytes(new string('a', conn.Settings.WriteBufferSize * 2));
+        if (async)
+            Assert.ThrowsAsync<NpgsqlException>(async () => await stream.WriteAsync(value));
+        else
+            Assert.Throws<NpgsqlException>(() => stream.Write(value));
+        Assert.That(conn.FullState, Is.EqualTo(ConnectionState.Broken));
+    }
+
     #endregion
 
     #region Utils
diff --git a/test/Npgsql.Tests/Support/PgServerMock.cs b/test/Npgsql.Tests/Support/PgServerMock.cs
@@ -328,12 +328,12 @@ internal PgServerMock WriteBackendKeyData(int processId, int secret)
     internal PgServerMock WriteCancellationResponse()
         => WriteErrorResponse(PostgresErrorCodes.QueryCanceled, "Cancellation", "Query cancelled");
 
-    internal PgServerMock WriteCopyInResponse()
+    internal PgServerMock WriteCopyInResponse(bool isBinary = false)
     {
         CheckDisposed();
         _writeBuffer.WriteByte((byte)BackendMessageCode.CopyInResponse);
         _writeBuffer.WriteInt32(5);
-        _writeBuffer.WriteByte(0);
+        _writeBuffer.WriteByte(isBinary ? (byte)1 : (byte)0);
         _writeBuffer.WriteInt16(1);
         _writeBuffer.WriteInt16(0);
         return this;
@@ -381,4 +381,4 @@ public void Dispose()
 
         _disposed = true;
     }
-}
+}
