huahaofeng
/
aspnetcore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Avoid zero-byte send in WebSockets (#3326)

This commit is contained in:
BrennanConroy 2018-11-21 11:47:39 -08:00 committed by GitHub
parent dde73d0b8a
commit c852bdcc33
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 42 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
src/Common/WebSocketExtensions.cs
View File

@ -39,22 +39,28 @@ namespace System.Net.WebSockets
private static async ValueTask SendMultiSegmentAsync(WebSocket webSocket, ReadOnlySequence<byte> buffer, WebSocketMessageType webSocketMessageType, CancellationToken cancellationToken = default)
{
var position = buffer.Start;
// Get a segment before the loop so we can be one segment behind while writing
// This allows us to do a non-zero byte write for the endOfMessage = true send
buffer.TryGet(ref position, out var prevSegment);
while (buffer.TryGet(ref position, out var segment))
{
#if NETCOREAPP3_0
await webSocket.SendAsync(segment, webSocketMessageType, endOfMessage: false, cancellationToken);
await webSocket.SendAsync(prevSegment, webSocketMessageType, endOfMessage: false, cancellationToken);
#else
var isArray = MemoryMarshal.TryGetArray(segment, out var arraySegment);
var isArray = MemoryMarshal.TryGetArray(prevSegment, out var arraySegment);
Debug.Assert(isArray);
await webSocket.SendAsync(arraySegment, webSocketMessageType, endOfMessage: false, cancellationToken);
#endif
prevSegment = segment;
}
// Empty end of message frame
// End of message frame
#if NETCOREAPP3_0
await webSocket.SendAsync(Memory<byte>.Empty, webSocketMessageType, endOfMessage: true, cancellationToken);
await webSocket.SendAsync(prevSegment, webSocketMessageType, endOfMessage: true, cancellationToken);
#else
await webSocket.SendAsync(new ArraySegment<byte>(Array.Empty<byte>()), webSocketMessageType, endOfMessage: true, cancellationToken);
var isArrayEnd = MemoryMarshal.TryGetArray(prevSegment, out var arraySegmentEnd);
Debug.Assert(isArrayEnd);
await webSocket.SendAsync(arraySegmentEnd, webSocketMessageType, endOfMessage: true, cancellationToken);
#endif
}
}

31
test/Microsoft.AspNetCore.Http.Connections.Tests/WebSocketsTests.cs
View File

@ -396,5 +396,36 @@ namespace Microsoft.AspNetCore.Http.Connections.Tests
}
}
}
[Fact]
public async Task MultiSegmentSendWillNotSendEmptyEndOfMessageFrame()
{
using (var feature = new TestWebSocketConnectionFeature())
{
var serverSocket = await feature.AcceptAsync();
var sequence = ReadOnlySequenceFactory.CreateSegments(new byte[] { 1 }, new byte[] { 15 });
Assert.False(sequence.IsSingleSegment);
await serverSocket.SendAsync(sequence, WebSocketMessageType.Text);
// Run the client socket
var client = feature.Client.ExecuteAndCaptureFramesAsync();
await serverSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "", default);
var messages = await client.OrTimeout();
Assert.Equal(2, messages.Received.Count);
// First message: 1 byte, endOfMessage false
Assert.Single(messages.Received[0].Buffer);
Assert.Equal(1, messages.Received[0].Buffer[0]);
Assert.False(messages.Received[0].EndOfMessage);
// Second message: 1 byte, endOfMessage true
Assert.Single(messages.Received[1].Buffer);
Assert.Equal(15, messages.Received[1].Buffer[0]);
Assert.True(messages.Received[1].EndOfMessage);
}
}
}
}