1 files changed, 75 insertions, 0 deletions

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