2 files changed, 20 insertions, 35 deletions

diff --git a/bin/bbackupd/BackupClientContext.cpp b/bin/bbackupd/BackupClientContext.cpp
index 26df04be..45c48a67 100644
--- a/bin/bbackupd/BackupClientContext.cpp
+++ b/bin/bbackupd/BackupClientContext.cpp
@@ -73,7 +73,8 @@ BackupClientContext::BackupClientContext
   mKeepAliveTimer(0, "KeepAliveTime"),
   mbIsManaged(false),
   mrProgressNotifier(rProgressNotifier),
-  mTcpNiceMode(TcpNiceMode)
+  mTcpNiceMode(TcpNiceMode),
+  mpNice(NULL)
 {
 }
 
@@ -113,13 +114,12 @@ BackupProtocolClient &BackupClientContext::GetConnection()
 	{
 		return *mapConnection;
 	}
-	
-	// there shouldn't be a connection open
-	ASSERT(mapSocket.get() == 0);
+
 	// Defensive. Must close connection before releasing any old socket.
 	mapConnection.reset();
-	mapSocket.reset(new SocketStreamTLS);
-	
+
+	std::auto_ptr<SocketStream> apSocket(new SocketStreamTLS);
+
 	try
 	{
 		// Defensive.
@@ -130,21 +130,22 @@ BackupProtocolClient &BackupClientContext::GetConnection()
 			mHostname << "'...");
 
 		// Connect!
-		((SocketStreamTLS *)(mapSocket.get()))->Open(mrTLSContext,
+		((SocketStreamTLS *)(apSocket.get()))->Open(mrTLSContext,
 			Socket::TypeINET, mHostname, mPort);
 		
 		if(mTcpNiceMode)
 		{
-			// Pass control of mapSocket to NiceSocketStream,
+			// Pass control of apSocket to NiceSocketStream,
 			// which will take care of destroying it for us.
-			mapNice.reset(new NiceSocketStream(mapSocket));
-			mapConnection.reset(new BackupProtocolClient(*mapNice));
+			// But we need to hang onto a pointer to the nice
+			// socket, so we can enable and disable nice mode.
+			// This is scary, it could be deallocated under us.
+			mpNice = new NiceSocketStream(apSocket);
+			apSocket.reset(mpNice);
 		}
-		else
-		{
-			mapConnection.reset(new BackupProtocolClient(*mapSocket));
-		}
-		
+
+		mapConnection.reset(new BackupProtocolClient(apSocket));
+
 		// Set logging option
 		mapConnection->SetLogToSysLog(mExtendedLogging);
 		
@@ -165,10 +166,10 @@ BackupProtocolClient &BackupClientContext::GetConnection()
 				mapConnection->SetLogToFile(mpExtendedLogFileHandle);
 			}
 		}
-		
+
 		// Handshake
 		mapConnection->Handshake();
-		
+
 		// Check the version of the server
 		{
 			std::auto_ptr<BackupProtocolVersion> serverVersion(
@@ -192,8 +193,6 @@ BackupProtocolClient &BackupClientContext::GetConnection()
 				try
 				{
 					mapConnection->QueryFinished();
-					mapNice.reset();
-					mapSocket.reset();
 				}
 				catch(...)
 				{
@@ -222,8 +221,6 @@ BackupProtocolClient &BackupClientContext::GetConnection()
 	{
 		// Clean up.
 		mapConnection.reset();
-		mapNice.reset();
-		mapSocket.reset();
 		throw;
 	}
 	
@@ -269,17 +266,6 @@ void BackupClientContext::CloseAnyOpenConnection()
 		mapConnection.reset();
 	}
 
-	try
-	{
-		// Be nice about closing the socket
-		mapNice.reset();
-		mapSocket.reset();
-	}
-	catch(...)
-	{
-		// Ignore errors
-	}
-
 	// Delete any pending list
 	if(mpDeleteList != 0)
 	{
diff --git a/bin/bbackupd/BackupClientContext.h b/bin/bbackupd/BackupClientContext.h
index 1fcc6ede..7e081e2d 100644
--- a/bin/bbackupd/BackupClientContext.h
+++ b/bin/bbackupd/BackupClientContext.h
@@ -214,7 +214,7 @@ public:
 	{
 		if(mTcpNiceMode)
 		{
-			mapNice->SetEnabled(enabled);
+			mpNice->SetEnabled(enabled);
 		}
 	}
 
@@ -226,8 +226,6 @@ private:
 	std::string mHostname;
 	int mPort;
 	uint32_t mAccountNumber;
-	std::auto_ptr<SocketStream> mapSocket;
-	std::auto_ptr<NiceSocketStream> mapNice;
 	std::auto_ptr<BackupProtocolClient> mapConnection;
 	bool mExtendedLogging;
 	bool mExtendedLogToFile;
@@ -246,6 +244,7 @@ private:
 	int mMaximumDiffingTime;
 	ProgressNotifier &mrProgressNotifier;
 	bool mTcpNiceMode;
+	NiceSocketStream *mpNice;
 };
 
 #endif // BACKUPCLIENTCONTEXT__H