From 8e3f9aa68104fefc959d2d6af7cba3c11bfde3fb Mon Sep 17 00:00:00 2001
From: Fabio Alessandrelli <fabio.alessandrelli@gmail.com>
Date: Sat, 6 Jun 2020 16:09:28 +0200
Subject: [PATCH] Implement RSA encryption/decryption.

---
 core/crypto/crypto.cpp             |  2 ++
 core/crypto/crypto.h               |  2 ++
 modules/mbedtls/crypto_mbedtls.cpp | 27 +++++++++++++++++++++++++++
 modules/mbedtls/crypto_mbedtls.h   |  2 ++
 4 files changed, 33 insertions(+)

diff --git a/core/crypto/crypto.cpp b/core/crypto/crypto.cpp
index f1d13b06337..29d02e11dff 100644
--- a/core/crypto/crypto.cpp
+++ b/core/crypto/crypto.cpp
@@ -88,6 +88,8 @@ void Crypto::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("generate_self_signed_certificate", "key", "issuer_name", "not_before", "not_after"), &Crypto::generate_self_signed_certificate, DEFVAL("CN=myserver,O=myorganisation,C=IT"), DEFVAL("20140101000000"), DEFVAL("20340101000000"));
 	ClassDB::bind_method(D_METHOD("sign", "hash_type", "hash", "key"), &Crypto::sign);
 	ClassDB::bind_method(D_METHOD("verify", "hash_type", "hash", "signature", "key"), &Crypto::verify);
+	ClassDB::bind_method(D_METHOD("encrypt", "key", "plaintext"), &Crypto::encrypt);
+	ClassDB::bind_method(D_METHOD("decrypt", "key", "ciphertext"), &Crypto::decrypt);
 }
 
 /// Resource loader/saver
diff --git a/core/crypto/crypto.h b/core/crypto/crypto.h
index 472ad8ab9d6..916f7798eb9 100644
--- a/core/crypto/crypto.h
+++ b/core/crypto/crypto.h
@@ -85,6 +85,8 @@ public:
 
 	virtual Vector<uint8_t> sign(HashingContext::HashType p_hash_type, Vector<uint8_t> p_hash, Ref<CryptoKey> p_key) = 0;
 	virtual bool verify(HashingContext::HashType p_hash_type, Vector<uint8_t> p_hash, Vector<uint8_t> p_signature, Ref<CryptoKey> p_key) = 0;
+	virtual Vector<uint8_t> encrypt(Ref<CryptoKey> p_key, Vector<uint8_t> p_plaintext) = 0;
+	virtual Vector<uint8_t> decrypt(Ref<CryptoKey> p_key, Vector<uint8_t> p_ciphertext) = 0;
 
 	Crypto() {}
 };
diff --git a/modules/mbedtls/crypto_mbedtls.cpp b/modules/mbedtls/crypto_mbedtls.cpp
index a432a88fd10..501bfff0753 100644
--- a/modules/mbedtls/crypto_mbedtls.cpp
+++ b/modules/mbedtls/crypto_mbedtls.cpp
@@ -362,3 +362,30 @@ bool CryptoMbedTLS::verify(HashingContext::HashType p_hash_type, Vector<uint8_t>
 	ERR_FAIL_COND_V_MSG(!key.is_valid(), false, "Invalid key provided.");
 	return mbedtls_pk_verify(&(key->pkey), type, p_hash.ptr(), size, p_signature.ptr(), p_signature.size()) == 0;
 }
+
+Vector<uint8_t> CryptoMbedTLS::encrypt(Ref<CryptoKey> p_key, Vector<uint8_t> p_plaintext) {
+	Ref<CryptoKeyMbedTLS> key = static_cast<Ref<CryptoKeyMbedTLS>>(p_key);
+	ERR_FAIL_COND_V_MSG(!key.is_valid(), Vector<uint8_t>(), "Invalid key provided.");
+	uint8_t buf[1024];
+	size_t size;
+	Vector<uint8_t> out;
+	int ret = mbedtls_pk_encrypt(&(key->pkey), p_plaintext.ptr(), p_plaintext.size(), buf, &size, sizeof(buf), mbedtls_ctr_drbg_random, &ctr_drbg);
+	ERR_FAIL_COND_V_MSG(ret, out, "Error while encrypting: " + itos(ret));
+	out.resize(size);
+	copymem(out.ptrw(), buf, size);
+	return out;
+}
+
+Vector<uint8_t> CryptoMbedTLS::decrypt(Ref<CryptoKey> p_key, Vector<uint8_t> p_ciphertext) {
+	Ref<CryptoKeyMbedTLS> key = static_cast<Ref<CryptoKeyMbedTLS>>(p_key);
+	ERR_FAIL_COND_V_MSG(!key.is_valid(), Vector<uint8_t>(), "Invalid key provided.");
+	ERR_FAIL_COND_V_MSG(key->is_public_only(), Vector<uint8_t>(), "Invalid key provided. Cannot decrypt using a public_only key.");
+	uint8_t buf[2048];
+	size_t size;
+	Vector<uint8_t> out;
+	int ret = mbedtls_pk_decrypt(&(key->pkey), p_ciphertext.ptr(), p_ciphertext.size(), buf, &size, sizeof(buf), mbedtls_ctr_drbg_random, &ctr_drbg);
+	ERR_FAIL_COND_V_MSG(ret, out, "Error while decrypting: " + itos(ret));
+	out.resize(size);
+	copymem(out.ptrw(), buf, size);
+	return out;
+}
diff --git a/modules/mbedtls/crypto_mbedtls.h b/modules/mbedtls/crypto_mbedtls.h
index c22ddcdb422..2a446f9d484 100644
--- a/modules/mbedtls/crypto_mbedtls.h
+++ b/modules/mbedtls/crypto_mbedtls.h
@@ -120,6 +120,8 @@ public:
 	virtual Ref<X509Certificate> generate_self_signed_certificate(Ref<CryptoKey> p_key, String p_issuer_name, String p_not_before, String p_not_after);
 	virtual Vector<uint8_t> sign(HashingContext::HashType p_hash_type, Vector<uint8_t> p_hash, Ref<CryptoKey> p_key);
 	virtual bool verify(HashingContext::HashType p_hash_type, Vector<uint8_t> p_hash, Vector<uint8_t> p_signature, Ref<CryptoKey> p_key);
+	virtual Vector<uint8_t> encrypt(Ref<CryptoKey> p_key, Vector<uint8_t> p_plaintext);
+	virtual Vector<uint8_t> decrypt(Ref<CryptoKey> p_key, Vector<uint8_t> p_ciphertext);
 
 	CryptoMbedTLS();
 	~CryptoMbedTLS();