dotnet · stephentoub · Oct 15, 2016 · Oct 14, 2016 · Oct 14, 2016 · jamesqo
diff --git a/src/Common/src/Interop/Windows/Winsock/SafeNativeOverlapped.cs b/src/Common/src/Interop/Windows/Winsock/SafeNativeOverlapped.cs
@@ -2,23 +2,17 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using Microsoft.Win32.SafeHandles;
-
 using System.Diagnostics;
-using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Threading;
 
 namespace System.Net.Sockets
 {
-    internal class SafeNativeOverlapped : SafeHandle
+    internal sealed class SafeNativeOverlapped : SafeHandle
     {
-        private static readonly SafeNativeOverlapped s_zero = new SafeNativeOverlapped();
-        private SafeCloseSocket _safeCloseSocket;
-
-        internal static SafeNativeOverlapped Zero { get { return s_zero; } }
+        internal static SafeNativeOverlapped Zero { get; } = new SafeNativeOverlapped();
 
-        protected SafeNativeOverlapped()
+        private SafeNativeOverlapped()
             : this(IntPtr.Zero)
         {
             if (GlobalLog.IsEnabled)
@@ -27,7 +21,7 @@ protected SafeNativeOverlapped()
             }
         }
 
-        protected SafeNativeOverlapped(IntPtr handle)
+        private SafeNativeOverlapped(IntPtr handle)
             : base(IntPtr.Zero, true)
         {
             SetHandle(handle);
@@ -36,32 +30,26 @@ protected SafeNativeOverlapped(IntPtr handle)
         public unsafe SafeNativeOverlapped(SafeCloseSocket socketHandle, NativeOverlapped* handle)
             : this((IntPtr)handle)
         {
-            _safeCloseSocket = socketHandle;
+            SocketHandle = socketHandle;
 
             if (GlobalLog.IsEnabled)
             {
                 GlobalLog.Print("SafeNativeOverlapped#" + LoggingHash.HashString(this) + "::ctor(socket#" + LoggingHash.HashString(socketHandle) + ")");
             }
 
 #if DEBUG
-            _safeCloseSocket.AddRef();
+            SocketHandle.AddRef();
 #endif
         }
 
-        protected override void Dispose(bool disposing)
+        internal unsafe void ReplaceHandle(NativeOverlapped* overlapped)
         {
-            if (disposing)
-            {
-                // It is important that the boundHandle is released immediately to allow new overlapped operations.
-                if (GlobalLog.IsEnabled)
-                {
-                    GlobalLog.Print("SafeNativeOverlapped#" + LoggingHash.HashString(this) + "::Dispose(true)");
-                }
-
-                FreeNativeOverlapped();
-            }
+            Debug.Assert(handle == IntPtr.Zero, "We should only be replacing the handle when it's already been freed.");
+            SetHandle((IntPtr)overlapped);
         }
 
+        internal SafeCloseSocket SocketHandle { get; }
+
         public override bool IsInvalid
         {
             get { return handle == IntPtr.Zero; }
@@ -75,37 +63,31 @@ protected override bool ReleaseHandle()
             }
 
             FreeNativeOverlapped();
+
+#if DEBUG
+            SocketHandle.Release();
+#endif
             return true;
         }
 
-        private void FreeNativeOverlapped()
+        internal void FreeNativeOverlapped()
         {
-            IntPtr oldHandle = Interlocked.Exchange(ref handle, IntPtr.Zero);
-
             // Do not call free during AppDomain shutdown, there may be an outstanding operation.
             // Overlapped will take care calling free when the native callback completes.
+            IntPtr oldHandle = Interlocked.Exchange(ref handle, IntPtr.Zero);
             if (oldHandle != IntPtr.Zero && !Environment.HasShutdownStarted)
             {
                 unsafe
                 {
-                    Debug.Assert(_safeCloseSocket != null, "m_SafeCloseSocket is null.");
-
-                    ThreadPoolBoundHandle boundHandle = _safeCloseSocket.IOCPBoundHandle;
-                    Debug.Assert(boundHandle != null, "SafeNativeOverlapped::ImmediatelyFreeNativeOverlapped - boundHandle is null");
+                    Debug.Assert(SocketHandle != null, "SocketHandle is null.");
 
-                    if (boundHandle != null)
-                    {
-                        // FreeNativeOverlapped will be called even if boundHandle was previously disposed.
-                        boundHandle.FreeNativeOverlapped((NativeOverlapped*)oldHandle);
-                    }
+                    ThreadPoolBoundHandle boundHandle = SocketHandle.IOCPBoundHandle;
+                    Debug.Assert(boundHandle != null, "SafeNativeOverlapped::FreeNativeOverlapped - boundHandle is null");
 
-#if DEBUG
-                    _safeCloseSocket.Release();
-#endif
-                    _safeCloseSocket = null;
+                    // FreeNativeOverlapped will be called even if boundHandle was previously disposed.
+                    boundHandle?.FreeNativeOverlapped((NativeOverlapped*)oldHandle);
                 }
             }
-            return;
         }
     }
 }
