Add WebSocketStream

src/libraries/System.Net.WebSockets/ref/System.Net.WebSockets.cs

@@ -46,6 +46,32 @@ public static void RegisterPrefixes() { }
         public virtual System.Threading.Tasks.ValueTask SendAsync(System.ReadOnlyMemory<byte> buffer, System.Net.WebSockets.WebSocketMessageType messageType, System.Net.WebSockets.WebSocketMessageFlags messageFlags, System.Threading.CancellationToken cancellationToken) { throw null; }
         protected static void ThrowOnInvalidState(System.Net.WebSockets.WebSocketState state, params System.Net.WebSockets.WebSocketState[] validStates) { }
     }
+    public sealed partial class WebSocketStream : System.IO.Stream
+    {
+        public WebSocketStream(System.Net.WebSockets.WebSocket webSocket, bool ownsWebSocket = false) { }
+        public override System.IAsyncResult BeginWrite(byte[] buffer, int offset, int count, System.AsyncCallback? callback, object? state) { throw null; }
+        public override System.IAsyncResult BeginRead(byte[] buffer, int offset, int count, System.AsyncCallback? callback, object? state) { throw null; }
+        public override bool CanRead { get { throw null; } }
+        public override bool CanSeek { get { throw null; } }
+        public override bool CanWrite { get { throw null; } }
+        protected override void Dispose(bool disposing) { }
+        public override System.Threading.Tasks.ValueTask DisposeAsync() { throw null; }
+        public override int EndRead(System.IAsyncResult asyncResult) { throw null; }
+        public override void EndWrite(IAsyncResult asyncResult) { }
+        public override void Flush() { }
+        public override System.Threading.Tasks.Task FlushAsync(System.Threading.CancellationToken cancellationToken) { throw null; }
+        public override long Length { get { throw null; } }
+        public override long Position { get { throw null; } set { } }
+        public override int Read(byte[] buffer, int offset, int count) { throw null; }
+        public override System.Threading.Tasks.Task<int> ReadAsync(byte[] buffer, int offset, int count, System.Threading.CancellationToken cancellationToken) { throw null; }
+        public override System.Threading.Tasks.ValueTask<int> ReadAsync(System.Memory<byte> buffer, System.Threading.CancellationToken cancellationToken = default) { throw null; }
+        public override long Seek(long offset, System.IO.SeekOrigin origin) { throw null; }
+        public override void SetLength(long value) { }
+        public System.Net.WebSockets.WebSocket WebSocket { get { throw null; } }
+        public override void Write(byte[] buffer, int offset, int count) { throw null; }
+        public override System.Threading.Tasks.Task WriteAsync(byte[] buffer, int offset, int count, System.Threading.CancellationToken cancellationToken) { throw null; }
+        public override System.Threading.Tasks.ValueTask WriteAsync(System.ReadOnlyMemory<byte> buffer, System.Threading.CancellationToken cancellationToken = default) { throw null; }
+    }
     public enum WebSocketCloseStatus
     {
         NormalClosure = 1000,

src/libraries/System.Net.WebSockets/src/System.Net.WebSockets.csproj

@@ -33,6 +33,7 @@
     <Compile Include="System\Net\WebSockets\WebSocketReceiveResult.cs" />
     <Compile Include="System\Net\WebSockets\WebSocketState.cs" />
     <Compile Include="System\Net\WebSockets\WebSocketStateHelper.cs" />
+    <Compile Include="System\Net\WebSockets\WebSocketStream.cs" />
     <Compile Include="$(CommonPath)System\Net\WebSockets\WebSocketDefaults.cs"
              Link="Common\System\Net\WebSockets\WebSocketDefaults.cs" />
     <Compile Include="$(CommonPath)System\Net\WebSockets\WebSocketValidate.cs"

src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/WebSocketStream.cs

@@ -0,0 +1,196 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace System.Net.WebSockets
+{
+    /// <summary>Provides a <see cref="Stream"/> that delegates to a wrapped <see cref="WebSocket"/>.</summary>
+    public sealed class WebSocketStream : Stream
+    {
+        /// <summary>Whether disposing this instance should dispose <see cref="WebSocket"/>.</summary>
+        private readonly bool _ownsWebSocket;
+        /// <summary>Whether the instance has been disposed.</summary>
+        private bool _disposed;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="WebSocketStream"/> class using a specified <see cref="WebSocket"/> instance.
+        /// </summary>
+        /// <param name="webSocket">The <see cref="WebSocket"/> wrapped by this instance.</param>
+        /// <param name="ownsWebSocket">
+        /// <see langword="true"/> to indicate that the <see cref="WebSocketStream"/> takes ownership of the <see cref="WebSocket"/>,
+        /// such that disposing of this <see cref="Stream"/> will dispose of the <see cref="WebSocket"/>; otherwise, <see langword="false"/>.
+        /// When <see langword="true"/>, disposing of this instance doesn't initiate a close handshake; it merely delegates to
+        /// <see cref="WebSocket.Dispose"/>.
+        /// </param>
+        public WebSocketStream(WebSocket webSocket, bool ownsWebSocket = false)
+        {
+            ArgumentNullException.ThrowIfNull(webSocket);
+
+            WebSocket = webSocket;
+            _ownsWebSocket = ownsWebSocket;
+        }
+
+        /// <inheritdoc />
+        protected override void Dispose(bool disposing)
+        {
+            if (_disposed)
+            {
+                return;
+            }
+            _disposed = true;
+
+            if (disposing && _ownsWebSocket)
+            {
+                if (WebSocket.State is WebSocketState.Open)
+                {
+                    // There's no synchronous close, so we're forced to do sync-over-async.
+                    WebSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, null, CancellationToken.None).GetAwaiter().GetResult();
+                }
+
+                WebSocket.Dispose();
+            }
+
+            base.Dispose(disposing);
+        }
+
+        /// <inheritdoc />
+        public override ValueTask DisposeAsync()
+        {
+            if (_disposed)
+            {
+                return default;
+            }
+            _disposed = true;
+
+            if (_ownsWebSocket && WebSocket.State is WebSocketState.Open)
+            {
+                return CloseAndDisposeAsync();
+            }
+            else
+            {
+                return base.DisposeAsync();
+            }
+
+            async ValueTask CloseAndDisposeAsync()
+            {
+                if (WebSocket.State is WebSocketState.Open)
+                {
+                    await WebSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, null, CancellationToken.None).ConfigureAwait(false);
+                }
+                WebSocket.Dispose();
+                await base.DisposeAsync().ConfigureAwait(false);
+            }
+        }
+
+        /// <summary>Gets the underlying <see cref="WebSocket"/> wrapped by this <see cref="WebSocketStream"/>.</summary>
+        /// <remarks>The <see cref="WebSocket"/> used to construct this instance.</remarks>
+        public WebSocket WebSocket { get; }
+
+        /// <inheritdoc />
+        public override bool CanRead => WebSocket.State is WebSocketState.Open or WebSocketState.CloseSent;
+
+        /// <inheritdoc />
+        public override bool CanWrite => WebSocket.State is WebSocketState.Open or WebSocketState.CloseReceived;
+
+        /// <inheritdoc />
+        public override bool CanSeek => false;
+
+        /// <inheritdoc />
+        public override void Flush() { }
+
+        /// <inheritdoc />
+        public override Task FlushAsync(CancellationToken cancellationToken) => Task.CompletedTask;
+
+        /// <inheritdoc />
+        public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        {
+            ValidateBufferArguments(buffer, offset, count);
+
+            return ReadAsync(buffer.AsMemory(offset, count), cancellationToken).AsTask();
+        }
+
+        /// <inheritdoc />
+        public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
+        {
+            ObjectDisposedException.ThrowIf(_disposed, this);
+            cancellationToken.ThrowIfCancellationRequested();
+
+            while (WebSocket.State < WebSocketState.CloseReceived)
+            {
+                ValueWebSocketReceiveResult result = await WebSocket.ReceiveAsync(buffer, cancellationToken).ConfigureAwait(false);
+                if (result.MessageType == WebSocketMessageType.Close)
+                {
+                    break;
+                }
+
+                if (result.Count > 0 || buffer.IsEmpty)
+                {
+                    return result.Count;
+                }
+            }
+
+            return 0;
+        }
+
+        /// <inheritdoc />
+        public override IAsyncResult BeginRead(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state) =>
+            TaskToAsyncResult.Begin(ReadAsync(buffer, offset, count), callback, state);
+
+        /// <inheritdoc />
+        public override int EndRead(IAsyncResult asyncResult) =>
+            TaskToAsyncResult.End<int>(asyncResult);
+
+        /// <inheritdoc />
+        public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        {
+            ValidateBufferArguments(buffer, offset, count);
+
+            return WriteAsync(buffer.AsMemory(offset, count), cancellationToken).AsTask();
+        }
+
+        /// <inheritdoc />
+        public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
+        {
+            if (_disposed)
+            {
+                return ValueTask.FromException(new ObjectDisposedException(GetType().FullName));
+            }
+
+            if (cancellationToken.IsCancellationRequested)
+            {
+                return ValueTask.FromCanceled(cancellationToken);
+            }
+
+            return WebSocket.SendAsync(buffer, WebSocketMessageType.Binary, endOfMessage: true, cancellationToken);
+        }
+
+        /// <inheritdoc />
+        public override IAsyncResult BeginWrite(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state) =>
+            TaskToAsyncResult.Begin(WriteAsync(buffer, offset, count), callback, state);
+
+        /// <inheritdoc />
+        public override void EndWrite(IAsyncResult asyncResult) =>
+            TaskToAsyncResult.End(asyncResult);
+
+        /// <inheritdoc />
+        public override int Read(byte[] buffer, int offset, int count) => ReadAsync(buffer, offset, count).GetAwaiter().GetResult();
+
+        /// <inheritdoc />
+        public override void Write(byte[] buffer, int offset, int count) => WriteAsync(buffer, offset, count).GetAwaiter().GetResult();
+
+        /// <inheritdoc />
+        public override long Length => throw new NotSupportedException();
+
+        /// <inheritdoc />
+        public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); }
+
+        /// <inheritdoc />
+        public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
+
+        /// <inheritdoc />
+        public override void SetLength(long value) => throw new NotSupportedException();
+    }
+}

