tpm2-device 0.5.6

// SPDX-License-Identifier: GPL-3-0-or-later
// Copyright (c) 2025 Opinsys Oy
// Copyright (c) 2024-2025 Jarkko Sakkinen

#![deny(clippy::all)]
#![deny(clippy::pedantic)]

use nix::{
    fcntl,
    poll::{poll, PollFd, PollFlags},
};
use rand::{thread_rng, RngCore};
use std::{
    cell::RefCell,
    fs::{File, OpenOptions},
    io::{Read, Write},
    os::fd::{AsFd, AsRawFd},
    path::{Path, PathBuf},
    rc::Rc,
    time::{Duration, Instant},
};

use thiserror::Error;
use tpm2_crypto::TpmHash;
use tpm2_protocol::{
    basic::{TpmHandle, TpmUint32},
    constant::{MAX_HANDLES, TPM_MAX_COMMAND_SIZE},
    data::{
        Tpm2bEncryptedSecret, Tpm2bName, Tpm2bNonce, TpmAlgId, TpmCap, TpmCc, TpmEccCurve, TpmHt,
        TpmPt, TpmRc, TpmRcBase, TpmRh, TpmSe, TpmSt, TpmaSession, TpmsAlgProperty,
        TpmsAuthCommand, TpmsCapabilityData, TpmsContext, TpmsPcrSelect, TpmsPcrSelection,
        TpmtPublic, TpmtSymDefObject, TpmuCapabilities,
    },
    frame::{
        tpm_marshal_command, tpm_unmarshal_response, TpmAuthCommands, TpmAuthResponses, TpmCommand,
        TpmContextLoadCommand, TpmContextSaveCommand, TpmFlushContextCommand, TpmFrame,
        TpmGetCapabilityCommand, TpmGetCapabilityResponse, TpmReadPublicCommand, TpmResponse,
        TpmStartAuthSessionCommand,
    },
    TpmWriter,
};
use tracing::{debug, trace};

/// Errors that can occur when talking to a TPM device.
#[derive(Debug, Error)]
pub enum TpmDeviceError {
    /// The TPM device is already mutably borrowed.
    #[error("device is already borrowed")]
    AlreadyBorrowed,

    /// The requested capability is not available from the TPM.
    #[error("capability not found: {0}")]
    CapabilityMissing(TpmCap),

    #[error("operation interrupted by user")]
    Interrupted,

    /// An invalid command code was used.
    #[error("invalid CC: {0}")]
    InvalidCc(tpm2_protocol::data::TpmCc),

    /// The TPM returned an invalid or malformed response.
    #[error("invalid response")]
    InvalidResponse,

