Blocktree/crates/btlib/src/crypto/tpm.rs

use super::*;

use btserde::read_from;
use log::error;
use openssl::{
    bn::{BigNum, BigNumRef},
    nid::Nid,
    pkcs5::pbkdf2_hmac,
};
use serde::ser;
use std::{
    ffi::CStr,
    fs::OpenOptions,
    io::{BufReader, BufWriter, Read, Write},
    mem::size_of,
    ops::Deref,
    os::{raw::c_char, unix::fs::PermissionsExt},
    path::{Path, PathBuf},
    sync::{Arc, RwLock, RwLockWriteGuard},
    time::Duration,
};
use tss_esapi::{
    attributes::{object::ObjectAttributes, SessionAttributesBuilder},
    constants::{
        response_code::Tss2ResponseCode,
        session_type::SessionType,
        tss::{TPM2_PERSISTENT_FIRST, TPM2_RH_NULL},
        CapabilityType, CommandCode, Tss2ResponseCodeKind,
    },
    handles::{KeyHandle, ObjectHandle, PersistentTpmHandle, TpmHandle},
    interface_types::{
        algorithm::HashingAlgorithm,
        dynamic_handles::Persistent,
        key_bits::RsaKeyBits,
        resource_handles::{Hierarchy, Provision},
        session_handles::{AuthSession, PolicySession},
    },
    structures::{
        Auth, CapabilityData, Data, Digest, EncryptedSecret, HashScheme, HashcheckTicket, Private,
        Public, PublicKeyRsa, PublicRsaParameters, RsaDecryptionScheme, RsaScheme, SignatureScheme,
        SymmetricDefinition, SymmetricDefinitionObject, Ticket,
    },
    tcti_ldr::{TabrmdConfig, TctiNameConf},
    traits::{Marshall, UnMarshall},
    Context,
};
use tss_esapi_sys::{TPM2_CAP, TPM2_HANDLE, TPM2_MAX_CAP_BUFFER, TPMT_TK_HASHCHECK, TSS2_RC};
use zeroize::Zeroizing;

impl From<tss_esapi::Error> for Error {
    fn from(err: tss_esapi::Error) -> Self {
        match err {
            tss_esapi::Error::WrapperError(err) => Error::custom(err),
            tss_esapi::Error::Tss2Error(err) => {
                let rc = err.tss2_rc();
                let text = tss2_rc_decode(err);
                let string = format!("response code: {}, response text: {}", rc, text);
                Error::custom(string)
            }
        }
    }
}

impl<Guard> From<std::sync::PoisonError<Guard>> for Error {
    fn from(error: std::sync::PoisonError<Guard>) -> Self {
        Error::custom(error.to_string())
    }
}

/// Trait which extends the `tss_esapi::Context`.
trait ContextExt {
    fn persistent_handles(&mut self) -> Result<Vec<PersistentTpmHandle>>;

    /// Returns the first free persistent handle available in the TPM.
    fn unused_persistent_primary_key(&mut self) -> Result<Persistent>;

    fn persist_key(&mut self, key_handle: KeyHandle) -> Result<TPM2_HANDLE>;

    fn evict_key(&mut self, tpm_handle: TPM2_HANDLE, key_handle: Option<KeyHandle>) -> Result<()>;

    fn key_handle(&mut self, tpm_handle: TpmHandle) -> Result<KeyHandle>;

    fn set_encrypt_decrypt(&mut self, session: AuthSession) -> Result<()>;

    fn start_default_auth_session(&mut self) -> Result<AuthSession>;

    fn start_policy_session(&mut self, is_trial: IsTrial) -> Result<PolicySession>;

    fn dup_with_password_policy(&mut self) -> Result<Digest>;

    fn export_key<S: Scheme>(
        &mut self,
        key_pair: &KeyPair<S>,
        new_parent: KeyHandle,
        key_auth: Auth,
        policy_session: PolicySession,
        encryption_key: &AeadKey,
    ) -> Result<ExportedKeyPair<S>>;

    fn import_key<S: Scheme>(
        &mut self,
        exported: ExportedKeyPair<S>,
        new_parent: KeyHandle,
        auth: Auth,
        encryption_key: &AeadKey,
    ) -> Result<KeyPair<S>>;

    fn set_auth(&mut self, cred_data: &CredData, password: &str) -> Result<()>;
}

impl ContextExt for Context {
    fn persistent_handles(&mut self) -> Result<Vec<PersistentTpmHandle>> {
        let capability = CapabilityType::Handles;
        let property = TPM2_PERSISTENT_FIRST;
        let max = usize::try_from(TPM2_MAX_CAP_BUFFER).unwrap();
        // This calculation was copied from tpm2_getcap.
        let count = (max - size_of::<TPM2_CAP>() - size_of::<u32>()) / size_of::<TPM2_HANDLE>();
        let property_count = u32::try_from(count).unwrap();
        let mut all_handles = Vec::new();
        let more = false;
        for _ in 0..255 {
            let (handles, more) = self.get_capability(capability, property, property_count)?;
            let list = match handles {
                CapabilityData::Handles(list) => list.into_inner(),
                _ => {
                    return Err(Error::custom(format!(
                        "Unexpected capability type returned by TPM: {:?}",
                        handles
                    )))
                }
            };
            all_handles.reserve(list.len());
            for handle in list {
                all_handles.push(PersistentTpmHandle::try_from(handle)?);
            }
            if !more {
                break;
            }
        }
        if more {
            return Err(Error::custom(
                "hit iteration limit before retrieving all persistent handles from the TPM",
            ));
        }
        all_handles.sort_unstable_by_key(|handle| TPM2_HANDLE::from(*handle));
        Ok(all_handles)
    }

    fn unused_persistent_primary_key(&mut self) -> Result<Persistent> {
        let mut used_handles = self.persistent_handles()?.into_iter();
        // These address regions are defined in Registry of Reserved TPM 2.0 Handles and Localities
        // The first regions for both are for primary keys and the second is marked as "Available".
        let storage = (0x81_00_00_00..0x81_00_01_00u32).chain(0x81_00_80_00..0x81_01_00_00u32);
        let endorsement = (0x81_01_00_00..0x81_01_01_00u32).chain(0x81_01_80_00..0x81_02_00_00u32);
        let possible_handles = storage
            .chain(endorsement)
            .map(|handle| PersistentTpmHandle::new(handle).unwrap());
        let mut unused: Option<PersistentTpmHandle> = None;
        // We simultaneously iterate over the possible handles and the used handles, breaking when
        // we find a handle which is not equal to the current used handle, or when
        // there are no more used handles. This works because both `used_handles` and
        // `possible_handles` are sorted.
        for handle in possible_handles {
            let used = match used_handles.next() {
                Some(used) => used,
                None => {
                    unused = Some(handle);
                    break;
                }
            };
            if used != handle {
                unused = Some(handle);
                break;
            }
        }
        match unused {
            Some(unused) => Ok(Persistent::Persistent(unused)),
            None => Err(Error::custom("failed to find an unused persistent handle")),
        }
    }

    fn persist_key(&mut self, key_handle: KeyHandle) -> Result<TPM2_HANDLE> {
        let mut persistent: Option<Persistent> = None;
        for _ in 0..255 {
            let handle =
                self.execute_with_session(None, |ctx| ctx.unused_persistent_primary_key())?;
            match self.evict_control(Provision::Owner, key_handle.into(), handle) {
                Ok(_) => {
                    persistent = Some(handle);
                    break;
                }
                Err(error) => {
                    if let tss_esapi::Error::Tss2Error(rc) = error {
                        if Some(Tss2ResponseCodeKind::NvDefined) == rc.kind() {
                            // This type of error occurs if another application used the persistent
                            // handle we found before we could, so we try again with a different
                            // handle.
                            continue;
                        }
                    }
                    return Err(error.into());
                }
            }
        }
        match persistent {
            Some(Persistent::Persistent(inner)) => Ok(inner.into()),
            None => Err(Error::custom("retry limit reached")),
        }
    }

    fn evict_key(&mut self, persistent: TPM2_HANDLE, key_handle: Option<KeyHandle>) -> Result<()> {
        let tpm_handle = TpmHandle::try_from(persistent)?;
        let key_handle = match key_handle {
            Some(key_handle) => key_handle,
            None => self.tr_from_tpm_public(tpm_handle)?.into(),
        };
        let persistent = Persistent::Persistent(tpm_handle.try_into()?);
        self.evict_control(Provision::Owner, key_handle.into(), persistent)?;
        Ok(())
    }

