Merge pull request #8699 from minosgalanakis/update/rsa_context_in_programs_5014

[MBEDTLS_PRIVATE] Update rsa context in programs 5014
diff --git a/programs/pkey/dh_client.c b/programs/pkey/dh_client.c
index 0cb1562..165cee2 100644
--- a/programs/pkey/dh_client.c
+++ b/programs/pkey/dh_client.c
@@ -13,14 +13,13 @@
 
 #if defined(MBEDTLS_AES_C) && defined(MBEDTLS_DHM_C) && \
     defined(MBEDTLS_ENTROPY_C) && defined(MBEDTLS_NET_C) && \
-    defined(MBEDTLS_RSA_C) && defined(MBEDTLS_MD_CAN_SHA256) && \
-    defined(MBEDTLS_FS_IO) && defined(MBEDTLS_CTR_DRBG_C) && \
-    defined(MBEDTLS_MD_CAN_SHA1)
+    defined(MBEDTLS_RSA_C) && defined(MBEDTLS_SHA256_C) && \
+    defined(MBEDTLS_FS_IO) && defined(MBEDTLS_CTR_DRBG_C)
 #include "mbedtls/net_sockets.h"
 #include "mbedtls/aes.h"
 #include "mbedtls/dhm.h"
 #include "mbedtls/rsa.h"
-#include "mbedtls/sha1.h"
+#include "mbedtls/sha256.h"
 #include "mbedtls/entropy.h"
 #include "mbedtls/ctr_drbg.h"
 
@@ -33,9 +32,8 @@
 
 #if !defined(MBEDTLS_AES_C) || !defined(MBEDTLS_DHM_C) ||     \
     !defined(MBEDTLS_ENTROPY_C) || !defined(MBEDTLS_NET_C) ||  \
-    !defined(MBEDTLS_RSA_C) || !defined(MBEDTLS_MD_CAN_SHA256) ||    \
-    !defined(MBEDTLS_FS_IO) || !defined(MBEDTLS_CTR_DRBG_C) || \
-    !defined(MBEDTLS_SHA1_C)
+    !defined(MBEDTLS_RSA_C) || !defined(MBEDTLS_SHA256_C) ||    \
+    !defined(MBEDTLS_FS_IO) || !defined(MBEDTLS_CTR_DRBG_C)
 int main(void)
 {
     mbedtls_printf("MBEDTLS_AES_C and/or MBEDTLS_DHM_C and/or MBEDTLS_ENTROPY_C "
@@ -60,12 +58,14 @@
 
     int ret = 1;
     int exit_code = MBEDTLS_EXIT_FAILURE;
+    unsigned int mdlen;
     size_t n, buflen;
     mbedtls_net_context server_fd;
 
     unsigned char *p, *end;
     unsigned char buf[2048];
-    unsigned char hash[32];
+    unsigned char hash[MBEDTLS_MD_MAX_SIZE];
+    mbedtls_mpi N, E;
     const char *pers = "dh_client";
 
     mbedtls_entropy_context entropy;
@@ -78,6 +78,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 +108,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,18 +181,24 @@
 
     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;
     }
 
-    if ((ret = mbedtls_sha1(buf, (int) (p - 2 - buf), hash)) != 0) {
-        mbedtls_printf(" failed\n  ! mbedtls_sha1 returned %d\n\n", ret);
+    mdlen = (unsigned int) mbedtls_md_get_size(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256));
+    if (mdlen == 0) {
+        mbedtls_printf(" failed\n  ! Invalid digest type\n\n");
+        goto exit;
+    }
+
+    if ((ret = mbedtls_sha256(buf, (int) (p - 2 - buf), hash, 0)) != 0) {
+        mbedtls_printf(" failed\n  ! mbedtls_sha256 returned %d\n\n", ret);
         goto exit;
     }
 
     if ((ret = mbedtls_rsa_pkcs1_verify(&rsa, MBEDTLS_MD_SHA256,
-                                        32, hash, p)) != 0) {
+                                        mdlen, hash, p)) != 0) {
         mbedtls_printf(" failed\n  ! mbedtls_rsa_pkcs1_verify returned %d\n\n", ret);
         goto exit;
     }
