Ensure disconnect() method always disconnect the underlying sockets.

Added additional checks after weak pointer locks.

7 files changed, 96 insertions, 24 deletions

diff --git a/src/vmime/net/imap/IMAPConnection.cpp b/src/vmime/net/imap/IMAPConnection.cpp
index df3da1c3..0e4020f1 100644
--- a/src/vmime/net/imap/IMAPConnection.cpp
+++ b/src/vmime/net/imap/IMAPConnection.cpp
@@ -51,14 +51,13 @@
 
 #include <sstream>
 
-
 // Helpers for service properties
 #define GET_PROPERTY(type, prop) \
-	(m_store.lock()->getInfos().getPropertyValue <type>(getSession(), \
-		dynamic_cast <const IMAPServiceInfos&>(m_store.lock()->getInfos()).getProperties().prop))
+	(getStoreOrThrow()->getInfos().getPropertyValue <type>(getSession(), \
+		dynamic_cast <const IMAPServiceInfos&>(getStoreOrThrow()->getInfos()).getProperties().prop))
 #define HAS_PROPERTY(prop) \
-	(m_store.lock()->getInfos().hasProperty(getSession(), \
-		dynamic_cast <const IMAPServiceInfos&>(m_store.lock()->getInfos()).getProperties().prop))
+	(getStoreOrThrow()->getInfos().hasProperty(getSession(), \
+		dynamic_cast <const IMAPServiceInfos&>(getStoreOrThrow()->getInfos()).getProperties().prop))
 
 
 namespace vmime {
@@ -127,6 +126,10 @@ void IMAPConnection::connect() {
 
 	shared_ptr <IMAPStore> store = m_store.lock();
 
+	if (!store) {
+		throw exceptions::illegal_state("Store not found");
+	}
+
 	// Create the time-out handler
 	if (store->getTimeoutHandlerFactory()) {
 		m_timeoutHandler = store->getTimeoutHandlerFactory()->create();
@@ -244,6 +247,17 @@ void IMAPConnection::connect() {
 	setState(STATE_AUTHENTICATED);
 }
 
+shared_ptr <IMAPStore> IMAPConnection::getStoreOrThrow() {
+
+	auto store = m_store.lock();
+
+	if (!store) {
+		throw exceptions::illegal_state("Store not found");
+	}
+
+	return store;
+}
+
 
 void IMAPConnection::authenticate() {
 
@@ -529,9 +543,15 @@ void IMAPConnection::startTLS() {
 			throw exceptions::command_error("STARTTLS", resp->getErrorLog(), "bad response");
 		}
 
+		auto store = m_store.lock();
+
+		if (!store) {
+			throw exceptions::illegal_state("Store not found");
+		}
+
 		shared_ptr <tls::TLSSession> tlsSession = tls::TLSSession::create(
-			m_store.lock()->getCertificateVerifier(),
-			m_store.lock()->getSession()->getTLSProperties()
+			store->getCertificateVerifier(),
+			store->getSession()->getTLSProperties()
 		);
 
 		shared_ptr <tls::TLSSocket> tlsSocket = tlsSession->getSocket(m_socket);
@@ -701,11 +721,11 @@ shared_ptr <connectionInfos> IMAPConnection::getConnectionInfos() const {
 
 void IMAPConnection::disconnect() {
 
-	if (!isConnected()) {
+	bool wasConnected = isConnected();
+	internalDisconnect();
+	if (!wasConnected) {
 		throw exceptions::not_connected();
 	}
-
-	internalDisconnect();
 }
 
 
@@ -714,7 +734,9 @@ void IMAPConnection::internalDisconnect() {
 	if (isConnected()) {
 
 		IMAPCommand::LOGOUT()->send(dynamicCast <IMAPConnection>(shared_from_this()));
+	}
 
+	if (m_socket) {
 		m_socket->disconnect();
 		m_socket = null;
 	}
@@ -836,7 +858,13 @@ shared_ptr <IMAPStore> IMAPConnection::getStore() {
 
 shared_ptr <session> IMAPConnection::getSession() {
 
-	return m_store.lock()->getSession();
+	auto store = m_store.lock();
+
+	if (!store) {
+		throw exceptions::illegal_state("Store not found");
+	}
+
+	return store->getSession();
 }
 
 
diff --git a/src/vmime/net/imap/IMAPConnection.hpp b/src/vmime/net/imap/IMAPConnection.hpp
index 99750d49..d229d0f7 100644
--- a/src/vmime/net/imap/IMAPConnection.hpp
+++ b/src/vmime/net/imap/IMAPConnection.hpp
@@ -114,6 +114,8 @@ public:
 
 private:
 
+	shared_ptr <IMAPStore> getStoreOrThrow();
+
 	void authenticate();
 #if VMIME_HAVE_SASL_SUPPORT
 	void authenticateSASL();
diff --git a/src/vmime/net/imap/IMAPFolder.cpp b/src/vmime/net/imap/IMAPFolder.cpp
index 98bc05d8..ed219448 100644
--- a/src/vmime/net/imap/IMAPFolder.cpp
+++ b/src/vmime/net/imap/IMAPFolder.cpp
@@ -290,9 +290,17 @@ void IMAPFolder::close(const bool expunge) {
 	shared_ptr <IMAPStore> store = m_store.lock();
 
 	if (!store) {
+		if (m_connection && m_connection->isConnected())
+			m_connection->disconnect();
+		m_connection = nullptr;
+		m_open = false;
 		throw exceptions::illegal_state("Store disconnected");
 	}
 
+	if (!m_connection || !m_connection->isConnected()) {
+		throw exceptions::illegal_state("Folder not connected");
+	}
+
 	if (!isOpen()) {
 		throw exceptions::illegal_state("Folder not open");
 	}
@@ -314,7 +322,7 @@ void IMAPFolder::close(const bool expunge) {
 	oldConnection->disconnect();
 
 	// Now use default store connection
-	m_connection = m_store.lock()->connection();
+	m_connection = store->connection();
 
 	m_open = false;
 	m_mode = -1;
@@ -546,7 +554,6 @@ bool IMAPFolder::isOpen() const {
 	return m_open;
 }
 
-
 shared_ptr <message> IMAPFolder::getMessage(const size_t num) {
 
 	if (!isOpen()) {
@@ -1421,6 +1428,10 @@ void IMAPFolder::noop() {
 		throw exceptions::illegal_state("Store disconnected");
 	}
 
+	if (!m_connection || !m_connection->isConnected()) {
+		throw exceptions::illegal_state("Store disconnected");
+	}
+
 	IMAPCommand::NOOP()->send(m_connection);
 
 	scoped_ptr <IMAPParser::response> resp(m_connection->readResponse());
diff --git a/src/vmime/net/imap/IMAPMessage.cpp b/src/vmime/net/imap/IMAPMessage.cpp
index f74a4a4b..7f155469 100644
--- a/src/vmime/net/imap/IMAPMessage.cpp
+++ b/src/vmime/net/imap/IMAPMessage.cpp
@@ -243,7 +243,7 @@ void IMAPMessage::extract(
 	shared_ptr <const IMAPFolder> folder = m_folder.lock();
 
 	if (!folder) {
-		throw exceptions::folder_not_found();
+		throw exceptions::illegal_state("Folder not open");
 	}
 
 	extractImpl(
@@ -318,6 +318,14 @@ size_t IMAPMessage::extractImpl(
 
 	shared_ptr <const IMAPFolder> folder = m_folder.lock();
 
+	if (!folder) {
+		throw exceptions::folder_not_found();
+	}
+	
+	if (!folder->isOpen()) {
+		throw exceptions::illegal_state("Folder not open");
+	}
+	
 	IMAPMessage_literalHandler literalHandler(os, progress);
 
 	if (length == 0) {
@@ -474,7 +482,6 @@ size_t IMAPMessage::extractImpl(
 	return literalHandler.getTarget()->getBytesWritten();
 }
 
-
 int IMAPMessage::processFetchResponse(
 	const fetchAttributes& options,
 	const IMAPParser::message_data& msgData
@@ -717,10 +724,16 @@ shared_ptr <vmime::message> IMAPMessage::getParsedMessage() {
 
 	} catch (exceptions::unfetched_object&) {
 
+		auto folder = m_folder.lock();
+
+		if (!folder) {
+			throw exceptions::folder_not_found();
+		}
+
 		std::vector <shared_ptr <message> > msgs;
 		msgs.push_back(dynamicCast <IMAPMessage>(shared_from_this()));
 
-		m_folder.lock()->fetchMessages(
+		folder->fetchMessages(
 			msgs, fetchAttributes(fetchAttributes::STRUCTURE), /* progress */ NULL
 		);
 
diff --git a/src/vmime/net/imap/IMAPMessagePartContentHandler.cpp b/src/vmime/net/imap/IMAPMessagePartContentHandler.cpp
index 85f6cb6e..040d2ab0 100644
--- a/src/vmime/net/imap/IMAPMessagePartContentHandler.cpp
+++ b/src/vmime/net/imap/IMAPMessagePartContentHandler.cpp
@@ -36,7 +36,6 @@
 #include "vmime/utility/outputStreamAdapter.hpp"
 #include "vmime/utility/inputStreamStringAdapter.hpp"
 
-
 namespace vmime {
 namespace net {
 namespace imap {
@@ -71,6 +70,9 @@ void IMAPMessagePartContentHandler::generate(
 	shared_ptr <IMAPMessage> msg = m_message.lock();
 	shared_ptr <messagePart> part = m_part.lock();
 
+	if (!msg || !part)
+		throw exceptions::message_not_found();
+
 	// Data is already encoded
 	if (isEncoded()) {
 
@@ -140,6 +142,9 @@ void IMAPMessagePartContentHandler::extract(
 	shared_ptr <IMAPMessage> msg = m_message.lock();
 	shared_ptr <messagePart> part = m_part.lock();
 
+	if (!msg || !part)
+		throw exceptions::message_not_found();
+
 	// No decoding to perform
 	if (!isEncoded()) {
 
@@ -172,12 +177,17 @@ void IMAPMessagePartContentHandler::extractRaw(
 	shared_ptr <IMAPMessage> msg = m_message.lock();
 	shared_ptr <messagePart> part = m_part.lock();
 
+	if (!msg || !part)
+		throw exceptions::message_not_found();
+
 	msg->extractPart(part, os, progress);
 }
 
 
 size_t IMAPMessagePartContentHandler::getLength() const {
-
+	shared_ptr <messagePart> part = m_part.lock();
+	if (!part)
+		throw exceptions::message_not_found();
 	return m_part.lock()->getSize();
 }
 
diff --git a/src/vmime/net/imap/IMAPParser.hpp b/src/vmime/net/imap/IMAPParser.hpp
index 409362da..2a37b9cd 100644
--- a/src/vmime/net/imap/IMAPParser.hpp
+++ b/src/vmime/net/imap/IMAPParser.hpp
@@ -4912,6 +4912,9 @@ public:
 		shared_ptr <timeoutHandler> toh = m_timeoutHandler.lock();
 		shared_ptr <socket> sok = m_socket.lock();
 
+		if (!sok)
+			throw exceptions::illegal_state("Store disconnected");
+
 		if (toh) {
 			toh->resetTimeOut();
 		}
@@ -4957,6 +4960,9 @@ public:
 		shared_ptr <timeoutHandler> toh = m_timeoutHandler.lock();
 		shared_ptr <socket> sok = m_socket.lock();
 
+		if (!sok)
+			throw exceptions::illegal_state("Store disconnected");
+
 		if (m_progress) {
 			m_progress->start(count);
 		}
diff --git a/src/vmime/net/imap/IMAPStore.cpp b/src/vmime/net/imap/IMAPStore.cpp
index eafa4442..be5c4de7 100644
--- a/src/vmime/net/imap/IMAPStore.cpp
+++ b/src/vmime/net/imap/IMAPStore.cpp
@@ -179,9 +179,7 @@ shared_ptr <IMAPConnection> IMAPStore::getConnection() {
 
 void IMAPStore::disconnect() {
 
-	if (!isConnected()) {
-		throw exceptions::not_connected();
-	}
+	bool wasConnected = isConnected();
 
 	for (std::list <IMAPFolder*>::iterator it = m_folders.begin() ;
 	     it != m_folders.end() ; ++it) {
@@ -191,10 +189,14 @@ void IMAPStore::disconnect() {
 
 	m_folders.clear();
 
+	if (m_connection) {
+		m_connection->disconnect();
+		m_connection = null;
+	}
 
-	m_connection->disconnect();
-
-	m_connection = null;
+	if (!wasConnected) {
+		throw exceptions::not_connected();
+	}
 }