4 files changed, 43 insertions, 12 deletions

diff --git a/common/connectionbase.cpp b/common/connectionbase.cpp
index 165614f..a113e6a 100644
--- a/common/connectionbase.cpp
+++ b/common/connectionbase.cpp
@@ -1,6 +1,6 @@
 #include "connectionbase.h"
 
-#include <stdexcept>
+#include <exception>
 
 #include <iostream>
 
@@ -11,7 +11,7 @@ void ConnectionBase::start_recv() {
 void ConnectionBase::got_data(uint8_t* data, std::size_t bytes) {
 	if(pending_size == 0) {
 		if(bytes != 4) {
-			throw std::runtime_error("Deserialization header error.");
+			error("Deserialization header error.");
 		}
 		
 		uint16_t* header = (uint16_t*)data;
@@ -29,7 +29,7 @@ void ConnectionBase::got_data(uint8_t* data, std::size_t bytes) {
 	
 	if(bytes != pending_size) {
 		std::cout << "Bytes: " << bytes << "  Pending size: " << pending_size << std::endl;
-		throw std::runtime_error("Deserialization attempted with incomplete data.");
+		error("Deserialization attempted with incomplete data.");
 		// TODO: Calling got_data() with incomplete data should be allowed.
 	}
 	
@@ -75,11 +75,16 @@ void ConnectionBase::got_data(uint8_t* data, std::size_t bytes) {
 			break;
 		
 		default:
-			throw std::runtime_error("Deserialization attempted on unknown message type.");
+			error("Deserialization attempted on unknown message type.");
 	}
 	
 	if(bytes) {
-		m->deserialize(data, bytes);
+		try {
+			m->deserialize(data, bytes);
+		
+		} catch(std::exception& e) {
+			error(std::string("Deserialization failed: ") + e.what());
+		}
 	}
 	
 	got_message(m);
diff --git a/common/connectionbase.h b/common/connectionbase.h
index 4604896..f86a9ab 100644
--- a/common/connectionbase.h
+++ b/common/connectionbase.h
@@ -3,6 +3,7 @@
 
 #include <stdint.h>
 #include <cstdlib>
+#include <string>
 
 #include "message.h"
 
@@ -33,6 +34,10 @@ class ConnectionBase {
 		//! \param msg Received message.
 		virtual void got_message(const Message::p& msg) = 0;
 		
+		//! Called upon error (lost connection or failed deserialization).
+		//! \param msg Error message.
+		virtual void error(const std::string& msg) = 0;
+		
 	public:
 		ConnectionBase();
 		virtual ~ConnectionBase();
diff --git a/server/connection.cpp b/server/connection.cpp
index fd2d257..fc6cadb 100644
--- a/server/connection.cpp
+++ b/server/connection.cpp
@@ -8,7 +8,7 @@ Connection::Connection(boost::asio::io_service& io_service) : socket(io_service)
 
 void Connection::handle_read(uint8_t* data, std::size_t bytes, const boost::system::error_code& error) {
 	if(error) {
-		recv_callback.clear();
+		error();
 		return;
 	}
 	
@@ -26,12 +26,26 @@ void Connection::request_data(std::size_t bytes) {
 	
 	boost::asio::async_read(socket, boost::asio::buffer(buf, bytes),
 		boost::bind(&Connection::handle_read, shared_from_this(), buf, bytes, boost::asio::placeholders::error));
-	
-	// boost::asio::placeholders::error
 }
 
 void Connection::got_message(const Message::p& msg) {
-	recv_callback(msg);
+	// TODO: Implement message queueing. For now, unexpected messages will cause an error.
+	
+	boost::function<void (Message::p)> f = recv_callback;
+	recv_callback.clear();
+	error_callback.clear();
+	
+	f(msg);
+}
+
+void Connection::error(const std::string& msg) {
+	if(error_callback) {
+		boost::function<void (const std::string&)> f = error_callback;
+		recv_callback.clear();
+		error_callback.clear();
+		
+		f(msg);
+	}
 }
 
 void Connection::write_data(uint8_t* data, std::size_t bytes) {
@@ -43,8 +57,9 @@ Connection::p Connection::create(boost::asio::io_service& io_service) {
 	return Connection::p(new Connection(io_service));
 }
 
-void Connection::recv(boost::function<void (Message::p)> f) {
-	recv_callback = f;
+void Connection::recv(boost::function<void (Message::p)> callback, boost::function<void (const std::string&)> error) {
+	recv_callback = callback;
+	error_callback = error;
 	
 	start_recv();
 }
\ No newline at end of file
diff --git a/server/connection.h b/server/connection.h
index e6680c9..0f16022 100644
--- a/server/connection.h
+++ b/server/connection.h
@@ -15,6 +15,7 @@ class Connection : public ConnectionBase, public boost::enable_shared_from_this<
 		boost::asio::ip::tcp::socket socket;
 		
 		boost::function<void (Message::p)> recv_callback;
+		boost::function<void (const std::string&)> error_callback;
 		
 		Connection(boost::asio::io_service& io_service);
 		
@@ -33,6 +34,9 @@ class Connection : public ConnectionBase, public boost::enable_shared_from_this<
 		
 		//! Implements got_message().
 		virtual void got_message(const Message::p& msg);
+		
+		//! Implements error().
+		virtual void error(const std::string& msg);
 	
 	public:
 		typedef boost::shared_ptr<Connection> p;
@@ -41,7 +45,9 @@ class Connection : public ConnectionBase, public boost::enable_shared_from_this<
 		static p create(boost::asio::io_service& io_service);
 		
 		//! Initiates an asynchronous message receive.
-		void recv(boost::function<void (Message::p)> f);
+		//! \param callback Callback for received message.
+		//! \param error Callback for error.
+		void recv(boost::function<void (Message::p)> callback, boost::function<void (const std::string&)> error = 0);
 };
 
 #endif