diff --git a/src/System.Net.Sockets/src/Resources/Strings.resx b/src/System.Net.Sockets/src/Resources/Strings.resx
@@ -291,4 +291,16 @@
   <data name="net_sockets_disconnect_notsupported" xml:space="preserve">
     <value>This platform does not support disconnecting a Socket.  Instead, close the Socket and create a new one.</value>
   </data>
+  <data name="NotSupported_UnreadableStream" xml:space="preserve">
+    <value>Stream does not support reading.</value>
+  </data>
+  <data name="NotSupported_UnwritableStream" xml:space="preserve">
+    <value>Stream does not support writing.</value>
+  </data>
+  <data name="ObjectDisposed_StreamClosed" xml:space="preserve">
+    <value>Can not access a closed Stream.</value>
+  </data>
+  <data name="ArgumentOutOfRange_NeedPosNum" xml:space="preserve">
+    <value>Positive number required.</value>
+  </data>
 </root>
diff --git a/src/System.Net.Sockets/src/System.Net.Sockets.csproj b/src/System.Net.Sockets/src/System.Net.Sockets.csproj
@@ -75,6 +75,12 @@
     <Compile Include="System\Net\Sockets\MultipleConnectAsync.cs" />
     <Compile Include="System\Net\Sockets\OverlappedAsyncResult.cs" />
     <Compile Include="System\Net\Sockets\ReceiveMessageOverlappedAsyncResult.cs" />
+    <Compile Include="$(CommonPath)\System\IO\StreamHelpers.ArrayPoolCopy.cs">
+      <Link>Common\System\IO\StreamHelpers.ArrayPoolCopy.cs</Link>
+    </Compile>
+    <Compile Include="$(CommonPath)\System\IO\StreamHelpers.CopyValidation.cs">
+      <Link>Common\System\IO\StreamHelpers.CopyValidation.cs</Link>
+    </Compile>
     <!-- Logging -->
     <Compile Include="$(CommonPath)\System\Net\Logging\GlobalLog.cs">
       <Link>Common\System\Net\Logging\GlobalLog.cs</Link>

diff --git a/src/System.Net.Sockets/src/System/Net/Sockets/NetworkStream.cs b/src/System.Net.Sockets/src/System/Net/Sockets/NetworkStream.cs
@@ -2,7 +2,10 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Buffers;
+using System.Diagnostics;
 using System.IO;
+using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Threading;
 using System.Threading.Tasks;
@@ -1025,6 +1028,63 @@ public override Task WriteAsync(byte[] buffer, int offset, int size, Cancellatio
                 this);
         }
 
