1 files changed, 201 insertions, 0 deletions

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);
+}
+
+
+
+