From 486762c8b16cc802aaffd0ea68498e6086b745e3 Mon Sep 17 00:00:00 2001
From: Xie Han <63350856@qq.com>
Date: Fri, 10 Feb 2023 03:35:54 +0800
Subject: [PATCH] Add MySQL RSA encrypt.

---
 src/factory/MySQLTaskImpl.cc  | 13 ++++++---
 src/protocol/MySQLMessage.cc  | 55 +++++++++++++++++++++++++++++++++++
 src/protocol/MySQLMessage.inl |  2 ++
 3 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/src/factory/MySQLTaskImpl.cc b/src/factory/MySQLTaskImpl.cc
index 2e52cf75..98449218 100644
--- a/src/factory/MySQLTaskImpl.cc
+++ b/src/factory/MySQLTaskImpl.cc
@@ -293,10 +293,14 @@ CommMessageIn *ComplexMySQLTask::message_in()
 
 	case ST_AUTH_REQUEST:
 	case ST_AUTH_SWITCH_REQUEST:
-	case ST_CLEAR_PASSWORD_REQUEST:
 		resp = new MySQLAuthResponse;
 		break;
 
+	case ST_CLEAR_PASSWORD_REQUEST:
+	case ST_RSA_AUTH_REQUEST:
+		resp = new MySQLResponse;
+		break;
+
 	case ST_SHA256_PUBLIC_KEY_REQUEST:
 	case ST_CSHA2_PUBLIC_KEY_REQUEST:
 		resp = new MySQLPublicKeyResponse;
@@ -437,6 +441,7 @@ int ComplexMySQLTask::keep_alive_timeout()
 	case ST_AUTH_REQUEST:
 	case ST_AUTH_SWITCH_REQUEST:
 	case ST_CLEAR_PASSWORD_REQUEST:
+	case ST_RSA_AUTH_REQUEST:
 		if (resp->is_ok_packet())
 		{
 			if (!res_charset_.empty())
@@ -447,7 +452,9 @@ int ComplexMySQLTask::keep_alive_timeout()
 			break;
 		}
 
-		if (resp->is_error_packet() || conn->state == ST_CLEAR_PASSWORD_REQUEST)
+		if (resp->is_error_packet() ||
+			conn->state == ST_CLEAR_PASSWORD_REQUEST ||
+			conn->state == ST_RSA_AUTH_REQUEST)
 		{
 			this->resp = std::move(*resp);
 			state_ = WFT_STATE_TASK_ERROR;
@@ -474,8 +481,6 @@ int ComplexMySQLTask::keep_alive_timeout()
 		conn->state = ST_RSA_AUTH_REQUEST;
 		break;
 
-	case ST_RSA_AUTH_REQUEST:
-		break;
 
 	case ST_CHARSET_REQUEST:
 		if (!resp->is_ok_packet())
diff --git a/src/protocol/MySQLMessage.cc b/src/protocol/MySQLMessage.cc
index e6edc54b..4e8a9b72 100644
--- a/src/protocol/MySQLMessage.cc
+++ b/src/protocol/MySQLMessage.cc
@@ -535,8 +535,63 @@ int MySQLPublicKeyResponse::encode(struct iovec vectors[], int max)
 	return MySQLMessage::encode(vectors, max);
 }
 
+int MySQLRSAAuthRequest::rsa_encrypt(void *ctx)
+{
+	EVP_PKEY_CTX *pkey_ctx = (EVP_PKEY_CTX *)ctx;
+	std::string pass = password_;
+	unsigned char *p = (unsigned char *)pass.c_str();
+	unsigned char out[256];
+	size_t outlen = 256;
+
+	for (size_t i = 0; i < pass.size(); i++)
+		p[i] ^= seed_[i % 20];
+
+	if (EVP_PKEY_encrypt_init(pkey_ctx) > 0 &&
+		EVP_PKEY_CTX_set_rsa_padding(pkey_ctx, RSA_PKCS1_OAEP_PADDING) > 0 &&
+		EVP_PKEY_encrypt(pkey_ctx, out, &outlen, p, pass.size()) > 0)
+	{
+		buf_.assign((char *)out, 256);
+		return 0;
+	}
+
+	return -1;
+}
+
 int MySQLRSAAuthRequest::encode(struct iovec vectors[], int max)
 {
+	RSA *rsa;
+	BIO *bio;
+
+	bio = BIO_new_mem_buf(public_key_.c_str(), public_key_.size());
+	if (!bio)
+		return -1;
+
+	rsa = PEM_read_bio_RSA_PUBKEY(bio, NULL, NULL, NULL);
+	BIO_free(bio);
+
+	EVP_PKEY_CTX *ctx;
+	EVP_PKEY *pkey;
+	int ret = -1;
+
+	pkey = EVP_PKEY_new();
+	if (pkey)
+	{
+		EVP_PKEY_assign_RSA(pkey, rsa);
+		ctx = EVP_PKEY_CTX_new(pkey, NULL);
+		if (ctx)
+		{
+			ret = rsa_encrypt(ctx);
+			EVP_PKEY_CTX_free(ctx);
+		}
+
+		EVP_PKEY_free(pkey);
+	}
+	else
+		RSA_free(rsa);
+
+	if (ret < 0)
+		return -1;
+
 	return MySQLMessage::encode(vectors, max);
 }
 
diff --git a/src/protocol/MySQLMessage.inl b/src/protocol/MySQLMessage.inl
index be2c148b..b4caf495 100644
--- a/src/protocol/MySQLMessage.inl
+++ b/src/protocol/MySQLMessage.inl
@@ -307,6 +307,8 @@ private:
 		return -2;
 	}
 
+	int rsa_encrypt(void *ctx);
+
 	std::string password_;
 	std::string public_key_;
 	unsigned char seed_[20];