    fn set_encrypt_decrypt(&mut self, session: AuthSession) -> Result<()> {
        let (attributes, mask) = SessionAttributesBuilder::new()
            .with_decrypt(true)
            .with_encrypt(true)
            .build();
        self.tr_sess_set_attributes(session, attributes, mask)?;
        Ok(())
    }

    fn start_default_auth_session(&mut self) -> Result<AuthSession> {
        let session = self
            .start_auth_session(
                None,
                None,
                None,
                SessionType::Hmac,
                SymmetricDefinition::AES_256_CFB,
                HashingAlgorithm::Sha256,
            )?
            .ok_or_else(|| Error::custom("empty session handle received from TPM"))?;
        self.set_encrypt_decrypt(session)?;
        Ok(session)
    }

    fn start_policy_session(&mut self, is_trial: IsTrial) -> Result<PolicySession> {
        let session = self
            .start_auth_session(
                None,
                None,
                None,
                is_trial.into(),
                SymmetricDefinition::AES_256_CFB,
                HashingAlgorithm::Sha256,
            )?
            .ok_or_else(|| Error::custom("empty session handle received from TPM"))?;
        self.set_encrypt_decrypt(session)?;
        Ok(session.try_into()?)
    }

    /// Loads the public information from the persistent object stored in the TPM referred to by
    /// `tpm_handle` and returns the `KeyHandle` for this object. If the TPM handle refers to an
    /// object which is not actually a key, then an error is returned.
    fn key_handle(&mut self, tpm_handle: TpmHandle) -> Result<KeyHandle> {
        let obj_handle =
            self.execute_with_session(None, |ctx| ctx.tr_from_tpm_public(tpm_handle))?;
        Ok(obj_handle.into())
    }

    /// Returns the digest for the policy which allows a key to be duplicated provided that the
    /// caller supplies the correct password (authValue) for the key being duplicated.
    fn dup_with_password_policy(&mut self) -> Result<Digest> {
        let policy_session = self.start_policy_session(IsTrial::True)?;
        self.execute_with_session(None, |ctx| {
            ctx.policy_password(policy_session)?;
            ctx.policy_command_code(policy_session, CommandCode::Duplicate)?;
            Ok(ctx.policy_get_digest(policy_session)?)
        })
    }

    fn export_key<S: Scheme>(
        &mut self,
        key_pair: &KeyPair<S>,
        new_parent: KeyHandle,
        key_auth: Auth,
        policy_session: PolicySession,
        encryption_key: &AeadKey,
    ) -> Result<ExportedKeyPair<S>> {
        let (public, ..) = self.read_public(key_pair.private)?;
        let result: Result<()> = self.execute_with_session(None, |ctx| {
            ctx.policy_password(policy_session)?;
            ctx.policy_command_code(policy_session, CommandCode::Duplicate)?;
            Ok(())
        });
        result?;
        let result: Result<(tss_esapi::structures::Private, EncryptedSecret)> = self
            .execute_with_session(Some(policy_session.into()), |ctx| {
                let obj_handle = ObjectHandle::from(key_pair.private);
                ctx.tr_set_auth(obj_handle, key_auth)?;
                let (_, private, secret) = ctx.duplicate(
                    obj_handle,
                    new_parent.into(),
                    None,
                    SymmetricDefinitionObject::Null,
                )?;
                Ok((private, secret))
            });
        let (private, secret) = result?;
        let tagged_ct =
            encryption_key.encrypt(PublicWrapper(public), &TpmBlobs { private, secret })?;
        Ok(ExportedKeyPair {
            scheme: key_pair.public.scheme,
            tagged_ct,
        })
    }

    fn import_key<S: Scheme>(
        &mut self,
        exported: ExportedKeyPair<S>,
        new_parent: KeyHandle,
        auth: Auth,
        encryption_key: &AeadKey,
    ) -> Result<KeyPair<S>> {
        let blobs = encryption_key.decrypt(&exported.tagged_ct)?;
        let public = exported.tagged_ct.aad;
        let private = self.import(
            new_parent.into(),
            None,
            public.clone(),
            blobs.private,
            blobs.secret,
            SymmetricDefinitionObject::Null,
        )?;
        let key_handle = self.load(new_parent, private, public.clone())?;
        self.tr_set_auth(key_handle.into(), auth)?;
        let public = AsymKeyPub::try_from(public.0, exported.scheme)?;
        Ok(KeyPair {
            public,
            private: key_handle,
        })
    }

    fn set_auth(&mut self, creds: &CredData, password: &str) -> Result<()> {
        let auth = Auth::try_from(password.as_bytes())?;
        self.tr_set_auth(creds.sign.private.into(), auth.clone())?;
        self.tr_set_auth(creds.enc.private.into(), auth)?;
        Ok(())
    }
}

enum IsTrial {
    True,
    False,
}

impl From<IsTrial> for SessionType {
    fn from(is_trial: IsTrial) -> Self {
        match is_trial {
            IsTrial::True => SessionType::Trial,
            IsTrial::False => SessionType::Policy,
        }
    }
}

trait DigestExt {
    fn empty() -> Digest;
}

impl DigestExt for Digest {
    fn empty() -> Digest {
        let empty = [0u8; 0];
        Digest::try_from(empty.as_slice()).unwrap()
    }
}

/// An authentication cookie which grants access to a node's credentials.
#[derive(Serialize, Deserialize, Clone)]
struct Cookie(#[serde(with = "BigArray")] [u8; Self::LEN]);

impl Cookie {
    const LEN: usize = TpmCredStore::SIGN_SCHEME.key_len_const() as usize;

    fn random() -> Result<Cookie> {
        Ok(Cookie(rand_array()?))
    }

    fn as_slice(&self) -> &[u8] {
        self.0.as_slice()
    }

    /// Returns the `Auth` value associated with this cookie.
    fn auth(&self) -> Auth {
        // This shouldn't fail because the given slice is only 64 bytes long.
        Auth::try_from(&self.as_slice()[..64]).unwrap()
    }
}

#[derive(Serialize, Deserialize)]
struct PublicWrapper(
    #[serde(serialize_with = "serialize_marshall")]
    #[serde(deserialize_with = "deserialize_unmarshall")]
    Public,
);

impl Deref for PublicWrapper {
    type Target = Public;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

fn serialize_marshall<S: Serializer, T: Marshall>(
    value: &T,
    ser: S,
) -> std::result::Result<S::Ok, S::Error> {
    let vec = value.marshall().map_err(ser::Error::custom)?;
    vec.as_slice().serialize(ser)
}

fn deserialize_unmarshall<'de, D: Deserializer<'de>, T: UnMarshall>(
    de: D,
) -> std::result::Result<T, D::Error> {
    let vec: Vec<u8> = Deserialize::deserialize(de)?;
    T::unmarshall(&vec).map_err(de::Error::custom)
}

#[derive(Serialize, Deserialize)]
struct TpmBlobs {
    #[serde(serialize_with = "serialize_as_vec")]
    #[serde(deserialize_with = "deserialize_as_vec")]
    private: Private,
    #[serde(serialize_with = "serialize_as_vec")]
    #[serde(deserialize_with = "deserialize_as_vec")]
    secret: EncryptedSecret,
}

fn serialize_as_vec<S: Serializer, T: Deref<Target = Vec<u8>>>(
    value: &T,
    ser: S,
) -> std::result::Result<S::Ok, S::Error> {
    value.deref().serialize(ser)
}

fn deserialize_as_vec<'de, D: Deserializer<'de>, T: TryFrom<Vec<u8>>>(
    de: D,
) -> std::result::Result<T, D::Error> {
    T::try_from(Vec::deserialize(de)?)
        .map_err(|_| de::Error::custom("Unable to convert from vector"))
}

#[derive(Serialize, Deserialize)]
pub struct ExportedKeyPair<S> {
    scheme: S,
    tagged_ct: TaggedCiphertext<TpmBlobs, PublicWrapper>,
}

#[derive(Serialize, Deserialize)]
pub struct ExportedCreds {
    sign: ExportedKeyPair<Sign>,
    enc: ExportedKeyPair<Encrypt>,
    params: DerivationParams,
    writecap: Option<Writecap>,
}

/// A public/private key pair in a form which can be serialized and deserialized.
#[derive(Serialize, Clone)]
struct StoredKeyPair<S: Scheme> {
    public: AsymKeyPub<S>,
    private: TPM2_HANDLE,
}

impl<'de, S: Scheme> Deserialize<'de> for StoredKeyPair<S> {
    fn deserialize<D: Deserializer<'de>>(d: D) -> std::result::Result<Self, D::Error> {
        const FIELDS: &[&str] = &["public", "private", "writecap"];

        struct StructVisitor<S: Scheme>(PhantomData<S>);

        impl<'de, S: Scheme> Visitor<'de> for StructVisitor<S> {
            type Value = StoredKeyPair<S>;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_fmt(format_args!("struct {}", stringify!(StoredKeyPair)))
            }

            fn visit_seq<A: SeqAccess<'de>>(
                self,
                mut seq: A,
            ) -> std::result::Result<Self::Value, A::Error> {
                let public: AsymKeyPub<S> = seq
                    .next_element()?
                    .ok_or_else(|| de::Error::missing_field(FIELDS[0]))?;
                let private: TPM2_HANDLE = seq
                    .next_element()?
                    .ok_or_else(|| de::Error::missing_field(FIELDS[1]))?;
                let pair = StoredKeyPair { public, private };
                Ok(pair)
            }
        }

        d.deserialize_struct(
            stringify!(StoredKeyPair),
            FIELDS,
            StructVisitor(PhantomData),
        )
    }
}

