Fix an issue that the socket may not be released.

- New wrapper to handle tcp socket connection.
- Fix some memory leak.

Signed-off-by: noisyfox <timemanager.rick@gmail.com>

shadowsocks-csharp/Controller/Service/PortForwarder.cs

@@ -1,7 +1,7 @@
 using System;
 using System.Net;
 using System.Net.Sockets;
-using Shadowsocks.Util;
+using Shadowsocks.Util.Sockets;
 
 namespace Shadowsocks.Controller
 {
@@ -29,7 +29,7 @@ namespace Shadowsocks.Controller
             private byte[] _firstPacket;
             private int _firstPacketLength;
             private Socket _local;
-            private Socket _remote;
+            private WrappedSocket _remote;
             private bool _closed = false;
             private bool _localShutdown = false;
             private bool _remoteShutdown = false;
@@ -49,7 +49,8 @@ namespace Shadowsocks.Controller
                     EndPoint remoteEP = SocketUtil.GetEndPoint("localhost", targetPort);
 
                     // Connect to the remote endpoint.
-                    SocketUtil.BeginConnectTcp(remoteEP, ConnectCallback, null);
+                    _remote = new WrappedSocket();
+                    _remote.BeginConnect(remoteEP, ConnectCallback, null);
                 }
                 catch (Exception e)
                 {
@@ -66,7 +67,8 @@ namespace Shadowsocks.Controller
                 }
                 try
                 {
-                    _remote = SocketUtil.EndConnectTcp(ar);
+                    _remote.EndConnect(ar);
+                    _remote.SetSocketOption(SocketOptionLevel.Tcp, SocketOptionName.NoDelay, true);
                     HandshakeReceive();
                 }
                 catch (Exception e)
@@ -243,7 +245,7 @@ namespace Shadowsocks.Controller
                     try
                     {
                         _remote.Shutdown(SocketShutdown.Both);
-                        _remote.Close();
+                        _remote.Dispose();
                     }
                     catch (SocketException e)
                     {

shadowsocks-csharp/Controller/Service/TCPRelay.cs

@@ -9,7 +9,7 @@ using Shadowsocks.Controller.Strategy;
 using Shadowsocks.Encryption;
 using Shadowsocks.Model;
 using Shadowsocks.Proxy;
-using Shadowsocks.Util;
+using Shadowsocks.Util.Sockets;
 
 namespace Shadowsocks.Controller
 {
@@ -40,7 +40,6 @@ namespace Shadowsocks.Controller
             handler.controller = _controller;
             handler.tcprelay = this;
 
-            handler.Start(firstPacket, length);
             IList<TCPHandler> handlersToClose = new List<TCPHandler>();
             lock (Handlers)
             {
@@ -59,6 +58,15 @@ namespace Shadowsocks.Controller
                 Logging.Debug("Closing timed out TCP connection.");
                 handler1.Close();
             }
+
+            /*
+             * Start after we put it into Handlers set. Otherwise if it failed in handler.Start()
+             * then it will call handler.Close() before we add it into the set.
+             * Then the handler will never release until the next Handle call. Sometimes it will
+             * cause odd problems (especially during memory profiling).
+             */
+            handler.Start(firstPacket, length);
+
             return true;
         }
 
@@ -455,7 +463,7 @@ namespace Shadowsocks.Controller
             timer.Dispose();
 
 
-            if (_proxyConnected || _destConnected)
+            if (_proxyConnected || _destConnected || _closed)
             {
                 return;
             }
@@ -527,7 +535,7 @@ namespace Shadowsocks.Controller
             timer.Enabled = false;
             timer.Dispose();
 
-            if (_destConnected)
+            if (_destConnected || _closed)
             {
                 return;
             }

shadowsocks-csharp/Proxy/DirectConnect.cs

@@ -2,7 +2,7 @@
 using System.Net;
 using System.Net.Sockets;
 using System.Threading;
-using Shadowsocks.Util;
+using Shadowsocks.Util.Sockets;
 
 namespace Shadowsocks.Proxy
 {
@@ -31,7 +31,7 @@ namespace Shadowsocks.Proxy
             }
         }
 