    /// An I/O error occurred when accessing the TPM device.
    #[error("I/O: {0}")]
    Io(#[from] std::io::Error),

    /// Marshaling a TPM protocol encoded object failed.
    #[error("marshal: {0}")]
    Marshal(tpm2_protocol::TpmProtocolError),

    /// No TPM device is available.
    #[error("device not available")]
    NotAvailable,

    /// The requested operation could not be completed.
    #[error("operation failed")]
    OperationFailed,

    /// No PCR banks are available on the TPM.
    #[error("PCR banks not available")]
    PcrBanksNotAvailable,

    /// The PCR selection masks differ between active banks.
    #[error("PCR bank selection mismatch")]
    PcrBankSelectionMismatch,

    /// The TPM response did not match the expected command code.
    #[error("response mismatch: {0}")]
    ResponseMismatch(TpmCc),

    /// The TPM command timed out.
    #[error("TPM command timed out")]
    Timeout,

    /// The TPM returned an error code.
    #[error("TPM return code: {0}")]
    TpmRc(TpmRc),

    /// Trailing data after the response.
    #[error("trailing data")]
    TrailingData,

    /// Unmarshaling a TPM protocol encoded object failed.
    #[error("unmarshal: {0}")]
    Unmarshal(tpm2_protocol::TpmProtocolError),

    /// An unexpected end-of-file was encountered.
    #[error("unexpected EOF")]
    UnexpectedEof,
}

impl From<TpmRc> for TpmDeviceError {
    fn from(rc: TpmRc) -> Self {
        Self::TpmRc(rc)
    }
}

impl From<nix::Error> for TpmDeviceError {
    fn from(err: nix::Error) -> Self {
        Self::Io(std::io::Error::from_raw_os_error(err as i32))
    }
}

/// Executes a closure with a mutable reference to a `TpmDevice`.
///
/// This helper function centralizes the boilerplate for safely acquiring a
/// mutable borrow of a `TpmDevice` from the shared `Rc<RefCell<...>>`.
///
/// # Errors
///
/// Returns [`NotAvailable`](crate::TpmDeviceError::NotAvailable) when no device
/// is present.
/// Returns [`AlreadyBorrowed`](crate::TpmDeviceError::AlreadyBorrowed) when the
/// device is already mutably borrowed.
/// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants depending
/// on function.
pub fn with_device<F, T, E>(device: Option<Rc<RefCell<TpmDevice>>>, function: F) -> Result<T, E>
where
    F: FnOnce(&mut TpmDevice) -> Result<T, E>,
    E: From<TpmDeviceError>,
{
    let device_rc = device.ok_or(TpmDeviceError::NotAvailable)?;
    let mut device_guard = device_rc
        .try_borrow_mut()
        .map_err(|_| TpmDeviceError::AlreadyBorrowed)?;
    function(&mut device_guard)
}

/// A builder for constructing a `TpmDevice`.
pub struct TpmDeviceBuilder {
    path: PathBuf,
    timeout: Duration,
    interrupted: Box<dyn Fn() -> bool>,
}

impl Default for TpmDeviceBuilder {
    fn default() -> Self {
        Self {
            path: PathBuf::from("/dev/tpmrm0"),
            timeout: Duration::from_secs(120),
            interrupted: Box::new(|| false),
        }
    }
}

impl TpmDeviceBuilder {
    /// Sets the device file path.
    #[must_use]
    pub fn with_path<P: AsRef<Path>>(mut self, path: P) -> Self {
        self.path = path.as_ref().to_path_buf();
        self
    }

    /// Sets the operation timeout.
    #[must_use]
    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        self
    }

    /// Sets the interruption check callback.
    #[must_use]
    pub fn with_interrupted<F>(mut self, handler: F) -> Self
    where
        F: Fn() -> bool + 'static,
    {
        self.interrupted = Box::new(handler);
        self
    }

    /// Opens the TPM device file and constructs the `TpmDevice`.
    ///
    /// # Errors
    ///
    /// Returns [`Io`](crate::TpmDeviceError::Io) when the device file cannot be
    /// opened or when configuring the file descriptor flags fails.
    pub fn build(self) -> Result<TpmDevice, TpmDeviceError> {
        let file = OpenOptions::new()
            .read(true)
            .write(true)
            .open(&self.path)
            .map_err(TpmDeviceError::Io)?;

        let fd = file.as_raw_fd();
        let flags = fcntl::fcntl(fd, fcntl::FcntlArg::F_GETFL)?;
        let mut oflags = fcntl::OFlag::from_bits_truncate(flags);
        oflags.insert(fcntl::OFlag::O_NONBLOCK);
        fcntl::fcntl(fd, fcntl::FcntlArg::F_SETFL(oflags))?;

        Ok(TpmDevice {
            file,
            interrupted: self.interrupted,
            timeout: self.timeout,
            command: Vec::with_capacity(TPM_MAX_COMMAND_SIZE),
            response: Vec::with_capacity(TPM_MAX_COMMAND_SIZE),
        })
    }
}

pub struct TpmDevice {
    file: File,
    interrupted: Box<dyn Fn() -> bool>,
    timeout: Duration,
    command: Vec<u8>,
    response: Vec<u8>,
}

