9 files changed, 335 insertions, 7 deletions

diff --git a/src/Bindings/LuaState.cpp b/src/Bindings/LuaState.cpp
index 73b114599..81770058c 100644
--- a/src/Bindings/LuaState.cpp
+++ b/src/Bindings/LuaState.cpp
@@ -343,6 +343,18 @@ bool cLuaState::PushFunction(const cTableRef & a_TableRef)
 
 
 
+void cLuaState::PushNil(void)
+{
+	ASSERT(IsValid());
+
+	lua_pushnil(m_LuaState);
+	m_NumCurrentFunctionArgs += 1;
+}
+
+
+
+
+
 void cLuaState::Push(const AString & a_String)
 {
 	ASSERT(IsValid());
diff --git a/src/Bindings/LuaState.h b/src/Bindings/LuaState.h
index f68b022ea..7fc3197eb 100644
--- a/src/Bindings/LuaState.h
+++ b/src/Bindings/LuaState.h
@@ -184,6 +184,8 @@ public:
 	/** Returns true if a_FunctionName is a valid Lua function that can be called */
 	bool HasFunction(const char * a_FunctionName);
 	
+	void PushNil(void);
+
 	// Push a const value onto the stack (keep alpha-sorted):
 	void Push(const AString & a_String);
 	void Push(const AStringVector & a_Vector);
diff --git a/src/Bindings/LuaTCPLink.cpp b/src/Bindings/LuaTCPLink.cpp
index 6b8395806..7e2c10e13 100644
--- a/src/Bindings/LuaTCPLink.cpp
+++ b/src/Bindings/LuaTCPLink.cpp
@@ -64,6 +64,13 @@ cLuaTCPLink::~cLuaTCPLink()
 
 bool cLuaTCPLink::Send(const AString & a_Data)
 {
+	// If running in SSL mode, push the data into the SSL context instead:
+	if (m_SslContext != nullptr)
+	{
+		m_SslContext->Send(a_Data);
+		return true;
+	}
+
 	// Safely grab a copy of the link:
 	cTCPLinkPtr Link = m_Link;
 	if (Link == nullptr)
@@ -179,6 +186,58 @@ void cLuaTCPLink::Close(void)
 
 
 
+AString cLuaTCPLink::StartTLSClient(
+	const AString & a_OwnCertData,
+	const AString & a_OwnPrivKeyData,
+	const AString & a_OwnPrivKeyPassword
+)
+{
+	// Check preconditions:
+	if (m_SslContext != nullptr)
+	{
+		return "TLS is already active on this link";
+	}
+	if (
+		(a_OwnCertData.empty()  && !a_OwnPrivKeyData.empty()) ||
+		(!a_OwnCertData.empty() && a_OwnPrivKeyData.empty())
+	)
+	{
+		return "Either provide both the certificate and private key, or neither";
+	}
+
+	// Create the SSL context:
+	m_SslContext = std::make_unique<cLinkSslContext>(*this);
+	m_SslContext->Initialize(true);
+
+	// Create the peer cert, if required:
+	if (!a_OwnCertData.empty() && !a_OwnPrivKeyData.empty())
+	{
+		auto OwnCert = std::make_shared<cX509Cert>();
+		int res = OwnCert->Parse(a_OwnCertData.data(), a_OwnCertData.size());
+		if (res != 0)
+		{
+			m_SslContext.reset();
+			return Printf("Cannot parse peer certificate: -0x%x", res);
+		}
+		auto OwnPrivKey = std::make_shared<cCryptoKey>();
+		res = OwnPrivKey->ParsePrivate(a_OwnPrivKeyData.data(), a_OwnPrivKeyData.size(), a_OwnPrivKeyPassword);
+		if (res != 0)
+		{
+			m_SslContext.reset();
+			return Printf("Cannot parse peer private key: -0x%x", res);
+		}
+		m_SslContext->SetOwnCert(OwnCert, OwnPrivKey);
+	}
+
+	// Start the handshake:
+	m_SslContext->Handshake();
+	return "";
+}
+
+
+
+
+
 void cLuaTCPLink::Terminated(void)
 {
 	// Disable the callbacks:
@@ -207,6 +266,26 @@ void cLuaTCPLink::Terminated(void)
 
 
 
+void cLuaTCPLink::ReceivedCleartextData(const char * a_Data, size_t a_NumBytes)
+{
+	// Check if we're still valid:
+	if (!m_Callbacks.IsValid())
+	{
+		return;
+	}
+
+	// Call the callback:
+	cPluginLua::cOperation Op(m_Plugin);
+	if (!Op().Call(cLuaState::cTableRef(m_Callbacks, "OnReceivedData"), this, AString(a_Data, a_NumBytes)))
+	{
+		LOGINFO("cTCPLink OnReceivedData callback failed in plugin %s.", m_Plugin.GetName().c_str());
+	}
+}
+
+
+
+
+
 void cLuaTCPLink::OnConnected(cTCPLink & a_Link)
 {
 	// Check if we're still valid:
@@ -269,6 +348,13 @@ void cLuaTCPLink::OnReceivedData(const char * a_Data, size_t a_Length)
 		return;
 	}
 
+	// If we're running in SSL mode, put the data into the SSL decryptor:
+	if (m_SslContext != nullptr)
+	{
+		m_SslContext->StoreReceivedData(a_Data, a_Length);
+		return;
+	}
+
 	// Call the callback:
 	cPluginLua::cOperation Op(m_Plugin);
 	if (!Op().Call(cLuaState::cTableRef(m_Callbacks, "OnReceivedData"), this, AString(a_Data, a_Length)))
@@ -302,3 +388,118 @@ void cLuaTCPLink::OnRemoteClosed(void)
 
 
 
+
+////////////////////////////////////////////////////////////////////////////////
+// cLuaTCPLink::cLinkSslContext:
+
+cLuaTCPLink::cLinkSslContext::cLinkSslContext(cLuaTCPLink & a_Link):
+	m_Link(a_Link)
+{
+}
+
+
+
+
+
+void cLuaTCPLink::cLinkSslContext::StoreReceivedData(const char * a_Data, size_t a_NumBytes)
+{
+	m_EncryptedData.append(a_Data, a_NumBytes);
+
+	// Try to finish a pending handshake:
+	TryFinishHandshaking();
+
+	// Flush any cleartext data that can be "received":
+	FlushBuffers();
+}
+
+
+
+
+
+void cLuaTCPLink::cLinkSslContext::FlushBuffers(void)
+{
+	// 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));
+	}
+}
+
+
+
+
+
+void cLuaTCPLink::cLinkSslContext::TryFinishHandshaking(void)
+{
+	// If the handshake hasn't finished yet, retry:
+	if (!HasHandshaken())
+	{
+		Handshake();
+	}
+
+	// If the handshake succeeded, write all the queued plaintext data:
+	if (HasHandshaken())
+	{
+		WritePlain(m_CleartextData.data(), m_CleartextData.size());
+		m_CleartextData.clear();
+	}
+}
+
+
+
+
+
+void cLuaTCPLink::cLinkSslContext::Send(const AString & a_Data)
+{
+	// If the handshake hasn't completed yet, queue the data:
+	if (!HasHandshaken())
+	{
+		m_CleartextData.append(a_Data);
+		TryFinishHandshaking();
+		return;
+	}
+
+	// The connection is all set up, write the cleartext data into the SSL context:
+	WritePlain(a_Data.data(), a_Data.size());
+	FlushBuffers();
+}
+
+
+
+
+
+int cLuaTCPLink::cLinkSslContext::ReceiveEncrypted(unsigned char * a_Buffer, size_t a_NumBytes)
+{
+	// 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 cLuaTCPLink::cLinkSslContext::SendEncrypted(const unsigned char * a_Buffer, size_t a_NumBytes)
+{
+	m_Link.m_Link->Send(a_Buffer, a_NumBytes);
+	return static_cast<int>(a_NumBytes);
+}
+
+
+
+
diff --git a/src/Bindings/LuaTCPLink.h b/src/Bindings/LuaTCPLink.h
index f2af911ec..9536c052b 100644
--- a/src/Bindings/LuaTCPLink.h
+++ b/src/Bindings/LuaTCPLink.h
@@ -11,6 +11,7 @@
 
 #include "../OSSupport/Network.h"
 #include "PluginLua.h"
+#include "../PolarSSL++/SslContext.h"
 
 
 
@@ -62,7 +63,53 @@ public:
 	Sends the RST packet, queued outgoing and incoming data is lost. */
 	void Close(void);
 
+	/** 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. */
+	AString StartTLSClient(
+		const AString & a_OwnCertData,
+		const AString & a_OwnPrivKeyData,
+		const AString & a_OwnPrivKeyPassword
+	);
+
 protected:
+	/** Wrapper around cSslContext that is used when this link is being encrypted by SSL. */
+	class cLinkSslContext :
+		public cSslContext
+	{
+		cLuaTCPLink & 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;
+
+	public:
+		cLinkSslContext(cLuaTCPLink & a_Link);
+
+		/** 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 AString & a_Data);
+
+		// 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;
+	};
+
+
 	/** The plugin for which the link is created. */
 	cPluginLua & m_Plugin;
 
@@ -76,11 +123,19 @@ protected:
 	/** The server that is responsible for this link, if any. */
 	cLuaServerHandleWPtr m_Server;
 
+	/** The SSL context used for encryption, if this link uses SSL.
+	If valid, the link uses encryption through this context. */
+	UniquePtr<cLinkSslContext> m_SslContext;
+
 
 	/** Common code called when the link is considered as terminated.
 	Releases m_Link, m_Callbacks and this from m_Server, each when applicable. */
 	void Terminated(void);
 
+	/** Called by the SSL context when there's incoming data available in the cleartext.
+	Reports the data via the Lua callback function. */
+	void ReceivedCleartextData(const char * a_Data, size_t a_NumBytes);
+
 	// cNetwork::cConnectCallbacks overrides:
 	virtual void OnConnected(cTCPLink & a_Link) override;
 	virtual void OnError(int a_ErrorCode, const AString & a_ErrorMsg) override;
diff --git a/src/Bindings/ManualBindings_Network.cpp b/src/Bindings/ManualBindings_Network.cpp
index 902f687c8..ff0f3568c 100644
--- a/src/Bindings/ManualBindings_Network.cpp
+++ b/src/Bindings/ManualBindings_Network.cpp
@@ -389,6 +389,51 @@ static int tolua_cTCPLink_GetRemotePort(lua_State * L)
 
 
 
+/** Binds cLuaTCPLink::StartTLSClient */
+static int tolua_cTCPLink_StartTLSClient(lua_State * L)
+{
+	// Function signature:
+	// LinkInstance:StartTLSClient(OwnCert, OwnPrivKey, OwnPrivKeyPassword) -> [true] or [nil, ErrMsg]
+
+	cLuaState S(L);
+	if (
+		!S.CheckParamUserType(1, "cTCPLink") ||
+		!S.CheckParamString(2, 4) ||
+		!S.CheckParamEnd(5)
+	)
+	{
+		return 0;
+	}
+	
+	// Get the link:
+	cLuaTCPLink * Link;
+	if (lua_isnil(L, 1))
+	{
+		LOGWARNING("cTCPLink:StartTLSClient(): invalid link object. Stack trace:");
+		S.LogStackTrace();
+		return 0;
+	}
+	Link = *static_cast<cLuaTCPLink **>(lua_touserdata(L, 1));
+
+	// Read the params:
+	AString OwnCert, OwnPrivKey, OwnPrivKeyPassword;
+	S.GetStackValues(2, OwnCert, OwnPrivKey, OwnPrivKeyPassword);
+
+	// Start the TLS handshake:
+	AString res = Link->StartTLSClient(OwnCert, OwnPrivKey, OwnPrivKeyPassword);
+	if (!res.empty())
+	{
+		S.PushNil();
+		S.Push(Printf("Cannot start TLS on link to %s:%d: %s", Link->GetRemoteIP().c_str(), Link->GetRemotePort(), res.c_str()));
+		return 2;
+	}
+	return 1;
+}
+
+
+
+
+
 ////////////////////////////////////////////////////////////////////////////////
 // cServerHandle bindings (routed through cLuaServerHandle):
 
@@ -495,11 +540,12 @@ void ManualBindings::BindNetwork(lua_State * tolua_S)
 	tolua_endmodule(tolua_S);
 
 	tolua_beginmodule(tolua_S, "cTCPLink");
-		tolua_function(tolua_S, "Send",          tolua_cTCPLink_Send);
-		tolua_function(tolua_S, "GetLocalIP",    tolua_cTCPLink_GetLocalIP);
-		tolua_function(tolua_S, "GetLocalPort",  tolua_cTCPLink_GetLocalPort);
-		tolua_function(tolua_S, "GetRemoteIP",   tolua_cTCPLink_GetRemoteIP);
-		tolua_function(tolua_S, "GetRemotePort", tolua_cTCPLink_GetRemotePort);
+		tolua_function(tolua_S, "Send",           tolua_cTCPLink_Send);
+		tolua_function(tolua_S, "GetLocalIP",     tolua_cTCPLink_GetLocalIP);
+		tolua_function(tolua_S, "GetLocalPort",   tolua_cTCPLink_GetLocalPort);
+		tolua_function(tolua_S, "GetRemoteIP",    tolua_cTCPLink_GetRemoteIP);
+		tolua_function(tolua_S, "GetRemotePort",  tolua_cTCPLink_GetRemotePort);
+		tolua_function(tolua_S, "StartTLSClient", tolua_cTCPLink_StartTLSClient);
 	tolua_endmodule(tolua_S);
 
 	tolua_beginmodule(tolua_S, "cServerHandle");
diff --git a/src/Globals.h b/src/Globals.h
index 29eaac871..7c2ab38d8 100644
--- a/src/Globals.h
+++ b/src/Globals.h
@@ -379,9 +379,10 @@ void inline LOG(const char * a_Format, ...)
 	#define assert_test(x) ( !!(x) || (assert(!#x), exit(1), 0))
 #endif
 
-// Unified shared ptr from before C++11. Also no silly undercores.
+// Unified ptr types from before C++11. Also no silly undercores.
 #define SharedPtr std::shared_ptr
 #define WeakPtr std::weak_ptr
+#define UniquePtr std::unique_ptr
 
 
 
diff --git a/src/HTTPServer/SslHTTPConnection.cpp b/src/HTTPServer/SslHTTPConnection.cpp
index f09daac8f..f8dea0731 100644
--- a/src/HTTPServer/SslHTTPConnection.cpp
+++ b/src/HTTPServer/SslHTTPConnection.cpp
@@ -25,6 +25,15 @@ cSslHTTPConnection::cSslHTTPConnection(cHTTPServer & a_HTTPServer, const cX509Ce
 
 
 
+cSslHTTPConnection::~cSslHTTPConnection()
+{
+	m_Ssl.NotifyClose();
+}
+
+
+
+
+
 void cSslHTTPConnection::OnReceivedData(const char * a_Data, size_t a_Size)
 {
 	// Process the received data:
diff --git a/src/HTTPServer/SslHTTPConnection.h b/src/HTTPServer/SslHTTPConnection.h
index dc54b1eff..c461a3a24 100644
--- a/src/HTTPServer/SslHTTPConnection.h
+++ b/src/HTTPServer/SslHTTPConnection.h
@@ -25,6 +25,8 @@ public:
 	/** Creates a new connection on the specified server.
 	Sends the specified cert as the server certificate, uses the private key for decryption. */
 	cSslHTTPConnection(cHTTPServer & a_HTTPServer, const cX509CertPtr & a_Cert, const cCryptoKeyPtr & a_PrivateKey);
+
+	~cSslHTTPConnection();
 	
 protected:
 	cBufferedSslContext m_Ssl;
diff --git a/src/PolarSSL++/CryptoKey.cpp b/src/PolarSSL++/CryptoKey.cpp
index 7c4f021b3..9354ddf50 100644
--- a/src/PolarSSL++/CryptoKey.cpp
+++ b/src/PolarSSL++/CryptoKey.cpp
@@ -45,7 +45,7 @@ cCryptoKey::cCryptoKey(const AString & a_PrivateKeyData, const AString & a_Passw
 	if (res != 0)
 	{
 		LOGWARNING("Failed to parse private key: -0x%x", res);
-		ASSERT(!"Cannot parse PubKey");
+		ASSERT(!"Cannot parse PrivKey");
 		return;
 	}
 }