From b6a96195fb188d9e0e0bc4b4e2c70a78db116d03 Mon Sep 17 00:00:00 2001
From: Minos Galanakis <minos.galanakis@arm.com>
Date: Fri, 12 Jan 2024 14:34:14 +0000
Subject: [PATCH] programs_dh_client/server: Updated programs to use
 `mbedtls_rsa_get_len()`

Signed-off-by: Minos Galanakis <minos.galanakis@arm.com>
---
 programs/pkey/dh_client.c | 16 +++++++++-------
 programs/pkey/dh_server.c |  7 ++++---
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/programs/pkey/dh_client.c b/programs/pkey/dh_client.c
index d8fc86fa0b..774051c5d6 100644
--- a/programs/pkey/dh_client.c
+++ b/programs/pkey/dh_client.c
@@ -66,6 +66,7 @@ int main(void)
     unsigned char *p, *end;
     unsigned char buf[2048];
     unsigned char hash[MBEDTLS_MD_CAN_SHA256_MAX_SIZE];
+    mbedtls_mpi N, E;
     const char *pers = "dh_client";
 
     mbedtls_entropy_context entropy;
@@ -78,6 +79,8 @@ int main(void)
     mbedtls_dhm_init(&dhm);
     mbedtls_aes_init(&aes);
     mbedtls_ctr_drbg_init(&ctr_drbg);
+    mbedtls_mpi_init(&N);
+    mbedtls_mpi_init(&E);
 
     /*
      * 1. Setup the RNG
@@ -106,16 +109,13 @@ int main(void)
     }
 
     mbedtls_rsa_init(&rsa);
-
-    if ((ret = mbedtls_mpi_read_file(&rsa.MBEDTLS_PRIVATE(N), 16, f)) != 0 ||
-        (ret = mbedtls_mpi_read_file(&rsa.MBEDTLS_PRIVATE(E), 16, f)) != 0) {
+    if ((ret = mbedtls_mpi_read_file(&N, 16, f)) != 0 ||
+        (ret = mbedtls_mpi_read_file(&E, 16, f)) != 0 ||
+        (ret = mbedtls_rsa_import(&rsa, &N, NULL, NULL, NULL, &E) != 0)) {
         mbedtls_printf(" failed\n  ! mbedtls_mpi_read_file returned %d\n\n", ret);
         fclose(f);
         goto exit;
     }
-
-    rsa.MBEDTLS_PRIVATE(len) = (mbedtls_mpi_bitlen(&rsa.MBEDTLS_PRIVATE(N)) + 7) >> 3;
-
     fclose(f);
 
     /*
@@ -182,7 +182,7 @@ int main(void)
 
     p += 2;
 
-    if ((n = (size_t) (end - p)) != rsa.MBEDTLS_PRIVATE(len)) {
+    if ((n = (size_t) (end - p)) != mbedtls_rsa_get_len(&rsa)) {
         mbedtls_printf(" failed\n  ! Invalid RSA signature size\n\n");
         goto exit;
     }
@@ -273,6 +273,8 @@ exit:
     mbedtls_dhm_free(&dhm);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+    mbedtls_mpi_free(&N);
+    mbedtls_mpi_free(&E);
 
     mbedtls_exit(exit_code);
 }
diff --git a/programs/pkey/dh_server.c b/programs/pkey/dh_server.c
index 7d7618be18..0c6cebc103 100644
--- a/programs/pkey/dh_server.c
+++ b/programs/pkey/dh_server.c
@@ -190,8 +190,9 @@ int main(void)
         goto exit;
     }
 
-    buf[n] = (unsigned char) (rsa.MBEDTLS_PRIVATE(len) >> 8);
-    buf[n + 1] = (unsigned char) (rsa.MBEDTLS_PRIVATE(len));
+    const size_t rsa_key_len = mbedtls_rsa_get_len(&rsa);
+    buf[n] = (unsigned char) (rsa_key_len >> 8);
+    buf[n + 1] = (unsigned char) (rsa_key_len);
 
     if ((ret = mbedtls_rsa_pkcs1_sign(&rsa, mbedtls_ctr_drbg_random, &ctr_drbg,
                                       MBEDTLS_MD_SHA256, MBEDTLS_MD_CAN_SHA256_MAX_SIZE,
@@ -200,7 +201,7 @@ int main(void)
         goto exit;
     }
 
-    buflen = n + 2 + rsa.MBEDTLS_PRIVATE(len);
+    buflen = n + 2 + rsa_key_len;
     buf2[0] = (unsigned char) (buflen >> 8);
     buf2[1] = (unsigned char) (buflen);