1 files changed, 31 insertions, 4 deletions

diff --git a/src/Net/Connection.h b/src/Net/Connection.h
index 51f40b0..add10b7 100644
--- a/src/Net/Connection.h
+++ b/src/Net/Connection.h
@@ -44,7 +44,7 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
     friend class Listener;
 
     enum State {
-      DISCONNECTED, CONNECT, CONNECTED, DISCONNECT, SHUTDOWN
+      DISCONNECTED, HANDSHAKE, CONNECT, CONNECTED, DISCONNECT, SHUTDOWN
     };
 
   private:
@@ -69,6 +69,8 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
 
     State state;
 
+    bool dontStart;
+
     boost::scoped_ptr<boost::array<boost::uint8_t, 1024*1024> > receiveBuffer;
     std::size_t received;
 
@@ -80,7 +82,7 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
 
     bool receiving;
     unsigned long sending;
-    
+
     void enterReceiveLoop();
 
     void handleHeaderReceive(const boost::shared_array<boost::uint8_t> &data);
@@ -107,7 +109,7 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
 
     bool _isConnected() const {return (state != DISCONNECTED);}
     bool _isConnecting() const {
-      return (state == CONNECT);
+      return (state == HANDSHAKE || state == CONNECT);
     }
 
     bool _isDisconnecting() const {
@@ -133,7 +135,9 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
     }
 
     Connection(Core::Application *application0, boost::shared_ptr<boost::asio::ssl::context> context0) :
-    application(application0), state(DISCONNECTED), receiveBuffer(new boost::array<boost::uint8_t, 1024*1024>), receiveSignal(application), connectedSignal(application),
+      application(application0), state(DISCONNECTED), dontStart(false),
+      receiveBuffer(new boost::array<boost::uint8_t, 1024*1024>),
+      receiveSignal(application), connectedSignal(application),
       disconnectedSignal(application), context(context0), socket(application->getIOService(), *context) {}
 
     static boost::shared_ptr<Connection> create(Core::Application *application, boost::shared_ptr<boost::asio::ssl::context> context) {
@@ -192,6 +196,29 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
       return peer;
     }
 
+    void setStart(bool start = true) {
+      boost::lock_guard<boost::shared_mutex> lock(connectionLock);
+
+      dontStart = !start;
+    }
+
+    void unsetStart() {
+      setStart(false);
+    }
+
+    void startReceive() {
+      {
+        boost::lock_guard<boost::shared_mutex> lock(connectionLock);
+
+        if(state != CONNECT)
+          return;
+
+        _setState(CONNECTED);
+      }
+
+      enterReceiveLoop();
+    }
+
     void disconnect();
 
     bool send(const Packet &packet);