1 files changed, 111 insertions, 6 deletions

diff --git a/src/Bindings/LuaTCPLink.cpp b/src/Bindings/LuaTCPLink.cpp
index c533456ad..40371d6da 100644
--- a/src/Bindings/LuaTCPLink.cpp
+++ b/src/Bindings/LuaTCPLink.cpp
@@ -160,10 +160,14 @@ void cLuaTCPLink::Shutdown(void)
 	cTCPLinkPtr Link = m_Link;
 	if (Link != nullptr)
 	{
+		if (m_SslContext != nullptr)
+		{
+			m_SslContext->NotifyClose();
+			m_SslContext->ResetSelf();
+			m_SslContext.reset();
+		}
 		Link->Shutdown();
 	}
-
-	Terminated();
 }
 
 
@@ -176,6 +180,12 @@ void cLuaTCPLink::Close(void)
 	cTCPLinkPtr Link = m_Link;
 	if (Link != nullptr)
 	{
+		if (m_SslContext != nullptr)
+		{
+			m_SslContext->NotifyClose();
+			m_SslContext->ResetSelf();
+			m_SslContext.reset();
+		}
 		Link->Close();
 	}
 
@@ -228,6 +238,58 @@ AString cLuaTCPLink::StartTLSClient(
 		}
 		m_SslContext->SetOwnCert(OwnCert, OwnPrivKey);
 	}
+	m_SslContext->SetSelf(cLinkSslContextWPtr(m_SslContext));
+
+	// Start the handshake:
+	m_SslContext->Handshake();
+	return "";
+}
+
+
+
+
+
+AString cLuaTCPLink::StartTLSServer(
+	const AString & a_OwnCertData,
+	const AString & a_OwnPrivKeyData,
+	const AString & a_OwnPrivKeyPassword,
+	const AString & a_StartTLSData
+)
+{
+	// Check preconditions:
+	if (m_SslContext != nullptr)
+	{
+		return "TLS is already active on this link";
+	}
+	if (a_OwnCertData.empty()  || a_OwnPrivKeyData.empty())
+	{
+		return "Provide the server certificate and private key";
+	}
+
+	// Create the SSL context:
+	m_SslContext.reset(new cLinkSslContext(*this));
+	m_SslContext->Initialize(false);
+
+	// Create the peer cert:
+	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 server 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 server private key: -0x%x", res);
+	}
+	m_SslContext->SetOwnCert(OwnCert, OwnPrivKey);
+	m_SslContext->SetSelf(cLinkSslContextWPtr(m_SslContext));
+
+	// Push the initial data:
+	m_SslContext->StoreReceivedData(a_StartTLSData.data(), a_StartTLSData.size());
 
 	// Start the handshake:
 	m_SslContext->Handshake();
@@ -254,12 +316,17 @@ void cLuaTCPLink::Terminated(void)
 	}
 
 	// If the link is still open, close it:
-	cTCPLinkPtr Link = m_Link;
-	if (Link != nullptr)
 	{
-		Link->Close();
-		m_Link.reset();
+		cTCPLinkPtr Link = m_Link;
+		if (Link != nullptr)
+		{
+			Link->Close();
+			m_Link.reset();
+		}
 	}
+
+	// If the SSL context still exists, free it:
+	m_SslContext.reset();
 }
 
 
@@ -401,8 +468,29 @@ cLuaTCPLink::cLinkSslContext::cLinkSslContext(cLuaTCPLink & a_Link):
 
 
 
+void cLuaTCPLink::cLinkSslContext::SetSelf(cLinkSslContextWPtr & a_Self)
+{
+	m_Self = a_Self;
+}
+
+
+
+
+
+void cLuaTCPLink::cLinkSslContext::ResetSelf(void)
+{
+	m_Self.reset();
+}
+
+
+
+
+
 void cLuaTCPLink::cLinkSslContext::StoreReceivedData(const char * a_Data, size_t a_NumBytes)
 {
+	// Hold self alive for the duration of this function
+	cLinkSslContextPtr Self(m_Self);
+
 	m_EncryptedData.append(a_Data, a_NumBytes);
 
 	// Try to finish a pending handshake:
@@ -418,6 +506,9 @@ void cLuaTCPLink::cLinkSslContext::StoreReceivedData(const char * a_Data, size_t
 
 void cLuaTCPLink::cLinkSslContext::FlushBuffers(void)
 {
+	// Hold self alive for the duration of this function
+	cLinkSslContextPtr Self(m_Self);
+
 	// If the handshake didn't complete yet, bail out:
 	if (!HasHandshaken())
 	{
@@ -429,6 +520,11 @@ void cLuaTCPLink::cLinkSslContext::FlushBuffers(void)
 	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;
+		}
 	}
 }
 
@@ -438,6 +534,9 @@ void cLuaTCPLink::cLinkSslContext::FlushBuffers(void)
 
 void cLuaTCPLink::cLinkSslContext::TryFinishHandshaking(void)
 {
+	// Hold self alive for the duration of this function
+	cLinkSslContextPtr Self(m_Self);
+
 	// If the handshake hasn't finished yet, retry:
 	if (!HasHandshaken())
 	{
@@ -458,6 +557,9 @@ void cLuaTCPLink::cLinkSslContext::TryFinishHandshaking(void)
 
 void cLuaTCPLink::cLinkSslContext::Send(const AString & a_Data)
 {
+	// Hold self alive for the duration of this function
+	cLinkSslContextPtr Self(m_Self);
+
 	// If the handshake hasn't completed yet, queue the data:
 	if (!HasHandshaken())
 	{
@@ -477,6 +579,9 @@ void cLuaTCPLink::cLinkSslContext::Send(const AString & a_Data)
 
 int cLuaTCPLink::cLinkSslContext::ReceiveEncrypted(unsigned char * a_Buffer, size_t a_NumBytes)
 {
+	// Hold self alive for the duration of this function
+	cLinkSslContextPtr Self(m_Self);
+
 	// If there's nothing queued in the buffer, report empty buffer:
 	if (m_EncryptedData.empty())
 	{