From 698c5ecf6311a5f81ef0b7f2c2c3f9a93df17074 Mon Sep 17 00:00:00 2001
From: David Garske <david@wolfssl.com>
Date: Tue, 24 Oct 2017 08:33:41 -0700
Subject: [PATCH] Cleanups to signature example. Added return code checking and
 clarified the malloc values used.

---
 signature/signature.c | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/signature/signature.c b/signature/signature.c
index 473dbb04..946c725d 100644
--- a/signature/signature.c
+++ b/signature/signature.c
@@ -32,6 +32,7 @@
 #include <stdio.h>
 
 #define RSA_KEY_SIZE    2048
+#define DER_FILE_BUFFER 2048 /* max DER size */
 
 void hexdump(const void *buffer, word32 len, byte cols)
 {
@@ -173,14 +174,15 @@ int rsa_load_der_file(const char* derFile, RsaKey *rsaKey)
     int ret = EXIT_FAILURE;
     FILE *file;
     byte *buffer = NULL;
+    word32 bufSize = DER_FILE_BUFFER;
     word32 bytes = 0;
     word32 idx = 0;
 
     file = fopen(derFile, "rb");
     if (file) {
-        buffer = malloc(RSA_KEY_SIZE);
+        buffer = malloc(bufSize);
         if(buffer) {
-            bytes = fread(buffer, 1, RSA_KEY_SIZE, file);
+            bytes = fread(buffer, 1, bufSize, file);
             fclose(file);
         }
     }
@@ -210,8 +212,19 @@ int rsa_sign_verify_test(enum wc_HashType hash_type, enum wc_SignatureType sig_t
 #endif
 
     /* Init */
-    wc_InitRng(&rng);
-    wc_InitRsaKey(&rsaKey, NULL);
+    ret = wc_InitRng(&rng);
+    if (ret != 0) {
+        printf("Init RNG failed! %d\n", ret);
+        ret = EXIT_FAILURE;
+        goto exit;
+    }
+
+    ret = wc_InitRsaKey(&rsaKey, NULL);
+    if (ret != 0) {
+        printf("Init RSA key failed! %d\n", ret);
+        ret = EXIT_FAILURE;
+        goto exit;
+    }
 
     printf("RSA Key Size %d\n", RSA_KEY_SIZE);
 
@@ -223,7 +236,7 @@ int rsa_sign_verify_test(enum wc_HashType hash_type, enum wc_SignatureType sig_t
     }
 
     /* Display key data */
-    rsaKeyLen = RSA_KEY_SIZE;
+    rsaKeyLen = DER_FILE_BUFFER;
     rsaKeyBuf = malloc(rsaKeyLen);
     ret = wc_RsaKeyToDer(&rsaKey, rsaKeyBuf, rsaKeyLen);
     if (ret <= 0) {
@@ -235,7 +248,7 @@ int rsa_sign_verify_test(enum wc_HashType hash_type, enum wc_SignatureType sig_t
     printf("RSA Key: Len %d\n", rsaKeyLen);
     hexdump(rsaKeyBuf, rsaKeyLen, 16);
 
-    rsaPubKeyLen = RSA_KEY_SIZE;
+    rsaPubKeyLen = DER_FILE_BUFFER;
     rsaPubKeyBuf = malloc(rsaPubKeyLen);
     ret = wc_RsaKeyToPublicDer(&rsaKey, rsaPubKeyBuf, rsaPubKeyLen);
     if (ret <= 0) {