1 files changed, 12 insertions, 6 deletions

diff --git a/host/lib/signature_digest.c b/host/lib/signature_digest.c
index dcc2cf26..b56233a0 100644
--- a/host/lib/signature_digest.c
+++ b/host/lib/signature_digest.c
@@ -9,13 +9,17 @@
 #include <stdlib.h>
 #include <unistd.h>
 
+#include "2sysincludes.h"
+
+#include "2common.h"
+#include "2sha.h"
 #include "cryptolib.h"
 #include "host_common.h"
 #include "signature_digest.h"
 
 
 uint8_t* PrependDigestInfo(unsigned int algorithm, uint8_t* digest) {
-  const int digest_size = hash_size_map[algorithm];
+  const int digest_size = vb2_digest_size(vb2_crypto_to_hash(algorithm));
   const int digestinfo_size = digestinfo_size_map[algorithm];
   const uint8_t* digestinfo = hash_digestinfo_map[algorithm];
   uint8_t* p = malloc(digestinfo_size + digest_size);
@@ -27,14 +31,16 @@ uint8_t* PrependDigestInfo(unsigned int algorithm, uint8_t* digest) {
 uint8_t* SignatureDigest(const uint8_t* buf, uint64_t len,
                          unsigned int algorithm) {
   uint8_t* info_digest  = NULL;
-  uint8_t* digest = NULL;
+
+  uint8_t digest[VB2_SHA512_DIGEST_SIZE];  /* Longest digest */
 
   if (algorithm >= kNumAlgorithms) {
     VBDEBUG(("SignatureDigest() called with invalid algorithm!\n"));
-  } else if ((digest = DigestBuf(buf, len, algorithm))) {
+  } else if (VB2_SUCCESS == vb2_digest_buffer(buf, len,
+					      vb2_crypto_to_hash(algorithm),
+					      digest, sizeof(digest))) {
     info_digest = PrependDigestInfo(algorithm, digest);
   }
-  free(digest);
   return info_digest;
 }
 
@@ -44,8 +50,8 @@ uint8_t* SignatureBuf(const uint8_t* buf, uint64_t len, const char* key_file,
   RSA* key = NULL;
   uint8_t* signature = NULL;
   uint8_t* signature_digest = SignatureDigest(buf, len, algorithm);
-  int signature_digest_len = (hash_size_map[algorithm] +
-                              digestinfo_size_map[algorithm]);
+  const int digest_size = vb2_digest_size(vb2_crypto_to_hash(algorithm));
+  int signature_digest_len = (digest_size + digestinfo_size_map[algorithm]);
   key_fp  = fopen(key_file, "r");
   if (!key_fp) {
     VBDEBUG(("SignatureBuf(): Couldn't open key file: %s\n", key_file));