2 files changed, 35 insertions, 21 deletions

diff --git a/src/security/tpm/tss/tcg-1.2/tss.c b/src/security/tpm/tss/tcg-1.2/tss.c
index 52bc2722b2..6b79aabe87 100644
--- a/src/security/tpm/tss/tcg-1.2/tss.c
+++ b/src/security/tpm/tss/tcg-1.2/tss.c
@@ -331,25 +331,20 @@ uint32_t tlcl_set_global_lock(void)
 	return tlcl_write(TPM_NV_INDEX0, NULL, 0);
 }
 
-uint32_t tlcl_extend(int pcr_num, const uint8_t *in_digest,
-		     uint8_t *out_digest)
+uint32_t tlcl_extend(int pcr_num, const uint8_t *digest_data,
+		     enum vb2_hash_algorithm digest_algo)
 {
 	struct s_tpm_extend_cmd cmd;
 	uint8_t response[kTpmResponseHeaderLength + kPcrDigestLength];
-	uint32_t result;
+
+	if (digest_algo != VB2_HASH_SHA1)
+		return TPM_E_INVALID_ARG;
 
 	memcpy(&cmd, &tpm_extend_cmd, sizeof(cmd));
 	to_tpm_uint32(cmd.buffer + tpm_extend_cmd.pcrNum, pcr_num);
-	memcpy(cmd.buffer + cmd.inDigest, in_digest, kPcrDigestLength);
-
-	result = tlcl_send_receive(cmd.buffer, response, sizeof(response));
-	if (result != TPM_SUCCESS)
-		return result;
+	memcpy(cmd.buffer + cmd.inDigest, digest_data, kPcrDigestLength);
 
-	if (out_digest)
-		memcpy(out_digest, response + kTpmResponseHeaderLength,
-		       kPcrDigestLength);
-	return result;
+	return tlcl_send_receive(cmd.buffer, response, sizeof(response));
 }
 
 uint32_t tlcl_get_permissions(uint32_t index, uint32_t *permissions)
diff --git a/src/security/tpm/tss/tcg-2.0/tss.c b/src/security/tpm/tss/tcg-2.0/tss.c
index 5d6cbf89b6..d228c7f6a8 100644
--- a/src/security/tpm/tss/tcg-2.0/tss.c
+++ b/src/security/tpm/tss/tcg-2.0/tss.c
@@ -118,21 +118,40 @@ uint32_t tlcl_assert_physical_presence(void)
 	return TPM_SUCCESS;
 }
 
-/*
- * The caller will provide the digest in a 32 byte buffer, let's consider it a
- * sha256 digest.
- */
-uint32_t tlcl_extend(int pcr_num, const uint8_t *in_digest,
-		     uint8_t *out_digest)
+static TPM_ALG_ID tpmalg_from_vb2_hash(enum vb2_hash_algorithm hash_type)
+{
+	switch (hash_type) {
+	case VB2_HASH_SHA1:
+		return TPM_ALG_SHA1;
+	case VB2_HASH_SHA256:
+		return TPM_ALG_SHA256;
+	case VB2_HASH_SHA384:
+		return TPM_ALG_SHA384;
+	case VB2_HASH_SHA512:
+		return TPM_ALG_SHA512;
+
+	default:
+		return TPM_ALG_ERROR;
+	}
+}
+
+uint32_t tlcl_extend(int pcr_num, const uint8_t *digest_data,
+		     enum vb2_hash_algorithm digest_type)
 {
 	struct tpm2_pcr_extend_cmd pcr_ext_cmd;
 	struct tpm2_response *response;
+	TPM_ALG_ID alg;
+
+	alg = tpmalg_from_vb2_hash(digest_type);
+	if (alg == TPM_ALG_ERROR)
+		return TPM_E_HASH_ERROR;
 
 	pcr_ext_cmd.pcrHandle = HR_PCR + pcr_num;
 	pcr_ext_cmd.digests.count = 1;
-	pcr_ext_cmd.digests.digests[0].hashAlg = TPM_ALG_SHA256;
-	memcpy(pcr_ext_cmd.digests.digests[0].digest.sha256, in_digest,
-	       sizeof(pcr_ext_cmd.digests.digests[0].digest.sha256));
+	pcr_ext_cmd.digests.digests[0].hashAlg = alg;
+	/* Always copying to sha512 as it's the largest one */
+	memcpy(pcr_ext_cmd.digests.digests[0].digest.sha512, digest_data,
+	       vb2_digest_size(digest_type));
 
 	response = tpm_process_command(TPM2_PCR_Extend, &pcr_ext_cmd);