Add SocketAsyncEventArgs to handle local connections

shadowsocks-csharp/Controller/AsyncSocketUserToken.cs

@@ -0,0 +1,37 @@
+using System;
+using System.Collections.Generic;
+using System.Text;
+using System.Net.Sockets;
+
+namespace Shadowsocks.Controller
+{
+    public class AsyncSocketUserToken
+    {
+        protected SocketAsyncEventArgs m_receiveEventArgs;
+        public SocketAsyncEventArgs ReceiveEventArgs { get { return m_receiveEventArgs; } set { m_receiveEventArgs = value; } }
+        protected byte[] m_asyncReceiveBuffer;
+
+        protected Socket m_connectSocket;
+        public Socket ConnectSocket
+        {
+            get
+            {
+                return m_connectSocket;
+            }
+            set
+            {
+                m_connectSocket = value;          
+                m_receiveEventArgs.AcceptSocket = m_connectSocket;
+            }
+        }
+
+        public AsyncSocketUserToken(int asyncReceiveBufferSize)
+        {
+            m_connectSocket = null;
+            m_receiveEventArgs = new SocketAsyncEventArgs();
+            m_receiveEventArgs.UserToken = this;
+            m_asyncReceiveBuffer = new byte[asyncReceiveBufferSize];
+            m_receiveEventArgs.SetBuffer(m_asyncReceiveBuffer, 0, m_asyncReceiveBuffer.Length);
+        }
+    }
+}

shadowsocks-csharp/Controller/AsyncSocketUserTokenPool.cs

@@ -0,0 +1,76 @@
+using System;
+using System.Collections.Generic;
+using System.Text;
+
+namespace Shadowsocks.Controller
+{
+    public class AsyncSocketUserTokenPool
+    {
+        private Stack<AsyncSocketUserToken> m_pool;
+
+        public AsyncSocketUserTokenPool(int capacity)
+        {
+            m_pool = new Stack<AsyncSocketUserToken>(capacity);
+        }
+
+        public void Push(AsyncSocketUserToken item)
+        {
+            if (item == null)
+            {
+                throw new ArgumentException("Items added to a AsyncSocketUserToken cannot be null");
+            }
+            lock (m_pool)
+            {
+                m_pool.Push(item);
+            }
+        }
+
+        public AsyncSocketUserToken Pop()
+        {
+            lock (m_pool)
+            {
+                return m_pool.Pop();
+            }
+        }
+
+        public int Count
+        {
+            get { return m_pool.Count; }
+        }
+    }
+
+    public class AsyncSocketUserTokenList : Object
+    {
+        private List<AsyncSocketUserToken> m_list;
+
+        public AsyncSocketUserTokenList()
+        {
+            m_list = new List<AsyncSocketUserToken>();
+        }
+
+        public void Add(AsyncSocketUserToken userToken)
+        {
+            lock(m_list)
+            {
+                m_list.Add(userToken);
+            }
+        }
+
+        public void Remove(AsyncSocketUserToken userToken)
+        {
+            lock (m_list)
+            {
+                m_list.Remove(userToken);
+            }
+        }
+
+        public void CopyList(ref AsyncSocketUserToken[] array)
+        {
+            lock (m_list)
+            {
+                array = new AsyncSocketUserToken[m_list.Count];
+                m_list.CopyTo(array);
+            }
+        }
+    }
+}

shadowsocks-csharp/Controller/Listener.cs

@@ -5,6 +5,7 @@ using System.Net;
 using System.Net.NetworkInformation;
 using System.Net.Sockets;
 using System.Text;
+using System.Threading;
 
 namespace Shadowsocks.Controller
 {
@@ -20,9 +21,29 @@ namespace Shadowsocks.Controller
         Socket _socket;
         IList<Service> _services;
 
+        private const int m_numConnections = 8000; 
+        private const int m_receiveBufferSize = 4096; 
+        private Semaphore m_maxNumberAcceptedClients;
+
+        private AsyncSocketUserTokenPool m_asyncSocketUserTokenPool;
+        private AsyncSocketUserTokenList m_asyncSocketUserTokenList;
+        public AsyncSocketUserTokenList AsyncSocketUserTokenList { get { return m_asyncSocketUserTokenList; } }
+
         public Listener(IList<Service> services)
         {
             this._services = services;
+
+            m_asyncSocketUserTokenPool = new AsyncSocketUserTokenPool(m_numConnections);
+            m_asyncSocketUserTokenList = new AsyncSocketUserTokenList();
+            m_maxNumberAcceptedClients = new Semaphore(m_numConnections, m_numConnections);
+
+            AsyncSocketUserToken userToken;
+            for (int i = 0; i < m_numConnections; i++)
+            {
+                userToken = new AsyncSocketUserToken(m_receiveBufferSize);
+                userToken.ReceiveEventArgs.Completed += new EventHandler<SocketAsyncEventArgs>(IO_Completed);
+                m_asyncSocketUserTokenPool.Push(userToken);
+            }
         }
 
         private bool CheckIfPortInUse(int port)
@@ -70,9 +91,8 @@ namespace Shadowsocks.Controller
 
                 // Start an asynchronous socket to listen for connections.
                 Console.WriteLine("Shadowsocks started");
-                _socket.BeginAccept(
-                    new AsyncCallback(AcceptCallback),
-                    _socket);
+
+                StartAccept(null);
             }
             catch (SocketException)
             {
@@ -87,77 +107,130 @@ namespace Shadowsocks.Controller
             {
                 _socket.Close();
                 _socket = null;
-            }
+            }         
         }
 