/// Credentials in a form which can be serialized and deserialized.
#[derive(Serialize, Deserialize, Clone)]
struct StoredCredData {
    sign: StoredKeyPair<Sign>,
    enc: StoredKeyPair<Encrypt>,
    writecap: Option<Writecap>,
}

impl StoredCredData {
    fn to_cred_data(&self, context: &mut Context) -> Result<CredData> {
        let sign = KeyPair::from_stored(context, &self.sign)?;
        let enc = KeyPair::from_stored(context, &self.enc)?;
        Ok(CredData {
            sign,
            enc,
            writecap: self.writecap.clone(),
        })
    }
}

#[derive(Serialize, Deserialize)]
struct Storage {
    cookie: Cookie,
    node: Option<StoredCredData>,
    root: Option<StoredCredData>,
    storage_key: Option<StoredKeyPair<Encrypt>>,
}

impl Storage {
    fn new(cookie: Cookie) -> Storage {
        Storage {
            cookie,
            node: None,
            root: None,
            storage_key: None,
        }
    }

    fn save<W: Write>(&self, to: &mut W) -> Result<()> {
        Ok(write_to(self, to)?)
    }

    fn load<R: Read>(from: &mut R) -> Result<Storage> {
        Ok(read_from(from)?)
    }

    fn init<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        let file = OpenOptions::new()
            .write(true)
            .create_new(true)
            .open(&path)?;
        // Only allow access from the current user.
        let metadata = file.metadata()?;
        let mut permissions = metadata.permissions();
        permissions.set_mode(0o600);
        file.set_permissions(permissions)?;
        let mut reader = BufWriter::new(file);
        self.save(&mut reader)?;
        reader.flush()?;
        Ok(())
    }

    fn load_or_init<P: AsRef<Path>>(path: P) -> Result<Storage> {
        match OpenOptions::new().read(true).open(&path) {
            Ok(file) => {
                let mut reader = BufReader::new(file);
                Self::load(&mut reader)
            }
            Err(error) => {
                if std::io::ErrorKind::NotFound != error.kind() {
                    return Err(error.into());
                }
                let state = Self::new(Cookie::random()?);
                state.init(path)?;
                Ok(state)
            }
        }
    }
}

impl From<HashKind> for HashingAlgorithm {
    fn from(kind: HashKind) -> Self {
        match kind {
            HashKind::Sha2_256 => HashingAlgorithm::Sha256,
            HashKind::Sha2_512 => HashingAlgorithm::Sha512,
        }
    }
}

impl TryInto<RsaScheme> for SchemeKind {
    type Error = Error;
    fn try_into(self) -> Result<RsaScheme> {
        match self {
            SchemeKind::Sign(sign) => match sign {
                Sign::RsaSsaPss(inner) => {
                    Ok(RsaScheme::RsaPss(HashScheme::new(inner.hash_kind.into())))
                }
            },
            SchemeKind::Encrypt(encrypt) => match encrypt {
                Encrypt::RsaEsOaep(inner) => {
                    Ok(RsaScheme::Oaep(HashScheme::new(inner.hash_kind.into())))
                }
            },
        }
    }
}

impl TryFrom<KeyLen> for RsaKeyBits {
    type Error = Error;

    fn try_from(len: KeyLen) -> Result<RsaKeyBits> {
        match len {
            KeyLen::Bits2048 => Ok(RsaKeyBits::Rsa2048),
            KeyLen::Bits3072 => Ok(RsaKeyBits::Rsa3072),
            KeyLen::Bits4096 => Ok(RsaKeyBits::Rsa4096),
            _ => Err(Error::custom(format!("unsupported key len: {}", len))),
        }
    }
}

enum Parent {
    Seed(Hierarchy),
    Key(KeyHandle),
}

struct KeyBuilder<'a, S: Scheme> {
    scheme: S,
    allow_dup: bool,
    unique: &'a [u8],
    auth_value: Option<Auth>,
    name_hash: HashingAlgorithm,
    policy_digest: Digest,
    parent: Parent,
    restricted: bool,
    symmetric: SymmetricDefinitionObject,
    rsa_exponent: u32,
}

impl<'a, S: Scheme> KeyBuilder<'a, S> {
    fn new(scheme: S, unique: &'a [u8]) -> KeyBuilder<'a, S> {
        KeyBuilder {
            scheme,
            allow_dup: false,
            unique,
            auth_value: None,
            name_hash: HashingAlgorithm::Sha256,
            policy_digest: Digest::empty(),
            parent: Parent::Seed(Hierarchy::Owner),
            restricted: false,
            symmetric: SymmetricDefinitionObject::Null,
            rsa_exponent: Rsa::EXP,
        }
    }

    fn with_allow_dup(mut self, allow_dup: bool) -> Self {
        self.allow_dup = allow_dup;
        self
    }

    fn with_auth(mut self, auth: Auth) -> Self {
        self.auth_value = Some(auth);
        self
    }

    fn with_policy_digest(mut self, policy_digest: Digest) -> Self {
        self.policy_digest = policy_digest;
        self
    }

    fn with_parent(mut self, parent: Parent) -> Self {
        self.parent = parent;
        self
    }

    fn with_restricted(mut self, restricted: bool) -> Self {
        self.restricted = restricted;
        self
    }

    fn with_symmetric(mut self, symmetric: SymmetricDefinitionObject) -> Self {
        self.symmetric = symmetric;
        self
    }

    fn with_rsa_exponent(mut self, rsa_exponent: u32) -> Self {
        self.rsa_exponent = rsa_exponent;
        self
    }

    fn rsa_template(&self) -> Result<Public> {
        let scheme_enum = self.scheme.as_enum();
        let decrypt = match scheme_enum {
            SchemeKind::Sign(_) => false,
            SchemeKind::Encrypt(_) => true,
        };
        let object_attributes = ObjectAttributes::builder()
            .with_fixed_tpm(!self.allow_dup)
            .with_fixed_parent(!self.allow_dup)
            .with_sensitive_data_origin(true)
            .with_user_with_auth(true)
            .with_decrypt(decrypt)
            .with_sign_encrypt(!decrypt)
            .with_restricted(self.restricted)
            .build()?;
        let name_hashing_algorithm = self.name_hash;
        let parameters = PublicRsaParameters::new(
            self.symmetric,
            if self.restricted {
                RsaScheme::Null
            } else {
                scheme_enum.try_into()?
            },
            self.scheme.key_len().try_into()?,
            self.rsa_exponent.try_into()?,
        );
        let unique = PublicKeyRsa::try_from(self.unique)?;
        let public = Public::Rsa {
            object_attributes,
            name_hashing_algorithm,
            auth_policy: self.policy_digest.clone(),
            parameters,
            unique,
        };
        Ok(public)
    }

    fn template(&self) -> Result<Public> {
        match self.scheme.as_enum() {
            SchemeKind::Encrypt(encrypt) => match encrypt {
                Encrypt::RsaEsOaep(_) => self.rsa_template(),
            },
            SchemeKind::Sign(sign) => match sign {
                Sign::RsaSsaPss(_) => self.rsa_template(),
            },
        }
    }

    fn build(self, context: &mut Context) -> Result<KeyPair<S>> {
        let (public, private) = match self.parent {
            Parent::Seed(hierarchy) => {
                let result = context.create_primary(
                    hierarchy,
                    self.template()?,
                    self.auth_value,
                    None,
                    None,
                    None,
                )?;
                (result.out_public, result.key_handle)
            }
            Parent::Key(parent) => {
                let result =
                    context.create(parent, self.template()?, self.auth_value, None, None, None)?;
                let private =
                    context.load(parent, result.out_private, result.out_public.clone())?;
                (result.out_public, private)
            }
        };
        let public = AsymKey::try_from(public, self.scheme)?;
        Ok(KeyPair { public, private })
    }
}

impl<'a> KeyBuilder<'a, Encrypt> {
    fn for_storage_key(scheme: Encrypt, unique: &'a [u8]) -> KeyBuilder<'a, Encrypt> {
        KeyBuilder::new(scheme, unique)
            .with_allow_dup(false)
            .with_restricted(true)
            .with_symmetric(SymmetricDefinitionObject::AES_128_CFB)
    }
}

/// A public/private key pair.
#[derive(Clone)]
struct KeyPair<S: Scheme> {
    public: AsymKeyPub<S>,
    /// A rust struct which wraps an `ESYS_TR` from the ESAPI.
    private: KeyHandle,
}

impl<S: Scheme> KeyPair<S> {
    fn from_stored(context: &mut Context, stored: &StoredKeyPair<S>) -> Result<KeyPair<S>> {
        let tpm_handle = TpmHandle::try_from(stored.private)?;
        let key_handle = context.key_handle(tpm_handle)?;
        Ok(KeyPair {
            public: stored.public.clone(),
            private: key_handle,
        })
    }

