3 files changed, 29 insertions, 27 deletions

diff --git a/src/Net/Connection.h b/src/Net/Connection.h
index 625fc94..6678bd9 100644
--- a/src/Net/Connection.h
+++ b/src/Net/Connection.h
@@ -34,10 +34,12 @@
 namespace Mad {
 namespace Net {
 
+class Listener;
 class ThreadManager;
 
 class Connection : boost::noncopyable {
   protected:
+    friend class Listener;
     friend class ThreadManager;
 
     enum State {
diff --git a/src/Net/Listener.cpp b/src/Net/Listener.cpp
index 16689db..be227fd 100644
--- a/src/Net/Listener.cpp
+++ b/src/Net/Listener.cpp
@@ -28,42 +28,48 @@
 namespace Mad {
 namespace Net {
 
-void Listener::handleAccept(const boost::system::error_code &error, boost::shared_ptr<ServerConnection> con) {
+void Listener::handleAccept(const boost::system::error_code &error, boost::shared_ptr<Connection> con) {
   if(error)
     return;
 
+  boost::lock_guard<boost::mutex> lock(mutex);
+
   {
-    boost::lock_guard<boost::shared_mutex> lock(con->connectionLock);
+    boost::lock_guard<boost::shared_mutex> conLock(con->connectionLock);
 
-    con->_setState(ServerConnection::CONNECT);
+    con->_setState(Connection::CONNECT);
 
     boost::signals::connection con1 = con->signalConnected().connect(boost::bind(&Listener::handleConnect, this, con));
     boost::signals::connection con2 = con->signalDisconnected().connect(boost::bind(&Listener::handleDisconnect, this, con));
 
     connections.insert(std::make_pair(con, std::make_pair(con1, con2)));
 
-    con->socket.async_handshake(boost::asio::ssl::stream_base::server, boost::bind(&ServerConnection::handleHandshake, con, boost::asio::placeholders::error));
+    con->socket.async_handshake(boost::asio::ssl::stream_base::server, boost::bind(&Connection::handleHandshake, con, boost::asio::placeholders::error));
   }
 
-  con.reset(new ServerConnection(sslContext));
+  con.reset(new Connection(sslContext));
   acceptor.async_accept(con->socket.lowest_layer(), boost::bind(&Listener::handleAccept, this, boost::asio::placeholders::error, con));
 }
 
-void Listener::handleConnect(boost::shared_ptr<ServerConnection> con) {
-  std::map<boost::shared_ptr<ServerConnection>, std::pair<boost::signals::connection, boost::signals::connection> >::iterator it = connections.find(con);
+void Listener::handleConnect(boost::shared_ptr<Connection> con) {
+  {
+    boost::lock_guard<boost::mutex> lock(mutex);
+    std::map<boost::shared_ptr<Connection>, std::pair<boost::signals::connection, boost::signals::connection> >::iterator it = connections.find(con);
 
-  if(it == connections.end())
-    return;
+    if(it == connections.end())
+      return;
 
-  // Disconnect signal handlers
-  it->second.first.disconnect();
-  it->second.second.disconnect();
-  connections.erase(it);
+    // Disconnect signal handlers
+    it->second.first.disconnect();
+    it->second.second.disconnect();
+    connections.erase(it);
+  }
 
   signal(con);
 }
 
-void Listener::handleDisconnect(boost::shared_ptr<ServerConnection> con) {
+void Listener::handleDisconnect(boost::shared_ptr<Connection> con) {
+  boost::lock_guard<boost::mutex> lock(mutex);
   connections.erase(con);
 }
 
@@ -81,12 +87,12 @@ sslContext(Connection::ioService, boost::asio::ssl::context::sslv23)
 
 
 
-  boost::shared_ptr<ServerConnection> con(new ServerConnection(sslContext));
+  boost::shared_ptr<Connection> con(new Connection(sslContext));
   acceptor.async_accept(con->socket.lowest_layer(), boost::bind(&Listener::handleAccept, this, boost::asio::placeholders::error, con));
 }
 
 Listener::~Listener() {
-  for(std::map<boost::shared_ptr<ServerConnection>,std::pair<boost::signals::connection, boost::signals::connection> >::iterator con = connections.begin(); con != connections.end(); ++con) {
+  for(std::map<boost::shared_ptr<Connection>,std::pair<boost::signals::connection, boost::signals::connection> >::iterator con = connections.begin(); con != connections.end(); ++con) {
     con->first->disconnect();
     // TODO wait...
   }
diff --git a/src/Net/Listener.h b/src/Net/Listener.h
index 0833cdf..903c907 100644
--- a/src/Net/Listener.h
+++ b/src/Net/Listener.h
@@ -29,29 +29,23 @@
 namespace Mad {
 namespace Net {
 
-// TODO XXX Thread-safeness XXX
 class Listener : boost::noncopyable {
   private:
-    class ServerConnection : public Connection {
-      public:
-        friend class Listener;
-
-        ServerConnection(boost::asio::ssl::context &sslContext) : Connection(sslContext) {}
-    };
+    boost::mutex mutex;
 
     std::string x905CertFile, x905KeyFile;
     boost::asio::ip::tcp::endpoint address;
     boost::asio::ip::tcp::acceptor acceptor;
     boost::asio::ssl::context sslContext;
 
-    std::map<boost::shared_ptr<ServerConnection>, std::pair<boost::signals::connection, boost::signals::connection> > connections;
+    std::map<boost::shared_ptr<Connection>, std::pair<boost::signals::connection, boost::signals::connection> > connections;
 
     boost::signal1<void, boost::shared_ptr<Connection> > signal;
 
-    void handleAccept(const boost::system::error_code &error, boost::shared_ptr<ServerConnection> con);
+    void handleAccept(const boost::system::error_code &error, boost::shared_ptr<Connection> con);
 
-    void handleConnect(boost::shared_ptr<ServerConnection> con);
-    void handleDisconnect(boost::shared_ptr<ServerConnection> con);
+    void handleConnect(boost::shared_ptr<Connection> con);
+    void handleDisconnect(boost::shared_ptr<Connection> con);
 
   public:
     Listener(const std::string &x905CertFile0, const std::string &x905KeyFile0,