-        public void AcceptCallback(IAsyncResult ar)
+        public void StartAccept(SocketAsyncEventArgs acceptEventArgs)
         {
-            Socket listener = (Socket)ar.AsyncState;
-            try
+            if (acceptEventArgs == null)
             {
-                Socket conn = listener.EndAccept(ar);
-
-                byte[] buf = new byte[4096];
-                object[] state = new object[] {
-                    conn,
-                    buf
-                };
+                acceptEventArgs = new SocketAsyncEventArgs();
+                acceptEventArgs.Completed += new EventHandler<SocketAsyncEventArgs>(AcceptEventArg_Completed);
+            }
+            else
+            {
+                acceptEventArgs.AcceptSocket = null; 
+            }
 
-                conn.BeginReceive(buf, 0, buf.Length, 0,
-                    new AsyncCallback(ReceiveCallback), state);
+            m_maxNumberAcceptedClients.WaitOne(); 
+            bool willRaiseEvent = _socket.AcceptAsync(acceptEventArgs);
+            if (!willRaiseEvent)
+            {
+                ProcessAccept(acceptEventArgs);
             }
-            catch (ObjectDisposedException)
+        }
+
+
+        void AcceptEventArg_Completed(object sender, SocketAsyncEventArgs acceptEventArgs)
+        {
+            try
             {
+                ProcessAccept(acceptEventArgs);
             }
             catch (Exception e)
             {
-                Console.WriteLine(e);
+                Logging.LogUsefulException(e);
             }
-            finally
+        }
+
+        private void ProcessAccept(SocketAsyncEventArgs acceptEventArgs)
+        {          
+            AsyncSocketUserToken userToken = m_asyncSocketUserTokenPool.Pop();
+            m_asyncSocketUserTokenList.Add(userToken);
+            userToken.ConnectSocket = acceptEventArgs.AcceptSocket;
+
+            try
             {
-                try
+                bool willRaiseEvent = userToken.ConnectSocket.ReceiveAsync(userToken.ReceiveEventArgs);  
+                if (!willRaiseEvent)
                 {
-                    listener.BeginAccept(
-                        new AsyncCallback(AcceptCallback),
-                        listener);
-                }
-                catch (ObjectDisposedException)
-                {
-                    // do nothing
+                    lock (userToken)
+                    {
+                        ProcessReceive(userToken.ReceiveEventArgs);
+                    }
                 }
-                catch (Exception e)
+            }
+            catch (Exception e)
+            {
+                Logging.LogUsefulException(e);
+            }
+
+            StartAccept(acceptEventArgs);
+        }
+
+
+        void IO_Completed(object sender, SocketAsyncEventArgs asyncEventArgs)
+        {
+            AsyncSocketUserToken userToken = asyncEventArgs.UserToken as AsyncSocketUserToken;
+            try
+            {
+                lock (userToken)
                 {
-                    Logging.LogUsefulException(e);
+                    if (asyncEventArgs.LastOperation == SocketAsyncOperation.Receive)
+                        ProcessReceive(asyncEventArgs);
                 }
             }
+            catch (Exception e)
+            {
+                Console.WriteLine(e);
+            }
         }
 
-
-        private void ReceiveCallback(IAsyncResult ar)
+        private void ProcessReceive(SocketAsyncEventArgs receiveEventArgs)
         {
-            object[] state = (object[])ar.AsyncState;
+            AsyncSocketUserToken userToken = receiveEventArgs.UserToken as AsyncSocketUserToken;
+            if (userToken.ConnectSocket == null)
+                return;
 
-            Socket conn = (Socket)state[0];
-            byte[] buf = (byte[])state[1];
-            try
+            if (userToken.ReceiveEventArgs.BytesTransferred > 0 && userToken.ReceiveEventArgs.SocketError == SocketError.Success)
             {
-                int bytesRead = conn.EndReceive(ar);
                 foreach (Service service in _services)
                 {
-                    if (service.Handle(buf, bytesRead, conn))
+                    if (service.Handle(userToken.ReceiveEventArgs.Buffer, userToken.ReceiveEventArgs.BytesTransferred, userToken.ConnectSocket))
                     {
                         return;
                     }
                 }
-                // no service found for this
-                // shouldn't happen
-                conn.Close();
+
+                bool willRaiseEvent = userToken.ConnectSocket.ReceiveAsync(userToken.ReceiveEventArgs);
+                if (!willRaiseEvent)
+                    ProcessReceive(userToken.ReceiveEventArgs);
+            }
+            else
+            {
+                CloseClientSocket(userToken);
+            }
+        }
+
+        public void CloseClientSocket(AsyncSocketUserToken userToken)
+        {
+            if (userToken.ConnectSocket == null)
+                return;
+            try
+            {
+                userToken.ConnectSocket.Shutdown(SocketShutdown.Both);
             }
             catch (Exception e)
             {
                 Console.WriteLine(e);
-                conn.Close();
             }
+            userToken.ConnectSocket.Close();
+            userToken.ConnectSocket = null;
+
+            m_maxNumberAcceptedClients.Release();
+            m_asyncSocketUserTokenPool.Push(userToken);
+            m_asyncSocketUserTokenList.Remove(userToken);
         }
     }
+
 }

shadowsocks-csharp/shadowsocks-csharp.csproj

@@ -124,6 +124,8 @@
     <Compile Include="3rd\zxing\ResultPoint.cs" />
     <Compile Include="3rd\zxing\ResultPointCallback.cs" />
     <Compile Include="3rd\zxing\WriterException.cs" />
+    <Compile Include="Controller\AsyncSocketUserToken.cs" />
+    <Compile Include="Controller\AsyncSocketUserTokenPool.cs" />
     <Compile Include="Controller\AutoStartup.cs" />
     <Compile Include="Controller\FileManager.cs" />
     <Compile Include="Controller\GFWListUpdater.cs" />