3 files changed, 402 insertions, 2 deletions

diff --git a/src/OSSupport/Network.h b/src/OSSupport/Network.h
index 1162d7fc6..78c5e92f0 100644
--- a/src/OSSupport/Network.h
+++ b/src/OSSupport/Network.h
@@ -20,6 +20,11 @@ typedef std::vector<cTCPLinkPtr> cTCPLinkPtrs;
 class cServerHandle;
 typedef SharedPtr<cServerHandle> cServerHandlePtr;
 typedef std::vector<cServerHandlePtr> cServerHandlePtrs;
+class cCryptoKey;
+typedef SharedPtr<cCryptoKey> cCryptoKeyPtr;
+class cX509Cert;
+typedef SharedPtr<cX509Cert> cX509CertPtr;
+
 
 
 
@@ -49,6 +54,10 @@ public:
 		Sending data on the link is not an error, but the data won't be delivered. */
 		virtual void OnRemoteClosed(void) = 0;
 
+		/** Called when the TLS handshake has been completed and communication can continue regularly.
+		Has an empty default implementation, so that link callback descendants don't need to specify TLS handlers when they don't use TLS at all. */
+		virtual void OnTlsHandshakeCompleted(void) {}
+
 		/** Called when an error is detected on the connection. */
 		virtual void OnError(int a_ErrorCode, const AString & a_ErrorMsg) = 0;
 	};
@@ -90,6 +99,30 @@ public:
 	Sends the RST packet, queued outgoing and incoming data is lost. */
 	virtual void Close(void) = 0;
 
+	/** Starts a TLS handshake as a client connection.
+	If a client certificate should be used for the connection, set the certificate into a_OwnCertData and
+	its corresponding private key to a_OwnPrivKeyData. If both are empty, no client cert is presented.
+	a_OwnPrivKeyPassword is the password to be used for decoding PrivKey, empty if not passworded.
+	Returns empty string on success, non-empty error description on failure. */
+	virtual AString StartTLSClient(
+		cX509CertPtr a_OwnCert,
+		cCryptoKeyPtr a_OwnPrivKey
+	) = 0;
+
+	/** Starts a TLS handshake as a server connection.
+	Set the server certificate into a_CertData and its corresponding private key to a_OwnPrivKeyData.
+	a_OwnPrivKeyPassword is the password to be used for decoding PrivKey, empty if not passworded.
+	a_StartTLSData is any data that should be pushed into the TLS before reading more data from the remote.
+	This is used mainly for protocols starting TLS in the middle of communication, when the TLS start command
+	can be received together with the TLS Client Hello message in one OnReceivedData() call, to re-queue the
+	Client Hello message into the TLS handshake buffer.
+	Returns empty string on success, non-empty error description on failure. */
+	virtual AString StartTLSServer(
+		cX509CertPtr a_OwnCert,
+		cCryptoKeyPtr a_OwnPrivKey,
+		const AString & a_StartTLSData
+	) = 0;
+
 	/** Returns the callbacks that are used. */
 	cCallbacksPtr GetCallbacks(void) const { return m_Callbacks; }
 
diff --git a/src/OSSupport/TCPLinkImpl.cpp b/src/OSSupport/TCPLinkImpl.cpp
index b15b6282f..d55dc9da1 100644
--- a/src/OSSupport/TCPLinkImpl.cpp
+++ b/src/OSSupport/TCPLinkImpl.cpp
@@ -48,6 +48,9 @@ cTCPLinkImpl::cTCPLinkImpl(evutil_socket_t a_Socket, cTCPLink::cCallbacksPtr a_L
 
 cTCPLinkImpl::~cTCPLinkImpl()
 {
+	// If the TLS context still exists, free it:
+	m_TlsContext.reset();
+
 	bufferevent_free(m_BufferEvent);
 }
 
@@ -129,7 +132,16 @@ bool cTCPLinkImpl::Send(const void * a_Data, size_t a_Length)
 		LOGD("%s: Cannot send data, the link is already shut down.", __FUNCTION__);
 		return false;
 	}