    fn to_stored(&self, private: TPM2_HANDLE) -> StoredKeyPair<S> {
        let public = self.public.clone();
        StoredKeyPair { public, private }
    }
}

/// Credentials which can be used for signing and encrypting.
struct CredData {
    sign: KeyPair<Sign>,
    enc: KeyPair<Encrypt>,
    writecap: Option<Writecap>,
}

struct State {
    context: Context,
    storage: Storage,
    node_creds: Option<TpmCreds>,
    storage_key: Option<KeyPair<Encrypt>>,
}

impl State {
    fn new(mut context: Context, storage: Storage) -> Result<State> {
        let storage_key = match &storage.storage_key {
            Some(storage_key) => Some(KeyPair::from_stored(&mut context, storage_key)?),
            None => None,
        };
        Ok(State {
            context,
            storage,
            node_creds: None,
            storage_key,
        })
    }

    fn init_node_creds(&mut self, state: &Arc<RwLock<State>>) -> Result<()> {
        let tpm_handles = match &self.storage.node {
            Some(handles) => handles,
            None => return Ok(()),
        };
        let key_handles = tpm_handles.to_cred_data(&mut self.context)?;
        let auth = self.storage.cookie.auth();
        self.context
            .tr_set_auth(key_handles.enc.private.into(), auth.clone())?;
        self.context
            .tr_set_auth(key_handles.sign.private.into(), auth)?;
        self.node_creds = Some(TpmCreds::new(key_handles, state));
        Ok(())
    }
}

pub struct TpmCredStore {
    state: Arc<RwLock<State>>,
    storage_path: PathBuf,
    cookie: Cookie,
}

impl TpmCredStore {
    const SIGN_SCHEME: Sign = Sign::RSA_PSS_2048_SHA_256;
    const ENCRYPT_SCHEME: Encrypt = Encrypt::RSA_OAEP_2048_SHA_256;
    const DEFAULT_WRITECAP_EXP: Duration = Duration::from_secs(60 * 60 * 24 * 365 * 10);

    pub fn new<P: AsRef<Path>>(mut context: Context, state_path: P) -> Result<TpmCredStore> {
        let storage = Storage::load_or_init(state_path.as_ref())?;
        let session = context.start_default_auth_session()?;
        context.set_sessions((Some(session), None, None));
        let cookie = storage.cookie.clone();
        let state = Arc::new(RwLock::new(State::new(context, storage)?));
        {
            let mut guard = state.write()?;
            guard.init_node_creds(&state)?;
        }
        Ok(TpmCredStore {
            state,
            storage_path: state_path.as_ref().to_owned(),
            cookie,
        })
    }

    fn save_storage(&self, guard: &mut RwLockWriteGuard<State>) -> Result<()> {
        let file = OpenOptions::new().write(true).open(&self.storage_path)?;
        let mut writer = BufWriter::new(file);
        guard.storage.save(&mut writer)?;
        writer.flush()?;
        Ok(())
    }

    fn persist<F: FnOnce(&mut Storage, StoredCredData)>(
        &self,
        creds: &TpmCreds,
        update_storage: F,
    ) -> Result<()> {
        let mut guard = self.state.write()?;
        let sign_handle = guard.context.persist_key(creds.sign.private)?;
        let enc_handle = match guard.context.persist_key(creds.enc.private) {
            Ok(handle) => handle,
            Err(error) => {
                guard
                    .context
                    .evict_key(sign_handle, Some(creds.sign.private))?;
                return Err(error);
            }
        };
        let handles = StoredCredData {
            sign: creds.sign.to_stored(sign_handle),
            enc: creds.enc.to_stored(enc_handle),
            writecap: creds.writecap.clone(),
        };
        update_storage(&mut guard.storage, handles);
        match self.save_storage(&mut guard) {
            Ok(_) => Ok(()),
            Err(error) => {
                let result = guard
                    .context
                    .evict_key(sign_handle, Some(creds.sign.private));
                if let Err(error) = result {
                    error!("failed to evict signing key due to error: {:?}", error)
                }
                let result = guard.context.evict_key(enc_handle, Some(creds.enc.private));
                if let Err(error) = result {
                    error!("failed to evict encryption key due to error: {:?}", error)
                }
                Err(error)
            }
        }
    }

    fn get_or_init_storage_key(&self) -> Result<KeyPair<Encrypt>> {
        {
            let guard = self.state.read()?;
            if let Some(storage_key) = &guard.storage_key {
                // We take this path if the storage key was generated and loaded.
                return Ok(storage_key.clone());
            }
        }
        {
            let mut guard = self.state.write()?;
            if let Some(storage_key) = guard.storage.storage_key.take() {
                let result = KeyPair::from_stored(&mut guard.context, &storage_key);
                guard.storage.storage_key = Some(storage_key);
                // We take this path if the storage key was generated but not loaded.
                return result;
            }
        }
        let params = KeyBuilder::for_storage_key(Self::ENCRYPT_SCHEME, self.cookie.as_slice())
            .with_parent(Parent::Seed(Hierarchy::Owner))
            .with_auth(self.cookie.auth());
        let mut guard = self.state.write()?;
        let storage_key = params.build(&mut guard.context)?;
        guard.storage_key = Some(storage_key.clone());
        let tpm_handle = guard.context.persist_key(storage_key.private)?;
        guard.storage.storage_key = Some(storage_key.to_stored(tpm_handle));
        if let Err(err) = self.save_storage(&mut guard) {
            guard.storage_key = None;
            guard.storage.storage_key = None;
            guard
                .context
                .evict_key(tpm_handle, Some(storage_key.private))?;
            return Err(err);
        }
        Ok(storage_key)
    }

    fn gen_key<S: Scheme>(&self, params: KeyBuilder<S>) -> Result<KeyPair<S>> {
        let mut guard = self.state.write()?;
        params.build(&mut guard.context)
    }

    fn gen_node_sign_key(&self) -> Result<KeyPair<Sign>> {
        let params = KeyBuilder::new(Self::SIGN_SCHEME, self.cookie.as_slice())
            .with_parent(Parent::Seed(Hierarchy::Owner))
            .with_allow_dup(false)
            .with_auth(self.cookie.auth());
        self.gen_key(params)
    }

    fn gen_node_enc_key(&self) -> Result<KeyPair<Encrypt>> {
        let params = KeyBuilder::new(Self::ENCRYPT_SCHEME, self.cookie.as_slice())
            .with_parent(Parent::Seed(Hierarchy::Owner))
            .with_allow_dup(false)
            .with_auth(self.cookie.auth());
        self.gen_key(params)
    }

    fn gen_node_creds(&self) -> Result<TpmCreds> {
        let sign = self.gen_node_sign_key()?;
        let enc = self.gen_node_enc_key()?;
        let cred_data = CredData {
            sign,
            enc,
            writecap: None,
        };
        let creds = TpmCreds::new(cred_data, &self.state);
        self.persist(&creds, |storage, handles| storage.node = Some(handles))?;
        Ok(creds)
    }

