1 files changed, 60 insertions, 19 deletions

diff --git a/extra/usb_updater/gsctool.c b/extra/usb_updater/gsctool.c
index 0046635e03..d91b1d6036 100644
--- a/extra/usb_updater/gsctool.c
+++ b/extra/usb_updater/gsctool.c
@@ -204,6 +204,17 @@ struct option_container {
 	const char *help_text;
 };
 
+/* SHA context used with our local sha_* abstraction functions */
+union sha_ctx {
+	SHA_CTX sha1;
+	SHA256_CTX sha256;
+};
+
+static void sha_init(union sha_ctx *ctx);
+static void sha_update(union sha_ctx *ctx, const void *data, size_t len);
+static void sha_final_into_block_digest(union sha_ctx *ctx, void *block_digest,
+					size_t size);
+
 /*
  * This by far exceeds the largest vendor command response size we ever
  * expect.
@@ -812,8 +823,7 @@ static void transfer_section(struct transfer_descriptor *td,
 	printf("sending 0x%zx bytes to %#x\n", data_len, section_addr);
 	while (data_len) {
 		size_t payload_size;
-		SHA_CTX ctx;
-		uint8_t digest[SHA_DIGEST_LENGTH];
+		union sha_ctx ctx;
 		int max_retries;
 		struct update_pdu updu;
 
@@ -825,15 +835,13 @@ static void transfer_section(struct transfer_descriptor *td,
 		updu.cmd.block_base = htobe32(section_addr);
 
 		/* Calculate the digest. */
-		SHA1_Init(&ctx);
-		SHA1_Update(&ctx, &updu.cmd.block_base,
-			    sizeof(updu.cmd.block_base));
-		SHA1_Update(&ctx, data_ptr, payload_size);
-		SHA1_Final(digest, &ctx);
-
-		/* Copy the first few bytes. */
-		memcpy(&updu.cmd.block_digest, digest,
-		       sizeof(updu.cmd.block_digest));
+		sha_init(&ctx);
+		sha_update(&ctx, &updu.cmd.block_base,
+			   sizeof(updu.cmd.block_base));
+		sha_update(&ctx, data_ptr, payload_size);
+		sha_final_into_block_digest(&ctx, &updu.cmd.block_digest,
+					    sizeof(updu.cmd.block_digest));
+
 		if (td->ep_type == usb_xfer) {
 			for (max_retries = 10; max_retries; max_retries--)
 				if (!transfer_block(&td->uep, &updu,
@@ -1282,8 +1290,7 @@ static int ext_cmd_over_usb(struct usb_endpoint *uep, uint16_t subcommand,
 	struct update_frame_header *ufh;
 	uint16_t *frame_ptr;
 	size_t usb_msg_size;
-	SHA_CTX ctx;
-	uint8_t digest[SHA_DIGEST_LENGTH];
+	union sha_ctx ctx;
 
 	usb_msg_size = sizeof(struct update_frame_header) +
 		sizeof(subcommand) + body_size;
@@ -1304,12 +1311,12 @@ static int ext_cmd_over_usb(struct usb_endpoint *uep, uint16_t subcommand,
 		memcpy(frame_ptr + 1, cmd_body, body_size);
 
 	/* Calculate the digest. */
-	SHA1_Init(&ctx);
-	SHA1_Update(&ctx, &ufh->cmd.block_base,
-		    usb_msg_size -
-		    offsetof(struct update_frame_header, cmd.block_base));
-	SHA1_Final(digest, &ctx);
-	memcpy(&ufh->cmd.block_digest, digest, sizeof(ufh->cmd.block_digest));
+	sha_init(&ctx);
+	sha_update(&ctx, &ufh->cmd.block_base,
+		   usb_msg_size -
+		   offsetof(struct update_frame_header, cmd.block_base));
+	sha_final_into_block_digest(&ctx, &ufh->cmd.block_digest,
+				    sizeof(ufh->cmd.block_digest));
 
 	do_xfer(uep, ufh, usb_msg_size, resp,
 		resp_size ? *resp_size : 0, 1, resp_size);
@@ -1601,6 +1608,40 @@ static void generate_reset_request(struct transfer_descriptor *td)
 	printf("reboot %s\n", reset_type);
 }
 
+/* Forward to correct SHA implementation based on image type */
+static void sha_init(union sha_ctx *ctx)
+{
+	if (image_magic == MAGIC_HAVEN)
+		SHA1_Init(&ctx->sha1);
+	else if (image_magic == MAGIC_DAUNTLESS)
+		SHA256_Init(&ctx->sha256);
+}
+
+/* Forward to correct SHA implementation based on image type */
+static void sha_update(union sha_ctx *ctx, const void *data, size_t len)
+{
+	if (image_magic == MAGIC_HAVEN)
+		SHA1_Update(&ctx->sha1, data, len);
+	else if (image_magic == MAGIC_DAUNTLESS)
+		SHA256_Update(&ctx->sha256, data, len);
+}
+
+/* Forward to correct SHA implementation based on image type */
+static void sha_final_into_block_digest(union sha_ctx *ctx, void *block_digest,
+					size_t size)
+{
+	/* Big enough for either hash algo */
+	uint8_t full_digest[SHA256_DIGEST_LENGTH];
+
+	if (image_magic == MAGIC_HAVEN)
+		SHA1_Final(full_digest, &ctx->sha1);
+	else if (image_magic == MAGIC_DAUNTLESS)
+		SHA256_Final(full_digest, &ctx->sha256);
+
+	/* Don't try to copy out more than the smallest (SHA1) digest */
+	memcpy(block_digest, full_digest, MIN(size, SHA_DIGEST_LENGTH));
+}
+
 /*
  * Machine output is formatted as "key=value", one key-value pair per line, and
  * parsed by other programs (e.g., debugd). The value part should be specified