+        public override Task CopyToAsync(Stream destination, int bufferSize, CancellationToken cancellationToken)
+        {
+            // Validate arguments as would the base CopyToAsync
+            StreamHelpers.ValidateCopyToArgs(this, destination, bufferSize);
+
+            // And bail early if cancellation has already been requested
+            if (cancellationToken.IsCancellationRequested)
+            {
+                return Task.FromCanceled(cancellationToken);
+            }
+
+            // Then do additional checks as ReadAsync would.
+
+            if (_cleanedUp)
+            {
+                throw new ObjectDisposedException(this.GetType().FullName);
+            }
+
+            Socket streamSocket = _streamSocket;
+            if (streamSocket == null)
+            {
+                throw new IOException(SR.Format(SR.net_io_readfailure, SR.net_io_connectionclosed));
+            }
+
+            // Do the copy.  We get a copy buffer from the shared pool, and we pass both it and the
+            // socket into the copy as part of the event args so as to avoid additional fields in
+            // the async method's state machine.
+            return CopyToAsyncCore(
+                destination,
+                new AwaitableSocketAsyncEventArgs(streamSocket, ArrayPool<byte>.Shared.Rent(bufferSize)),
+                cancellationToken);
+        }
+
+        private static async Task CopyToAsyncCore(Stream destination, AwaitableSocketAsyncEventArgs ea, CancellationToken cancellationToken)
+        {
+            try
+            {
+                while (true)
+                {
+                    cancellationToken.ThrowIfCancellationRequested();
+
+                    int bytesRead = await ea.ReceiveAsync();
+                    if (bytesRead == 0)
+                    {
+                        break;
+                    }
+
+                    await destination.WriteAsync(ea.Buffer, 0, bytesRead, cancellationToken).ConfigureAwait(false);
+                }
+            }
+            finally
+            {
+                ArrayPool<byte>.Shared.Return(ea.Buffer, clearArray: true);
+                ea.Dispose();
+            }
+        }
+
         // Flushes data from the stream.  This is meaningless for us, so it does nothing.
         public override void Flush()
         {
@@ -1092,5 +1152,114 @@ internal void DebugMembers()
                 _streamSocket.DebugMembers();
             }
         }
+
+        /// <summary>A SocketAsyncEventArgs that can be awaited to get the result of an operation.</summary>
+        internal sealed class AwaitableSocketAsyncEventArgs : SocketAsyncEventArgs, ICriticalNotifyCompletion
+        {
+            /// <summary>Sentinal object used to indicate that the operation has completed prior to OnCompleted being called.</summary>
+            private static readonly Action s_completedSentinel = () => { };
+            /// <summary>
+            /// null if the operation has not completed, <see cref="s_completedSentinel"/> if it has, and another object
+            /// if OnCompleted was called before the operation could complete, in which case it's the delegate to invoke
+            /// when the operation does complete.
+            /// </summary>
+            private Action _continuation;
+
+            /// <summary>Initializes the event args.</summary>
+            /// <param name="socket">The associated socket.</param>
+            /// <param name="buffer">The buffer to use for all operations.</param>
+            public AwaitableSocketAsyncEventArgs(Socket socket, byte[] buffer)
+            {
+                Debug.Assert(socket != null);
+                Debug.Assert(buffer != null && buffer.Length > 0);
+
+                // Store the socket into the base's UserToken.  This avoids the need for an extra field, at the expense
+                // of an object=>Socket cast when we need to access it, which is only once per operation.
+                UserToken = socket;
+
+                // Store the buffer for use by all operations with this instance.
+                SetBuffer(buffer, 0, buffer.Length);
+
+                // Hook up the completed event.
+                Completed += delegate
+                {
+                    // When the operation completes, see if OnCompleted was already called to hook up a continuation.
+                    // If it was, invoke the continuation.
+                    Action c = _continuation;
+                    if (c != null)
+                    {
+                        c();
+                    }
+                    else
+                    {
+                        // We may be racing with OnCompleted, so check with synchronization, trying to swap in our
+                        // completion sentinel.  If we lose the race and OnCompleted did hook up a continuation,
+                        // invoke it.  Otherwise, there's nothing more to be done.
+                        Interlocked.CompareExchange(ref _continuation, s_completedSentinel, null)?.Invoke();
+                    }
+                };
+            }
+
+            /// <summary>Initiates a receive operation on the associated socket.</summary>
+            /// <returns>This instance.</returns>
+            public AwaitableSocketAsyncEventArgs ReceiveAsync()
+            {
+                if (!Socket.ReceiveAsync(this))
+                {
+                    _continuation = s_completedSentinel;
+                }
+                return this;
+            }
+
+            /// <summary>Gets this instance.</summary>
+            public AwaitableSocketAsyncEventArgs GetAwaiter() => this;
+
+            /// <summary>Gets whether the operation has already completed.</summary>
+            /// <remarks>
+            /// This is not a generically usable IsCompleted operation that suggests the whole operation has completed.
+            /// Rather, it's specifically used as part of the await pattern, and is only usable to determine whether the
+            /// operation has completed by the time the instance is awaited.
+            /// </remarks>
+            public bool IsCompleted => _continuation != null;
+
+            /// <summary>Same as <see cref="OnCompleted(Action)"/> </summary>
+            public void UnsafeOnCompleted(Action continuation) => OnCompleted(continuation);
+
+            /// <summary>Queues the provided continuation to be executed once the operation has completed.</summary>
+            public void OnCompleted(Action continuation)
+            {
+                if (_continuation == s_completedSentinel || Interlocked.CompareExchange(ref _continuation, continuation, null) == s_completedSentinel)
+                {
+                    Task.Run(continuation);
+                }
+            }
+
+            /// <summary>Gets the result of the completion operation.</summary>
+            /// <returns>Number of bytes transferred.</returns>
+            /// <remarks>
+            /// Unlike Task's awaiter's GetResult, this does not block until the operation completes: it must only
+            /// be used once the operation has completed.  This is handled implicitly by await.
+            /// </remarks>
+            public int GetResult()
+            {
+                _continuation = null;
+                if (SocketError != SocketError.Success)
+                {
+                    ThrowIOSocketException();
+                }
+                return BytesTransferred;
+            }
+
+            /// <summary>Gets the associated socket.</summary>
+            internal Socket Socket => (Socket)UserToken; // stored in the base's UserToken to avoid an extra field in the object
+
+            /// <summary>Throws an IOException wrapping a SocketException using the current <see cref="SocketError"/>.</summary>
+            [MethodImpl(MethodImplOptions.NoInlining)]
+            private void ThrowIOSocketException()
+            {
+                var se = new SocketException((int)SocketError);
+                throw new IOException(SR.Format(SR.net_io_readfailure, se.Message), se);
+            }
+        }
     }
 }