    fn gen_root_sign_key(&self, password: &str, policy: Digest) -> Result<KeyPair<Sign>> {
        let scheme = Self::SIGN_SCHEME;
        let unique = rand_vec(scheme.key_len() as usize)?;
        let storage_key = self.get_or_init_storage_key()?;
        let params = KeyBuilder::new(scheme, unique.as_slice())
            .with_parent(Parent::Key(storage_key.private))
            .with_allow_dup(true)
            .with_auth(password.as_bytes().try_into()?)
            .with_policy_digest(policy);
        self.gen_key(params)
    }

    fn gen_root_enc_key(&self, password: &str, policy: Digest) -> Result<KeyPair<Encrypt>> {
        let scheme = Self::ENCRYPT_SCHEME;
        let unique = rand_vec(scheme.key_len() as usize)?;
        let storage_key = self.get_or_init_storage_key()?;
        let params = KeyBuilder::new(scheme, unique.as_slice())
            .with_parent(Parent::Key(storage_key.private))
            .with_allow_dup(true)
            .with_auth(password.as_bytes().try_into()?)
            .with_policy_digest(policy);
        self.gen_key(params)
    }
}

#[derive(Serialize, Deserialize)]
struct DerivationParams {
    iter: usize,
    hash: HashKind,
    kind: AeadKeyKind,
    salt: Vec<u8>,
    iv: Vec<u8>,
}

impl DerivationParams {
    const PBKDF2_ITER: usize = 1000000;
    const PBKDF2_HASH: HashKind = HashKind::Sha2_256;
    const EXPORT_KEY_KIND: AeadKeyKind = AeadKeyKind::AesGcm256;

    fn new() -> Result<DerivationParams> {
        const_assert!(
            DerivationParams::PBKDF2_HASH.len() == DerivationParams::EXPORT_KEY_KIND.key_len()
        );
        Ok(DerivationParams {
            iter: Self::PBKDF2_ITER,
            hash: Self::PBKDF2_HASH,
            kind: Self::EXPORT_KEY_KIND,
            salt: rand_vec(Self::PBKDF2_HASH.len())?,
            iv: rand_vec(Self::EXPORT_KEY_KIND.iv_len())?,
        })
    }

    fn derive_key(&self, password: &str) -> Result<AeadKey> {
        let mut key = Zeroizing::new([0u8; Self::EXPORT_KEY_KIND.key_len()]);
        pbkdf2_hmac(
            password.as_bytes(),
            self.salt.as_slice(),
            self.iter,
            self.hash.into(),
            key.as_mut_slice(),
        )?;
        AeadKey::copy_components(self.kind, key.as_slice(), &self.iv)
    }
}

impl CredStore for TpmCredStore {
    type CredHandle = TpmCreds;
    type ExportedCreds = ExportedCreds;

    fn node_creds(&self) -> Result<TpmCreds> {
        {
            let guard = self.state.read()?;
            if let Some(creds) = &guard.node_creds {
                return Ok(creds.clone());
            }
        }
        self.gen_node_creds()
    }

    fn root_creds(&self, password: &str) -> Result<Self::CredHandle> {
        let root_handles = {
            let guard = self.state.read()?;
            guard
                .storage
                .root
                .as_ref()
                .ok_or_else(|| Error::custom("root creds have not yet been generated"))?
                .clone()
        };
        let mut guard = self.state.write()?;
        let key_handles = root_handles.to_cred_data(&mut guard.context)?;
        guard.context.set_auth(&key_handles, password)?;
        Ok(TpmCreds::new(key_handles, &self.state))
    }

    fn gen_root_creds(&self, password: &str) -> Result<Self::CredHandle> {
        {
            let guard = self.state.read()?;
            if guard.storage.root.is_some() {
                return Err(Error::custom("root creds have already been generated"));
            }
        }
        let policy = {
            let mut guard = self.state.write()?;
            guard.context.dup_with_password_policy()?
        };
        let sign = self.gen_root_sign_key(password, policy.clone())?;
        let enc = self.gen_root_enc_key(password, policy)?;
        let cred_data = CredData {
            sign,
            enc,
            writecap: None,
        };
        {
            let mut guard = self.state.write()?;
            guard.context.set_auth(&cred_data, password)?;
        }
        let mut creds = TpmCreds::new(cred_data, &self.state);
        creds.init_root_writecap(Epoch::now() + Self::DEFAULT_WRITECAP_EXP)?;
        self.persist(&creds, |storage, handles| storage.root = Some(handles))?;
        Ok(creds)
    }

    fn storage_key(&self) -> Result<AsymKeyPub<Encrypt>> {
        let pair = self.get_or_init_storage_key()?;
        Ok(pair.public)
    }

    fn export_root_creds(
        &self,
        root_creds: &TpmCreds,
        password: &str,
        new_parent: &AsymKeyPub<Encrypt>,
    ) -> Result<ExportedCreds> {
        let params = DerivationParams::new()?;
        let aead_key = params.derive_key(password)?;
        let new_parent = new_parent.storage_key_public()?;
        let mut guard = self.state.write()?;
        if let Some(storage_key) = guard.storage_key.take() {
            // Save memory by flushing the storage key from the TPM's RAM.
            guard.context.flush_context(storage_key.private.into())?;
        }
        let new_parent_handle = guard
            .context
            .load_external_public(new_parent, Hierarchy::Null)?;
        let policy_session = guard.context.start_policy_session(IsTrial::False)?;
        let auth = Auth::try_from(password.as_bytes())?;
        let sign = guard.context.export_key(
            &root_creds.sign,
            new_parent_handle,
            auth.clone(),
            policy_session,
            &aead_key,
        )?;
        let enc = guard.context.export_key(
            &root_creds.enc,
            new_parent_handle,
            auth,
            policy_session,
            &aead_key,
        )?;
        Ok(ExportedCreds {
            sign,
            enc,
            params,
            writecap: root_creds.writecap.clone(),
        })
    }

    fn import_root_creds(&self, password: &str, exported: ExportedCreds) -> Result<TpmCreds> {
        let aead_key = exported.params.derive_key(password)?;
        let auth = Auth::try_from(password.as_bytes())?;
        let storage_key = self.get_or_init_storage_key()?;
        let creds = {
            let mut guard = self.state.write()?;
            let sign = guard.context.import_key(
                exported.sign,
                storage_key.private,
                auth.clone(),
                &aead_key,
            )?;
            let enc =
                guard
                    .context
                    .import_key(exported.enc, storage_key.private, auth, &aead_key)?;
            let cred_data = CredData {
                sign,
                enc,
                writecap: exported.writecap,
            };
            TpmCreds::new(cred_data, &self.state)
        };
        self.persist(&creds, |storage, handles| storage.root = Some(handles))?;
        Ok(creds)
    }
}

impl<S: Scheme> AsymKeyPub<S> {
    fn try_from(public: Public, scheme: S) -> Result<AsymKeyPub<S>> {
        match public {
            Public::Rsa {
                parameters, unique, ..
            } => {
                let exponent_value = parameters.exponent().value();
                let exponent = BigNum::from_u32(exponent_value)?;
                let modulus = BigNum::from_slice(unique.as_slice())?;
                let rsa = OsslRsa::from_public_components(modulus, exponent)?;
                let pkey = PKey::from_rsa(rsa)?.conv_pub();
                Ok(AsymKey { pkey, scheme })
            }
            _ => Err(Error::custom("Unsupported key type returned by TPM")),
        }
    }
}

impl AsymKeyPub<Encrypt> {
    fn storage_key_public(&self) -> Result<Public> {
        fn from_rsa(scheme: Encrypt, rsa: openssl::rsa::Rsa<super::Public>) -> Result<Public> {
            let exponent = rsa.e().try_into_u32()?;
            let modulus = rsa.n().to_vec();
            let builder = KeyBuilder::for_storage_key(scheme, &modulus).with_rsa_exponent(exponent);
            builder.rsa_template()
        }
        match self.scheme {
            Encrypt::RsaEsOaep(_) => from_rsa(self.scheme, self.pkey.rsa()?),
        }
    }
}

trait NidExt {
    fn sha3_256() -> Nid {
        Nid::from_raw(1097)
    }

    fn sha3_384() -> Nid {
        Nid::from_raw(1098)
    }