-        private Socket _remote;
+        private WrappedSocket _remote = new WrappedSocket();
 
         public EndPoint LocalEndPoint => _remote.LocalEndPoint;
 
@@ -55,12 +55,13 @@ namespace Shadowsocks.Proxy
         {
             DestEndPoint = destEndPoint;
 
-            SocketUtil.BeginConnectTcp(destEndPoint, callback, state);
+            _remote.BeginConnect(destEndPoint, callback, state);
         }
 
         public void EndConnectDest(IAsyncResult asyncResult)
         {
-            _remote = SocketUtil.EndConnectTcp(asyncResult);
+            _remote.EndConnect(asyncResult);
+            _remote.SetSocketOption(SocketOptionLevel.Tcp, SocketOptionName.NoDelay, true);
         }
 
         public void BeginSend(byte[] buffer, int offset, int size, SocketFlags socketFlags, AsyncCallback callback,
@@ -92,7 +93,7 @@ namespace Shadowsocks.Proxy
 
         public void Close()
         {
-            _remote?.Close();
+            _remote?.Dispose();
         }
     }
 }

shadowsocks-csharp/Proxy/Socks5Proxy.cs

@@ -1,12 +1,10 @@
 using System;
-using System.Collections.Generic;
-using System.Linq;
 using System.Net;
 using System.Net.Sockets;
 using System.Text;
 using System.Threading;
 using Shadowsocks.Controller;
-using Shadowsocks.Util;
+using Shadowsocks.Util.Sockets;
 
 namespace Shadowsocks.Proxy
 {
@@ -41,7 +39,7 @@ namespace Shadowsocks.Proxy
             public Exception ex { get; set; }
         }
 
-        private Socket _remote;
+        private WrappedSocket _remote = new WrappedSocket();
 
         private const int Socks5PktMaxSize = 4 + 16 + 2;
         private readonly byte[] _receiveBuffer = new byte[Socks5PktMaxSize];
@@ -58,7 +56,7 @@ namespace Shadowsocks.Proxy
 
             ProxyEndPoint = remoteEP;
 
-            SocketUtil.BeginConnectTcp(remoteEP, ConnectCallback, st);
+            _remote.BeginConnect(remoteEP, ConnectCallback, st);
         }
 
         public void EndConnectProxy(IAsyncResult asyncResult)
@@ -168,7 +166,7 @@ namespace Shadowsocks.Proxy
 
         public void Close()
         {
-            _remote?.Close();
+            _remote?.Dispose();
         }
         
 
