From 449df1022da409ecd49b6f6763c37ddd51fed9cd Mon Sep 17 00:00:00 2001
From: clowwindy <clowwindy42@gmail.com>
Date: Sat, 25 Jul 2015 17:40:08 +0800
Subject: [PATCH] add connect timeout

---
 shadowsocks-csharp/Controller/Logging.cs      |  3 +
 .../Controller/Service/TCPRelay.cs            | 76 +++++++++++++++----
 .../Strategy/HighAvailabilityStrategy.cs      |  6 +-
 3 files changed, 70 insertions(+), 15 deletions(-)

diff --git a/shadowsocks-csharp/Controller/Logging.cs b/shadowsocks-csharp/Controller/Logging.cs
index 145547bb..354df7a5 100755
--- a/shadowsocks-csharp/Controller/Logging.cs
+++ b/shadowsocks-csharp/Controller/Logging.cs
@@ -63,6 +63,9 @@ namespace Shadowsocks.Controller
                     Console.WriteLine(e);
                 }
             }
+            else if (e is ObjectDisposedException)
+            {
+            }
             else
             {
                 Console.WriteLine(e);
diff --git a/shadowsocks-csharp/Controller/Service/TCPRelay.cs b/shadowsocks-csharp/Controller/Service/TCPRelay.cs
index 3094bc76..15c3bc34 100644
--- a/shadowsocks-csharp/Controller/Service/TCPRelay.cs
+++ b/shadowsocks-csharp/Controller/Service/TCPRelay.cs
@@ -6,6 +6,7 @@ using System.Net;
 using Shadowsocks.Encryption;
 using Shadowsocks.Model;
 using Shadowsocks.Controller.Strategy;
+using System.Timers;
 
 namespace Shadowsocks.Controller
 {
@@ -48,6 +49,7 @@ namespace Shadowsocks.Controller
         public Socket connection;
         public ShadowsocksController controller;
         private int retryCount = 0;
+        private bool connected;
 
         private byte command;
         private byte[] _firstPacket;
@@ -129,7 +131,7 @@ namespace Shadowsocks.Controller
                     remote.Shutdown(SocketShutdown.Both);
                     remote.Close();
                 }
-                catch (SocketException e)
+                catch (Exception e)
                 {
                     Logging.LogUsefulException(e);
                 }
@@ -315,6 +317,15 @@ namespace Shadowsocks.Controller
             }
         }
 
+        private class ServerTimer : Timer
+        {
+            public Server Server;
+
+            public ServerTimer(int p) :base(p)
+            {
+            }
+        }
+
         private void StartConnect()
         {
             try
@@ -337,10 +348,16 @@ namespace Shadowsocks.Controller
                 remote.SetSocketOption(SocketOptionLevel.Tcp, SocketOptionName.NoDelay, true);
 
                 _startConnectTime = DateTime.Now;
+                ServerTimer connectTimer = new ServerTimer(3000);
+                connectTimer.AutoReset = false;
+                connectTimer.Elapsed += connectTimer_Elapsed;
+                connectTimer.Enabled = true;
+                connectTimer.Server = server;
 
+                connected = false;
                 // Connect to the remote endpoint.
                 remote.BeginConnect(remoteEP,
-                    new AsyncCallback(ConnectCallback), null);
+                    new AsyncCallback(ConnectCallback), connectTimer);
             }
             catch (Exception e)
             {
@@ -349,39 +366,72 @@ namespace Shadowsocks.Controller
             }
         }
 
+        private void connectTimer_Elapsed(object sender, ElapsedEventArgs e)
+        {
+            if (connected)
+            {
+                return;
+            }
+            Server server = ((ServerTimer)sender).Server;
+            controller.GetCurrentStrategy().SetFailure(server);
+            Console.WriteLine(String.Format("{0} timed out", server.FriendlyName()));
+            remote.Close();
+            RetryConnect();
+        }
+
+        private void RetryConnect()
+        {
+            if (retryCount < 4)
+            {
+                Console.WriteLine("Connection failed, retrying");
+                StartConnect();
+                retryCount++;
+            }
+            else
+            {
+                this.Close();
+            }
+        }
+
         private void ConnectCallback(IAsyncResult ar)
         {
+            Server server = null;
             if (closed)
             {
                 return;
             }
             try
             {
+                ServerTimer timer = (ServerTimer)ar.AsyncState;
+                server = timer.Server;
+                timer.Elapsed -= connectTimer_Elapsed;
+                timer.Enabled = false;
+                timer.Dispose();
+
                 // Complete the connection.
                 remote.EndConnect(ar);
 
+                connected = true;
+
                 //Console.WriteLine("Socket connected to {0}",
                 //    remote.RemoteEndPoint.ToString());
 
                 var latency = DateTime.Now - _startConnectTime;
-                controller.GetCurrentStrategy().UpdateLatency(this.server, latency);
+                controller.GetCurrentStrategy().UpdateLatency(server, latency);
 
                 StartPipe();
             }
+            catch (ArgumentException e)
+            {
+            }
             catch (Exception e)
             {
-                controller.GetCurrentStrategy().SetFailure(this.server);
-                Logging.LogUsefulException(e);
-                if (retryCount < 3)
-                {
-                    Console.WriteLine("Connection failed, retrying");
-                    StartConnect();
-                    retryCount++;
-                }
-                else
+                if (server != null)
                 {
-                    this.Close();
+                    controller.GetCurrentStrategy().SetFailure(server);
                 }
+                Logging.LogUsefulException(e);
+                RetryConnect();
             }
         }
 
diff --git a/shadowsocks-csharp/Controller/Strategy/HighAvailabilityStrategy.cs b/shadowsocks-csharp/Controller/Strategy/HighAvailabilityStrategy.cs
index d8c8dfb8..a4173da2 100644
--- a/shadowsocks-csharp/Controller/Strategy/HighAvailabilityStrategy.cs
+++ b/shadowsocks-csharp/Controller/Strategy/HighAvailabilityStrategy.cs
@@ -75,8 +75,10 @@ namespace Shadowsocks.Controller.Strategy
 
         public Server GetAServer(IStrategyCallerType type, System.Net.IPEndPoint localIPEndPoint)
         {
-            // TODO don't choose new server too frequently
-            ChooseNewServer();
+            if (type == IStrategyCallerType.TCP)
+            {
+                ChooseNewServer();
+            }
             return _currentServer;
         }