From 1304af47480295707ce93ab8c445eccd4cf7b6bd Mon Sep 17 00:00:00 2001
From: akallabeth <akallabeth@posteo.net>
Date: Fri, 25 Nov 2022 10:12:25 +0100
Subject: [PATCH] [core,rdp] Refactor rdp security encryption

Unify rc4 encryption key handling, use common free and reset functions
---
 libfreerdp/core/connection.c | 27 ++++++++++------------
 libfreerdp/core/rdp.c        | 43 +++++++++++++++++++++++++++++++-----
 libfreerdp/core/rdp.h        |  6 +++++
 libfreerdp/core/security.c   | 14 ++----------
 4 files changed, 58 insertions(+), 32 deletions(-)

diff --git a/libfreerdp/core/connection.c b/libfreerdp/core/connection.c
index 53ce1f6dd..2a3a5c566 100644
--- a/libfreerdp/core/connection.c
+++ b/libfreerdp/core/connection.c
@@ -790,10 +790,9 @@ static BOOL rdp_client_establish_keys(rdpRdp* rdp)
 		goto end;
 	}
 
-	rdp->rc4_decrypt_key = winpr_RC4_New(rdp->decrypt_key, rdp->rc4_key_len);
-	rdp->rc4_encrypt_key = winpr_RC4_New(rdp->encrypt_key, rdp->rc4_key_len);
-
-	if (!rdp->rc4_decrypt_key || !rdp->rc4_encrypt_key)
+	if (!rdp_reset_rc4_encrypt_keys(rdp))
+		goto end;
+	if (!rdp_reset_rc4_decrypt_keys(rdp))
 		goto end;
 
 	ret = TRUE;
@@ -804,12 +803,11 @@ end:
 	{
 		winpr_Cipher_Free(rdp->fips_decrypt);
 		winpr_Cipher_Free(rdp->fips_encrypt);
-		winpr_RC4_Free(rdp->rc4_decrypt_key);
-		winpr_RC4_Free(rdp->rc4_encrypt_key);
 		rdp->fips_decrypt = NULL;
 		rdp->fips_encrypt = NULL;
-		rdp->rc4_decrypt_key = NULL;
-		rdp->rc4_encrypt_key = NULL;
+
+		rdp_free_rc4_decrypt_keys(rdp);
+		rdp_free_rc4_encrypt_keys(rdp);
 	}
 
 	return ret;
@@ -922,10 +920,10 @@ BOOL rdp_server_establish_keys(rdpRdp* rdp, wStream* s)
 		goto end;
 	}
 
-	rdp->rc4_decrypt_key = winpr_RC4_New(rdp->decrypt_key, rdp->rc4_key_len);
-	rdp->rc4_encrypt_key = winpr_RC4_New(rdp->encrypt_key, rdp->rc4_key_len);
+	if (!rdp_reset_rc4_encrypt_keys(rdp))
+		goto end;
 
-	if (!rdp->rc4_decrypt_key || !rdp->rc4_encrypt_key)
+	if (!rdp_reset_rc4_decrypt_keys(rdp))
 		goto end;
 
 	ret = tpkt_ensure_stream_consumed(s, length);
@@ -936,12 +934,11 @@ end:
 	{
 		winpr_Cipher_Free(rdp->fips_encrypt);
 		winpr_Cipher_Free(rdp->fips_decrypt);
-		winpr_RC4_Free(rdp->rc4_encrypt_key);
-		winpr_RC4_Free(rdp->rc4_decrypt_key);
 		rdp->fips_encrypt = NULL;
 		rdp->fips_decrypt = NULL;
-		rdp->rc4_encrypt_key = NULL;
-		rdp->rc4_decrypt_key = NULL;
+
+		rdp_free_rc4_encrypt_keys(rdp);
+		rdp_free_rc4_decrypt_keys(rdp);
 	}
 
 	return ret;
diff --git a/libfreerdp/core/rdp.c b/libfreerdp/core/rdp.c
index 001cbb84f..7166f2baf 100644
--- a/libfreerdp/core/rdp.c
+++ b/libfreerdp/core/rdp.c
@@ -2211,13 +2211,12 @@ fail:
 static void rdp_reset_free(rdpRdp* rdp)
 {
 	WINPR_ASSERT(rdp);
-	winpr_RC4_Free(rdp->rc4_decrypt_key);
-	winpr_RC4_Free(rdp->rc4_encrypt_key);
+
+	rdp_free_rc4_decrypt_keys(rdp);
+	rdp_free_rc4_encrypt_keys(rdp);
+
 	winpr_Cipher_Free(rdp->fips_encrypt);
 	winpr_Cipher_Free(rdp->fips_decrypt);
-
-	rdp->rc4_decrypt_key = NULL;
-	rdp->rc4_encrypt_key = NULL;
 	rdp->fips_encrypt = NULL;
 	rdp->fips_decrypt = NULL;
 
@@ -2442,3 +2441,37 @@ BOOL rdp_finalize_is_flag_set(rdpRdp* rdp, UINT32 flag)
 	WINPR_ASSERT(rdp);
 	return (rdp->finalize_sc_pdus & flag) == flag;
 }