    fn sha3_512() -> Nid {
        Nid::from_raw(1099)
    }
}

impl NidExt for Nid {}

trait BigNumRefExt {
    fn try_into_u32(self) -> Result<u32>;
}

impl BigNumRefExt for &BigNumRef {
    fn try_into_u32(self) -> Result<u32> {
        let data = self.to_vec();
        if data.len() > 4 {
            return Err(Error::custom(format!("data was too long: {}", data.len())));
        }
        let mut buf = [0u8; 4];
        // Note that BigNum data is stored in big endian format, so padding zeros go at the
        // beginning of the buffer.
        let subslice = &mut buf[4 - data.len()..];
        subslice.copy_from_slice(data.as_slice());
        let int = u32::from_be_bytes(buf);
        Ok(int)
    }
}

trait MessageDigestExt {
    fn hash_algo(&self) -> Result<HashingAlgorithm>;
    fn hash_scheme(&self) -> Result<HashScheme> {
        Ok(HashScheme::new(self.hash_algo()?))
    }
}

impl MessageDigestExt for MessageDigest {
    fn hash_algo(&self) -> Result<HashingAlgorithm> {
        let nid = self.type_();
        let algo = if Nid::SHA1 == nid {
            HashingAlgorithm::Sha1
        } else if Nid::SHA256 == nid {
            HashingAlgorithm::Sha256
        } else if Nid::SHA384 == nid {
            HashingAlgorithm::Sha384
        } else if Nid::SHA512 == nid {
            HashingAlgorithm::Sha512
        } else if Nid::sha3_256() == nid {
            HashingAlgorithm::Sha3_256
        } else if Nid::sha3_384() == nid {
            HashingAlgorithm::Sha3_384
        } else if Nid::sha3_512() == nid {
            HashingAlgorithm::Sha3_512
        } else {
            return Err(Error::custom(format!(
                "Unsupported hash algorithm with NID: {:?}",
                nid
            )));
        };
        Ok(algo)
    }
}

trait HashcheckTicketExt {
    fn null() -> HashcheckTicket;
}

impl HashcheckTicketExt for HashcheckTicket {
    /// Returns the NULL Ticket of the hashcheck type, as defined in part 1 of the TPM spec,
    /// clause 4.47.
    fn null() -> HashcheckTicket {
        TPMT_TK_HASHCHECK {
            tag: HashcheckTicket::POSSIBLE_TAGS[0].into(),
            digest: Default::default(),
            hierarchy: TPM2_RH_NULL,
        }
        .try_into()
        .unwrap()
    }
}

#[derive(Clone)]
pub struct TpmCreds {
    state: Arc<RwLock<State>>,
    sign: KeyPair<Sign>,
    enc: KeyPair<Encrypt>,
    writecap: Option<Writecap>,
}

impl TpmCreds {
    fn new(key_handles: CredData, state: &Arc<RwLock<State>>) -> TpmCreds {
        TpmCreds {
            sign: key_handles.sign,
            enc: key_handles.enc,
            state: state.clone(),
            writecap: key_handles.writecap,
        }
    }

    fn init_root_writecap(&mut self, expires: Epoch) -> Result<()> {
        let writecap = self.issue_writecap(self.principal(), Vec::new(), expires)?;
        self.writecap = Some(writecap);
        Ok(())
    }
}

impl Principaled for TpmCreds {
    fn principal_of_kind(&self, kind: HashKind) -> Principal {
        self.sign.public.principal_of_kind(kind)
    }
}

impl Verifier for TpmCreds {
    type Op<'v> = OsslVerifyOp<'v>;

    fn init_verify<'v>(&'v self) -> Result<Self::Op<'v>> {
        self.sign.public.init_verify()
    }

    fn verify<'a, I: Iterator<Item = &'a [u8]>>(&self, parts: I, signature: &[u8]) -> Result<()> {
        self.sign.public.verify(parts, signature)
    }
}

impl Encrypter for TpmCreds {
    fn encrypt(&self, slice: &[u8]) -> Result<Vec<u8>> {
        self.enc.public.encrypt(slice)
    }
}

impl CredsPub for TpmCreds {
    fn public_sign(&self) -> &AsymKeyPub<Sign> {
        &self.sign.public
    }
}

pub struct TpmSignOp<'a> {
    op: VarHashOp,
    creds: &'a TpmCreds,
}

impl<'a> Op for TpmSignOp<'a> {
    type Arg = &'a TpmCreds;

    fn init(arg: Self::Arg) -> Result<Self> {
        let op = VarHashOp::init(arg.sign.public.scheme.hash_kind())?;
        Ok(TpmSignOp { op, creds: arg })
    }

    fn update(&mut self, data: &[u8]) -> Result<()> {
        self.op.update(data)
    }

    fn finish_into(self, buf: &mut [u8]) -> Result<usize> {
        let hash = self.op.finish()?;
        let digest = Digest::try_from(hash.as_ref())?;

        let validation = HashcheckTicket::null();
        // Null means the scheme used during key creation will be used.
        let scheme = SignatureScheme::Null;

        let sig = {
            let mut guard = self.creds.state.write()?;
            guard
                .context
                .sign(self.creds.sign.private, digest, scheme, validation)?
        };
        let slice: &[u8] = match &sig {
            tss_esapi::structures::Signature::RsaSsa(inner) => inner.signature(),
            tss_esapi::structures::Signature::RsaPss(inner) => inner.signature(),
            _ => {
                return Err(Error::custom(format!(
                    "Unexpected signature type: {:?}",
                    sig.algorithm()
                )))
            }
        };

        if buf.len() < slice.len() {
            return Err(Error::IncorrectSize { expected: slice.len(), actual: buf.len() });
        }
        buf.copy_from_slice(slice);
        Ok(slice.len())
    }
}

impl<'a> SignOp for TpmSignOp<'a> {
    fn scheme(&self) -> Sign {
        self.creds.sign.public.scheme
    }
}

impl Signer for TpmCreds {
    type Op<'s> = TpmSignOp<'s>;

    fn init_sign<'s>(&'s self) -> Result<Self::Op<'s>> {
        TpmSignOp::init(self)
    }

    fn sign<'a, I: Iterator<Item = &'a [u8]>>(&self, parts: I) -> Result<Signature> {
        let mut op = self.init_sign()?;
        for part in parts {
            op.update(part)?;
        }
        op.finish()
    }
}

impl Decrypter for TpmCreds {
    fn decrypt(&self, slice: &[u8]) -> Result<Vec<u8>> {
        let cipher_text = PublicKeyRsa::try_from(slice)?;
        // Null means the scheme used during key creation will be used.
        let in_scheme = RsaDecryptionScheme::Null;
        let empty = [0u8; 0];
        let label = Data::try_from(empty.as_slice())?;
        let plain_text = {
            let mut guard = self.state.write()?;
            guard
                .context
                .rsa_decrypt(self.enc.private, cipher_text, in_scheme, label)?
        };
        Ok(Vec::from(plain_text.value()))
    }
}

impl CredsPriv for TpmCreds {
    fn writecap(&self) -> Option<&Writecap> {
        self.writecap.as_ref()
    }

    fn set_writecap(&mut self, writecap: Writecap) {
        self.writecap = Some(writecap)
    }
}

impl Creds for TpmCreds {}

trait TctiNameConfExt {
    fn default() -> TctiNameConf;
}

impl TctiNameConfExt for TctiNameConf {
    /// Returns a configuration which specifies that Tabrmd should be connected to using the system
    /// DBus.
    fn default() -> TctiNameConf {
        TctiNameConf::Tabrmd(TabrmdConfig::default())
    }
}

#[link(name = "tss2-rc")]
extern "C" {
    fn Tss2_RC_Decode(rc: TSS2_RC) -> *const c_char;
}

/// Interface for types which can be converted to TSS2 response codes.
trait HasResponseCode {
    /// Returns the TSS2 response code associated with this instance.
    fn tss2_rc(&self) -> TSS2_RC;
}

/// Returns the error message associated with the given TSS2 response code.
fn tss2_rc_decode<E: HasResponseCode>(err: E) -> &'static str {
    let c_str = unsafe {
        let ptr = Tss2_RC_Decode(err.tss2_rc());
        CStr::from_ptr(ptr)
    };
    // We're relying on Tss2_RC_Decode to return valid C strings.
    c_str.to_str().unwrap()
}

impl HasResponseCode for Tss2ResponseCode {
    fn tss2_rc(&self) -> TSS2_RC {
        match self {
            Tss2ResponseCode::Success => 0,
            Tss2ResponseCode::FormatZero(code) => code.0,
            Tss2ResponseCode::FormatOne(code) => code.0,
        }
    }
}

impl HasResponseCode for TSS2_RC {
    fn tss2_rc(&self) -> TSS2_RC {
        *self
    }
}

#[cfg(test)]
mod test {
    use crate::test_helpers::{BtCursor, SwtpmHarness};

    use super::*;

    use ctor::ctor;
    use std::{fs::File, io::SeekFrom};
    use tss_esapi::{
        interface_types::ecc::EccCurve,
        structures::{EccPoint, EccScheme, KeyDerivationFunctionScheme, PublicEccParameters},
    };

