1 files changed, 27 insertions, 18 deletions

diff --git a/src/Net/Connection.h b/src/Net/Connection.h
index 19ee826..51f40b0 100644
--- a/src/Net/Connection.h
+++ b/src/Net/Connection.h
@@ -38,15 +38,13 @@ namespace Mad {
 namespace Net {
 
 class Listener;
-class ThreadManager;
 
 class MAD_NET_EXPORT Connection : boost::noncopyable {
   protected:
     friend class Listener;
-    friend class ThreadManager;
 
     enum State {
-      DISCONNECTED, CONNECT, CONNECTED, DISCONNECT
+      DISCONNECTED, CONNECT, CONNECTED, DISCONNECT, SHUTDOWN
     };
 
   private:
@@ -83,10 +81,6 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
     bool receiving;
     unsigned long sending;
     
-    void _initSocket() {
-      socket.reset(new boost::asio::ssl::stream<boost::asio::ip::tcp::socket>(application->getIOService(), context));
-    }
-    
     void enterReceiveLoop();
 
     void handleHeaderReceive(const boost::shared_array<boost::uint8_t> &data);
@@ -101,10 +95,12 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
     void rawSend(const boost::uint8_t *data, std::size_t length);
 
   protected:
+    boost::weak_ptr<Connection> thisPtr;
+
     boost::shared_mutex connectionLock;
     
-    boost::asio::ssl::context context;
-    boost::scoped_ptr<boost::asio::ssl::stream<boost::asio::ip::tcp::socket> > socket;
+    boost::shared_ptr<boost::asio::ssl::context> context;
+    boost::asio::ssl::stream<boost::asio::ip::tcp::socket> socket;
     boost::asio::ip::tcp::endpoint peer;
 
     void handleHandshake(const boost::system::error_code& error);
@@ -115,25 +111,38 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
     }
 
     bool _isDisconnecting() const {
-      return (state == DISCONNECT);
+      return (state == DISCONNECT || state == SHUTDOWN);
     }
 
     void _setState(State newState) {
       state = newState;
       
-      if(_isConnected() && !socket.get())
-        _initSocket();
-      else if(!_isConnected() && socket.get())
-         socket.reset();
-      
       stateChanged.notify_all();
     }
 
-    void doDisconnect();
+    void doDisconnect() {
+      boost::unique_lock<boost::shared_mutex> lock(connectionLock);
+
+      if(_isConnected() && state != SHUTDOWN) {
+        _setState(SHUTDOWN);
+        boost::system::error_code error;
+        socket.lowest_layer().cancel(error);
 
-    Connection(Core::Application *application0) :
+        socket.async_shutdown(boost::bind(&Connection::handleShutdown, thisPtr.lock(), boost::asio::placeholders::error));
+      }
+    }
+
+    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),
-      disconnectedSignal(application), context(application->getIOService(), boost::asio::ssl::context::sslv23) {}
+      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) {
+      boost::shared_ptr<Connection> connection(new Connection(application, context));
+
+      connection->thisPtr = connection;
+
+      return connection;
+    }
 
   public:
     virtual ~Connection();