+
+BOOL rdp_reset_rc4_encrypt_keys(rdpRdp* rdp)
+{
+	WINPR_ASSERT(rdp);
+	rdp_free_rc4_encrypt_keys(rdp);
+	rdp->rc4_encrypt_key = winpr_RC4_New(rdp->encrypt_key, rdp->rc4_key_len);
+
+	rdp->encrypt_use_count = 0;
+	return rdp->rc4_encrypt_key != NULL;
+}
+
+void rdp_free_rc4_encrypt_keys(rdpRdp* rdp)
+{
+	WINPR_ASSERT(rdp);
+	winpr_RC4_Free(rdp->rc4_encrypt_key);
+	rdp->rc4_encrypt_key = NULL;
+}
+
+void rdp_free_rc4_decrypt_keys(rdpRdp* rdp)
+{
+	WINPR_ASSERT(rdp);
+	winpr_RC4_Free(rdp->rc4_decrypt_key);
+	rdp->rc4_decrypt_key = NULL;
+}
+
+BOOL rdp_reset_rc4_decrypt_keys(rdpRdp* rdp)
+{
+	WINPR_ASSERT(rdp);
+	rdp_free_rc4_decrypt_keys(rdp);
+	rdp->rc4_decrypt_key = winpr_RC4_New(rdp->decrypt_key, rdp->rc4_key_len);
+
+	rdp->decrypt_use_count = 0;
+	return rdp->rc4_decrypt_key != NULL;
+}
diff --git a/libfreerdp/core/rdp.h b/libfreerdp/core/rdp.h
index 937b917b1..1582fb7cf 100644
--- a/libfreerdp/core/rdp.h
+++ b/libfreerdp/core/rdp.h
@@ -285,4 +285,10 @@ BOOL rdp_decrypt(rdpRdp* rdp, wStream* s, UINT16* pLength, UINT16 securityFlags)
 BOOL rdp_set_error_info(rdpRdp* rdp, UINT32 errorInfo);
 BOOL rdp_send_error_info(rdpRdp* rdp);
 
+void rdp_free_rc4_encrypt_keys(rdpRdp* rdp);
+BOOL rdp_reset_rc4_encrypt_keys(rdpRdp* rdp);
+
+void rdp_free_rc4_decrypt_keys(rdpRdp* rdp);
+BOOL rdp_reset_rc4_decrypt_keys(rdpRdp* rdp);
+
 #endif /* FREERDP_LIB_CORE_RDP_H */
diff --git a/libfreerdp/core/security.c b/libfreerdp/core/security.c
index 79267bee5..c147f10a8 100644
--- a/libfreerdp/core/security.c
+++ b/libfreerdp/core/security.c
@@ -741,13 +741,8 @@ BOOL security_encrypt(BYTE* data, size_t length, rdpRdp* rdp)
 		if (!security_key_update(rdp->encrypt_key, rdp->encrypt_update_key, rdp->rc4_key_len, rdp))
 			goto fail;
 
-		winpr_RC4_Free(rdp->rc4_encrypt_key);
-		rdp->rc4_encrypt_key = winpr_RC4_New(rdp->encrypt_key, rdp->rc4_key_len);
-
-		if (!rdp->rc4_encrypt_key)
+		if (!rdp_reset_rc4_encrypt_keys(rdp))
 			goto fail;
-
-		rdp->encrypt_use_count = 0;
 	}
 
 	if (!winpr_RC4_Update(rdp->rc4_encrypt_key, length, data, data))
@@ -780,13 +775,8 @@ BOOL security_decrypt(BYTE* data, size_t length, rdpRdp* rdp)
 		if (!security_key_update(rdp->decrypt_key, rdp->decrypt_update_key, rdp->rc4_key_len, rdp))
 			goto fail;
 
-		winpr_RC4_Free(rdp->rc4_decrypt_key);
-		rdp->rc4_decrypt_key = winpr_RC4_New(rdp->decrypt_key, rdp->rc4_key_len);
-
-		if (!rdp->rc4_decrypt_key)
+		if (!rdp_reset_rc4_decrypt_keys(rdp))
 			goto fail;
-
-		rdp->decrypt_use_count = 0;
 	}
 
 	if (!winpr_RC4_Update(rdp->rc4_decrypt_key, length, data, data))