impl std::fmt::Debug for TpmDevice {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Device")
            .field("file", &self.file)
            .field("timeout", &self.timeout)
            .finish_non_exhaustive()
    }
}

impl TpmDevice {
    const NO_SESSIONS: &'static [TpmsAuthCommand] = &[];

    /// Creates a new builder for `TpmDevice`.
    #[must_use]
    pub fn builder() -> TpmDeviceBuilder {
        TpmDeviceBuilder::default()
    }

    fn receive(&mut self, buf: &mut [u8]) -> Result<usize, TpmDeviceError> {
        let fd = self.file.as_fd();
        let mut fds = [PollFd::new(fd, PollFlags::POLLIN)];

        let num_events = match poll(&mut fds, 100u16) {
            Ok(num) => num,
            Err(nix::Error::EINTR) => return Ok(0),
            Err(e) => return Err(e.into()),
        };

        if num_events == 0 {
            return Ok(0);
        }

        let revents = fds[0].revents().unwrap_or(PollFlags::empty());

        if revents.intersects(PollFlags::POLLERR | PollFlags::POLLNVAL) {
            return Err(TpmDeviceError::UnexpectedEof);
        }

        if revents.contains(PollFlags::POLLIN) {
            match self.file.read(buf) {
                Ok(0) => Err(TpmDeviceError::UnexpectedEof),
                Ok(n) => Ok(n),
                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => Ok(0),
                Err(e) if e.kind() == std::io::ErrorKind::Interrupted => Ok(0),
                Err(e) => Err(e.into()),
            }
        } else if revents.contains(PollFlags::POLLHUP) {
            Err(TpmDeviceError::UnexpectedEof)
        } else {
            Ok(0)
        }
    }

    /// Performs the whole TPM command transmission process.
    ///
    /// # Errors
    ///
    /// Returns [`Interrupted`](crate::TpmDeviceError::Interrupted) when the
    /// interrupt callback requests cancellation.
    /// Returns [`Io`](crate::TpmDeviceError::Io) when a write, flush, or read
    /// operation on the device file fails, or when polling the device file
    /// descriptor fails.
    /// Returns [`Marshal`](crate::TpmDeviceError::Marshal) when marshal
    /// operation on TPM protocol compliant data fails.
    /// Returns [`Timeout`](crate::TpmDeviceError::Timeout) when the TPM does
    /// not respond within the configured timeout.
    /// Returns [`TpmRc`](crate::TpmDeviceError::TpmRc) when the TPM returns an
    /// error code.
    /// Returns [`Unmarshal`](crate::TpmDeviceError::Unmarshal) when unmarshal
    /// operation on TPM protocol compliant data fails.
    pub fn transmit<C: TpmFrame>(
        &mut self,
        command: &C,
        sessions: &[TpmsAuthCommand],
    ) -> Result<(TpmResponse, TpmAuthResponses), TpmDeviceError> {
        self.prepare_command(command, sessions)?;
        let cc = command.cc();

        self.file.write_all(&self.command)?;
        self.file.flush()?;

        let start_time = Instant::now();
        self.response.clear();
        let mut total_size: Option<usize> = None;
        let mut temp_buf = [0u8; 1024];

        loop {
            if (self.interrupted)() {
                return Err(TpmDeviceError::Interrupted);
            }
            if start_time.elapsed() > self.timeout {
                return Err(TpmDeviceError::Timeout);
            }

            let n = self.receive(&mut temp_buf)?;
            if n > 0 {
                self.response.extend_from_slice(&temp_buf[..n]);
            }

            if total_size.is_none() && self.response.len() >= 10 {
                let Ok(size_bytes): Result<[u8; 4], _> = self.response[2..6].try_into() else {
                    return Err(TpmDeviceError::OperationFailed);
                };
                let size = u32::from_be_bytes(size_bytes) as usize;
                if !(10..={ TPM_MAX_COMMAND_SIZE }).contains(&size) {
                    return Err(TpmDeviceError::OperationFailed);
                }
                total_size = Some(size);
            }

            if let Some(size) = total_size {
                if self.response.len() == size {
                    break;
                }
                if self.response.len() > size {
                    return Err(TpmDeviceError::TrailingData);
                }
            }
        }

        let result = tpm_unmarshal_response(cc, &self.response).map_err(TpmDeviceError::Unmarshal);
        trace!("{} R: {}", cc, hex::encode(&self.response));
        Ok(result??)
    }

