commit b6a96195fb188d9e0e0bc4b4e2c70a78db116d03
author Minos Galanakis <minos.galanakis@arm.com> Fri Jan 12 14:34:14 2024 +0000
committer Minos Galanakis <minos.galanakis@arm.com> Thu Jan 18 14:29:41 2024 +0000
tree 2407051c192cd956cff60c3fe55ec8a656237180
parent f4dfd1c8a5282ea0a9d0641d2fd6dd0649a5c92f

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

diff --git a/programs/pkey/dh_client.c b/programs/pkey/dh_client.c
index d8fc86f..774051c 100644
--- a/programs/pkey/dh_client.c
+++ b/programs/pkey/dh_client.c
@@ -66,6 +66,7 @@
     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 @@
     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 @@
     }
 
     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 @@
 
     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 @@
     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);
 }

programs/pkey/dh_server.c

diff --git a/programs/pkey/dh_server.c b/programs/pkey/dh_server.c
index 7d7618b..0c6cebc 100644
--- a/programs/pkey/dh_server.c
+++ b/programs/pkey/dh_server.c
@@ -190,8 +190,9 @@
         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 @@
         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);