From b93613e835be2c2316006ddef9bb3108646a534d Mon Sep 17 00:00:00 2001
From: JoelKatz <DavidJoelSchwartz@GMail.com>
Date: Thu, 19 Jan 2012 18:38:31 -0800
Subject: [PATCH] Fix crash on bad resolve. Fix peer shutdown logic. Fix peer
 tracking logic.

---
 ConnectionPool.cpp | 35 +++++++++++++++++++++--------------
 ConnectionPool.h   |  2 +-
 Peer.cpp           | 30 +++++++++++++++++++++++++++---
 Peer.h             |  3 ++-
 4 files changed, 51 insertions(+), 19 deletions(-)

diff --git a/ConnectionPool.cpp b/ConnectionPool.cpp
index c56c36dcd0..48ea049671 100644
--- a/ConnectionPool.cpp
+++ b/ConnectionPool.cpp
@@ -53,7 +53,7 @@ bool ConnectionPool::addToMap(const uint160& hanko, Peer::pointer peer)
 	return peerMap.insert(std::make_pair(hanko, peer)).second;
 }
 
-bool ConnectionPool::delFromMap(const uint160& hanko, Peer::pointer peer)
+bool ConnectionPool::delFromMap(const uint160& hanko)
 {
 	boost::mutex::scoped_lock sl(peerLock);
 	std::map<uint160, Peer::pointer>::iterator it=peerMap.find(hanko);
@@ -84,20 +84,27 @@ std::map<uint160, Peer::pointer> ConnectionPool::getAllConnected()
 
 bool ConnectionPool::connectTo(const std::string& host, const std::string& port)
 {
-	boost::asio::ip::tcp::resolver res(theApp->getIOService());
-	boost::asio::ip::tcp::resolver::query query(host.c_str(), port.c_str());
-	boost::asio::ip::tcp::resolver::iterator it(res.resolve(query)), end;
-
-	Peer::pointer peer(Peer::create(theApp->getIOService()));
-	boost::system::error_code error = boost::asio::error::host_not_found;
-	while (error && (it!=end))
+	try
 	{
-		peer->getSocket().close();
-		peer->getSocket().connect(*it++, error);
+		boost::asio::ip::tcp::resolver res(theApp->getIOService());
+		boost::asio::ip::tcp::resolver::query query(host.c_str(), port.c_str());
+		boost::asio::ip::tcp::resolver::iterator it(res.resolve(query)), end;
+
+		Peer::pointer peer(Peer::create(theApp->getIOService()));
+		boost::system::error_code error = boost::asio::error::host_not_found;
+		while (error && (it!=end))
+		{
+			peer->getSocket().close();
+			peer->getSocket().connect(*it++, error);
+		}
+		if(error) return false;
+		boost::mutex::scoped_lock sl(peerLock);
+		mPeers.push_back(peer);
+		peer->connected(boost::system::error_code());
+	}
+	catch (...)
+	{
+		return false;
 	}
-	if(error) return false;
-	boost::mutex::scoped_lock sl(peerLock);
-	mPeers.push_back(peer);
-	peer->connected(boost::system::error_code());
 	return true;
 }
diff --git a/ConnectionPool.h b/ConnectionPool.h
index 6b9f1daf9c..0d39a37a4f 100644
--- a/ConnectionPool.h
+++ b/ConnectionPool.h
@@ -29,7 +29,7 @@ public:
 	// hanko->peer mapping functions
 	bool inMap(const uint160& hanko);
 	bool addToMap(const uint160& hanko, Peer::pointer peer);
-	bool delFromMap(const uint160& hanko, Peer::pointer peer);
+	bool delFromMap(const uint160& hanko);
 	Peer::pointer findInMap(const uint160& hanko);
 	std::map<uint160, Peer::pointer> getAllConnected();
 
diff --git a/Peer.cpp b/Peer.cpp
index bc20c2ec27..1b2455c946 100644
--- a/Peer.cpp
+++ b/Peer.cpp
@@ -43,6 +43,12 @@ void Peer::handle_write(const boost::system::error_code& error, size_t bytes_tra
 	}
 }
 
+void Peer::detach()
+{
+	mSocket.close();
+	if(!!mHanko) theApp->getConnectionPool().delFromMap(mHanko);
+}
+
 void Peer::connected(const boost::system::error_code& error)
 {
 	if(!error)
@@ -52,7 +58,11 @@ void Peer::connected(const boost::system::error_code& error)
 		sendHello();
 		start_read_header();
 	}
-	else cout  << "Peer::connected Error: " << error << endl; //else BOOST_LOG_TRIVIAL(info) << "Error: " << error;
+	else
+	{
+		detach();
+		cout  << "Peer::connected Error: " << error << endl; //else BOOST_LOG_TRIVIAL(info) << "Error: " << error;
+	}
 	
 }
 
@@ -103,9 +113,19 @@ void Peer::handle_read_header(const boost::system::error_code& error)
 	if(!error)
 	{
 		unsigned msg_len = PackedMessage::getLength(mReadbuf);
+		// WRITEME: Compare to maximum message length, abort if too large
+		if(msg_len>(32*1024*1024))
+		{
+			detach();
+			return;
+		}
 		start_read_body(msg_len);
 	}
-	else cout  << "Peer::connected Error: " << error << endl; //else BOOST_LOG_TRIVIAL(info) << "Error: " << error;
+	else
+	{
+		detach();
+		cout  << "Peer::connected Error: " << error << endl; //else BOOST_LOG_TRIVIAL(info) << "Error: " << error;
+	}
 }
 
 void Peer::handle_read_body(const boost::system::error_code& error)
@@ -115,7 +135,11 @@ void Peer::handle_read_body(const boost::system::error_code& error)
 		processReadBuffer();
 		start_read_header();
 	}
-	else cout  << "Peer::connected Error: " << error << endl; //else BOOST_LOG_TRIVIAL(info) << "Error: " << error;
+	else
+	{
+		detach();
+		cout  << "Peer::connected Error: " << error << endl; //else BOOST_LOG_TRIVIAL(info) << "Error: " << error;
+	}
 }
 
 
diff --git a/Peer.h b/Peer.h
index 6a98dc4217..51657b14bc 100644
--- a/Peer.h
+++ b/Peer.h
@@ -24,7 +24,7 @@ protected:
 	std::list<PackedMessage::pointer> mSendQ;
 	PackedMessage::pointer mSendingPacket;
 	std::bitset<32> mPeerBits;
-	uint160 hanko;
+	uint160 mHanko;
 
 	Peer(boost::asio::io_service& io_service);
 
@@ -76,6 +76,7 @@ public:
 	}
 
 	void connected(const boost::system::error_code& error);
+	void detach();
 
 	void sendPacket(PackedMessage::pointer packet);
 	void sendLedgerProposal(Ledger::pointer ledger);