diff --git a/src/Frame.cs b/src/Frame.cs index f57c23a..c822e22 100644 --- a/src/Frame.cs +++ b/src/Frame.cs @@ -37,9 +37,7 @@ public byte[] ToArray() } public static async Task FromStream(Stream stream, CancellationToken token) { - var header = new byte[2]; - var read = await stream.ReadAsync(header, token); - if (read == 0) throw new EndOfStreamException(); + var header = await stream.ReadExactly(2, token); var isMasked = (header[1] & 0x80) == 0x80; var isFinal = (header[0] & 0x80) == 0x80; var opCode = (FrameOpCode)(header[0] & 0x0F); @@ -49,29 +47,15 @@ public static async Task FromStream(Stream stream, CancellationToken toke 127 => new byte[8], _ => new byte[0], }; - if (payloadLength.Length != 0) - { - read = await stream.ReadAsync(payloadLength, token); - if (read == 0) throw new EndOfStreamException(); - } + await stream.ReadExactly(payloadLength, token); var realLength = payloadLength.Length switch { 2 => BitConverter.ToUInt16(payloadLength), 8 => BitConverter.ToUInt64(payloadLength), _ => (ulong)(header[1] & 0x7F), }; - var maskingKey = isMasked ? new byte[4] : null; - if (maskingKey != null) - { - read = await stream.ReadAsync(maskingKey, token); - if (read == 0) throw new EndOfStreamException(); - } - var payload = new byte[realLength]; - if (realLength != 0) - { - read = await stream.ReadAsync(payload, token); - if (read == 0) throw new EndOfStreamException(); - } + var maskingKey = isMasked ? await stream.ReadExactly(4, token) : null; + var payload = await stream.ReadExactly((int)realLength, token); return new Frame(isFinal, opCode, maskingKey, payload); } diff --git a/src/StreamHelper.cs b/src/StreamHelper.cs new file mode 100644 index 0000000..019524e --- /dev/null +++ b/src/StreamHelper.cs @@ -0,0 +1,32 @@ + +using System.IO; +using System.Threading; +using System.Threading.Tasks; + +namespace SwebSocket; + +internal static class StreamHelper +{ + public static async Task ReadExactly(this Stream stream, byte[] buffer, CancellationToken token) + { + if (buffer.Length == 0) return; + + int alreadyRead = 0; + int toBeRead = buffer.Length; + + while (alreadyRead != toBeRead) + { + var actuallyRead = await stream.ReadAsync(buffer, alreadyRead, toBeRead - alreadyRead, token); + if (actuallyRead == 0) throw new EndOfStreamException(); + alreadyRead += actuallyRead; + } + } + + public static async Task ReadExactly(this Stream stream, int bytes, CancellationToken token) + { + var buffer = new byte[bytes]; + await ReadExactly(stream, buffer, token); + + return buffer; + } +} \ No newline at end of file