4 files changed, 43 insertions, 11 deletions

diff --git a/src/Net/ClientConnection.cpp b/src/Net/ClientConnection.cpp
index 10a03f1..3d218e7 100644
--- a/src/Net/ClientConnection.cpp
+++ b/src/Net/ClientConnection.cpp
@@ -43,7 +43,7 @@ void ClientConnection::connect(const boost::asio::ip::tcp::endpoint &address) th
   }
 
   peer = address;
-  _setState(CONNECT);
+  _setState(HANDSHAKE);
 
   socket.lowest_layer().async_connect(address, boost::bind(&ClientConnection::handleConnect, boost::dynamic_pointer_cast<ClientConnection>(thisPtr.lock()), boost::asio::placeholders::error));
 }
diff --git a/src/Net/Connection.cpp b/src/Net/Connection.cpp
index fc917d4..94b5493 100644
--- a/src/Net/Connection.cpp
+++ b/src/Net/Connection.cpp
@@ -44,17 +44,20 @@ void Connection::handleHandshake(const boost::system::error_code& error) {
 
   {
     boost::lock_guard<boost::shared_mutex> lock(connectionLock);
-    _setState(CONNECTED);
+    _setState(CONNECT);
 
     receiving = false;
     sending = 0;
 
     received = 0;
-  }
 
-  connectedSignal.emit();
+    connectedSignal.emit();
 
-  enterReceiveLoop();
+    if(dontStart)
+      return;
+  }
+
+  startReceive();
 }
 
 void Connection::handleShutdown(const boost::system::error_code& error) {
@@ -73,7 +76,7 @@ void Connection::enterReceiveLoop() {
   {
     boost::lock_guard<boost::shared_mutex> lock(connectionLock);
 
-    if(!_isConnected() || _isDisconnecting())
+    if(state != CONNECTED)
       return;
   }
 
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);
diff --git a/src/Net/Listener.cpp b/src/Net/Listener.cpp
index b3974d1..a67e608 100644
--- a/src/Net/Listener.cpp
+++ b/src/Net/Listener.cpp
@@ -29,6 +29,8 @@ namespace Net {
 void Listener::accept() {
   boost::shared_ptr<Connection> con(Connection::create(application, context));
   
+  con->unsetStart();
+
   acceptor.async_accept(con->socket.lowest_layer(), boost::bind(&Listener::handleAccept, this, boost::asio::placeholders::error, con));
 }
 
@@ -41,7 +43,7 @@ void Listener::handleAccept(const boost::system::error_code &error, boost::share
   {
     boost::lock_guard<boost::shared_mutex> conLock(con->connectionLock);
 
-    con->_setState(Connection::CONNECT);
+    con->_setState(Connection::HANDSHAKE);
 
     Core::Signals::Connection con1 = con->connectSignalConnected(boost::bind(&Listener::handleConnect, this, con));
     Core::Signals::Connection con2 = con->connectSignalDisconnected(boost::bind(&Listener::handleDisconnect, this, con));