From 84a79ae2d7e092bc42a3fde7a5b53116fc81d4fe Mon Sep 17 00:00:00 2001
From: Ionut Mihalcea <ionut.mihalcea@arm.com>
Date: Tue, 2 Nov 2021 10:27:18 +0000
Subject: [PATCH] Refactor attesting key setup

Split up the setup for the attesting key into different methods so it's
clearer what is being done.

Signed-off-by: Ionut Mihalcea <ionut.mihalcea@arm.com>
---
 .../abstraction/transient/key_attestation.rs  | 128 +++++++++---------
 1 file changed, 65 insertions(+), 63 deletions(-)

diff --git a/tss-esapi/src/abstraction/transient/key_attestation.rs b/tss-esapi/src/abstraction/transient/key_attestation.rs
index 5bf4163b..dd5cfd45 100644
--- a/tss-esapi/src/abstraction/transient/key_attestation.rs
+++ b/tss-esapi/src/abstraction/transient/key_attestation.rs
@@ -4,10 +4,10 @@ use super::{ObjectWrapper, TransientKeyContext};
 use crate::{
     abstraction::ek,
     constants::SessionType,
-    handles::{AuthHandle, SessionHandle},
+    handles::{AuthHandle, KeyHandle, SessionHandle},
     interface_types::{
         algorithm::{AsymmetricAlgorithm, HashingAlgorithm},
-        session_handles::PolicySession,
+        session_handles::{AuthSession, PolicySession},
     },
     structures::{EncryptedSecret, IDObject, SymmetricDefinition},
     tss2_esys::{TPM2B_PUBLIC, TPMT_PUBLIC},
@@ -118,67 +118,14 @@ impl TransientKeyContext {
         let credential_blob = IDObject::try_from(credential_blob)?;
         let secret = EncryptedSecret::try_from(secret)?;
         let object_handle = self.load_key(object.params, object.material, object.auth)?;
-        let session_2;
-        let key_handle = match key {
-            None => {
-                // No key was given, use the EK. This requires using a Policy session
-                session_2 = self
-                    .context
-                    .start_auth_session(
-                        None,
-                        None,
-                        None,
-                        SessionType::Policy,
-                        SymmetricDefinition::AES_128_CFB,
-                        HashingAlgorithm::Sha256,
-                    )
-                    .or_else(|e| {
-                        self.context.flush_context(object_handle.into())?;
-                        Err(e)
-                    })?;
-                let _ = self.context.policy_secret(
-                    PolicySession::try_from(session_2.unwrap())
-                        .expect("Failed to convert auth session to policy session"),
-                    AuthHandle::Endorsement,
-                    Default::default(),
-                    Default::default(),
-                    Default::default(),
-                    None,
-                );
-                ek::create_ek_object(&mut self.context, AsymmetricAlgorithm::Rsa, None).or_else(
-                    |e| {
-                        self.context.flush_context(object_handle.into())?;
-                        self.context
-                            .flush_context(SessionHandle::from(session_2).into())?;
-                        Err(e)
-                    },
-                )?
-            }
-            Some(key) => {
-                // Load key and create a HMAC session for it
-                session_2 = self
-                    .context
-                    .start_auth_session(
-                        None,
-                        None,
-                        None,
-                        SessionType::Hmac,
-                        SymmetricDefinition::AES_128_CFB,
-                        HashingAlgorithm::Sha256,
-                    )
-                    .or_else(|e| {
-                        self.context.flush_context(object_handle.into())?;
-                        Err(e)
-                    })?;
-                self.load_key(key.params, key.material, key.auth)
-                    .or_else(|e| {
-                        self.context.flush_context(object_handle.into())?;
-                        self.context
-                            .flush_context(SessionHandle::from(session_2).into())?;
-                        Err(e)
-                    })?
-            }
-        };
+        let (key_handle, session_2) = match key {
+            Some(key) => self.prepare_key_activate_cred(key),
+            None => self.prepare_ek_activate_cred(),
+        }
+        .or_else(|e| {
+            self.context.flush_context(object_handle.into())?;
+            Err(e)
+        })?;
 
         let (session_1, _, _) = self.context.sessions();
         let credential = self
@@ -200,4 +147,59 @@ impl TransientKeyContext {
             .flush_context(SessionHandle::from(session_2).into())?;
         Ok(credential.value().to_vec())
     }
+
+    // No key was given, use the EK. This requires using a Policy session
+    fn prepare_ek_activate_cred(&mut self) -> Result<(KeyHandle, Option<AuthSession>)> {
+        let session = self.context.start_auth_session(
+            None,
+            None,
+            None,
+            SessionType::Policy,
+            SymmetricDefinition::AES_128_CFB,
+            HashingAlgorithm::Sha256,
+        )?;
+        let _ = self.context.policy_secret(
+            PolicySession::try_from(session.unwrap())
+                .expect("Failed to convert auth session to policy session"),
+            AuthHandle::Endorsement,
+            Default::default(),
+            Default::default(),
+            Default::default(),
+            None,
+        );
+        Ok((
+            ek::create_ek_object(&mut self.context, AsymmetricAlgorithm::Rsa, None).or_else(
+                |e| {
+                    self.context
+                        .flush_context(SessionHandle::from(session).into())?;
+                    Err(e)
+                },
+            )?,
+            session,
+        ))
+    }
+
+    // Load key and create a HMAC session for it
+    fn prepare_key_activate_cred(
+        &mut self,
+        key: ObjectWrapper,
+    ) -> Result<(KeyHandle, Option<AuthSession>)> {
+        let session = self.context.start_auth_session(
+            None,
+            None,
+            None,
+            SessionType::Hmac,
+            SymmetricDefinition::AES_128_CFB,
+            HashingAlgorithm::Sha256,
+        )?;
+        Ok((
+            self.load_key(key.params, key.material, key.auth)
+                .or_else(|e| {
+                    self.context
+                        .flush_context(SessionHandle::from(session).into())?;
+                    Err(e)
+                })?,
+            session,
+        ))
+    }
 }