3 files changed, 54 insertions, 36 deletions

diff --git a/src/HTTP/UrlClient.cpp b/src/HTTP/UrlClient.cpp
index 29ee7e18d..c6ac69c32 100644
--- a/src/HTTP/UrlClient.cpp
+++ b/src/HTTP/UrlClient.cpp
@@ -54,12 +54,10 @@ public:
 		m_Callbacks->OnError(a_ErrorMessage);
 
 		// Terminate the request's TCP link:
-		auto link = m_Link;
-		if (link != nullptr)
+		if (auto link = m_Link.lock())
 		{
 			link->Close();
 		}
-		m_Self.reset();
 	}
 
 
@@ -76,7 +74,7 @@ public:
 		{
 			return nullptr;
 		}
-		cX509CertPtr cert;
+		cX509CertPtr cert = std::make_shared<cX509Cert>();
 		if (!cert->Parse(itr->second.data(), itr->second.size()))
 		{
 			LOGD("OwnCert failed to parse");
@@ -92,7 +90,7 @@ public:
 		{
 			return nullptr;
 		}
-		cCryptoKeyPtr key;
+		cCryptoKeyPtr key = std::make_shared<cCryptoKey>();
 		auto passItr = m_Options.find("OwnPrivKeyPassword");
 		auto pass = (passItr == m_Options.end()) ? AString() : passItr->second;
 		if (!key->ParsePrivate(itr->second.data(), itr->second.size(), pass))
@@ -126,15 +124,15 @@ protected:
 	/** Extra options to be used for the request. */
 	AStringMap m_Options;
 
-	/** SharedPtr to self, so that this object can keep itself alive for as long as it needs,
+	/** weak_ptr to self, so that this object can keep itself alive as needed by calling lock(),
 	and pass self as callbacks to cNetwork functions. */
-	std::shared_ptr<cUrlClientRequest> m_Self;
+	std::weak_ptr<cUrlClientRequest> m_Self;
 
 	/** The handler that "talks" the protocol specified in m_UrlScheme, handles all the sending and receiving. */
 	std::shared_ptr<cSchemeHandler> m_SchemeHandler;
 
 	/** The link handling the request. */
-	cTCPLinkPtr m_Link;
+	std::weak_ptr<cTCPLink> m_Link;
 
 	/** The number of redirect attempts that will still be followed.
 	If the response specifies a redirect and this is nonzero, the redirect is followed.
@@ -171,7 +169,6 @@ protected:
 	virtual void OnError(int a_ErrorCode, const AString & a_ErrorMsg) override
 	{
 		m_Callbacks->OnError(Printf("Network error %d (%s)", a_ErrorCode, a_ErrorMsg.c_str()));
-		m_Self.reset();
 	}
 
 
@@ -332,8 +329,8 @@ public:
 	// cHTTPResponseParser::cCallbacks overrides:
 	virtual void OnError(const AString & a_ErrorDescription) override
 	{
-		m_ParentRequest.CallErrorCallback(a_ErrorDescription);
 		m_Link = nullptr;
+		m_ParentRequest.CallErrorCallback(a_ErrorDescription);
 	}
 
 
@@ -430,6 +427,8 @@ public:
 		else
 		{
 			m_ParentRequest.GetCallbacks().OnBodyFinished();
+			// Finished recieving data, shutdown the link
+			m_Link->Shutdown();
 		}
 	}
 
@@ -493,15 +492,21 @@ void cUrlClientRequest::RedirectTo(const AString & a_RedirectUrl)
 		return;
 	}
 
+	// Keep ourself alive while the link drops ownership
+	auto Self = m_Self.lock();
+
 	// Do the actual redirect:
-	m_Link->Close();
+	if (auto Link = m_Link.lock())
+	{
+		Link->Close();
+	}
+
 	m_Url = a_RedirectUrl;
 	m_NumRemainingRedirects = m_NumRemainingRedirects - 1;
-	auto res = DoRequest(m_Self);
+	auto res = DoRequest(Self);
 	if (!res.first)
 	{
 		m_Callbacks->OnError(Printf("Redirection failed: %s", res.second.c_str()));
-		return;
 	}
 }
 
@@ -560,9 +565,6 @@ void cUrlClientRequest::OnRemoteClosed()
 	{
 		handler->OnRemoteClosed();
 	}
-
-	// Let ourselves be deleted
-	m_Self.reset();
 }
 
 
@@ -590,7 +592,8 @@ std::pair<bool, AString> cUrlClientRequest::DoRequest(std::shared_ptr<cUrlClient
 		return std::make_pair(false, Printf("Unknown Url scheme: %s", m_UrlScheme.c_str()));
 	}
 
-	if (!cNetwork::Connect(m_UrlHost, m_UrlPort, m_Self, m_Self))
+	// Connect and transfer ownership to the link
+	if (!cNetwork::Connect(m_UrlHost, m_UrlPort, a_Self, a_Self))
 	{
 		return std::make_pair(false, "Network connection failed");
 	}
diff --git a/src/OSSupport/NetworkSingleton.cpp b/src/OSSupport/NetworkSingleton.cpp
index af3a701a8..2f7ed3802 100644
--- a/src/OSSupport/NetworkSingleton.cpp
+++ b/src/OSSupport/NetworkSingleton.cpp
@@ -6,6 +6,7 @@
 
 #include "Globals.h"
 #include "NetworkSingleton.h"
+#include "OSSupport/Network.h"
 #include <event2/thread.h>
 #include <event2/bufferevent.h>
 #include <event2/listener.h>
@@ -102,11 +103,25 @@ void cNetworkSingleton::Terminate(void)
 	event_base_loopbreak(m_EventBase);
 	m_EventLoopThread.join();
 
-	// Remove all objects:
+	// Close all open connections:
 	{
 		cCSLock Lock(m_CS);
-		m_Connections.clear();
-		m_Servers.clear();
+		// Must take copies because Close will modify lists
+		auto Conns = m_Connections;
+		for (auto & Conn : Conns)
+		{
+			Conn->Close();
+		}
+
+		auto Servers = m_Servers;
+		for (auto & Server : Servers)
+		{
+			Server->Close();
+		}
+
+		// Closed handles should have removed themself
+		ASSERT(m_Connections.empty());
+		ASSERT(m_Servers.empty());
 	}
 
 	// Free the underlying LibEvent objects:
@@ -167,7 +182,7 @@ void cNetworkSingleton::SignalizeStartup(evutil_socket_t a_Socket, short a_Event
 
 
 
-void cNetworkSingleton::AddLink(cTCPLinkImplPtr a_Link)
+void cNetworkSingleton::AddLink(cTCPLinkPtr a_Link)
 {
 	ASSERT(!m_HasTerminated);
 	cCSLock Lock(m_CS);
@@ -178,7 +193,7 @@ void cNetworkSingleton::AddLink(cTCPLinkImplPtr a_Link)
 
 
 
-void cNetworkSingleton::RemoveLink(const cTCPLinkImpl * a_Link)
+void cNetworkSingleton::RemoveLink(const cTCPLink * a_Link)
 {
 	ASSERT(!m_HasTerminated);
 	cCSLock Lock(m_CS);
@@ -196,7 +211,7 @@ void cNetworkSingleton::RemoveLink(const cTCPLinkImpl * a_Link)
 
 
 
-void cNetworkSingleton::AddServer(cServerHandleImplPtr a_Server)
+void cNetworkSingleton::AddServer(cServerHandlePtr a_Server)
 {
 	ASSERT(!m_HasTerminated);
 	cCSLock Lock(m_CS);
@@ -207,7 +222,7 @@ void cNetworkSingleton::AddServer(cServerHandleImplPtr a_Server)
 
 
 
-void cNetworkSingleton::RemoveServer(const cServerHandleImpl * a_Server)
+void cNetworkSingleton::RemoveServer(const cServerHandle * a_Server)
 {
 	ASSERT(!m_HasTerminated);
 	cCSLock Lock(m_CS);
diff --git a/src/OSSupport/NetworkSingleton.h b/src/OSSupport/NetworkSingleton.h
index 902c62550..5f3c11720 100644
--- a/src/OSSupport/NetworkSingleton.h
+++ b/src/OSSupport/NetworkSingleton.h
@@ -24,12 +24,12 @@
 
 // fwd:
 struct event_base;
-class cTCPLinkImpl;
-typedef std::shared_ptr<cTCPLinkImpl> cTCPLinkImplPtr;
-typedef std::vector<cTCPLinkImplPtr> cTCPLinkImplPtrs;
-class cServerHandleImpl;
-typedef std::shared_ptr<cServerHandleImpl> cServerHandleImplPtr;
-typedef std::vector<cServerHandleImplPtr> cServerHandleImplPtrs;
+class cTCPLink;
+typedef std::shared_ptr<cTCPLink> cTCPLinkPtr;
+typedef std::vector<cTCPLinkPtr> cTCPLinkPtrs;
+class cServerHandle;
+typedef std::shared_ptr<cServerHandle> cServerHandlePtr;
+typedef std::vector<cServerHandlePtr> cServerHandlePtrs;
 
 
 
@@ -61,20 +61,20 @@ public:
 
 	/** Adds the specified link to m_Connections.
 	Used by the underlying link implementation when a new link is created. */
-	void AddLink(cTCPLinkImplPtr a_Link);
+	void AddLink(cTCPLinkPtr a_Link);
 
 	/** Removes the specified link from m_Connections.
 	Used by the underlying link implementation when the link is closed / errored. */
-	void RemoveLink(const cTCPLinkImpl * a_Link);
+	void RemoveLink(const cTCPLink * a_Link);
 
 	/** Adds the specified link to m_Servers.
 	Used by the underlying server handle implementation when a new listening server is created.
 	Only servers that succeed in listening are added. */
-	void AddServer(cServerHandleImplPtr a_Server);
+	void AddServer(cServerHandlePtr a_Server);
 
 	/** Removes the specified server from m_Servers.
 	Used by the underlying server handle implementation when the server is closed. */
-	void RemoveServer(const cServerHandleImpl * a_Server);
+	void RemoveServer(const cServerHandle * a_Server);
 
 protected:
 
@@ -82,10 +82,10 @@ protected:
 	event_base * m_EventBase;
 
 	/** Container for all client connections, including ones with pending-connect. */
-	cTCPLinkImplPtrs m_Connections;
+	cTCPLinkPtrs m_Connections;
 
 	/** Container for all servers that are currently active. */
-	cServerHandleImplPtrs m_Servers;
+	cServerHandlePtrs m_Servers;
 
 	/** Mutex protecting all containers against multithreaded access. */
 	cCriticalSection m_CS;