    fn prepare_command<C: TpmFrame>(
        &mut self,
        command: &C,
        sessions: &[TpmsAuthCommand],
    ) -> Result<(), TpmDeviceError> {
        let cc = command.cc();
        let tag = if sessions.is_empty() {
            TpmSt::NoSessions
        } else {
            TpmSt::Sessions
        };

        self.command.resize(TPM_MAX_COMMAND_SIZE, 0);

        let len = {
            let mut writer = TpmWriter::new(&mut self.command);
            tpm_marshal_command(command, tag, sessions, &mut writer)
                .map_err(TpmDeviceError::Marshal)?;
            writer.len()
        };
        self.command.truncate(len);

        trace!("{} C: {}", cc, hex::encode(&self.command));
        Ok(())
    }

    /// Fetches a complete list of capabilities from the TPM, handling
    /// pagination.
    ///
    /// # Errors
    ///
    /// Returns [`ResponseMismatch`](crate::TpmDeviceError::ResponseMismatch)
    /// when receiving unepected TPM response.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    fn get_capability<T, F, N>(
        &mut self,
        cap: TpmCap,
        property_start: u32,
        count: u32,
        mut extract: F,
        next_prop: N,
    ) -> Result<Vec<T>, TpmDeviceError>
    where
        T: Copy,
        F: for<'a> FnMut(&'a TpmuCapabilities) -> Result<&'a [T], TpmDeviceError>,
        N: Fn(&T) -> u32,
    {
        let mut results = Vec::new();
        let mut prop = property_start;
        loop {
            let (more_data, cap_data) =
                self.get_capability_page(cap, TpmUint32(prop), TpmUint32(count))?;
            let items: &[T] = extract(&cap_data.data)?;
            results.extend_from_slice(items);

            if more_data {
                if let Some(last) = items.last() {
                    prop = next_prop(last);
                } else {
                    break;
                }
            } else {
                break;
            }
        }
        Ok(results)
    }

    /// Retrieves all algorithm properties supported by the TPM.
    ///
    /// # Errors
    ///
    /// Returns [`ResponseMismatch`](crate::TpmDeviceError::ResponseMismatch)
    /// when receiving unepected TPM response.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn fetch_algorithm_properties(&mut self) -> Result<Vec<TpmsAlgProperty>, TpmDeviceError> {
        self.get_capability(
            TpmCap::Algs,
            0,
            u32::try_from(MAX_HANDLES).map_err(|_| TpmDeviceError::OperationFailed)?,
            |caps| match caps {
                TpmuCapabilities::Algs(algs) => Ok(algs),
                _ => Err(TpmDeviceError::CapabilityMissing(TpmCap::Algs)),
            },
            |last| last.alg as u32 + 1,
        )
    }

    /// Retrieves all handles of a specific type from the TPM.
    ///
    /// # Errors
    ///
    /// Returns [`ResponseMismatch`](crate::TpmDeviceError::ResponseMismatch)
    /// when receiving unepected TPM response.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn fetch_handles(&mut self, class: TpmHt) -> Result<Vec<TpmHandle>, TpmDeviceError> {
        self.get_capability(
            TpmCap::Handles,
            (class as u32) << 24,
            u32::try_from(MAX_HANDLES).map_err(|_| TpmDeviceError::OperationFailed)?,
            |caps| match caps {
                TpmuCapabilities::Handles(handles) => Ok(handles),
                _ => Err(TpmDeviceError::CapabilityMissing(TpmCap::Handles)),
            },
            |last| last.value() + 1,
        )
        .map(|handles| handles.into_iter().collect())
    }

    /// Retrieves all available ECC curves supported by the TPM.
    ///
    /// # Errors
    ///
    /// Returns [`ResponseMismatch`](crate::TpmDeviceError::ResponseMismatch)
    /// when receiving unepected TPM response.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn fetch_ecc_curves(&mut self) -> Result<Vec<TpmEccCurve>, TpmDeviceError> {
        self.get_capability(
            TpmCap::EccCurves,
            0,
            u32::try_from(MAX_HANDLES).map_err(|_| TpmDeviceError::OperationFailed)?,
            |caps| match caps {
                TpmuCapabilities::EccCurves(curves) => Ok(curves),
                _ => Err(TpmDeviceError::CapabilityMissing(TpmCap::EccCurves)),
            },
            |last| *last as u32 + 1,
        )
    }

    /// Retrieves the list of active PCR banks and the bank selection mask.
    ///
    /// # Errors
    ///
    /// Returns
    /// [`PcrBanksNotAvailable`](crate::TpmDeviceError::PcrBanksNotAvailable)
    /// when no PCR banks are available.
    /// Return
    /// [`PcrBankSelectionMismatch`](crate::TpmDeviceError::PcrBankSelectionMismatch)
    /// when the PCR selection masks differ between active banks.
    /// Returns [`ResponseMismatch`](crate::TpmDeviceError::ResponseMismatch)
    /// when receiving unepected TPM response.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn fetch_pcr_bank_list(
        &mut self,
    ) -> Result<(Vec<TpmAlgId>, TpmsPcrSelect), TpmDeviceError> {
        let pcrs: Vec<TpmsPcrSelection> = self.get_capability(
            TpmCap::Pcrs,
            0,
            u32::try_from(MAX_HANDLES).map_err(|_| TpmDeviceError::OperationFailed)?,
            |caps| match caps {
                TpmuCapabilities::Pcrs(pcrs) => Ok(pcrs),
                _ => Err(TpmDeviceError::CapabilityMissing(TpmCap::Pcrs)),
            },
            |last| last.hash as u32 + 1,
        )?;

        if pcrs.is_empty() {
            return Err(TpmDeviceError::PcrBanksNotAvailable);
        }

        let mut common_select: Option<TpmsPcrSelect> = None;
        let mut algs = Vec::with_capacity(pcrs.len());

        for bank in pcrs {
            if bank.pcr_select.iter().all(|&b| b == 0) {
                debug!(
                    "skipping unallocated bank {:?} (mask: {})",
                    bank.hash,
                    hex::encode(&*bank.pcr_select)
                );
                continue;
            }

            if let Some(ref select) = common_select {
                if bank.pcr_select != *select {
                    return Err(TpmDeviceError::PcrBankSelectionMismatch);
                }
            } else {
                common_select = Some(bank.pcr_select);
            }
            algs.push(bank.hash);
        }

        let select = common_select.ok_or(TpmDeviceError::PcrBanksNotAvailable)?;

        algs.sort();
        Ok((algs, select))
    }

    /// Fetches and returns one page of capabilities of a certain type from the
    /// TPM.
    ///
    /// # Errors
    ///
    /// Returns [`ResponseMismatch`](crate::TpmDeviceError::ResponseMismatch)
    /// when receiving unepected TPM response.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    fn get_capability_page(
        &mut self,
        cap: TpmCap,
        property: TpmUint32,
        property_count: TpmUint32,
    ) -> Result<(bool, TpmsCapabilityData), TpmDeviceError> {
        let cmd = TpmGetCapabilityCommand {
            cap,
            property,
            property_count,
            handles: [],
        };

        let (resp, _) = self.transmit(&cmd, Self::NO_SESSIONS)?;
        let TpmGetCapabilityResponse {
            more_data,
            capability_data,
            handles: [],
        } = resp
            .GetCapability()
            .map_err(|_| TpmDeviceError::ResponseMismatch(TpmCc::GetCapability))?;

        Ok((more_data.into(), capability_data))
    }

    /// Reads a specific TPM property.
    ///
    /// # Errors
    ///
    /// Returns [`ResponseMismatch`](crate::TpmDeviceError::ResponseMismatch)
    /// when receiving unepected TPM response.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn get_tpm_property(&mut self, property: TpmPt) -> Result<TpmUint32, TpmDeviceError> {
        let (_, cap_data) = self.get_capability_page(
            TpmCap::TpmProperties,
            TpmUint32(property as u32),
            TpmUint32(1),
        )?;

        let TpmuCapabilities::TpmProperties(props) = &cap_data.data else {
            return Err(TpmDeviceError::CapabilityMissing(TpmCap::TpmProperties));
        };

        let Some(prop) = props.iter().find(|prop| prop.property == property) else {
            return Err(TpmDeviceError::CapabilityMissing(TpmCap::TpmProperties));
        };

        Ok(prop.value)
    }

    /// Reads the public area of a TPM object.
    ///
    /// # Errors
    ///
    /// Returns [`ResponseMismatch`](crate::TpmDeviceError::ResponseMismatch)
    /// when receiving unepected TPM response.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn read_public(
        &mut self,
        handle: TpmHandle,
    ) -> Result<(TpmtPublic, Tpm2bName), TpmDeviceError> {
        let cmd = TpmReadPublicCommand { handles: [handle] };
        let (resp, _) = self.transmit(&cmd, Self::NO_SESSIONS)?;
        let read_public_resp = resp
            .ReadPublic()
            .map_err(|_| TpmDeviceError::ResponseMismatch(TpmCc::ReadPublic))?;
        let public = read_public_resp.out_public.inner;
        let name = read_public_resp.name;
        Ok((public, name))
    }

    /// Finds a persistent handle by its `Tpm2bName`.
    ///
    /// # Errors
    ///
    /// Returns [`ResponseMismatch`](crate::TpmDeviceError::ResponseMismatch)
    /// when receiving unepected TPM response.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn find_persistent(
        &mut self,
        target_name: &Tpm2bName,
    ) -> Result<Option<TpmHandle>, TpmDeviceError> {
        for handle in self.fetch_handles(TpmHt::Persistent)? {
            match self.read_public(handle) {
                Ok((_, name)) => {
                    if name == *target_name {
                        return Ok(Some(handle));
                    }
                }
                Err(TpmDeviceError::TpmRc(rc)) => {
                    let base = rc.base();
                    if base == TpmRcBase::ReferenceH0 || base == TpmRcBase::Handle {
                        continue;
                    }
                    return Err(TpmDeviceError::TpmRc(rc));
                }
                Err(e) => return Err(e),
            }
        }
        Ok(None)
    }

    /// Saves the context of a transient object or session.
    ///
    /// # Errors
    ///
    /// Returns [`ResponseMismatch`](crate::TpmDeviceError::ResponseMismatch)
    /// when receiving unepected TPM response.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn save_context(&mut self, save_handle: TpmHandle) -> Result<TpmsContext, TpmDeviceError> {
        let cmd = TpmContextSaveCommand {
            handles: [save_handle],
        };
        let (resp, _) = self.transmit(&cmd, Self::NO_SESSIONS)?;
        let save_resp = resp
            .ContextSave()
            .map_err(|_| TpmDeviceError::ResponseMismatch(TpmCc::ContextSave))?;
        Ok(save_resp.context)
    }

    /// Loads a TPM context and returns the handle.
    ///
    /// # Errors
    ///
    /// Returns [`ResponseMismatch`](crate::TpmDeviceError::ResponseMismatch)
    /// when receiving unepected TPM response.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn load_context(&mut self, context: TpmsContext) -> Result<TpmHandle, TpmDeviceError> {
        let cmd = TpmContextLoadCommand {
            context,
            handles: [],
        };
        let (resp, _) = self.transmit(&cmd, Self::NO_SESSIONS)?;
        let resp_inner = resp
            .ContextLoad()
            .map_err(|_| TpmDeviceError::ResponseMismatch(TpmCc::ContextLoad))?;
        Ok(resp_inner.handles[0])
    }

    /// Flushes a transient object or session from the TPM and removes it from
    /// the cache.
    ///
    /// # Errors
    ///
    /// Returns [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn flush_context(&mut self, handle: TpmHandle) -> Result<(), TpmDeviceError> {
        let cmd = TpmFlushContextCommand {
            flush_handle: handle,
            handles: [],
        };
        self.transmit(&cmd, Self::NO_SESSIONS)?;
        Ok(())
    }

    /// Loads a session context and then flushes the resulting handle.
    ///
    /// # Errors
    ///
    /// Returns [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn flush_session(&mut self, context: TpmsContext) -> Result<(), TpmDeviceError> {
        match self.load_context(context) {
            Ok(handle) => self.flush_context(handle),
            Err(TpmDeviceError::TpmRc(rc)) => {
                let base = rc.base();
                if base == TpmRcBase::ReferenceH0 || base == TpmRcBase::Handle {
                    Ok(())
                } else {
                    Err(TpmDeviceError::TpmRc(rc))
                }
            }
            Err(e) => Err(e),
        }
    }
}

