From 788f18086e01c67c817c1c095e88fd4cbbd3d345 Mon Sep 17 00:00:00 2001
From: Fabio Alessandrelli <fabio.alessandrelli@gmail.com>
Date: Fri, 5 Jun 2020 21:28:16 +0200
Subject: [PATCH] CryptoKey supports public keys.

---
 core/crypto/crypto.cpp             | 25 ++++++++++----
 core/crypto/crypto.h               |  7 ++--
 modules/mbedtls/crypto_mbedtls.cpp | 53 +++++++++++++++++++++++++++---
 modules/mbedtls/crypto_mbedtls.h   | 10 ++++--
 4 files changed, 79 insertions(+), 16 deletions(-)

diff --git a/core/crypto/crypto.cpp b/core/crypto/crypto.cpp
index a9a7cabee96..a6b1fa9f034 100644
--- a/core/crypto/crypto.cpp
+++ b/core/crypto/crypto.cpp
@@ -45,8 +45,11 @@ CryptoKey *CryptoKey::create() {
 }
 
 void CryptoKey::_bind_methods() {
-	ClassDB::bind_method(D_METHOD("save", "path"), &CryptoKey::save);
-	ClassDB::bind_method(D_METHOD("load", "path"), &CryptoKey::load);
+	ClassDB::bind_method(D_METHOD("save", "path", "public_only"), &CryptoKey::save, DEFVAL(false));
+	ClassDB::bind_method(D_METHOD("load", "path", "public_only"), &CryptoKey::load, DEFVAL(false));
+	ClassDB::bind_method(D_METHOD("is_public_only"), &CryptoKey::is_public_only);
+	ClassDB::bind_method(D_METHOD("save_to_string", "public_only"), &CryptoKey::save_to_string, DEFVAL(false));
+	ClassDB::bind_method(D_METHOD("load_from_string", "string_key", "public_only"), &CryptoKey::load_from_string, DEFVAL(false));
 }
 
 X509Certificate *(*X509Certificate::_create)() = nullptr;
@@ -98,9 +101,14 @@ RES ResourceFormatLoaderCrypto::load(const String &p_path, const String &p_origi
 	} else if (el == "key") {
 		CryptoKey *key = CryptoKey::create();
 		if (key) {
-			key->load(p_path);
+			key->load(p_path, false);
 		}
 		return key;
+	} else if (el == "pub") {
+		CryptoKey *key = CryptoKey::create();
+		if (key)
+			key->load(p_path, true);
+		return key;
 	}
 	return nullptr;
 }
