Skip to content

Commit 7dd2d4e

Browse files
committed
Fix possible nre while writing to RawCopyStream after break (#5213)
Fixes #5209 (cherry picked from commit d8c70e5) # Conflicts: # test/Npgsql.Tests/CopyTests.cs
1 parent 089442a commit 7dd2d4e

File tree

3 files changed

+43
-35
lines changed

3 files changed

+43
-35
lines changed

src/Npgsql/NpgsqlRawCopyStream.cs

Lines changed: 17 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -142,25 +142,17 @@ public override void Write(ReadOnlySpan<byte> buffer)
142142
return;
143143
}
144144

145-
try
146-
{
147-
// Value is too big, flush.
148-
Flush();
149-
150-
if (buffer.Length <= _writeBuf.WriteSpaceLeft)
151-
{
152-
_writeBuf.WriteBytes(buffer);
153-
return;
154-
}
145+
// Value is too big, flush.
146+
Flush();
155147

156-
// Value is too big even after a flush - bypass the buffer and write directly.
157-
_writeBuf.DirectWrite(buffer);
158-
}
159-
catch (Exception e)
148+
if (buffer.Length <= _writeBuf.WriteSpaceLeft)
160149
{
161-
_connector.Break(e);
162-
throw;
150+
_writeBuf.WriteBytes(buffer);
151+
return;
163152
}
153+
154+
// Value is too big even after a flush - bypass the buffer and write directly.
155+
_writeBuf.DirectWrite(buffer);
164156
}
165157

166158
#if NETSTANDARD2_0
@@ -187,25 +179,17 @@ async ValueTask WriteAsyncInternal(ReadOnlyMemory<byte> buffer, CancellationToke
187179
return;
188180
}
189181

190-
try
191-
{
192-
// Value is too big, flush.
193-
await FlushAsync(true, cancellationToken);
194-
195-
if (buffer.Length <= _writeBuf.WriteSpaceLeft)
196-
{
197-
_writeBuf.WriteBytes(buffer.Span);
198-
return;
199-
}
182+
// Value is too big, flush.
183+
await FlushAsync(true, cancellationToken);
200184

201-
// Value is too big even after a flush - bypass the buffer and write directly.
202-
await _writeBuf.DirectWrite(buffer, true, cancellationToken);
203-
}
204-
catch (Exception e)
185+
if (buffer.Length <= _writeBuf.WriteSpaceLeft)
205186
{
206-
_connector.Break(e);
207-
throw;
187+
_writeBuf.WriteBytes(buffer.Span);
188+
return;
208189
}
190+
191+
// Value is too big even after a flush - bypass the buffer and write directly.
192+
await _writeBuf.DirectWrite(buffer, true, cancellationToken);
209193
}
210194
}
211195

@@ -580,4 +564,4 @@ public ValueTask DisposeAsync()
580564
Dispose();
581565
return default;
582566
}
583-
}
567+
}

test/Npgsql.Tests/CopyTests.cs

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
using System.Threading;
99
using System.Threading.Tasks;
1010
using Npgsql.Internal;
11+
using Npgsql.Tests.Support;
1112
using NpgsqlTypes;
1213
using NUnit.Framework;
1314
using static Npgsql.Tests.TestUtil;
@@ -1113,6 +1114,29 @@ public async Task Copy_is_not_supported_in_regular_command_execution()
11131114
Assert.That(() => conn.ExecuteNonQuery($@"COPY {table} (foo) FROM stdin"), Throws.Exception.TypeOf<NotSupportedException>());
11141115
}
11151116

1117+
[Test, IssueLink("https://github.com/npgsql/npgsql/issues/5209")]
1118+
[Platform(Exclude = "MacOsX", Reason = "Write might not throw an exception")]
1119+
public async Task RawBinaryCopy_write_nre([Values] bool async)
1120+
{
1121+
await using var postmasterMock = PgPostmasterMock.Start(ConnectionString);
1122+
using var _ = CreateTempPool(postmasterMock.ConnectionString, out var connectionString);
1123+
await using var conn = await OpenConnectionAsync(connectionString);
1124+
1125+
var server = await postmasterMock.WaitForServerConnection();
1126+
await server
1127+
.WriteCopyInResponse(isBinary: true)
1128+
.FlushAsync();
1129+
1130+
await using var stream = await conn.BeginRawBinaryCopyAsync("COPY SomeTable (field_text, field_int4) FROM STDIN");
1131+
server.Close();
1132+
var value = Encoding.UTF8.GetBytes(new string('a', conn.Settings.WriteBufferSize * 2));
1133+
if (async)
1134+
Assert.ThrowsAsync<NpgsqlException>(async () => await stream.WriteAsync(value));
1135+
else
1136+
Assert.Throws<NpgsqlException>(() => stream.Write(value));
1137+
Assert.That(conn.FullState, Is.EqualTo(ConnectionState.Broken));
1138+
}
1139+
11161140
#endregion
11171141

11181142
#region Utils

test/Npgsql.Tests/Support/PgServerMock.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -328,12 +328,12 @@ internal PgServerMock WriteBackendKeyData(int processId, int secret)
328328
internal PgServerMock WriteCancellationResponse()
329329
=> WriteErrorResponse(PostgresErrorCodes.QueryCanceled, "Cancellation", "Query cancelled");
330330

331-
internal PgServerMock WriteCopyInResponse()
331+
internal PgServerMock WriteCopyInResponse(bool isBinary = false)
332332
{
333333
CheckDisposed();
334334
_writeBuffer.WriteByte((byte)BackendMessageCode.CopyInResponse);
335335
_writeBuffer.WriteInt32(5);
336-
_writeBuffer.WriteByte(0);
336+
_writeBuffer.WriteByte(isBinary ? (byte)1 : (byte)0);
337337
_writeBuffer.WriteInt16(1);
338338
_writeBuffer.WriteInt16(0);
339339
return this;

0 commit comments

Comments
 (0)