src/libraries/System.Net.WebSockets/tests/System.Net.WebSockets.Tests.csproj

@@ -8,6 +8,7 @@
   </ItemGroup>
   <ItemGroup>
     <Compile Include="WebSocketCloseTests.cs" />
+    <Compile Include="WebSocketStreamTests.cs" />
     <Compile Include="WebSocketTests.cs" />
     <Compile Include="WebSocketExceptionTests.cs" />
     <Compile Include="WebSocketKeepAliveTests.cs" />
@@ -24,4 +25,7 @@
     <Compile Include="$(CommonPath)System\Net\StreamBuffer.cs" Link="ProductionCode\Common\System\Net\StreamBuffer.cs" />
     <Compile Include="$(CommonPath)System\Net\MultiArrayBuffer.cs" Link="ProductionCode\Common\System\Net\MultiArrayBuffer.cs" />
   </ItemGroup>
+  <ItemGroup>
+    <ProjectReference Include="$(CommonTestPath)StreamConformanceTests\StreamConformanceTests.csproj" />
+  </ItemGroup>
 </Project>

src/libraries/System.Net.WebSockets/tests/WebSocketStreamTests.cs

@@ -0,0 +1,72 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.IO;
+using System.IO.Tests;
+using System.Threading;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace System.Net.WebSockets.Tests
+{
+    public sealed class WebSocketStreamTests : ConnectedStreamConformanceTests
+    {
+        protected override bool BlocksOnZeroByteReads => true;
+        protected override bool FlushRequiredToWriteData => false;
+        protected override bool ReadsReadUntilSizeOrEof => false;
+        protected override bool UsableAfterCanceledReads => false;
+        protected override Type UnsupportedConcurrentExceptionType => null;
+
+        protected override Task<StreamPair> CreateConnectedStreamsAsync()
+        {
+            (Stream stream1, Stream stream2) = ConnectedStreams.CreateBidirectional();
+
+            WebSocket webSocket1 = WebSocket.CreateFromStream(stream1, isServer: false, null, Timeout.InfiniteTimeSpan);
+            WebSocket webSocket2 = WebSocket.CreateFromStream(stream2, isServer: true, null, Timeout.InfiniteTimeSpan);
+
+            return Task.FromResult(new StreamPair(
+                new WebSocketStream(webSocket1, ownsWebSocket: true),
+                new WebSocketStream(webSocket2, ownsWebSocket: true)));
+        }
+
+        [Fact]
+        public void Ctor_InvalidArgs_Throws()
+        {
+            AssertExtensions.Throws<ArgumentNullException>("webSocket", () => new WebSocketStream(null));
+            AssertExtensions.Throws<ArgumentNullException>("webSocket", () => new WebSocketStream(null, ownsWebSocket: true));
+        }
+
+        [Fact]
+        public void Ctor_Roundtrips()
+        {
+            (Stream stream1, Stream stream2) = ConnectedStreams.CreateBidirectional();
+
+            WebSocket webSocket = WebSocket.CreateFromStream(stream1, isServer: false, null, Timeout.InfiniteTimeSpan);
+
+            WebSocketStream stream = new WebSocketStream(webSocket);
+            Assert.Same(webSocket, stream.WebSocket);
+            stream.Dispose();
+            Assert.Same(webSocket, stream.WebSocket);
+
+            stream2.Dispose();
+        }
+
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public void Dispose_ClosesWebSocketIfOwned(bool ownsWebSocket)
+        {
+            (Stream stream1, Stream stream2) = ConnectedStreams.CreateBidirectional();
+
+            WebSocket webSocket = WebSocket.CreateFromStream(stream1, isServer: false, null, Timeout.InfiniteTimeSpan);
+
+            WebSocketStream stream = new WebSocketStream(webSocket, ownsWebSocket);
+            Assert.Equal(WebSocketState.Open, webSocket.State);
+
+            stream.Dispose();
+            Assert.Equal(ownsWebSocket ? WebSocketState.Closed : WebSocketState.Open, webSocket.State);
+
+            stream2.Dispose();
+        }
+    }
+}