@@ -273,6 +278,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);
 }
diff --git a/programs/pkey/dh_server.c b/programs/pkey/dh_server.c
index adddbf2..91bac0e 100644
--- a/programs/pkey/dh_server.c
+++ b/programs/pkey/dh_server.c
@@ -13,14 +13,13 @@
 
 #if defined(MBEDTLS_AES_C) && defined(MBEDTLS_DHM_C) && \
     defined(MBEDTLS_ENTROPY_C) && defined(MBEDTLS_NET_C) && \
-    defined(MBEDTLS_RSA_C) && defined(MBEDTLS_MD_CAN_SHA256) && \
-    defined(MBEDTLS_FS_IO) && defined(MBEDTLS_CTR_DRBG_C) && \
-    defined(MBEDTLS_MD_CAN_SHA1)
+    defined(MBEDTLS_RSA_C) && defined(MBEDTLS_SHA256_C) && \
+    defined(MBEDTLS_FS_IO) && defined(MBEDTLS_CTR_DRBG_C)
 #include "mbedtls/net_sockets.h"
 #include "mbedtls/aes.h"
 #include "mbedtls/dhm.h"
 #include "mbedtls/rsa.h"
-#include "mbedtls/sha1.h"
+#include "mbedtls/sha256.h"
 #include "mbedtls/entropy.h"
 #include "mbedtls/ctr_drbg.h"
 
@@ -33,9 +32,8 @@
 
 #if !defined(MBEDTLS_AES_C) || !defined(MBEDTLS_DHM_C) ||     \
     !defined(MBEDTLS_ENTROPY_C) || !defined(MBEDTLS_NET_C) ||  \
-    !defined(MBEDTLS_RSA_C) || !defined(MBEDTLS_MD_CAN_SHA256) ||    \
-    !defined(MBEDTLS_FS_IO) || !defined(MBEDTLS_CTR_DRBG_C) || \
-    !defined(MBEDTLS_SHA1_C)
+    !defined(MBEDTLS_RSA_C) || !defined(MBEDTLS_SHA256_C) ||    \
+    !defined(MBEDTLS_FS_IO) || !defined(MBEDTLS_CTR_DRBG_C)
 int main(void)
 {
     mbedtls_printf("MBEDTLS_AES_C and/or MBEDTLS_DHM_C and/or MBEDTLS_ENTROPY_C "
@@ -53,11 +51,12 @@
 
     int ret = 1;
     int exit_code = MBEDTLS_EXIT_FAILURE;
+    unsigned int mdlen;
     size_t n, buflen;
     mbedtls_net_context listen_fd, client_fd;
 
     unsigned char buf[2048];
-    unsigned char hash[32];
+    unsigned char hash[MBEDTLS_MD_MAX_SIZE];
     unsigned char buf2[2];
     const char *pers = "dh_server";
 
@@ -186,21 +185,30 @@
     /*
      * 5. Sign the parameters and send them
      */
-    if ((ret = mbedtls_sha1(buf, n, hash)) != 0) {
-        mbedtls_printf(" failed\n  ! mbedtls_sha1 returned %d\n\n", ret);
+
+    mdlen = (unsigned int) mbedtls_md_get_size(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256));
+    if (mdlen == 0) {
+        mbedtls_printf(" failed\n  ! Invalid digest type\n\n");
         goto exit;
     }
 
-    buf[n] = (unsigned char) (rsa.MBEDTLS_PRIVATE(len) >> 8);
-    buf[n + 1] = (unsigned char) (rsa.MBEDTLS_PRIVATE(len));
+    if ((ret = mbedtls_sha256(buf, n, hash, 0)) != 0) {
+        mbedtls_printf(" failed\n  ! mbedtls_sha256 returned %d\n\n", ret);
+        goto exit;
+    }
 
-    if ((ret = mbedtls_rsa_pkcs1_sign(&rsa, NULL, NULL, MBEDTLS_MD_SHA256,
-                                      32, hash, buf + n + 2)) != 0) {
+    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, mdlen,
+                                      hash, buf + n + 2)) != 0) {
         mbedtls_printf(" failed\n  ! mbedtls_rsa_pkcs1_sign returned %d\n\n", ret);
         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);
 