@@ -108,6 +116,7 @@ RES ResourceFormatLoaderCrypto::load(const String &p_path, const String &p_origi
 void ResourceFormatLoaderCrypto::get_recognized_extensions(List<String> *p_extensions) const {
 	p_extensions->push_back("crt");
 	p_extensions->push_back("key");
+	p_extensions->push_back("pub");
 }
 
 bool ResourceFormatLoaderCrypto::handles_type(const String &p_type) const {
@@ -118,7 +127,7 @@ String ResourceFormatLoaderCrypto::get_resource_type(const String &p_path) const
 	String el = p_path.get_extension().to_lower();
 	if (el == "crt") {
 		return "X509Certificate";
-	} else if (el == "key") {
+	} else if (el == "key" || el == "pub") {
 		return "CryptoKey";
 	}
 	return "";
@@ -131,7 +140,8 @@ Error ResourceFormatSaverCrypto::save(const String &p_path, const RES &p_resourc
 	if (cert.is_valid()) {
 		err = cert->save(p_path);
 	} else if (key.is_valid()) {
-		err = key->save(p_path);
+		String el = p_path.get_extension().to_lower();
+		err = key->save(p_path, el == "pub");
 	} else {
 		ERR_FAIL_V(ERR_INVALID_PARAMETER);
 	}
@@ -146,7 +156,10 @@ void ResourceFormatSaverCrypto::get_recognized_extensions(const RES &p_resource,
 		p_extensions->push_back("crt");
 	}
 	if (key) {
-		p_extensions->push_back("key");
+		if (!key->is_public_only()) {
+			p_extensions->push_back("key");
+		}
+		p_extensions->push_back("pub");
 	}
 }
 
diff --git a/core/crypto/crypto.h b/core/crypto/crypto.h
index 6cc5f461645..c1bc2024bbf 100644
--- a/core/crypto/crypto.h
+++ b/core/crypto/crypto.h
@@ -45,8 +45,11 @@ protected:
 
 public:
 	static CryptoKey *create();
-	virtual Error load(String p_path) = 0;
-	virtual Error save(String p_path) = 0;
+	virtual Error load(String p_path, bool p_public_only = false) = 0;
+	virtual Error save(String p_path, bool p_public_only = false) = 0;
+	virtual String save_to_string(bool p_public_only = false) = 0;
+	virtual Error load_from_string(String p_string_key, bool p_public_only = false) = 0;
+	virtual bool is_public_only() const = 0;
 };
 
 class X509Certificate : public Resource {
diff --git a/modules/mbedtls/crypto_mbedtls.cpp b/modules/mbedtls/crypto_mbedtls.cpp
index fbb3a39eed6..0d1a1f0138a 100644
--- a/modules/mbedtls/crypto_mbedtls.cpp
+++ b/modules/mbedtls/crypto_mbedtls.cpp
@@ -50,7 +50,7 @@ CryptoKey *CryptoKeyMbedTLS::create() {
 	return memnew(CryptoKeyMbedTLS);
 }
 
-Error CryptoKeyMbedTLS::load(String p_path) {
+Error CryptoKeyMbedTLS::load(String p_path, bool p_public_only) {
 	ERR_FAIL_COND_V_MSG(locks, ERR_ALREADY_IN_USE, "Key is in use");
 
 	PackedByteArray out;
@@ -63,22 +63,33 @@ Error CryptoKeyMbedTLS::load(String p_path) {
 	out.write[flen] = 0; // string terminator
 	memdelete(f);
 
-	int ret = mbedtls_pk_parse_key(&pkey, out.ptr(), out.size(), nullptr, 0);
+	int ret = 0;
+	if (p_public_only) {
+		ret = mbedtls_pk_parse_public_key(&pkey, out.ptr(), out.size());
+	} else {
+		ret = mbedtls_pk_parse_key(&pkey, out.ptr(), out.size(), nullptr, 0);
+	}
 	// We MUST zeroize the memory for safety!
 	mbedtls_platform_zeroize(out.ptrw(), out.size());
-	ERR_FAIL_COND_V_MSG(ret, FAILED, "Error parsing private key '" + itos(ret) + "'.");
+	ERR_FAIL_COND_V_MSG(ret, FAILED, "Error parsing key '" + itos(ret) + "'.");
 
+	public_only = p_public_only;
 	return OK;
 }
 
-Error CryptoKeyMbedTLS::save(String p_path) {
+Error CryptoKeyMbedTLS::save(String p_path, bool p_public_only) {
 	FileAccess *f = FileAccess::open(p_path, FileAccess::WRITE);
 	ERR_FAIL_COND_V_MSG(!f, ERR_INVALID_PARAMETER, "Cannot save CryptoKeyMbedTLS file '" + p_path + "'.");
 
 	unsigned char w[16000];
 	memset(w, 0, sizeof(w));
 
-	int ret = mbedtls_pk_write_key_pem(&pkey, w, sizeof(w));
+	int ret = 0;
+	if (p_public_only) {
+		ret = mbedtls_pk_write_pubkey_pem(&pkey, w, sizeof(w));
+	} else {
+		ret = mbedtls_pk_write_key_pem(&pkey, w, sizeof(w));
+	}
 	if (ret != 0) {
 		memdelete(f);
 		mbedtls_platform_zeroize(w, sizeof(w)); // Zeroize anything we might have written.
@@ -92,6 +103,37 @@ Error CryptoKeyMbedTLS::save(String p_path) {
 	return OK;
 }
 
+Error CryptoKeyMbedTLS::load_from_string(String p_string_key, bool p_public_only) {
+	int ret = 0;
+	if (p_public_only) {
+		ret = mbedtls_pk_parse_public_key(&pkey, (unsigned char *)p_string_key.utf8().get_data(), p_string_key.utf8().size());
+	} else {
+		ret = mbedtls_pk_parse_key(&pkey, (unsigned char *)p_string_key.utf8().get_data(), p_string_key.utf8().size(), nullptr, 0);
+	}
+	ERR_FAIL_COND_V_MSG(ret, FAILED, "Error parsing key '" + itos(ret) + "'.");
+
+	public_only = p_public_only;
+	return OK;
+}
+
+String CryptoKeyMbedTLS::save_to_string(bool p_public_only) {
+	unsigned char w[16000];
+	memset(w, 0, sizeof(w));
+
+	int ret = 0;
+	if (p_public_only) {
+		ret = mbedtls_pk_write_pubkey_pem(&pkey, w, sizeof(w));
+	} else {
+		ret = mbedtls_pk_write_key_pem(&pkey, w, sizeof(w));
+	}
+	if (ret != 0) {
+		mbedtls_platform_zeroize(w, sizeof(w));
+		ERR_FAIL_V_MSG("", "Error saving key '" + itos(ret) + "'.");
+	}
+	String s = String::utf8((char *)w);
+	return s;
+}
+
 X509Certificate *X509CertificateMbedTLS::create() {
 	return memnew(X509CertificateMbedTLS);
 }
@@ -221,6 +263,7 @@ Ref<CryptoKey> CryptoMbedTLS::generate_rsa(int p_bytes) {
 	int ret = mbedtls_pk_setup(&(out->pkey), mbedtls_pk_info_from_type(MBEDTLS_PK_RSA));
 	ERR_FAIL_COND_V(ret != 0, nullptr);
 	ret = mbedtls_rsa_gen_key(mbedtls_pk_rsa(out->pkey), mbedtls_ctr_drbg_random, &ctr_drbg, p_bytes, 65537);
+	out->public_only = false;
 	ERR_FAIL_COND_V(ret != 0, nullptr);
 	return out;
 }
diff --git a/modules/mbedtls/crypto_mbedtls.h b/modules/mbedtls/crypto_mbedtls.h
index 48855d082a2..b3dd0e2a392 100644
--- a/modules/mbedtls/crypto_mbedtls.h
+++ b/modules/mbedtls/crypto_mbedtls.h
@@ -43,15 +43,19 @@ class SSLContextMbedTLS;
 class CryptoKeyMbedTLS : public CryptoKey {
 private:
 	mbedtls_pk_context pkey;
-	int locks;
+	int locks = 0;
+	bool public_only = true;
 
 public:
 	static CryptoKey *create();
 	static void make_default() { CryptoKey::_create = create; }
 	static void finalize() { CryptoKey::_create = nullptr; }
 
-	virtual Error load(String p_path);
-	virtual Error save(String p_path);
+	virtual Error load(String p_path, bool p_public_only);
+	virtual Error save(String p_path, bool p_public_only);
+	virtual String save_to_string(bool p_public_only);
+	virtual Error load_from_string(String p_string_key, bool p_public_only);
+	virtual bool is_public_only() const { return public_only; };
 
 	CryptoKeyMbedTLS() {
 		mbedtls_pk_init(&pkey);