10 files changed, 349 insertions, 128 deletions

diff --git a/src/PolarSSL++/BufferedSslContext.cpp b/src/PolarSSL++/BufferedSslContext.cpp
index 885b30c68..9f7caeb8a 100644
--- a/src/PolarSSL++/BufferedSslContext.cpp
+++ b/src/PolarSSL++/BufferedSslContext.cpp
@@ -20,6 +20,37 @@ cBufferedSslContext::cBufferedSslContext(size_t a_BufferSize):
 
 
 
+size_t cBufferedSslContext::WriteIncoming(const void * a_Data, size_t a_NumBytes)
+{
+	size_t NumBytes = std::min(m_IncomingData.GetFreeSpace(), a_NumBytes);
+	if (NumBytes > 0)
+	{
+		m_IncomingData.Write(a_Data, NumBytes);
+		return NumBytes;
+	}
+	return 0;
+}
+
+
+
+
+
+size_t cBufferedSslContext::ReadOutgoing(void * a_Data, size_t a_DataMaxSize)
+{
+	size_t NumBytes = std::min(m_OutgoingData.GetReadableSpace(), a_DataMaxSize);
+	if (NumBytes > 0)
+	{
+		m_OutgoingData.ReadBuf(a_Data, NumBytes);
+		m_OutgoingData.CommitRead();
+		return NumBytes;
+	}
+	return 0;
+}
+
+
+
+
+
 int cBufferedSslContext::ReceiveEncrypted(unsigned char * a_Buffer, size_t a_NumBytes)
 {
 	// Called when PolarSSL wants to read encrypted data from the SSL peer
diff --git a/src/PolarSSL++/CMakeLists.txt b/src/PolarSSL++/CMakeLists.txt
index b0a592760..9a59cdb2c 100644
--- a/src/PolarSSL++/CMakeLists.txt
+++ b/src/PolarSSL++/CMakeLists.txt
@@ -1,4 +1,3 @@
-
 cmake_minimum_required (VERSION 2.6)
 project (MCServer)
 
@@ -11,8 +10,8 @@ set(SOURCES
 	BufferedSslContext.cpp
 	CallbackSslContext.cpp
 	CtrDrbgContext.cpp
+	CryptoKey.cpp
 	EntropyContext.cpp
-	PublicKey.cpp
 	RsaPrivateKey.cpp
 	Sha1Checksum.cpp
 	SslContext.cpp
@@ -26,8 +25,8 @@ set(HEADERS
 	BufferedSslContext.h
 	CallbackSslContext.h
 	CtrDrbgContext.h
+	CryptoKey.h
 	EntropyContext.h
-	PublicKey.h
 	RsaPrivateKey.h
 	SslContext.h
 	Sha1Checksum.h
diff --git a/src/PolarSSL++/CryptoKey.cpp b/src/PolarSSL++/CryptoKey.cpp
new file mode 100644
index 000000000..0763c387b
--- /dev/null
+++ b/src/PolarSSL++/CryptoKey.cpp
@@ -0,0 +1,149 @@
+
+// CryptoKey.cpp
+
+// Implements the cCryptoKey class representing a RSA public key in PolarSSL
+
+#include "Globals.h"
+#include "CryptoKey.h"
+
+
+
+
+
+cCryptoKey::cCryptoKey(void)
+{
+	pk_init(&m_Pk);
+	m_CtrDrbg.Initialize("rsa_pubkey", 10);
+}
+
+
+
+
+
+cCryptoKey::cCryptoKey(const AString & a_PublicKeyData)
+{
+	pk_init(&m_Pk);
+	m_CtrDrbg.Initialize("rsa_pubkey", 10);
+	int res = ParsePublic(a_PublicKeyData.data(), a_PublicKeyData.size());
+	if (res != 0)
+	{
+		LOGWARNING("Failed to parse public key: -0x%x", res);
+		ASSERT(!"Cannot parse PubKey");
+		return;
+	}
+}
+
+
+
+
+
+cCryptoKey::cCryptoKey(const AString & a_PrivateKeyData, const AString & a_Password)
+{
+	pk_init(&m_Pk);
+	m_CtrDrbg.Initialize("rsa_privkey", 11);
+	int res = ParsePrivate(a_PrivateKeyData.data(), a_PrivateKeyData.size(), a_Password);
+	if (res != 0)
+	{
+		LOGWARNING("Failed to parse private key: -0x%x", res);
+		ASSERT(!"Cannot parse PubKey");
+		return;
+	}
+}
+
+
+
+
+
+cCryptoKey::~cCryptoKey()
+{
+	pk_free(&m_Pk);
+}
+
+
+
+
+
+int cCryptoKey::Decrypt(const Byte * a_EncryptedData, size_t a_EncryptedLength, Byte * a_DecryptedData, size_t a_DecryptedMaxLength)
+{
+	ASSERT(IsValid());
+	
+	size_t DecryptedLen = a_DecryptedMaxLength;
+	int res = pk_decrypt(&m_Pk,
+		a_EncryptedData, a_EncryptedLength,
+		a_DecryptedData, &DecryptedLen, a_DecryptedMaxLength,
+		ctr_drbg_random, m_CtrDrbg.GetInternal()
+	);
+	if (res != 0)
+	{
+		return res;
+	}
+	return (int)DecryptedLen;
+}
+
+
+
+
+
+int cCryptoKey::Encrypt(const Byte * a_PlainData, size_t a_PlainLength, Byte * a_EncryptedData, size_t a_EncryptedMaxLength)
+{
+	ASSERT(IsValid());
+	
+	size_t EncryptedLength = a_EncryptedMaxLength;
+	int res = pk_encrypt(&m_Pk,
+		a_PlainData, a_PlainLength, a_EncryptedData, &EncryptedLength, a_EncryptedMaxLength,
+		ctr_drbg_random, m_CtrDrbg.GetInternal()
+	);
+	if (res != 0)
+	{
+		return res;
+	}
+	return (int)EncryptedLength;
+}
+
+
+
+
+
+
+int cCryptoKey::ParsePublic(const void * a_Data, size_t a_NumBytes)
+{
+	ASSERT(!IsValid());  // Cannot parse a second key
+	
+	return pk_parse_public_key(&m_Pk, (const unsigned char *)a_Data, a_NumBytes);
+}
+
+
+
+
+
+
+int cCryptoKey::ParsePrivate(const void * a_Data, size_t a_NumBytes, const AString & a_Password)
+{
+	ASSERT(!IsValid());  // Cannot parse a second key
+	
+	if (a_Password.empty())
+	{
+		return pk_parse_key(&m_Pk, (const unsigned char *)a_Data, a_NumBytes, NULL, 0);
+	}
+	else
+	{
+		return pk_parse_key(
+			&m_Pk,
+			(const unsigned char *)a_Data, a_NumBytes,
+			(const unsigned char *)a_Password.c_str(), a_Password.size()
+		);
+	}
+}
+
+
+
+
+
+bool cCryptoKey::IsValid(void) const
+{
+	return (pk_get_type(&m_Pk) != POLARSSL_PK_NONE);
+}
+
+
+
+
diff --git a/src/PolarSSL++/CryptoKey.h b/src/PolarSSL++/CryptoKey.h
new file mode 100644
index 000000000..9c298e501
--- /dev/null
+++ b/src/PolarSSL++/CryptoKey.h
@@ -0,0 +1,76 @@
+
+// CryptoKey.h
+
+// Declares the cCryptoKey class representing a RSA public key in PolarSSL
+
+
+
+
+
+#pragma once
+
+#include "CtrDrbgContext.h"
+#include "polarssl/pk.h"
+
+
+
+
+
+class cCryptoKey
+{
+	friend class cSslContext;
+	
+public:
+	/** Constructs an empty key instance. Before use, it needs to be filled by ParsePublic() or ParsePrivate() */
+	cCryptoKey(void);
+	
+	/** Constructs the public key out of the DER- or PEM-encoded pubkey data */
+	cCryptoKey(const AString & a_PublicKeyData);
+	
+	/** Constructs the private key out of the DER- or PEM-encoded privkey data, with the specified password.
+	If a_Password is empty, no password is assumed. */
+	cCryptoKey(const AString & a_PrivateKeyData, const AString & a_Password);
+	
+	~cCryptoKey();
+	
+	/** Decrypts the data using the stored public key
+	Both a_EncryptedData and a_DecryptedData must be at least <KeySizeBytes> bytes large.
+	Returns the number of bytes decrypted, or negative number for error. */
+	int Decrypt(const Byte * a_EncryptedData, size_t a_EncryptedLength, Byte * a_DecryptedData, size_t a_DecryptedMaxLength);
+	
+	/** Encrypts the data using the stored public key
+	Both a_EncryptedData and a_DecryptedData must be at least <KeySizeBytes> bytes large.
+	Returns the number of bytes decrypted, or negative number for error. */
+	int Encrypt(const Byte * a_PlainData, size_t a_PlainLength, Byte * a_EncryptedData, size_t a_EncryptedMaxLength);
+	
+	/** Parses the specified data into a public key representation.
+	The key can be DER- or PEM-encoded.
+	Returns 0 on success, PolarSSL error code on failure. */
+	int ParsePublic(const void * a_Data, size_t a_NumBytes);
+
+	/** Parses the specified data into a private key representation.
+	If a_Password is empty, no password is assumed.
+	The key can be DER- or PEM-encoded.
+	Returns 0 on success, PolarSSL error code on failure. */
+	int ParsePrivate(const void * a_Data, size_t a_NumBytes, const AString & a_Password);
+	
+	/** Returns true if the contained key is valid. */
+	bool IsValid(void) const;
+
+protected:
+	/** The PolarSSL representation of the key data */
+	pk_context m_Pk;
+	
+	/** The random generator used in encryption and decryption */
+	cCtrDrbgContext m_CtrDrbg;
+	
+	
+	/** Returns the internal context ptr. Only use in PolarSSL API calls. */
+	pk_context * GetInternal(void) { return &m_Pk; }
+} ;
+
+typedef SharedPtr<cCryptoKey> cCryptoKeyPtr;
+
+
+
+
diff --git a/src/PolarSSL++/CtrDrbgContext.h b/src/PolarSSL++/CtrDrbgContext.h
index 65e9a2374..230db8753 100644
--- a/src/PolarSSL++/CtrDrbgContext.h
+++ b/src/PolarSSL++/CtrDrbgContext.h
@@ -26,7 +26,7 @@ class cCtrDrbgContext
 {
 	friend class cSslContext;
 	friend class cRsaPrivateKey;
-	friend class cPublicKey;
+	friend class cCryptoKey;
 	
 public:
 	/** Constructs the context with a new entropy context. */
diff --git a/src/PolarSSL++/PublicKey.cpp b/src/PolarSSL++/PublicKey.cpp
deleted file mode 100644
index 49794a0c8..000000000
--- a/src/PolarSSL++/PublicKey.cpp
+++ /dev/null
@@ -1,73 +0,0 @@
-
-// PublicKey.cpp
-
-// Implements the cPublicKey class representing a RSA public key in PolarSSL
-
-#include "Globals.h"
-#include "PublicKey.h"
-
-
-
-
-
-cPublicKey::cPublicKey(const AString & a_PublicKeyDER)
-{
-	pk_init(&m_Pk);
-	if (pk_parse_public_key(&m_Pk, (const Byte *)a_PublicKeyDER.data(), a_PublicKeyDER.size()) != 0)
-	{
-		ASSERT(!"Cannot parse PubKey");
-		return;
-	}
-	m_CtrDrbg.Initialize("rsa_pubkey", 10);
-}
-
-
-
-
-
-cPublicKey::~cPublicKey()
-{
-	pk_free(&m_Pk);
-}
-
-
-
-
-
-int cPublicKey::Decrypt(const Byte * a_EncryptedData, size_t a_EncryptedLength, Byte * a_DecryptedData, size_t a_DecryptedMaxLength)
-{
-	size_t DecryptedLen = a_DecryptedMaxLength;
-	int res = pk_decrypt(&m_Pk,
-		a_EncryptedData, a_EncryptedLength,
-		a_DecryptedData, &DecryptedLen, a_DecryptedMaxLength,
-		ctr_drbg_random, m_CtrDrbg.GetInternal()
-	);
-	if (res != 0)
-	{
-		return res;
-	}
-	return (int)DecryptedLen;
-}
-
-
-
-
-
-int cPublicKey::Encrypt(const Byte * a_PlainData, size_t a_PlainLength, Byte * a_EncryptedData, size_t a_EncryptedMaxLength)
-{
-	size_t EncryptedLength = a_EncryptedMaxLength;
-	int res = pk_encrypt(&m_Pk,
-		a_PlainData, a_PlainLength, a_EncryptedData, &EncryptedLength, a_EncryptedMaxLength,
-		ctr_drbg_random, m_CtrDrbg.GetInternal()
-	);
-	if (res != 0)
-	{
-		return res;
-	}
-	return (int)EncryptedLength;
-}
-
-
-
-
-
diff --git a/src/PolarSSL++/PublicKey.h b/src/PolarSSL++/PublicKey.h
deleted file mode 100644
index 5a0a57147..000000000
--- a/src/PolarSSL++/PublicKey.h
+++ /dev/null
@@ -1,48 +0,0 @@
-
-// PublicKey.h
-
-// Declares the cPublicKey class representing a RSA public key in PolarSSL
-
-
-
-
-
-#pragma once
-
-#include "CtrDrbgContext.h"
-#include "polarssl/pk.h"
-
-
-
-
-
-class cPublicKey
-{
-public:
-	/** Constructs the public key out of the DER-encoded pubkey data */
-	cPublicKey(const AString & a_PublicKeyDER);
-	
-	~cPublicKey();
-	
-	/** Decrypts the data using the stored public key
-	Both a_EncryptedData and a_DecryptedData must be at least <KeySizeBytes> bytes large.
-	Returns the number of bytes decrypted, or negative number for error. */
-	int Decrypt(const Byte * a_EncryptedData, size_t a_EncryptedLength, Byte * a_DecryptedData, size_t a_DecryptedMaxLength);
-	
-	/** Encrypts the data using the stored public key
-	Both a_EncryptedData and a_DecryptedData must be at least <KeySizeBytes> bytes large.
-	Returns the number of bytes decrypted, or negative number for error. */
-	int Encrypt(const Byte * a_PlainData, size_t a_PlainLength, Byte * a_EncryptedData, size_t a_EncryptedMaxLength);
-
-protected:
-	/** The public key PolarSSL representation */
-	pk_context m_Pk;
-	
-	/** The random generator used in encryption and decryption */
-	cCtrDrbgContext m_CtrDrbg;
-} ;
-
-
-
-
-
diff --git a/src/PolarSSL++/RsaPrivateKey.h b/src/PolarSSL++/RsaPrivateKey.h
index ffacde11b..4d03f27df 100644
--- a/src/PolarSSL++/RsaPrivateKey.h
+++ b/src/PolarSSL++/RsaPrivateKey.h
@@ -19,6 +19,8 @@
 /** Encapsulates an RSA private key used in PKI cryptography */
 class cRsaPrivateKey
 {
+	friend class cSslContext;
+	
 public:
 	/** Creates a new empty object, the key is not assigned */
 	cRsaPrivateKey(void);
@@ -51,8 +53,14 @@ protected:
 
 	/** The random generator used for generating the key and encryption / decryption */
 	cCtrDrbgContext m_CtrDrbg;
+	
+	
+	/** Returns the internal context ptr. Only use in PolarSSL API calls. */
+	rsa_context * GetInternal(void) { return &m_Rsa; }
 } ;
 
+typedef SharedPtr<cRsaPrivateKey> cRsaPrivateKeyPtr;
+
 
 
 
diff --git a/src/PolarSSL++/SslContext.cpp b/src/PolarSSL++/SslContext.cpp
index 1994cf844..bc397b655 100644
--- a/src/PolarSSL++/SslContext.cpp
+++ b/src/PolarSSL++/SslContext.cpp
@@ -40,7 +40,7 @@ int cSslContext::Initialize(bool a_IsClient, const SharedPtr<cCtrDrbgContext> &
 	if (m_IsValid)
 	{
 		LOGWARNING("SSL: Double initialization is not supported.");
-		return POLARSSL_ERR_SSL_MALLOC_FAILED;  // There is no return value well-suited for this, reuse this one.
+		return POLARSSL_ERR_SSL_BAD_INPUT_DATA;  // There is no return value well-suited for this, reuse this one.
 	}
 	
 	// Set the CtrDrbg context, create a new one if needed:
@@ -70,6 +70,18 @@ int cSslContext::Initialize(bool a_IsClient, const SharedPtr<cCtrDrbgContext> &
 		ssl_set_dbg(&m_Ssl, &SSLDebugMessage, this);
 		ssl_set_verify(&m_Ssl, &SSLVerifyCert, this);
 		*/
+		
+		/*
+		// Set ciphersuite to the easiest one to decode, so that the connection can be wireshark-decoded:
+		static const int CipherSuites[] =
+		{
+			TLS_RSA_WITH_RC4_128_MD5,
+			TLS_RSA_WITH_RC4_128_SHA,
+			TLS_RSA_WITH_AES_128_CBC_SHA,
+			0,  // Must be 0-terminated!
+		};
+		ssl_set_ciphersuites(&m_Ssl, CipherSuites);
+		*/
 	#endif
 	
 	m_IsValid = true;
@@ -80,8 +92,56 @@ int cSslContext::Initialize(bool a_IsClient, const SharedPtr<cCtrDrbgContext> &
 
 
 
+void cSslContext::SetOwnCert(const cX509CertPtr & a_OwnCert, const cRsaPrivateKeyPtr & a_OwnCertPrivKey)
+{
+	ASSERT(m_IsValid);  // Call Initialize() first
+	
+	// Check that both the cert and the key is valid:
+	if ((a_OwnCert.get() == NULL) || (a_OwnCertPrivKey.get() == NULL))
+	{
+		LOGWARNING("SSL: Own certificate is not valid, skipping the set.");
+		return;
+	}
+	
+	// Make sure we have the cert stored for later, PolarSSL only uses the cert later on
+	m_OwnCert = a_OwnCert;
+	m_OwnCertPrivKey = a_OwnCertPrivKey;
+	
+	// Set into the context:
+	ssl_set_own_cert_rsa(&m_Ssl, m_OwnCert->GetInternal(), m_OwnCertPrivKey->GetInternal());
+}
+
+
+
+
+
+void cSslContext::SetOwnCert(const cX509CertPtr & a_OwnCert, const cCryptoKeyPtr & a_OwnCertPrivKey)
+{
+	ASSERT(m_IsValid);  // Call Initialize() first
+	
+	// Check that both the cert and the key is valid:
+	if ((a_OwnCert.get() == NULL) || (a_OwnCertPrivKey.get() == NULL))
+	{
+		LOGWARNING("SSL: Own certificate is not valid, skipping the set.");
+		return;
+	}
+	
+	// Make sure we have the cert stored for later, PolarSSL only uses the cert later on
+	m_OwnCert = a_OwnCert;
+	m_OwnCertPrivKey2 = a_OwnCertPrivKey;
+	
+	// Set into the context:
+	ssl_set_own_cert(&m_Ssl, m_OwnCert->GetInternal(), m_OwnCertPrivKey2->GetInternal());
+}
+
+
+
+
+
 void cSslContext::SetCACerts(const cX509CertPtr & a_CACert, const AString & a_ExpectedPeerName)
 {
+	ASSERT(m_IsValid);  // Call Initialize() first
+	
 	// Store the data in our internal buffers, to avoid losing the pointers later on
 	// PolarSSL will need these after this call returns, and the caller may move / delete the data before that:
 	m_ExpectedPeerName = a_ExpectedPeerName;
diff --git a/src/PolarSSL++/SslContext.h b/src/PolarSSL++/SslContext.h
index 85add5f8b..6b4f2c1e7 100644
--- a/src/PolarSSL++/SslContext.h
+++ b/src/PolarSSL++/SslContext.h
@@ -11,6 +11,8 @@
 
 #include "polarssl/ssl.h"
 #include "../ByteBuffer.h"
+#include "CryptoKey.h"
+#include "RsaPrivateKey.h"
 #include "X509Cert.h"
 
 
@@ -38,7 +40,7 @@ public:
 	/** Creates a new uninitialized context */
 	cSslContext(void);
 	
-	~cSslContext();
+	virtual ~cSslContext();
 	
 	/** Initializes the context for use as a server or client.
 	Returns 0 on success, PolarSSL error on failure. */
@@ -47,7 +49,15 @@ public:
 	/** Returns true if the object has been initialized properly. */
 	bool IsValid(void) const { return m_IsValid; }
 	
-	/** Sets a cert chain as the trusted cert store for this context.
+	/** Sets the certificate to use as our own. Must be used when representing a server, optional when client.
+	Must be called after Initialize(). */
+	void SetOwnCert(const cX509CertPtr & a_OwnCert, const cRsaPrivateKeyPtr & a_OwnCertPrivKey);
+	
+	/** Sets the certificate to use as our own. Must be used when representing a server, optional when client.
+	Must be called after Initialize(). */
+	void SetOwnCert(const cX509CertPtr & a_OwnCert, const cCryptoKeyPtr & a_OwnCertPrivKey);
+	
+	/** Sets a cert chain as the trusted cert store for this context. Must be called after Initialize().
 	Calling this will switch the context into strict cert verification mode.
 	a_ExpectedPeerName is the CommonName that we expect the SSL peer to have in its cert,
 	if it is different, the verification will fail. An empty string will disable the CN check. */
@@ -93,6 +103,15 @@ protected:
 	/** The SSL context that PolarSSL uses. */
 	ssl_context m_Ssl;
 	
+	/** The certificate that we present to the peer. */
+	cX509CertPtr m_OwnCert;
+	
+	/** Private key for m_OwnCert, if initialized from a cRsaPrivateKey. */
+	cRsaPrivateKeyPtr m_OwnCertPrivKey;
+	
+	/** Private key for m_OwnCert, if initialized from a cCryptoKey. */
+	cCryptoKeyPtr m_OwnCertPrivKey2;
+	
 	/** True if the SSL handshake has been completed. */
 	bool m_HasHandshaken;