Connection: Allow setting a receive limit

6 files changed, 28 insertions, 50 deletions

diff --git a/src/Common/ClientConnection.cpp b/src/Common/ClientConnection.cpp
index 02abcd1..db6efb2 100644
--- a/src/Common/ClientConnection.cpp
+++ b/src/Common/ClientConnection.cpp
@@ -62,19 +62,5 @@ bool ClientConnection::disconnect() {
   return true;
 }
 
-/*void* ClientConnection::getCertificate(size_t *size) const {
-  const gnutls_datum_t *cert = connection->getCertificate();
-
-  *size = cert->size;
-  return cert->data;
-}
-
-void* ClientConnection::getPeerCertificate(size_t *size) const {
-  const gnutls_datum_t *cert = connection->getPeerCertificate();
-
-  *size = cert->size;
-  return cert->data;
-}*/
-
 }
 }
diff --git a/src/Common/ClientConnection.h b/src/Common/ClientConnection.h
index 302c938..c521a5a 100644
--- a/src/Common/ClientConnection.h
+++ b/src/Common/ClientConnection.h
@@ -59,9 +59,6 @@ class MAD_COMMON_EXPORT ClientConnection : public Connection {
     virtual bool isAuthenticated() const {
       return true;
     }
-
-    //virtual void* getCertificate(size_t *size) const;
-    //virtual void* getPeerCertificate(size_t *size) const;
 };
 
 }