/// A builder for creating a TPM policy session.
pub struct TpmPolicySessionBuilder {
    bind: TpmHandle,
    tpm_key: TpmHandle,
    nonce_caller: Option<Tpm2bNonce>,
    encrypted_salt: Option<Tpm2bEncryptedSecret>,
    session_type: TpmSe,
    symmetric: TpmtSymDefObject,
    auth_hash: TpmAlgId,
}

impl Default for TpmPolicySessionBuilder {
    fn default() -> Self {
        Self {
            bind: (TpmRh::Null as u32).into(),
            tpm_key: (TpmRh::Null as u32).into(),
            nonce_caller: None,
            encrypted_salt: None,
            session_type: TpmSe::Policy,
            symmetric: TpmtSymDefObject::default(),
            auth_hash: TpmAlgId::Sha256,
        }
    }
}

impl TpmPolicySessionBuilder {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    #[must_use]
    pub fn with_bind(mut self, bind: TpmHandle) -> Self {
        self.bind = bind;
        self
    }

    #[must_use]
    pub fn with_tpm_key(mut self, tpm_key: TpmHandle) -> Self {
        self.tpm_key = tpm_key;
        self
    }

    #[must_use]
    pub fn with_nonce_caller(mut self, nonce: Tpm2bNonce) -> Self {
        self.nonce_caller = Some(nonce);
        self
    }

