3 files changed, 188 insertions, 7 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())
 	{
diff --git a/src/Bindings/LuaTCPLink.h b/src/Bindings/LuaTCPLink.h
index 9536c052b..4e0d7dcec 100644
--- a/src/Bindings/LuaTCPLink.h
+++ b/src/Bindings/LuaTCPLink.h
@@ -74,7 +74,27 @@ public:
 		const AString & a_OwnPrivKeyPassword
 	);
 
+	/** 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. */
+	AString StartTLSServer(
+		const AString & a_OwnCertData,
+		const AString & a_OwnPrivKeyData,
+		const AString & a_OwnPrivKeyPassword,
+		const AString & a_StartTLSData
+	);
+
 protected:
+	// fwd:
+	class cLinkSslContext;
+	typedef SharedPtr<cLinkSslContext> cLinkSslContextPtr;
+	typedef WeakPtr<cLinkSslContext> cLinkSslContextWPtr;
+
 	/** Wrapper around cSslContext that is used when this link is being encrypted by SSL. */
 	class cLinkSslContext :
 		public cSslContext
@@ -87,9 +107,18 @@ protected:
 		/** 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. */
+		cLinkSslContextWPtr m_Self;
+
 	public:
 		cLinkSslContext(cLuaTCPLink & a_Link);
 
+		/** Shares ownership of self, so that this object can keep itself alive for as long as it needs. */
+		void SetSelf(cLinkSslContextWPtr & 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);
@@ -125,7 +154,7 @@ protected:
 
 	/** The SSL context used for encryption, if this link uses SSL.
 	If valid, the link uses encryption through this context. */
-	UniquePtr<cLinkSslContext> m_SslContext;
+	cLinkSslContextPtr m_SslContext;
 
 
 	/** Common code called when the link is considered as terminated.
diff --git a/src/Bindings/ManualBindings_Network.cpp b/src/Bindings/ManualBindings_Network.cpp
index 4a6b7bc0e..30a34815c 100644
--- a/src/Bindings/ManualBindings_Network.cpp
+++ b/src/Bindings/ManualBindings_Network.cpp
@@ -502,6 +502,52 @@ static int tolua_cTCPLink_StartTLSClient(lua_State * L)
 
 
 
+/** Binds cLuaTCPLink::StartTLSServer */
+static int tolua_cTCPLink_StartTLSServer(lua_State * L)
+{
+	// Function signature:
+	// LinkInstance:StartTLSServer(OwnCert, OwnPrivKey, OwnPrivKeyPassword, StartTLSData) -> [true] or [nil, ErrMsg]
+
+	cLuaState S(L);
+	if (
+		!S.CheckParamUserType(1, "cTCPLink") ||
+		!S.CheckParamString(2, 4) ||
+		// Param 5 is optional, don't check
+		!S.CheckParamEnd(6)
+	)
+	{
+		return 0;
+	}
+	
+	// Get the link:
+	cLuaTCPLink * Link;
+	if (lua_isnil(L, 1))
+	{
+		LOGWARNING("cTCPLink:StartTLSServer(): invalid link object. Stack trace:");
+		S.LogStackTrace();
+		return 0;
+	}
+	Link = *static_cast<cLuaTCPLink **>(lua_touserdata(L, 1));
+
+	// Read the params:
+	AString OwnCert, OwnPrivKey, OwnPrivKeyPassword, StartTLSData;
+	S.GetStackValues(2, OwnCert, OwnPrivKey, OwnPrivKeyPassword, StartTLSData);
+
+	// Start the TLS handshake:
+	AString res = Link->StartTLSServer(OwnCert, OwnPrivKey, OwnPrivKeyPassword, StartTLSData);
+	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):
 
@@ -616,6 +662,7 @@ void ManualBindings::BindNetwork(lua_State * tolua_S)
 		tolua_function(tolua_S, "Send",           tolua_cTCPLink_Send);
 		tolua_function(tolua_S, "Shutdown",       tolua_cTCPLink_Shutdown);
 		tolua_function(tolua_S, "StartTLSClient", tolua_cTCPLink_StartTLSClient);
+		tolua_function(tolua_S, "StartTLSServer", tolua_cTCPLink_StartTLSServer);
 	tolua_endmodule(tolua_S);
 
 	tolua_beginmodule(tolua_S, "cServerHandle");