-	return (bufferevent_write(m_BufferEvent, a_Data, a_Length) == 0);
+
+	// If running in TLS mode, push the data into the TLS context instead:
+	if (m_TlsContext != nullptr)
+	{
+		m_TlsContext->Send(a_Data, a_Length);
+		return true;
+	}
+
+	// Send the data:
+	return SendRaw(a_Data, a_Length);
 }
 
 
@@ -138,6 +150,14 @@ bool cTCPLinkImpl::Send(const void * a_Data, size_t a_Length)
 
 void cTCPLinkImpl::Shutdown(void)
 {
+	// If running in TLS mode, notify the TLS layer:
+	if (m_TlsContext != nullptr)
+	{
+		m_TlsContext->NotifyClose();
+		m_TlsContext->ResetSelf();
+		m_TlsContext.reset();
+	}
+
 	// If there's no outgoing data, shutdown the socket directly:
 	if (evbuffer_get_length(bufferevent_get_output(m_BufferEvent)) == 0)
 	{
@@ -155,6 +175,14 @@ void cTCPLinkImpl::Shutdown(void)
 
 void cTCPLinkImpl::Close(void)
 {
+	// If running in TLS mode, notify the TLS layer:
+	if (m_TlsContext != nullptr)
+	{
+		m_TlsContext->NotifyClose();
+		m_TlsContext->ResetSelf();
+		m_TlsContext.reset();
+	}
+
 	// Disable all events on the socket, but keep it alive:
 	bufferevent_disable(m_BufferEvent, EV_READ | EV_WRITE);
 	if (m_Server == nullptr)
@@ -173,18 +201,98 @@ void cTCPLinkImpl::Close(void)
 
 
 
+AString cTCPLinkImpl::StartTLSClient(
+	cX509CertPtr a_OwnCert,
+	cCryptoKeyPtr a_OwnPrivKey
+)
+{
+	// Check preconditions:
+	if (m_TlsContext != nullptr)
+	{
+		return "TLS is already active on this link";
+	}
+	if (
+		((a_OwnCert == nullptr) && (a_OwnPrivKey != nullptr)) ||
+		((a_OwnCert != nullptr) && (a_OwnPrivKey != nullptr))
+	)
+	{
+		return "Either provide both the certificate and private key, or neither";
+	}
+
+	// Create the TLS context:
+	m_TlsContext.reset(new cLinkTlsContext(*this));
+	m_TlsContext->Initialize(true);
+	if (a_OwnCert != nullptr)
+	{
+		m_TlsContext->SetOwnCert(a_OwnCert, a_OwnPrivKey);
+	}
+	m_TlsContext->SetSelf(cLinkTlsContextWPtr(m_TlsContext));
+
+	// Start the handshake:
+	m_TlsContext->Handshake();
+	return "";
+}
+
+
+
+
+
+AString cTCPLinkImpl::StartTLSServer(
+	cX509CertPtr a_OwnCert,
+	cCryptoKeyPtr a_OwnPrivKey,
+	const AString & a_StartTLSData
+)
+{
+	// Check preconditions:
+	if (m_TlsContext != nullptr)
+	{
+		return "TLS is already active on this link";
+	}
+	if ((a_OwnCert == nullptr)  || (a_OwnPrivKey == nullptr))
+	{
+		return "Provide the server certificate and private key";
+	}
+
+	// Create the TLS context:
+	m_TlsContext.reset(new cLinkTlsContext(*this));
+	m_TlsContext->Initialize(false);
+	m_TlsContext->SetOwnCert(a_OwnCert, a_OwnPrivKey);
+	m_TlsContext->SetSelf(cLinkTlsContextWPtr(m_TlsContext));
+
+	// Push the initial data:
+	m_TlsContext->StoreReceivedData(a_StartTLSData.data(), a_StartTLSData.size());
+
+	// Start the handshake:
+	m_TlsContext->Handshake();
+	return "";
+}
+
+
+
+
+
 void cTCPLinkImpl::ReadCallback(bufferevent * a_BufferEvent, void * a_Self)
 {
 	ASSERT(a_Self != nullptr);
 	cTCPLinkImpl * Self = static_cast<cTCPLinkImpl *>(a_Self);
+	ASSERT(Self->m_BufferEvent == a_BufferEvent);
 	ASSERT(Self->m_Callbacks != nullptr);
 
 	// Read all the incoming data, in 1024-byte chunks:
 	char data[1024];
 	size_t length;
+	auto tlsContext = Self->m_TlsContext;
 	while ((length = bufferevent_read(a_BufferEvent, data, sizeof(data))) > 0)
 	{
-		Self->m_Callbacks->OnReceivedData(data, length);
+		if (tlsContext != nullptr)
+		{
+			ASSERT(tlsContext->IsLink(Self));
+			tlsContext->StoreReceivedData(data, length);
+		}
+		else
+		{
+			Self->ReceivedCleartextData(data, length);
+		}
 	}
 }
 
@@ -214,6 +322,11 @@ void cTCPLinkImpl::EventCallback(bufferevent * a_BufferEvent, short a_What, void
 {
 	ASSERT(a_Self != nullptr);
 	cTCPLinkImplPtr Self = static_cast<cTCPLinkImpl *>(a_Self)->m_Self;
+	if (Self == nullptr)
+	{
+		// The link has already been freed
+		return;
+	}
 
 	// If an error is reported, call the error callback:
 	if (a_What & BEV_EVENT_ERROR)
@@ -262,6 +375,13 @@ void cTCPLinkImpl::EventCallback(bufferevent * a_BufferEvent, short a_What, void
 	// If the connection has been closed, call the link callback and remove the connection:
 	if (a_What & BEV_EVENT_EOF)
 	{
+		// If running in TLS mode and there's data left in the TLS contect, report it:
+		auto tlsContext = Self->m_TlsContext;
+		if (tlsContext != nullptr)
+		{
+			tlsContext->FlushBuffers();
+		}
+
 		Self->m_Callbacks->OnRemoteClosed();
 		if (Self->m_Server != nullptr)
 		{
@@ -357,6 +477,178 @@ void cTCPLinkImpl::DoActualShutdown(void)
 
 
 
+bool cTCPLinkImpl::SendRaw(const void * a_Data, size_t a_Length)
+{
+	return (bufferevent_write(m_BufferEvent, a_Data, a_Length) == 0);
+}
+
+
+
+
+
+void cTCPLinkImpl::ReceivedCleartextData(const char * a_Data, size_t a_Length)
+{
+	ASSERT(m_Callbacks != nullptr);
+	m_Callbacks->OnReceivedData(a_Data, a_Length);
+}
+
+
+
+
+
+////////////////////////////////////////////////////////////////////////////////
+// cTCPLinkImpl::cLinkTlsContext:
+
+cTCPLinkImpl::cLinkTlsContext::cLinkTlsContext(cTCPLinkImpl & a_Link):
+	m_Link(a_Link)
+{
+}
+
+
+
+
+
+void cTCPLinkImpl::cLinkTlsContext::SetSelf(cLinkTlsContextWPtr a_Self)
+{
+	m_Self = a_Self;
+}
+
+
+
+
+
+void cTCPLinkImpl::cLinkTlsContext::ResetSelf(void)
+{
+	m_Self.reset();
+}
+
+
+
+
+
+void cTCPLinkImpl::cLinkTlsContext::StoreReceivedData(const char * a_Data, size_t a_NumBytes)
+{
+	// Hold self alive for the duration of this function
+	cLinkTlsContextPtr Self(m_Self);
+
+	m_EncryptedData.append(a_Data, a_NumBytes);
+
+	// Try to finish a pending handshake:
+	TryFinishHandshaking();
+
+	// Flush any cleartext data that can be "received":
+	FlushBuffers();
+}
+
+
+
+
+
+void cTCPLinkImpl::cLinkTlsContext::FlushBuffers(void)
+{
+	// Hold self alive for the duration of this function
+	cLinkTlsContextPtr Self(m_Self);
+
+	// If the handshake didn't complete yet, bail out:
+	if (!HasHandshaken())
+	{
+		return;
+	}
+
+	char Buffer[1024];
+	int NumBytes;
+	while ((NumBytes = ReadPlain(Buffer, sizeof(Buffer))) > 0)
+	{
+		m_Link.ReceivedCleartextData(Buffer, static_cast<size_t>(NumBytes));
+		if (m_Self.expired())
+		{
+			// The callback closed the SSL context, bail out
+			return;
+		}
+	}
+}
+
+
+
+
+
+void cTCPLinkImpl::cLinkTlsContext::TryFinishHandshaking(void)
+{
+	// Hold self alive for the duration of this function
+	cLinkTlsContextPtr Self(m_Self);
+
+	// If the handshake hasn't finished yet, retry:
+	if (!HasHandshaken())
+	{
+		Handshake();
+		// If the handshake succeeded, write all the queued plaintext data:
+		if (HasHandshaken())
+		{
+			m_Link.GetCallbacks()->OnTlsHandshakeCompleted();
+			WritePlain(m_CleartextData.data(), m_CleartextData.size());
+			m_CleartextData.clear();
+		}
+	}
+}
+
+
+
+
+
+void cTCPLinkImpl::cLinkTlsContext::Send(const void * a_Data, size_t a_Length)
+{
+	// Hold self alive for the duration of this function
+	cLinkTlsContextPtr Self(m_Self);
+
+	// If the handshake hasn't completed yet, queue the data:
+	if (!HasHandshaken())
+	{
+		m_CleartextData.append(reinterpret_cast<const char *>(a_Data), a_Length);
+		TryFinishHandshaking();
+		return;
+	}
+
+	// The connection is all set up, write the cleartext data into the SSL context:
+	WritePlain(a_Data, a_Length);
+	FlushBuffers();
+}
+
+
+
+
+
+int cTCPLinkImpl::cLinkTlsContext::ReceiveEncrypted(unsigned char * a_Buffer, size_t a_NumBytes)
+{
+	// Hold self alive for the duration of this function
+	cLinkTlsContextPtr Self(m_Self);
+
+	// If there's nothing queued in the buffer, report empty buffer:
+	if (m_EncryptedData.empty())
+	{
+		return POLARSSL_ERR_NET_WANT_READ;
+	}
+
+	// Copy as much data as possible to the provided buffer:
+	size_t BytesToCopy = std::min(a_NumBytes, m_EncryptedData.size());
+	memcpy(a_Buffer, m_EncryptedData.data(), BytesToCopy);
+	m_EncryptedData.erase(0, BytesToCopy);
+	return static_cast<int>(BytesToCopy);
+}
+
+
+
+
+
+int cTCPLinkImpl::cLinkTlsContext::SendEncrypted(const unsigned char * a_Buffer, size_t a_NumBytes)
+{
+	m_Link.SendRaw(a_Buffer, a_NumBytes);
+	return static_cast<int>(a_NumBytes);
+}
+
+
+
+
+
 ////////////////////////////////////////////////////////////////////////////////
 // cNetwork API:
 
diff --git a/src/OSSupport/TCPLinkImpl.h b/src/OSSupport/TCPLinkImpl.h
index bea21aeff..b54c1a2cc 100644
--- a/src/OSSupport/TCPLinkImpl.h
+++ b/src/OSSupport/TCPLinkImpl.h
@@ -14,6 +14,7 @@
 #include "Network.h"
 #include <event2/event.h>
 #include <event2/bufferevent.h>
+#include "../PolarSSL++/SslContext.h"
 
 
 
@@ -64,9 +65,73 @@ public:
 	virtual UInt16 GetRemotePort(void) const override { return m_RemotePort; }
 	virtual void Shutdown(void) override;
 	virtual void Close(void) override;
+	virtual AString StartTLSClient(
+		cX509CertPtr a_OwnCert,
+		cCryptoKeyPtr a_OwnPrivKey
+	) override;
+	virtual AString StartTLSServer(
+		cX509CertPtr a_OwnCert,
+		cCryptoKeyPtr a_OwnPrivKey,
+		const AString & a_StartTLSData
+	) override;
 
 protected:
 
+	// fwd:
+	class cLinkTlsContext;
+	typedef SharedPtr<cLinkTlsContext> cLinkTlsContextPtr;
+	typedef WeakPtr<cLinkTlsContext> cLinkTlsContextWPtr;
+
+	/** Wrapper around cSslContext that is used when this link is being encrypted by SSL. */
+	class cLinkTlsContext :
+		public cSslContext
+	{
+		cTCPLinkImpl & m_Link;
+
+		/** Buffer for storing the incoming encrypted data until it is requested by the SSL decryptor. */
+		AString m_EncryptedData;
+
+		/** Buffer for storing the outgoing cleartext data until the link has finished handshaking. */
+		AString m_CleartextData;
+
+		/** Shared ownership of self, so that this object can keep itself alive for as long as it needs. */
+		cLinkTlsContextWPtr m_Self;
+
+	public:
+		cLinkTlsContext(cTCPLinkImpl & a_Link);
+
+		/** Shares ownership of self, so that this object can keep itself alive for as long as it needs. */
+		void SetSelf(cLinkTlsContextWPtr a_Self);
+
+		/** Removes the self ownership so that we can detect the SSL closure. */
+		void ResetSelf(void);
+
+		/** Stores the specified block of data into the buffer of the data to be decrypted (incoming from remote).
+		Also flushes the SSL buffers by attempting to read any data through the SSL context. */
+		void StoreReceivedData(const char * a_Data, size_t a_NumBytes);
+
+		/** Tries to read any cleartext data available through the SSL, reports it in the link. */
+		void FlushBuffers(void);
+
+		/** Tries to finish handshaking the SSL. */
+		void TryFinishHandshaking(void);
+
+		/** Sends the specified cleartext data over the SSL to the remote peer.
+		If the handshake hasn't been completed yet, queues the data for sending when it completes. */
+		void Send(const void * a_Data, size_t a_Length);
+
+		// cSslContext overrides:
+		virtual int ReceiveEncrypted(unsigned char * a_Buffer, size_t a_NumBytes) override;
+		virtual int SendEncrypted(const unsigned char * a_Buffer, size_t a_NumBytes) override;
+
+		/** Returns true if the context's associated TCP link is the same link as a_Link. */
+		bool IsLink(cTCPLinkImpl * a_Link)
+		{
+			return (a_Link == &m_Link);
+		}
+	};
+
+
 	/** Callbacks to call when the connection is established.
 	May be NULL if not used. Only used for outgoing connections (cNetwork::Connect()). */
 	cNetwork::cConnectCallbacksPtr m_ConnectCallbacks;
@@ -99,6 +164,10 @@ protected:
 	data is sent to the OS TCP stack, the socket gets shut down. */
 	bool m_ShouldShutdown;
 
+	/** The SSL context used for encryption, if this link uses SSL.
+	If valid, the link uses encryption through this context. */
+	cLinkTlsContextPtr m_TlsContext;
+
 
 	/** Creates a new link to be queued to connect to a specified host:port.
 	Used for outgoing connections created using cNetwork::Connect().
@@ -127,6 +196,12 @@ protected:
 	/** Calls shutdown on the link and disables LibEvent writing.
 	Called after all data from LibEvent buffers is sent to the OS TCP stack and shutdown() has been called before. */
 	void DoActualShutdown(void);
+
+	/** Sends the data directly to the socket (without the optional TLS). */
+	bool SendRaw(const void * a_Data, size_t a_Length);
+
+	/** Called by the TLS when it has decoded a piece of incoming cleartext data from the socket. */
+	void ReceivedCleartextData(const char * a_Data, size_t a_Length);
 };