    #[must_use]
    pub fn with_encrypted_salt(mut self, salt: Tpm2bEncryptedSecret) -> Self {
        self.encrypted_salt = Some(salt);
        self
    }

    #[must_use]
    pub fn with_session_type(mut self, session_type: TpmSe) -> Self {
        self.session_type = session_type;
        self
    }

    #[must_use]
    pub fn with_symmetric(mut self, symmetric: TpmtSymDefObject) -> Self {
        self.symmetric = symmetric;
        self
    }

    #[must_use]
    pub fn with_auth_hash(mut self, auth_hash: TpmAlgId) -> Self {
        self.auth_hash = auth_hash;
        self
    }

    /// Opens the policy session on the provided device.
    ///
    /// # Errors
    ///
    /// Returns [`ResponseMismatch`](crate::TpmDeviceError::ResponseMismatch) if
    /// the TPM response is unexpected.
    /// Returns [`Unmarshal`](crate::TpmDeviceError::Unmarshal) when unmarshal
    /// operation on TPM protocol compliant data fails.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants depending
    /// on function.
    pub fn open(self, device: &mut TpmDevice) -> Result<TpmPolicySession, TpmDeviceError> {
        let nonce_caller = if let Some(nonce) = self.nonce_caller {
            nonce
        } else {
            let digest_len = TpmHash::from(self.auth_hash).size();
            let mut nonce_bytes = vec![0; digest_len];
            thread_rng().fill_bytes(&mut nonce_bytes);
            Tpm2bNonce::try_from(nonce_bytes.as_slice()).map_err(TpmDeviceError::Unmarshal)?
        };

        let cmd = TpmStartAuthSessionCommand {
            nonce_caller,
            encrypted_salt: self.encrypted_salt.unwrap_or_default(),
            session_type: self.session_type,
            symmetric: self.symmetric,
            auth_hash: self.auth_hash,
            handles: [self.tpm_key, self.bind],
        };

        let (resp, _) = device.transmit(&cmd, TpmDevice::NO_SESSIONS)?;
        let start_resp = resp
            .StartAuthSession()
            .map_err(|_| TpmDeviceError::ResponseMismatch(TpmCc::StartAuthSession))?;

        Ok(TpmPolicySession {
            handle: start_resp.handles[0],
            attributes: TpmaSession::CONTINUE_SESSION,
            hash_alg: self.auth_hash,
            nonce_tpm: start_resp.nonce_tpm,
        })
    }
}

