2 files changed, 28 insertions, 7 deletions

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