diff --git a/programs/pkey/rsa_decrypt.c b/programs/pkey/rsa_decrypt.c
index 76bfddf..a84af50 100644
--- a/programs/pkey/rsa_decrypt.c
+++ b/programs/pkey/rsa_decrypt.c
@@ -133,7 +133,7 @@
 
     fclose(f);
 
-    if (i != rsa.MBEDTLS_PRIVATE(len)) {
+    if (i != mbedtls_rsa_get_len(&rsa)) {
         mbedtls_printf("\n  ! Invalid RSA signature format\n\n");
         goto exit;
     }
diff --git a/programs/pkey/rsa_encrypt.c b/programs/pkey/rsa_encrypt.c
index 4bbb54e..6538f8a 100644
--- a/programs/pkey/rsa_encrypt.c
+++ b/programs/pkey/rsa_encrypt.c
@@ -126,7 +126,7 @@
         goto exit;
     }
 
-    for (i = 0; i < rsa.MBEDTLS_PRIVATE(len); i++) {
+    for (i = 0; i < mbedtls_rsa_get_len(&rsa); i++) {
         mbedtls_fprintf(f, "%02X%s", buf[i],
                         (i + 1) % 16 == 0 ? "\r\n" : " ");
     }
diff --git a/programs/pkey/rsa_sign.c b/programs/pkey/rsa_sign.c
index 9d8ebe3..e14953b 100644
--- a/programs/pkey/rsa_sign.c
+++ b/programs/pkey/rsa_sign.c
@@ -131,7 +131,7 @@
         goto exit;
     }
 
-    for (i = 0; i < rsa.MBEDTLS_PRIVATE(len); i++) {
+    for (i = 0; i < mbedtls_rsa_get_len(&rsa); i++) {
         mbedtls_fprintf(f, "%02X%s", buf[i],
                         (i + 1) % 16 == 0 ? "\r\n" : " ");
     }
diff --git a/programs/pkey/rsa_verify.c b/programs/pkey/rsa_verify.c
index e7d72fd..4a9af77 100644
--- a/programs/pkey/rsa_verify.c
+++ b/programs/pkey/rsa_verify.c
@@ -37,11 +37,14 @@
     int exit_code = MBEDTLS_EXIT_FAILURE;
     size_t i;
     mbedtls_rsa_context rsa;
+    mbedtls_mpi N, E;
     unsigned char hash[32];
     unsigned char buf[MBEDTLS_MPI_MAX_SIZE];
     char filename[512];
 
     mbedtls_rsa_init(&rsa);
+    mbedtls_mpi_init(&N);
+    mbedtls_mpi_init(&E);
 
     if (argc != 2) {
         mbedtls_printf("usage: rsa_verify <filename>\n");
@@ -62,15 +65,13 @@
         goto exit;
     }
 
-    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);
 
     /*
@@ -91,7 +92,7 @@
 
     fclose(f);
 
-    if (i != rsa.MBEDTLS_PRIVATE(len)) {
+    if (i != mbedtls_rsa_get_len(&rsa)) {
         mbedtls_printf("\n  ! Invalid RSA signature format\n\n");
         goto exit;
     }
@@ -124,6 +125,8 @@
 exit:
 
     mbedtls_rsa_free(&rsa);
+    mbedtls_mpi_free(&N);
+    mbedtls_mpi_free(&E);
 
     mbedtls_exit(exit_code);
 }