From 9c6bff6c959994da280f8b3b5bad62a1600063d0 Mon Sep 17 00:00:00 2001
From: Magnus Kulke <magnuskulke@microsoft.com>
Date: Mon, 8 Jul 2024 15:15:58 +0200
Subject: [PATCH] vtpm: add extend functionality (#55)

This can be used to perform runtime measurements

Signed-off-by: Magnus Kulke <magnuskulke@microsoft.com>
---
 az-cvm-vtpm/Cargo.toml      |  2 +-
 az-cvm-vtpm/src/vtpm/mod.rs | 60 +++++++++++++++++++++++++++++++++++--
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/az-cvm-vtpm/Cargo.toml b/az-cvm-vtpm/Cargo.toml
index 9755105..fdc0edb 100644
--- a/az-cvm-vtpm/Cargo.toml
+++ b/az-cvm-vtpm/Cargo.toml
@@ -29,7 +29,7 @@ serde-big-array = "0.5.1"
 sev.workspace = true
 sha2 = "0.10.8"
 thiserror.workspace = true
-tss-esapi = "7.4"
+tss-esapi = "7.5"
 zerocopy.workspace = true
 
 [features]
diff --git a/az-cvm-vtpm/src/vtpm/mod.rs b/az-cvm-vtpm/src/vtpm/mod.rs
index 40df097..a4ab367 100644
--- a/az-cvm-vtpm/src/vtpm/mod.rs
+++ b/az-cvm-vtpm/src/vtpm/mod.rs
@@ -4,13 +4,13 @@
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use tss_esapi::abstraction::{nv, pcr, public::DecodedKey};
-use tss_esapi::handles::TpmHandle;
+use tss_esapi::handles::{PcrHandle, TpmHandle};
 use tss_esapi::interface_types::algorithm::HashingAlgorithm;
 use tss_esapi::interface_types::resource_handles::NvAuth;
 use tss_esapi::interface_types::session_handles::AuthSession;
 use tss_esapi::structures::pcr_selection_list::PcrSelectionListBuilder;
 use tss_esapi::structures::pcr_slot::PcrSlot;
-use tss_esapi::structures::{Attest, AttestInfo, Data, Signature, SignatureScheme};
+use tss_esapi::structures::{Attest, AttestInfo, Data, DigestValues, Signature, SignatureScheme};
 use tss_esapi::tcti_ldr::{DeviceConfig, TctiNameConf};
 use tss_esapi::traits::{Marshall, UnMarshall};
 use tss_esapi::Context;
@@ -50,6 +50,36 @@ const VTPM_QUOTE_PCR_SLOTS: [PcrSlot; 24] = [
     PcrSlot::Slot23,
 ];
 
+fn to_pcr_handle(pcr: u8) -> Result<PcrHandle, ExtendError> {
+    match pcr {
+        0 => Ok(PcrHandle::Pcr0),
+        1 => Ok(PcrHandle::Pcr1),
+        2 => Ok(PcrHandle::Pcr2),
+        3 => Ok(PcrHandle::Pcr3),
+        4 => Ok(PcrHandle::Pcr4),
+        5 => Ok(PcrHandle::Pcr5),
+        6 => Ok(PcrHandle::Pcr6),
+        7 => Ok(PcrHandle::Pcr7),
+        8 => Ok(PcrHandle::Pcr8),
+        9 => Ok(PcrHandle::Pcr9),
+        10 => Ok(PcrHandle::Pcr10),
+        11 => Ok(PcrHandle::Pcr11),
+        12 => Ok(PcrHandle::Pcr12),
+        13 => Ok(PcrHandle::Pcr13),
+        14 => Ok(PcrHandle::Pcr14),
+        15 => Ok(PcrHandle::Pcr15),
+        16 => Ok(PcrHandle::Pcr16),
+        17 => Ok(PcrHandle::Pcr17),
+        18 => Ok(PcrHandle::Pcr18),
+        19 => Ok(PcrHandle::Pcr19),
+        20 => Ok(PcrHandle::Pcr20),
+        21 => Ok(PcrHandle::Pcr21),
+        22 => Ok(PcrHandle::Pcr22),
+        23 => Ok(PcrHandle::Pcr23),
+        _ => Err(ExtendError::InvalidPcr),
+    }
+}
+
 #[derive(Error, Debug)]
 pub enum ReportError {
     #[error("tpm error")]
@@ -70,6 +100,32 @@ pub fn get_report() -> Result<Vec<u8>, ReportError> {
     Ok(report)
 }
 
+#[derive(Error, Debug)]
+pub enum ExtendError {
+    #[error("tpm error")]
+    Tpm(#[from] tss_esapi::Error),
+    #[error("invalid pcr number (expected 0-23)")]
+    InvalidPcr,
+}
+
+/// Extend a PCR register with a sha256 digest
+pub fn extend_pcr(pcr: u8, digest: &[u8; 32]) -> Result<(), ExtendError> {
+    let pcr_handle = to_pcr_handle(pcr)?;
+
+    let mut vals = DigestValues::new();
+    let sha256_digest = digest.to_vec().try_into()?;
+    vals.set(HashingAlgorithm::Sha256, sha256_digest);
+
+    let conf: TctiNameConf = TctiNameConf::Device(DeviceConfig::default());
+    let mut context = Context::new(conf)?;
+
+    let auth_session = AuthSession::Password;
+    context.set_sessions((Some(auth_session), None, None));
+    context.pcr_extend(pcr_handle, vals)?;
+
+    Ok(())
+}
+
 #[derive(Error, Debug)]
 pub enum AKPubError {
     #[error("tpm error")]