    #[ctor]
    fn ctor() {
        env_logger::init();
    }

    /// Displays the message associated with a TSS2 return code.
    #[test]
    fn print_error_message() {
        const RC: TSS2_RC = 2461;
        let msg = tss2_rc_decode(RC);
        println!("{}", msg);
    }

    #[test]
    fn create_context() {
        let harness = SwtpmHarness::new().unwrap();
        harness.context().unwrap().self_test(true).unwrap();
    }

    #[test]
    fn create_primary_key() {
        let harness = SwtpmHarness::new().unwrap();
        let mut context = harness.context().unwrap();

        let public = {
            let object_attributes = ObjectAttributes::builder()
                .with_fixed_tpm(true)
                .with_fixed_parent(true)
                .with_sensitive_data_origin(true)
                .with_user_with_auth(true)
                .with_decrypt(false)
                .with_sign_encrypt(true)
                .with_restricted(false)
                .build()
                .expect("ObjectAttributesBuilder failed");
            let name_hashing_algorithm = HashingAlgorithm::Sha256;
            let empty = [0u8; 0];
            let auth_policy = Digest::try_from(empty.as_slice()).unwrap();
            let parameters = PublicEccParameters::new(
                SymmetricDefinitionObject::Null,
                EccScheme::EcDsa(HashScheme::new(HashingAlgorithm::Sha256)),
                EccCurve::NistP256,
                KeyDerivationFunctionScheme::Null,
            );
            Public::Ecc {
                object_attributes,
                name_hashing_algorithm,
                auth_policy,
                parameters,
                unique: EccPoint::default(),
            }
        };

        let session = context
            .start_auth_session(
                None,
                None,
                None,
                SessionType::Hmac,
                SymmetricDefinition::AES_256_CFB,
                HashingAlgorithm::Sha256,
            )
            .expect("Failed to create session")
            .expect("Received invalid handle");

        context.execute_with_session(Some(session), |ctx| {
            let primary = ctx
                .create_primary(Hierarchy::Null, public, None, None, None, None)
                .expect("create_primary failed")
                .key_handle;

            ctx.flush_context(primary.into())
                .expect("flush_context failed");
        });
    }

    /// Tests that a TPM Credential Store can be created when a cookie does not already exist.
    #[test]
    fn tpm_cred_store_new() -> Result<()> {
        let harness = SwtpmHarness::new()?;
        let cookie_path = harness.dir_path().join("cookie.bin");
        let store = TpmCredStore::new(harness.context()?, &cookie_path)?;
        let cookie = File::open(&cookie_path)?;
        let metadata = cookie.metadata()?;
        let actual = metadata.permissions().mode();
        // Assert that the cookie can only be read by its owner.
        assert_eq!(0o600, 0o777 & actual);
        drop(store);
        Ok(())
    }

    #[test]
    fn gen_creds() -> Result<()> {
        let harness = SwtpmHarness::new()?;
        let cookie_path = harness.dir_path().join("cookie.bin");
        let store = TpmCredStore::new(harness.context()?, &cookie_path)?;
        store.gen_node_creds()?;
        Ok(())
    }

    /// Displays the numeric identifiers used by the supported hash algorithms.
    //#[test]
    #[allow(dead_code)]
    fn show_nids() {
        fn show_nid(digest: MessageDigest) {
            let nid = digest.type_();
            println!("{}: {:?}", nid.long_name().unwrap(), nid);
        }

        show_nid(MessageDigest::sha1());
        show_nid(MessageDigest::sha256());
        show_nid(MessageDigest::sha384());
        show_nid(MessageDigest::sha512());
        show_nid(MessageDigest::sha3_256());
        show_nid(MessageDigest::sha3_384());
        show_nid(MessageDigest::sha3_512());
    }

    /// Verifies that the NIDs returned by the supported hash algorithms are as expected.
    #[test]
    fn verify_expected_nids() {
        fn assert_eq(digest: MessageDigest, nid: Nid) {
            assert_eq!(digest.type_(), nid);
        }

        assert_eq(MessageDigest::sha1(), Nid::SHA1);
        assert_eq(MessageDigest::sha256(), Nid::SHA256);
        assert_eq(MessageDigest::sha384(), Nid::SHA384);
        assert_eq(MessageDigest::sha512(), Nid::SHA512);
        assert_eq(MessageDigest::sha3_256(), Nid::sha3_256());
        assert_eq(MessageDigest::sha3_384(), Nid::sha3_384());
        assert_eq(MessageDigest::sha3_512(), Nid::sha3_512());
    }

    /// Returns a SwtpmHarness and a TpmCredStore that uses it. Note that the order of the entries
    /// in the returned tuple is significant, as TpmCredStore must be dropped _before_ SwtpmHarness.
    fn test_store() -> Result<(SwtpmHarness, TpmCredStore)> {
        let harness = SwtpmHarness::new()?;
        let store = TpmCredStore::new(harness.context()?, &harness.state_path())?;
        Ok((harness, store))
    }

    fn sign_verify_test(creds: &TpmCreds) -> Result<()> {
        let data: [u8; 1024] = rand_array()?;
        let parts = [data.as_slice()];
        let sig = creds.sign(parts.into_iter())?;

        creds.verify(parts.into_iter(), sig.as_slice())
    }

    #[test]
    fn tpm_sign_verify() -> Result<()> {
        let (_harness, store) = test_store()?;
        let creds = store.gen_node_creds()?;
        sign_verify_test(&creds)
    }

    fn encrypt_decrypt_test(creds: &TpmCreds) -> Result<()> {
        let expected: [u8; Cookie::LEN / 2] = rand_array()?;
        let ct = creds.encrypt(expected.as_slice())?;
        let actual = creds.decrypt(&ct)?;
        if expected.as_slice() == actual {
            Ok(())
        } else {
            Err(Error::custom("decrypted data did not match input"))
        }
    }

    #[test]
    fn tpm_encrypt_decrypt() -> Result<()> {
        let (_harness, store) = test_store()?;
        let creds = store.gen_node_creds()?;
        encrypt_decrypt_test(&creds)
    }

    /// Tests that `HashcheckTicket::null` doesn't panic.
    #[test]
    fn hashcheck_null() {
        HashcheckTicket::null();
    }

    #[test]
    fn persistent_handles() -> Result<()> {
        let harness = SwtpmHarness::new()?;
        let mut context = harness.context()?;
        context.persistent_handles()?;
        Ok(())
    }

    #[test]
    fn first_free_persistent() -> Result<()> {
        let harness = SwtpmHarness::new()?;
        let mut context = harness.context()?;
        context.unused_persistent_primary_key()?;
        Ok(())
    }

    #[test]
    fn persist_key() -> Result<()> {
        let (_harness, store) = test_store()?;
        let cookie = Cookie::random()?;
        let params = KeyBuilder::new(Sign::RSA_PSS_3072_SHA_256, cookie.as_slice());
        let pair = store.gen_key(params)?;
        let mut guard = store.state.write()?;
        guard.context.persist_key(pair.private)?;
        Ok(())
    }

    /// Tests that the same node key is returned by after a cred store is dropped and recreated.
    #[test]
    fn node_key_persisted() -> Result<()> {
        let (harness, store) = test_store()?;
        let expected = {
            let creds = store.node_creds()?;
            creds.principal()
        };
        drop(store);
        let store = TpmCredStore::new(harness.context()?, harness.state_path())?;
        let creds = store.node_creds()?;
        let actual = creds.principal();
        assert_eq!(expected, actual);
        sign_verify_test(&creds)?;
        encrypt_decrypt_test(&creds)?;
        Ok(())
    }

    #[test]
    fn root_key_can_be_used_after_generation() {
        let (_harness, store) = test_store().expect("failed to make test store");
        let creds = store
            .gen_root_creds(&"TimeInvariant")
            .expect("failed to gen root creds");
        let data = [1u8; 32];
        creds
            .sign(std::iter::once(data.as_slice()))
            .expect("sign failed");
    }

    #[test]
    fn root_and_node_keys_generated() {
        let (_harness, store) = test_store().expect("failed to make test store");
        let _root_creds = store
            .gen_root_creds(&"TranslationInvariant")
            .expect("failed to gen root creds");
        let _node_creds = store.node_creds().expect("failed to gen node creds");
    }

    #[test]
    fn verify_root_writecap() {
        let (_harness, store) = test_store().expect("failed to make test store");
        let root_creds = store
            .gen_root_creds(&"TranslationInvariant")
            .expect("failed to gen root creds");
        let writecap = root_creds.writecap().expect("no root writecap was present");
        let path = crate::BlockPath::new(root_creds.principal(), Vec::new());
        writecap
            .assert_valid_for(&path)
            .expect("failed to verify root writecap");
    }