diff --git a/src/System.Net.Sockets/src/System/Net/Sockets/SocketAsyncEventArgs.Windows.cs b/src/System.Net.Sockets/src/System/Net/Sockets/SocketAsyncEventArgs.Windows.cs
@@ -155,9 +155,21 @@ private unsafe void PrepareIOCPOperation()
                         "). Returned = " + ((IntPtr)overlapped).ToString("x"));
                 }
             }
-
             Debug.Assert(overlapped != null, "NativeOverlapped is null.");
-            _ptrNativeOverlapped = new SafeNativeOverlapped(_currentSocket.SafeHandle, overlapped);
+
+            // If we already have a SafeNativeOverlapped SafeHandle and it's associated with the same
+            // socket (due to the last operation that used this SocketAsyncEventArgs using the same socket),
+            // then we can reuse the same SafeHandle object.  Otherwise, this is either the first operation
+            // or the last operation was with a different socket, so create a new SafeHandle.
+            if (_ptrNativeOverlapped?.SocketHandle == _currentSocket.SafeHandle)
+            {
+                _ptrNativeOverlapped.ReplaceHandle(overlapped);
+            }
+            else
+            {
+                _ptrNativeOverlapped?.Dispose();
+                _ptrNativeOverlapped = new SafeNativeOverlapped(_currentSocket.SafeHandle, overlapped);
+            }
         }
 
         private void CompleteIOCPOperation()
@@ -169,12 +181,10 @@ private void CompleteIOCPOperation()
             // it is guaranteed that the IOCP operation will be completed in the callback even if Socket.Success was 
             // returned by the Win32 API.
 
-            // Required to allow another IOCP operation for the same handle.
-            if (_ptrNativeOverlapped != null)
-            {
-                _ptrNativeOverlapped.Dispose();
-                _ptrNativeOverlapped = null;
-            }
+            // Required to allow another IOCP operation for the same handle.  We release the native overlapped
+            // in the safe handle, but keep the safe handle object around so as to be able to reuse it
+            // for other operations.
+            _ptrNativeOverlapped?.FreeNativeOverlapped();
         }
 
         private void InnerStartOperationAccept(bool userSuppliedBuffer)

diff --git a/src/System.Net.Sockets/src/netcore50/project.json b/src/System.Net.Sockets/src/netcore50/project.json
@@ -5,6 +5,7 @@
         "Microsoft.NETCore.Platforms": "1.0.1",
         "Microsoft.NETCore.Targets": "1.0.1",
         "Microsoft.TargetingPack.Private.WinRT": "1.0.1",
+        "System.Buffers": "4.0.0",
         "System.Collections": "4.0.0",
         "System.Diagnostics.Debug": "4.0.10",
         "System.Diagnostics.Tracing": "4.0.20",