1 files changed, 94 insertions, 55 deletions

diff --git a/src/Core/ConnectionManager.cpp b/src/Core/ConnectionManager.cpp
index 3ce52fc..3d252c2 100644
--- a/src/Core/ConnectionManager.cpp
+++ b/src/Core/ConnectionManager.cpp
@@ -45,12 +45,49 @@ namespace Core {
 
 ConnectionManager ConnectionManager::connectionManager;
 
+bool ConnectionManager::Connection::send(const Net::Packet &packet) {
+  return connection->send(packet);
+}
+
+ConnectionManager::Connection::Connection(Net::ServerConnection *connection0, ConnectionType type0)
+: connection(connection0), type(type0), hostInfo(0) {
+  connection->signalReceive().connect(sigc::mem_fun(this, &Connection::receive));
+}
+
+ConnectionManager::Connection::~Connection() {
+  delete connection;
+}
+
+bool ConnectionManager::Connection::isConnected() const {
+  return connection->isConnected();
+}
+
+bool ConnectionManager::Connection::disconnect() {
+  connection->disconnect();
+
+  return true;
+}
+
+void* ConnectionManager::Connection::getCertificate(size_t *size) const {
+  const gnutls_datum_t *cert = connection->getCertificate();
+
+  *size = cert->size;
+  return cert->data;
+}
+
+void* ConnectionManager::Connection::getPeerCertificate(size_t *size) const {
+  const gnutls_datum_t *cert = connection->getPeerCertificate();
+
+  *size = cert->size;
+  return cert->data;
+}
 
-void ConnectionManager::updateState(const std::string &name, Common::HostInfo::State state) {
-  daemonInfo[name].setState(state);
+void ConnectionManager::updateState(Common::HostInfo *hostInfo, Common::HostInfo::State state) {
+  hostInfo->setState(state);
 
-  for(std::list<Net::ServerConnection*>::iterator con = clientConnections.begin(); con != clientConnections.end(); ++con) {
-    Common::RequestManager::get()->sendRequest<Requests::DaemonStateUpdateRequest>(*con, Common::Request::slot_type(), name, state);
+  for(std::list<Connection*>::iterator con = connections.begin(); con != connections.end(); ++con) {
+    if((*con)->getConnectionType() == Connection::CLIENT)
+      Common::RequestManager::get()->sendRequest<Requests::DaemonStateUpdateRequest>(*con, Common::Request::slot_type(), hostInfo->getName(), state);
   }
 }
 
@@ -92,7 +129,6 @@ bool ConnectionManager::handleConfigEntry(const Common::ConfigEntry &entry, bool
     if(entry[0].getSize() == 1) {
       if(entry[1].empty()) {
         daemonInfo.insert(std::make_pair(entry[0][0], Common::HostInfo(entry[0][0])));
-        identifiedDaemonConnections.insert(std::make_pair<std::string,Net::ServerConnection*>(entry[0][0], 0));
 
         return true;
       }
@@ -146,11 +182,9 @@ void ConnectionManager::doInit() {
 }
 
 void ConnectionManager::doDeinit() {
-  for(std::list<Net::ServerConnection*>::iterator con = daemonConnections.begin(); con != daemonConnections.end(); ++con)
+  for(std::list<Connection*>::iterator con = connections.begin(); con != connections.end(); ++con)
     delete *con;
 
-  for(std::list<Net::ServerConnection*>::iterator con = clientConnections.begin(); con != clientConnections.end(); ++con)
-    delete *con;
 
   Common::RequestManager::get()->unregisterPacketType("AuthGSSAPI");
   Common::RequestManager::get()->unregisterPacketType("DaemonCommand");
@@ -166,19 +200,15 @@ void ConnectionManager::doDeinit() {
   Net::Connection::deinit();
 }
 
-void ConnectionManager::handleConnections(std::list<Net::ServerConnection*>& connections) {
-  for(std::list<Net::ServerConnection*>::iterator con = connections.begin(); con != connections.end();) {
+void ConnectionManager::run() {
+  // TODO Logging
+
+  Net::FdManager::get()->run();
+
+  for(std::list<Connection*>::iterator con = connections.begin(); con != connections.end();) {
     if(!(*con)->isConnected()) {
-      if((*con)->isIdentified()) {
-        for(std::map<std::string,Net::ServerConnection*>::iterator idCon = identifiedDaemonConnections.begin(); idCon != identifiedDaemonConnections.end(); ++idCon) {
-          if(idCon->second == *con) {
-            idCon->second = 0;
-
-            updateState(idCon->first, Common::HostInfo::INACTIVE);
-            break;
-          }
-        }
-      }
+      if((*con)->isIdentified())
+        updateState((*con)->getHostInfo(), Common::HostInfo::INACTIVE);
 
       Common::RequestManager::get()->unregisterConnection(*con);
       delete *con;
@@ -187,69 +217,78 @@ void ConnectionManager::handleConnections(std::list<Net::ServerConnection*>& con
     else
       ++con;
   }
-}
-
-void ConnectionManager::run() {
-  // TODO Logging
-
-  Net::FdManager::get()->run();
-
-  handleConnections(daemonConnections);
-  handleConnections(clientConnections);
 
   for(std::list<Net::Listener*>::iterator listener = listeners.begin(); listener != listeners.end(); ++listener) {
     Net::ServerConnection *con;
 
     while((con = (*listener)->getConnection()) != 0) {
-      (con->isDaemonConnection() ? daemonConnections : clientConnections).push_back(con);
-      Common::RequestManager::get()->registerConnection(con);
+      Connection *connection = new Connection(con,
+          con->isDaemonConnection() ? Connection::DAEMON : Connection::CLIENT);
+      connections.push_back(connection);
+      Common::RequestManager::get()->registerConnection(connection);
     }
   }
 }
 
-Net::Connection* ConnectionManager::getDaemonConnection(const std::string &name) const throw (Common::Exception&) {
-  std::map<std::string,Net::ServerConnection*>::const_iterator daemon = identifiedDaemonConnections.find(name);
-  if(daemon == identifiedDaemonConnections.end())
+Common::Connection* ConnectionManager::getDaemonConnection(const std::string &name) const throw (Common::Exception&) {
+  const Common::HostInfo *hostInfo;
+
+  try {
+    hostInfo = &daemonInfo.at(name);
+  }
+  catch(std::out_of_range&) {
     throw Common::Exception(Common::Exception::UNKNOWN_DAEMON);
+  }
 
-  if(!daemon->second)
-    throw Common::Exception(Common::Exception::NOT_AVAILABLE);
+  if(hostInfo->getState() != Common::HostInfo::INACTIVE) {
+    for(std::list<Connection*>::const_iterator it = connections.begin(); it != connections.end(); ++it) {
+      if((*it)->getHostInfo() == hostInfo) {
+        return *it;
+      }
+    }
+  }
 
-  return daemon->second;
+  throw(Common::Exception::NOT_AVAILABLE);
 }
 
-std::string ConnectionManager::getDaemonName(const Net::Connection *con) const throw (Common::Exception&) {
-  for(std::map<std::string,Net::ServerConnection*>::const_iterator daemon = identifiedDaemonConnections.begin(); daemon != identifiedDaemonConnections.end(); ++daemon) {
-    if(daemon->second == con)
-      return daemon->first;
+std::string ConnectionManager::getDaemonName(const Common::Connection *con) const throw (Common::Exception&) {
+  const Connection *connection = dynamic_cast<const Connection*>(con);
+
+  if(connection) {
+    if(connection->isIdentified()) {
+      return connection->getHostInfo()->getName();
+    }
   }
 
   throw Common::Exception(Common::Exception::UNKNOWN_DAEMON);
 }
 
-void ConnectionManager::identifyDaemonConnection(Net::Connection *connection, const std::string &name) throw (Common::Exception&) {
+void ConnectionManager::identifyDaemonConnection(Common::Connection *con, const std::string &name) throw (Common::Exception&) {
   // TODO Logging
 
-  if(connection->isIdentified())
-    throw Common::Exception(Common::Exception::ALREADY_IDENTIFIED);
+  Connection *connection = dynamic_cast<Connection*>(con);
 
-  std::list<Net::ServerConnection*>::iterator con = std::find(daemonConnections.begin(), daemonConnections.end(), connection);
-  if(con == daemonConnections.end())
+  if(!connection || (connection->getConnectionType() != Connection::DAEMON))
     throw Common::Exception(Common::Exception::INVALID_ACTION);
 
-  std::map<std::string,Net::ServerConnection*>::iterator idCon = identifiedDaemonConnections.find(name);
-  if(idCon == identifiedDaemonConnections.end()) {
+  if(connection->isIdentified())
+    throw Common::Exception(Common::Exception::ALREADY_IDENTIFIED);
+
+  if(daemonInfo.count(name) == 0)
     throw Common::Exception(Common::Exception::UNKNOWN_DAEMON);
-  }
 
-  if(idCon->second) {
-    idCon->second->disconnect();
-    Common::Logger::log(Common::Logger::WARNING, "Disconnecting old connection.");
+  Common::HostInfo *hostInfo = &daemonInfo[name];
+
+  if(hostInfo->getState() != Common::HostInfo::INACTIVE) {
+    try {
+      getDaemonConnection(name)->disconnect();
+      Common::Logger::log(Common::Logger::WARNING, "Disconnecting old connection.");
+    }
+    catch(Common::Exception&) {}
   }
 
-  idCon->second = *con;
-  updateState(idCon->first, Common::HostInfo::RUNNING);
-  connection->setIdentified();
+  connection->identify(hostInfo);
+  updateState(hostInfo, Common::HostInfo::RUNNING);
 
   Common::Logger::logf("Identified as '%s'.", name.c_str());
 }