    #[test]
    fn issue_writecap_to_node() {
        let (_harness, store) = test_store().expect("failed to make test store");
        let root_creds = store
            .gen_root_creds(&"TranslationInvariant")
            .expect("failed to gen root creds");
        let path = crate::BlockPath::new(
            root_creds.principal(),
            vec!["apps".to_string(), "comms".to_string()],
        );
        let node_creds = store.node_creds().expect("failed to gen node creds");
        let writecap = root_creds
            .issue_writecap(
                node_creds.principal(),
                path.components().map(|e| e.to_string()).collect(),
                Epoch::now() + Duration::from_secs(3600),
            )
            .expect("failed to issue writecap");
        writecap
            .assert_valid_for(&path)
            .expect("failed to verify writecap");
    }

    #[test]
    fn root_key_persisted() -> Result<()> {
        const PASSWORD: &str = "Scaramouch";
        let (harness, store) = test_store()?;
        let expected = {
            let creds = store.gen_root_creds(PASSWORD)?;
            creds.principal()
        };
        drop(store);
        let store = TpmCredStore::new(harness.context()?, harness.state_path())?;
        let creds = store.root_creds(PASSWORD)?;
        let actual = creds.principal();
        assert_eq!(expected, actual);
        sign_verify_test(&creds)?;
        encrypt_decrypt_test(&creds)?;
        Ok(())
    }

    #[test]
    fn root_key_unusable_when_password_wrong() -> Result<()> {
        let (harness, store) = test_store()?;
        store.gen_root_creds("Galileo")?;
        drop(store);
        let store = TpmCredStore::new(harness.context()?, harness.state_path())?;
        let creds = store.root_creds("Figaro")?;
        assert!(sign_verify_test(&creds).is_err());
        assert!(encrypt_decrypt_test(&creds).is_err());
        Ok(())
    }

    #[test]
    fn root_key_export_import() {
        const PASSWORD: &str = "Frobinate";
        let (_src_harness, src_store) = test_store().unwrap();
        let src_root_creds = src_store.gen_root_creds(PASSWORD).unwrap();

        let (_dest_harness, dest_store) = test_store().unwrap();
        let dest_storage_key = dest_store.storage_key().unwrap();

        let exported = src_store
            .export_root_creds(&src_root_creds, PASSWORD, &dest_storage_key)
            .unwrap();
        let vec = to_vec(&exported).unwrap();
        let mut slice = vec.as_slice();
        let exported = read_from(&mut slice).unwrap();
        let dest_root_creds = dest_store.import_root_creds(PASSWORD, exported).unwrap();

        let message = rand_vec(TpmCredStore::ENCRYPT_SCHEME.key_len() as usize / 2).unwrap();
        let sig = dest_root_creds
            .sign([message.as_slice()].into_iter())
            .unwrap();
        src_root_creds
            .verify([message.as_slice()].into_iter(), sig.as_slice())
            .unwrap();

        let ct = src_root_creds.encrypt(message.as_slice()).unwrap();
        let pt = dest_root_creds.decrypt(ct.as_slice()).unwrap();
        assert_eq!(message.as_slice(), pt.as_slice());
    }

    /// Tests that a signature computed using a `TpmSignOp` can be verified.
    #[test]
    fn sign_write_can_be_verified() {
        const LEN: usize = 512;
        let cursor = BtCursor::new([0u8; LEN]);
        let (_harness, store) = test_store().unwrap();
        let creds = store.gen_node_creds().expect("gen_node_creds failed");
        let sign_op = creds.init_sign().expect("init_sign failed");
        let mut sign_wrap = SignWrite::new(cursor, sign_op);
        let part_values = (1..9u8).map(|k| [k; LEN / 8]).collect::<Vec<_>>();
        let get_parts = || part_values.iter().map(|arr| arr.as_slice());

        for part in get_parts() {
            sign_wrap.write(part).expect("write failed");
        }
        let (sig, ..) = sign_wrap.finish().expect("finish failed");

        creds
            .verify(get_parts(), sig.as_ref())
            .expect("verify failed");
    }

    /// Tests that data written into a `SignWrite` using a `TpmSignOp` and later read from a
    /// `VerifyRead` using a `OsslVerifyOp` can be successfully verified.
    #[test]
    fn sign_write_then_verify_read() {
        use std::io::Seek;
        const LEN: usize = 512;
        let cursor = BtCursor::new([0u8; LEN]);
        let (_harness, store) = test_store().unwrap();
        let creds = store.gen_node_creds().expect("gen_node_creds failed");
        let sign_op = creds.init_sign().expect("init_sign failed");
        let mut sign_wrap = SignWrite::new(cursor, sign_op);

        for part in (1..9u8).map(|k| [k; LEN / 8]) {
            sign_wrap.write(part.as_slice()).expect("write failed");
        }
        let (sig, mut cursor) = sign_wrap.finish().expect("finish failed");

        cursor.seek(SeekFrom::Start(0)).expect("seek failed");

        let verify_op = creds.init_verify().expect("init_verify failed");
        let mut verify_read = VerifyRead::new(cursor, verify_op);
        let mut buf = Vec::with_capacity(LEN);
        verify_read
            .read_to_end(&mut buf)
            .expect("read_to_end failed");

        verify_read
            .finish(sig.as_ref())
            .expect("failed to verify signature");
    }

    /// This test is broken for unknown reasons. It passes if no inner encryption key is supplied.
    /// To work-around this issue I've chosen to wrap the data returned by `duplicate` with a
    /// symmetric key in software.
    //#[test]
    #[allow(dead_code)]
    fn key_export_import() -> Result<()> {
        let auth = Auth::try_from(vec![0u8; 32])?;

        let src_harness = SwtpmHarness::new()?;
        let mut src_ctx = src_harness.context()?;
        {
            let session = src_ctx.start_default_auth_session()?;
            src_ctx.set_sessions((Some(session), None, None));
        }

        let src_storage = {
            let unique = [0u8; 0];
            KeyBuilder::for_storage_key(Encrypt::RSA_OAEP_2048_SHA_256, unique.as_slice())
                .with_parent(Parent::Seed(Hierarchy::Owner))
                .with_auth(auth.clone())
                .build(&mut src_ctx)?
                .private
        };

        let dest_harness = SwtpmHarness::new()?;
        let mut dest_ctx = dest_harness.context()?;
        {
            let session = dest_ctx.start_default_auth_session()?;
            dest_ctx.set_sessions((Some(session), None, None));
        }

        let dest_storage = {
            let unique = [0u8; 0];
            KeyBuilder::for_storage_key(Encrypt::RSA_OAEP_2048_SHA_256, unique.as_slice())
                .with_parent(Parent::Seed(Hierarchy::Owner))
                .with_auth(auth.clone())
                .build(&mut dest_ctx)?
                .private
        };

        let key_handle = {
            let policy = src_ctx.dup_with_password_policy()?;
            let unique = rand_array::<256>()?;
            KeyBuilder::new(Sign::RSA_PSS_2048_SHA_256, unique.as_slice())
                .with_parent(Parent::Key(src_storage))
                .with_allow_dup(true)
                .with_policy_digest(policy)
                .build(&mut src_ctx)?
                .private
        };

        let new_parent = {
            let (public, ..) = dest_ctx.read_public(dest_storage)?;
            src_ctx.load_external_public(public, Hierarchy::Owner)?
        };

        let (public, ..) = src_ctx.read_public(key_handle)?;

        let encryption_key = Data::try_from(vec![7u8; 16])?;

        let (_, private, secret) = {
            let session = src_ctx.start_policy_session(IsTrial::False)?;
            let result: Result<()> = src_ctx.execute_with_session(None, |ctx| {
                ctx.policy_password(session)?;
                ctx.policy_command_code(session, CommandCode::Duplicate)?;
                Ok(())
            });
            result?;
            src_ctx
                .execute_with_session(Some(session.into()), |ctx| {
                    ctx.duplicate(
                        key_handle.into(),
                        new_parent.into(),
                        Some(encryption_key.clone()),
                        SymmetricDefinitionObject::AES_128_CFB,
                    )
                })
                .unwrap()
        };

        dest_ctx
            .import(
                dest_storage.into(),
                Some(encryption_key),
                public,
                private,
                secret,
                SymmetricDefinitionObject::AES_128_CFB,
            )
            .unwrap();
        Ok(())
    }
}