diff --git a/src/Net/Connection.cpp b/src/Net/Connection.cpp
index f1beb35..256bbfe 100644
--- a/src/Net/Connection.cpp
+++ b/src/Net/Connection.cpp
@@ -35,7 +35,7 @@ Connection::~Connection() {
 
 void Connection::handleHandshake(const boost::system::error_code& error) {
   if(error) {
-    application->logf("Error: %s", error.message().c_str());
+    application->logf(Core::Logger::LOG_NETWORK, "Error: %s", error.message().c_str());
 
     // TODO Error handling
     doDisconnect();
@@ -64,7 +64,7 @@ void Connection::handleShutdown(const boost::system::error_code& error) {
   boost::lock_guard<boost::shared_mutex> lock(connectionLock);
 
   if(error) {
-    application->logf(Core::Logger::LOG_VERBOSE, "Shutdown error: %s", error.message().c_str());
+    application->logf(Core::Logger::LOG_NETWORK, Core::Logger::LOG_VERBOSE, "Shutdown error: %s", error.message().c_str());
   }
 
   _setState(DISCONNECTED);
@@ -90,11 +90,17 @@ void Connection::handleHeaderReceive(const boost::shared_array<boost::uint8_t> &
     header = *reinterpret_cast<const Packet::Header*>(data.get());
   }
 
-  if(header.length == 0) {
+  boost::uint32_t length = ntohl(header.length);
+
+  if(length == 0) {
     receiveSignal.emit(boost::shared_ptr<Packet>(new Packet(ntohs(header.requestId))));
 
     enterReceiveLoop();
   }
+  else if(length > receiveLimit) {
+    application->log(Core::Logger::LOG_NETWORK, Core::Logger::LOG_WARNING, "Packet size limit exceeded. Disconnecting.");
+    doDisconnect();
+  }
   else {
     rawReceive(ntohl(header.length), boost::bind(&Connection::handleDataReceive, thisPtr.lock(), _1));
   }
@@ -115,7 +121,7 @@ void Connection::handleRead(const boost::system::error_code& error, std::size_t
     if(error == boost::system::errc::operation_canceled)
       return;
 
-    application->logf(Core::Logger::LOG_DEFAULT, "Read error: %s", error.message().c_str());
+    application->logf(Core::Logger::LOG_NETWORK, "Read error: %s", error.message().c_str());
 
     // TODO Error
     doDisconnect();
@@ -176,7 +182,7 @@ void Connection::rawReceive(std::size_t length, const boost::function1<void, con
 
 void Connection::handleWrite(const boost::system::error_code& error, std::size_t) {
   if(error)
-    application->logf(Core::Logger::LOG_VERBOSE, "Write error: %s", error.message().c_str());
+    application->logf(Core::Logger::LOG_NETWORK, Core::Logger::LOG_VERBOSE, "Write error: %s", error.message().c_str());
 
   {
     boost::unique_lock<boost::shared_mutex> lock(connectionLock);
diff --git a/src/Net/Connection.h b/src/Net/Connection.h
index add10b7..64b12c6 100644
--- a/src/Net/Connection.h
+++ b/src/Net/Connection.h
@@ -80,6 +80,8 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
     Core::Signals::Signal0 connectedSignal;
     Core::Signals::Signal0 disconnectedSignal;
 
+    boost::uint32_t receiveLimit;
+
     bool receiving;
     unsigned long sending;
 
@@ -138,7 +140,8 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
       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) {}
+      disconnectedSignal(application), receiveLimit(0xFFFF) /* 64K */, receiving(false), sending(0),
+      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));
@@ -180,17 +183,6 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
         stateChanged.wait(lock);
     }
 
-    /*const gnutls_datum_t* getCertificate() const {
-      // TODO Thread-safeness
-      return gnutls_certificate_get_ours(session);
-    }
-
-    const gnutls_datum_t* getPeerCertificate() const {
-      // TODO Thread-safeness
-      unsigned int n;
-      return gnutls_certificate_get_peers(session, &n);
-    }*/
-
     boost::asio::ip::tcp::endpoint getPeer() {
       boost::shared_lock<boost::shared_mutex> lock(connectionLock);
       return peer;
@@ -206,6 +198,16 @@ class MAD_NET_EXPORT Connection : boost::noncopyable {
       setStart(false);
     }
 
+    boost::uint32_t getReceiveLimit() {
+      boost::shared_lock<boost::shared_mutex> lock(connectionLock);
+      return receiveLimit;
+    }
+
+    void setReceiveLimit(boost::uint32_t limit) {
+      boost::lock_guard<boost::shared_mutex> lock(connectionLock);
+      receiveLimit = limit;
+    }
+
     void startReceive() {
       {
         boost::lock_guard<boost::shared_mutex> lock(connectionLock);
diff --git a/src/Server/ConnectionManager.cpp b/src/Server/ConnectionManager.cpp
index ad5a7d3..1ffd3f1 100644
--- a/src/Server/ConnectionManager.cpp
+++ b/src/Server/ConnectionManager.cpp
@@ -64,23 +64,12 @@ boost::shared_ptr<const Common::AuthContext> ConnectionManager::ServerConnection
 
   authContext = application->getAuthManager()->authenticate(method, subMethod, user, data, response, authContext);
 
-  return authContext;
-}
-
-/*void* ConnectionManager::ServerConnection::getCertificate(size_t *size) const {
-  const gnutls_datum_t *cert = connection->getCertificate();
+  if(authContext->isAuthenticated())
+    connection->setReceiveLimit(0xFFFFFFFF); // 4 Gigs
 
-  *size = cert->size;
-  return cert->data;
+  return authContext;
 }
 
-void* ConnectionManager::ServerConnection::getPeerCertificate(size_t *size) const {
-  const gnutls_datum_t *cert = connection->getPeerCertificate();
-
-  *size = cert->size;
-  return cert->data;
-}*/
-
 boost::asio::ip::tcp::endpoint ConnectionManager::parseAddress(const std::string &str) throw(Core::Exception) {
   try {
     if(str == "*")
diff --git a/src/Server/ConnectionManager.h b/src/Server/ConnectionManager.h
index 099072c..8c989b9 100644
--- a/src/Server/ConnectionManager.h
+++ b/src/Server/ConnectionManager.h
@@ -73,8 +73,6 @@ class MAD_SERVER_EXPORT ConnectionManager : public Core::Configurable, private b
         bool isConnected() const;
 
         virtual bool disconnect();
-        //virtual void* getCertificate(size_t *size) const;
-        //virtual void* getPeerCertificate(size_t *size) const;
 
         ConnectionType getConnectionType() const {
           return type;