/// Represents an active TPM policy session.
#[derive(Debug, Clone)]
pub struct TpmPolicySession {
    handle: TpmHandle,
    attributes: TpmaSession,
    hash_alg: TpmAlgId,
    nonce_tpm: Tpm2bNonce,
}

impl TpmPolicySession {
    /// Creates a new builder for `TpmPolicySession`.
    #[must_use]
    pub fn builder() -> TpmPolicySessionBuilder {
        TpmPolicySessionBuilder::new()
    }

    /// Returns the session handle.
    #[must_use]
    pub fn handle(&self) -> TpmHandle {
        self.handle
    }

    /// Returns the session attributes.
    #[must_use]
    pub fn attributes(&self) -> TpmaSession {
        self.attributes
    }

    /// Returns the hash algorithm used by the session.
    #[must_use]
    pub fn hash_alg(&self) -> TpmAlgId {
        self.hash_alg
    }

    /// Returns the nonce generated by the TPM.
    #[must_use]
    pub fn nonce_tpm(&self) -> &Tpm2bNonce {
        &self.nonce_tpm
    }

    /// Applies a list of policy commands to this session.
    ///
    /// This method iterates through the provided commands, updates the first handle
    /// of each command (or second for `PolicySecret`) to point to this session,
    /// and transmits them to the device.
    ///
    /// # Errors
    ///
    /// Returns [`InvalidCc`](crate::TpmDeviceError::InvalidCc) when a command is not
    /// a supported policy command.
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn run(
        &self,
        device: &mut TpmDevice,
        commands: Vec<(TpmCommand, TpmAuthCommands)>,
    ) -> Result<(), TpmDeviceError> {
        for (mut command_body, auth_sessions) in commands {
            match &mut command_body {
                TpmCommand::PolicyPcr(cmd) => cmd.handles[0] = self.handle,
                TpmCommand::PolicyOr(cmd) => cmd.handles[0] = self.handle,
                TpmCommand::PolicyRestart(cmd) => {
                    cmd.handles[0] = self.handle;
                }
                TpmCommand::PolicySecret(cmd) => {
                    cmd.handles[1] = self.handle;
                }
                _ => {
                    return Err(TpmDeviceError::InvalidCc(command_body.cc()));
                }
            }
            device.transmit(&command_body, auth_sessions.as_ref())?;
        }
        Ok(())
    }

    /// Flushes the session context from the TPM.
    ///
    /// # Errors
    ///
    /// Returns other [`TpmDeviceError`](crate::TpmDeviceError) variants when
    /// [`TpmDevice::transmit`](crate::TpmDevice::transmit) fails.
    pub fn flush(&self, device: &mut TpmDevice) -> Result<(), TpmDeviceError> {
        device.flush_context(self.handle)
    }
}