@@ -177,7 +175,9 @@ namespace Shadowsocks.Proxy
             var state = (Socks5State) ar.AsyncState;
             try
             {
-                _remote = SocketUtil.EndConnectTcp(ar);
+                _remote.EndConnect(ar);
+
+                _remote.SetSocketOption(SocketOptionLevel.Tcp, SocketOptionName.NoDelay, true);
 
                 byte[] handshake = {5, 1, 0};
                 _remote.BeginSend(handshake, 0, handshake.Length, 0, Socks5HandshakeSendCallback, state);

shadowsocks-csharp/Util/SocketUtil.cs

@@ -1,132 +0,0 @@
-using System;
-using System.Net;
-using System.Net.Sockets;
-using System.Threading;
-
-namespace Shadowsocks.Util
-{
-    public static class SocketUtil
-    {
-        private class DnsEndPoint2 : DnsEndPoint
-        {
-            public DnsEndPoint2(string host, int port) : base(host, port)
-            {
-            }
-
-            public DnsEndPoint2(string host, int port, AddressFamily addressFamily) : base(host, port, addressFamily)
-            {
-            }
-
-            public override string ToString()
-            {
-                return this.Host + ":" + this.Port;
-            }
-        }
-
-        public static EndPoint GetEndPoint(string host, int port)
-        {
-            IPAddress ipAddress;
-            bool parsed = IPAddress.TryParse(host, out ipAddress);
-            if (parsed)
-            {
-                return new IPEndPoint(ipAddress, port);
-            }
-
-            // maybe is a domain name
-            return new DnsEndPoint2(host, port);
-        }
-
-        private class AutoReleaseAsyncResult : IAsyncResult
-        {
-
-            public bool IsCompleted { get; } = true;
-            public WaitHandle AsyncWaitHandle { get; } = null;
-            public object AsyncState { get; set; }
-            public bool CompletedSynchronously { get; } = true;
-
-            public TcpUserToken UserToken { get; set; }
-
-            ~AutoReleaseAsyncResult()
-            {
-                UserToken.Dispose();
-            }
-        }
-
-        private class TcpUserToken
-        {
-            public AsyncCallback Callback { get; private set; }
-            public SocketAsyncEventArgs Args { get; private set; }
-            public object AsyncState { get; private set; }
-
-            public TcpUserToken(AsyncCallback callback, object state, SocketAsyncEventArgs args)
-            {
-                Callback = callback;
-                AsyncState = state;
-                Args = args;
-            }
-
-            public void Dispose()
-            {
-                Args?.Dispose();
-                Callback = null;
-                Args = null;
-                AsyncState = null;
-            }
-        }
-
-        private static void OnTcpConnectCompleted(object sender, SocketAsyncEventArgs args)
-        {
-            args.Completed -= OnTcpConnectCompleted;
-            TcpUserToken token = (TcpUserToken) args.UserToken;
-
-            AutoReleaseAsyncResult r = new AutoReleaseAsyncResult
-            {
-                AsyncState = token.AsyncState,
-                UserToken = token
-            };
-
-            token.Callback(r);
-        }
-
-        public static void BeginConnectTcp(EndPoint endPoint, AsyncCallback callback, object state)
-        {
-            var arg = new SocketAsyncEventArgs();
-            arg.RemoteEndPoint = endPoint;
-            arg.Completed += OnTcpConnectCompleted;
-            arg.UserToken = new TcpUserToken(callback, state, arg);
-
-
-            Socket.ConnectAsync(SocketType.Stream, ProtocolType.Tcp, arg);
-        }
-
-        public static Socket EndConnectTcp(IAsyncResult asyncResult)
-        {
-            var r = asyncResult as AutoReleaseAsyncResult;
-            if (r == null)
-            {
-                throw new ArgumentException("Invalid asyncResult.", nameof(asyncResult));
-            }
-
-            var tut = r.UserToken;
-
-            var arg = tut.Args;
-
-            if (arg.SocketError != SocketError.Success)
-            {
-                if (arg.ConnectByNameError != null)
-                {
-                    throw arg.ConnectByNameError;
-                }
-
-                var ex = new SocketException((int)arg.SocketError);
-                throw ex;
-            }
-
-            var so = tut.Args.ConnectSocket;
-
-            so.SetSocketOption(SocketOptionLevel.Tcp, SocketOptionName.NoDelay, true);
-
-            return so;
-        }
-    }
-}

shadowsocks-csharp/Util/Sockets/SocketUtil.cs

@@ -0,0 +1,72 @@
+using System;
+using System.Net;
+using System.Net.Sockets;
+
+namespace Shadowsocks.Util.Sockets
+{
+    public static class SocketUtil
+    {
+        private class DnsEndPoint2 : DnsEndPoint
+        {
+            public DnsEndPoint2(string host, int port) : base(host, port)
+            {
+            }
+
+            public DnsEndPoint2(string host, int port, AddressFamily addressFamily) : base(host, port, addressFamily)
+            {
+            }
+
+            public override string ToString()
+            {
+                return this.Host + ":" + this.Port;
+            }
+        }
+
+        public static EndPoint GetEndPoint(string host, int port)
+        {
+            IPAddress ipAddress;
+            bool parsed = IPAddress.TryParse(host, out ipAddress);
+            if (parsed)
+            {
+                return new IPEndPoint(ipAddress, port);
+            }
+
+            // maybe is a domain name
+            return new DnsEndPoint2(host, port);
+        }
+
+
+        public static void FullClose(this System.Net.Sockets.Socket s)
+        {
+            try
+            {
+                s.Shutdown(SocketShutdown.Both);
+            }
+            catch (Exception)
+            {
+            }
+            try
+            {
+                s.Disconnect(false);
+            }
+            catch (Exception)
+            {
+            }
+            try
+            {
+                s.Close();
+            }
+            catch (Exception)
+            {
+            }
+            try
+            {
+                s.Dispose();
+            }
+            catch (Exception)
+            {
+            }
+        }
+        
+    }
+}

shadowsocks-csharp/Util/Sockets/WrappedSocket.cs

@@ -0,0 +1,268 @@
+using System;
+using System.Net;
+using System.Net.Sockets;
+using System.Threading;
+
+namespace Shadowsocks.Util.Sockets
+{
+    /*
+     * A wrapped socket class which support both ipv4 and ipv6 based on the
+     * connected remote endpoint.
+     * 
+     * If the server address is host name, then it may have both ipv4 and ipv6 address
+     * after resolving. The main idea is we don't want to resolve and choose the address
+     * by ourself. Instead, Socket.ConnectAsync() do handle this thing internally by trying
+     * each address and returning an established socket connection. Such approach solves
+     * two problem:
+     * 1. Async DNS resolving.
+     * 2. 
+     */
+    public class WrappedSocket
+    {
+        public EndPoint LocalEndPoint => _activeSocket?.LocalEndPoint;
+
+        private SpinLock _socketSyncLock = new SpinLock();
+            // Only used during connection and close, so it won't cost too much.
+
+        private volatile bool _disposed;
+        private bool Connected => _activeSocket != null;
+        private Socket _activeSocket;
+
+
+        public void BeginConnect(EndPoint remoteEP, AsyncCallback callback, object state)
+        {
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(GetType().FullName);
+            }
+            if (Connected)
+            {
+                throw new SocketException((int) SocketError.IsConnected);
+            }
+
+            var arg = new SocketAsyncEventArgs();
+            arg.RemoteEndPoint = remoteEP;
+            arg.Completed += OnTcpConnectCompleted;
+            arg.UserToken = new TcpUserToken(callback, state);
+
+            Socket.ConnectAsync(SocketType.Stream, ProtocolType.Tcp, arg);
+        }
+
+        private class FakeAsyncResult : IAsyncResult
+        {
+            public bool IsCompleted { get; } = true;
+            public WaitHandle AsyncWaitHandle { get; } = null;
+            public object AsyncState { get; set; }
+            public bool CompletedSynchronously { get; } = true;
+            public Exception InternalException { get; set; } = null;
+        }
+
+        private class TcpUserToken
+        {
+            public AsyncCallback Callback { get; }
+            public object AsyncState { get; }
+
+            public TcpUserToken(AsyncCallback callback, object state)
+            {
+                Callback = callback;
+                AsyncState = state;
+            }
+        }
+
+        private void OnTcpConnectCompleted(object sender, SocketAsyncEventArgs args)
+        {
+            using (args)
+            {
+                args.Completed -= OnTcpConnectCompleted;
+                var token = (TcpUserToken) args.UserToken;
+
+                if (args.SocketError != SocketError.Success)
+                {
+                    var ex = args.ConnectByNameError ?? new SocketException((int) args.SocketError);
+
+                    var r = new FakeAsyncResult()
+                    {
+                        AsyncState = token.AsyncState,
+                        InternalException = ex
+                    };
+
+                    token.Callback(r);
+                }
+                else
+                {
+                    var lockTaken = false;
+                    if (!_socketSyncLock.IsHeldByCurrentThread)
+                    {
+                        _socketSyncLock.TryEnter(ref lockTaken);
+                    }
+                    try
+                    {
+                        if (Connected)
+                        {
+                            args.ConnectSocket.FullClose();
+                        }
+                        else
+                        {
+                            _activeSocket = args.ConnectSocket;
+                            if (_disposed)
+                            {
+                                _activeSocket.FullClose();
+                            }
+
+                            var r = new FakeAsyncResult()
+                            {
+                                AsyncState = token.AsyncState
+                            };
+                            token.Callback(r);
+                        }
+                    }
+                    finally
+                    {
+                        if (lockTaken)
+                        {
+                            _socketSyncLock.Exit();
+                        }
+                    }
+                }
+            }
+        }
+
+        public void EndConnect(IAsyncResult asyncResult)
+        {
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(GetType().FullName);
+            }
+
+            var r = asyncResult as FakeAsyncResult;
+            if (r == null)
+            {
+                throw new ArgumentException("Invalid asyncResult.", nameof(asyncResult));
+            }
+
+            if (r.InternalException != null)
+            {
+                throw r.InternalException;
+            }
+        }
+
+        public void Dispose()
+        {
+            if (_disposed)
+            {
+                return;
+            }
+            var lockTaken = false;
+            if (!_socketSyncLock.IsHeldByCurrentThread)
+            {
+                _socketSyncLock.TryEnter(ref lockTaken);
+            }
+            try
+            {
+                _disposed = true;
+                _activeSocket?.FullClose();
+            }
+            finally
+            {
+                if (lockTaken)
+                {
+                    _socketSyncLock.Exit();
+                }
+            }
+
+        }
+
+        public IAsyncResult BeginSend(byte[] buffer, int offset, int size, SocketFlags socketFlags,
+            AsyncCallback callback,
+            object state)
+        {
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(GetType().FullName);
+            }
+            if (!Connected)
+            {
+                throw new SocketException((int) SocketError.NotConnected);
+            }
+
+            return _activeSocket.BeginSend(buffer, offset, size, socketFlags, callback, state);
+        }
+
+        public int EndSend(IAsyncResult asyncResult)
+        {
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(GetType().FullName);
+            }
+            if (!Connected)
+            {
+                throw new SocketException((int) SocketError.NotConnected);
+            }
+
+            return _activeSocket.EndSend(asyncResult);
+        }
+
+        public IAsyncResult BeginReceive(byte[] buffer, int offset, int size, SocketFlags socketFlags,
+            AsyncCallback callback,
+            object state)
+        {
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(GetType().FullName);
+            }
+            if (!Connected)
+            {
+                throw new SocketException((int) SocketError.NotConnected);
+            }
+
+            return _activeSocket.BeginReceive(buffer, offset, size, socketFlags, callback, state);
+        }
+
+        public int EndReceive(IAsyncResult asyncResult)
+        {
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(GetType().FullName);
+            }
+            if (!Connected)
+            {
+                throw new SocketException((int) SocketError.NotConnected);
+            }
+
+            return _activeSocket.EndReceive(asyncResult);
+        }
+
+        public void Shutdown(SocketShutdown how)
+        {
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(GetType().FullName);
+            }
+            if (!Connected)
+            {
+                return;
+            }
+
+            _activeSocket.Shutdown(how);
+        }
+
+        public void SetSocketOption(SocketOptionLevel optionLevel, SocketOptionName optionName, bool optionValue)
+        {
+            SetSocketOption(optionLevel, optionName, optionValue ? 1 : 0);
+        }
+
+        public void SetSocketOption(SocketOptionLevel optionLevel, SocketOptionName optionName, int optionValue)
+        {
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(GetType().FullName);
+            }
+            if (!Connected)
+            {
+                throw new SocketException((int)SocketError.NotConnected);
+            }
+
+            _activeSocket.SetSocketOption(optionLevel, optionName, optionValue);
+        }
+    }
+}

shadowsocks-csharp/shadowsocks-csharp.csproj

@@ -179,7 +179,8 @@
     <Compile Include="StringEx.cs" />
     <Compile Include="Util\ProcessManagement\Job.cs" />
     <Compile Include="Util\ProcessManagement\ThreadUtil.cs" />
-    <Compile Include="Util\SocketUtil.cs" />
+    <Compile Include="Util\Sockets\SocketUtil.cs" />
+    <Compile Include="Util\Sockets\WrappedSocket.cs" />
     <Compile Include="Util\Util.cs" />
     <Compile Include="View\ConfigForm.cs">
       <SubType>Form</SubType>