Blocktree/crates/btlib/src/test_helpers.rs

/// Test data and functions to help with testing.
use super::*;
use btserde::{Error, Result};
use crypto::{self, *};
use nix::{
    sys::signal::{self, Signal},
    unistd::Pid,
};
use std::{
    cell::RefCell,
    fmt::Write as FmtWrite,
    fs::File,
    io::{Cursor, Write},
    path::PathBuf,
    process::{Child, Command, ExitStatus, Stdio},
    str::FromStr,
    sync::{
        atomic::{AtomicU16, Ordering},
        mpsc::{channel, Receiver, TryRecvError},
    },
};
use tempdir::TempDir;
use tss_esapi::{
    tcti_ldr::{TabrmdConfig, TctiNameConf},
    Context,
};

pub const PRINCIPAL: [u8; 32] = [
    0x75, 0x28, 0xA9, 0xE0, 0x9D, 0x24, 0xBA, 0xB3, 0x79, 0x56, 0x15, 0x68, 0xFD, 0xA4, 0xE2, 0xA4,
    0xCF, 0xB2, 0xC0, 0xE3, 0x96, 0xAE, 0xA2, 0x6E, 0x45, 0x15, 0x50, 0xED, 0xA6, 0xBE, 0x6D, 0xEC,
];

pub const PRINCIPAL2: [u8; 32] = [
    0x6E, 0x8A, 0x76, 0x1D, 0xC7, 0x3B, 0x50, 0x90, 0x51, 0x9E, 0x73, 0x9B, 0xEE, 0x7B, 0x02, 0xF9,
    0xB9, 0x17, 0x7C, 0xF6, 0xBB, 0xC8, 0xD5, 0x30, 0xBF, 0x2A, 0xB4, 0xDE, 0x1B, 0x38, 0xCC, 0xF6,
];

pub const SIGNATURE: [u8; 384] = [
    0x12, 0xB1, 0x09, 0x2F, 0x2B, 0x3C, 0x53, 0xE8, 0x1B, 0x2B, 0x6A, 0xE7, 0x97, 0x42, 0x9D, 0x83,
    0x71, 0x75, 0x65, 0x25, 0xFD, 0xB0, 0x0E, 0x2F, 0xAB, 0x53, 0xB7, 0x03, 0x03, 0x39, 0xEE, 0xE9,
    0x15, 0x19, 0xAB, 0x1A, 0xF1, 0x2C, 0x3C, 0xAB, 0x5C, 0x02, 0xEA, 0xD6, 0xF6, 0x94, 0x36, 0x80,
    0x5F, 0xD1, 0x8E, 0xC4, 0xB9, 0xB5, 0x04, 0x73, 0xBB, 0x77, 0x2C, 0x9C, 0xAF, 0xE6, 0x38, 0xB7,
    0xD8, 0xC1, 0xA3, 0x1B, 0xAD, 0xFB, 0xB2, 0x5E, 0x77, 0xBD, 0x79, 0xDC, 0x4E, 0xEC, 0xEA, 0xFA,
    0x80, 0x9F, 0x12, 0x60, 0xC4, 0xB3, 0x71, 0xAF, 0x1C, 0x28, 0xEE, 0x91, 0x1D, 0x6A, 0x69, 0x31,
    0xA4, 0x5A, 0xAA, 0xBE, 0x0B, 0x20, 0x6A, 0xE1, 0x61, 0x0E, 0x0B, 0x9F, 0x60, 0x73, 0xBE, 0x29,
    0x1F, 0x11, 0x0A, 0x8F, 0x42, 0x8F, 0x80, 0xC1, 0x06, 0x40, 0x29, 0xE9, 0xF1, 0x74, 0x44, 0x69,
    0xCC, 0xD3, 0x05, 0x22, 0x24, 0x87, 0x94, 0x49, 0x4B, 0x5B, 0x62, 0x48, 0x01, 0x7B, 0x05, 0x28,
    0xBE, 0x4D, 0x50, 0x06, 0xE5, 0x67, 0xA8, 0xFB, 0x9A, 0x19, 0x52, 0x57, 0x4B, 0xC4, 0xBB, 0x6A,
    0xF5, 0x0D, 0x67, 0xD0, 0x08, 0x93, 0x81, 0xD2, 0x9E, 0xE2, 0xC8, 0xFA, 0x25, 0xEE, 0x9D, 0xC5,
    0xD8, 0xF4, 0xD5, 0x17, 0xE4, 0xEF, 0xFC, 0x09, 0xB0, 0x09, 0xCA, 0xCA, 0x71, 0x36, 0x61, 0xBE,
    0xC1, 0xD0, 0x53, 0x6D, 0x0C, 0x0A, 0x4D, 0xF1, 0xE0, 0xAB, 0x2C, 0x89, 0x98, 0x81, 0xED, 0x25,
    0xF8, 0x81, 0x8E, 0x2C, 0x46, 0x74, 0x57, 0xCF, 0x5B, 0xE2, 0x14, 0xA5, 0xBF, 0x56, 0xED, 0xD3,
    0x11, 0xE7, 0xB4, 0x8D, 0x89, 0x3A, 0xB2, 0x78, 0xF1, 0xA9, 0x82, 0x9D, 0x3B, 0xE5, 0x6B, 0xD2,
    0xA5, 0xB6, 0xFB, 0x71, 0x0F, 0x98, 0x52, 0x54, 0x1E, 0x98, 0xD2, 0x3B, 0x78, 0x51, 0xD1, 0xE2,
    0x57, 0xAE, 0xB3, 0x34, 0x0E, 0x20, 0x26, 0xCF, 0x4B, 0x1F, 0x12, 0xCA, 0xB1, 0xD5, 0x67, 0x55,
    0xF9, 0xE2, 0xA9, 0x7B, 0xAE, 0x35, 0x35, 0x67, 0xA3, 0xCA, 0xBD, 0xF2, 0x41, 0x4E, 0x3F, 0x7E,
    0x20, 0x72, 0xC2, 0xCC, 0xE4, 0x90, 0x00, 0x04, 0x5D, 0xF7, 0xD9, 0xC7, 0x0C, 0x2C, 0xF8, 0x80,
    0xCB, 0xF2, 0x3E, 0x1A, 0xE9, 0x43, 0x23, 0xAF, 0x62, 0x57, 0xD1, 0x4C, 0xFD, 0x05, 0xB6, 0xB6,
    0x86, 0xB9, 0x86, 0x40, 0xC2, 0x56, 0x45, 0x52, 0xE2, 0x93, 0x62, 0x65, 0x23, 0xE7, 0x1D, 0x89,
    0xCA, 0x40, 0x23, 0x44, 0xDB, 0x56, 0x90, 0x4A, 0x0E, 0xB1, 0xB1, 0xE1, 0x9B, 0x5A, 0x5D, 0x53,
    0xEF, 0x4E, 0xE1, 0x0E, 0x12, 0xA9, 0x7C, 0x11, 0x0B, 0x6A, 0x4C, 0x3A, 0x8F, 0x2B, 0x9F, 0xC6,
    0x42, 0xD7, 0xFE, 0x7D, 0x10, 0x02, 0xBD, 0x31, 0x21, 0xCD, 0xD1, 0x32, 0x04, 0x22, 0xE2, 0x36,
];

/// The key of the parent of test blocks.
pub(crate) static PARENT_KEY: SymKey = {
    let key = [
        0x35, 0x3D, 0x8C, 0x95, 0x6C, 0x8D, 0xE6, 0xC0, 0xB0, 0xD5, 0x1C, 0xE9, 0x94, 0xB0, 0x58,
        0xD3, 0x80, 0x46, 0x12, 0x1C, 0xF3, 0x9B, 0x8A, 0xEC, 0x38, 0xD5, 0x8B, 0x05, 0x92, 0x8D,
        0xA1, 0x18,
    ];
    let iv = [
        0xDA, 0x96, 0x0E, 0xF5, 0x00, 0xDB, 0xC8, 0x4E, 0xFB, 0x83, 0x59, 0x99, 0x6C, 0x18, 0xD1,
        0x27,
    ];
    SymKey::Aes256Cbc { key, iv }
};

lazy_static! {
    static ref INHERIT: Ciphertext<SymKey> = PARENT_KEY
        .ser_encrypt(&BLOCK_KEY)
        .expect("failed to encrypt block key");
}

/// The key used to encrypt test blocks.
pub(crate) static BLOCK_KEY: SymKey = {
    let key = [
        0xB2, 0xB3, 0xDA, 0x5A, 0x1A, 0xF6, 0xB3, 0x78, 0x30, 0xAB, 0x1D, 0x33, 0x33, 0xE7, 0xE3,
        0x5B, 0xBB, 0xF9, 0xFE, 0xD0, 0xC1, 0xF7, 0x90, 0x34, 0x69, 0xB7, 0xE7, 0xC6, 0x1C, 0x46,
        0x85, 0x48,
    ];
    let iv = [
        0xEC, 0x19, 0x59, 0x3A, 0x1D, 0x1E, 0x4A, 0x58, 0x66, 0xC1, 0xD1, 0x9A, 0x61, 0x6E, 0xBA,
        0x16,
    ];
    SymKey::Aes256Cbc { key, iv }
};

lazy_static! {
    pub static ref ROOT_CREDS: ConcreteCreds =
        ConcreteCreds::generate().expect("root cred generation failed");
    pub static ref NODE_CREDS: ConcreteCreds =
        ConcreteCreds::generate().expect("node cred generation failed");
}

/// Converts the given error to a serde_block_tree error by turning it into a message.
fn convert_err<E: Display>(err: E) -> Error {
    Error::Message(err.to_string())
}

pub(crate) fn make_principal() -> Principal {
    Principal(Hash::Sha2_256(PRINCIPAL))
}

pub(crate) fn make_path_with_root(root: Principal, rel_components: Vec<&str>) -> BlockPath {
    let mut components = Vec::with_capacity(rel_components.len() + 1);
    components.push(root.0.to_string());
    for component in rel_components {
        components.push(component.to_string());
    }
    BlockPath::new(root, components)
}

pub(crate) fn make_path(rel_components: Vec<&str>) -> BlockPath {
    make_path_with_root(make_principal(), rel_components)
}

pub(crate) fn make_writecap_and_creds(rel_components: Vec<&str>) -> (Writecap, impl Creds) {
    let (root_writecap, root_key) = make_self_signed_writecap();
    let issued_to = Principal(Hash::Sha2_256(PRINCIPAL));
    (
        make_writecap_trusted_by(root_writecap, &root_key, issued_to, rel_components),
        root_key,
    )
}

pub(crate) fn make_writecap(rel_components: Vec<&str>) -> Writecap {
    let (writecap, ..) = make_writecap_and_creds(rel_components);
    writecap
}

pub(crate) fn make_writecap_trusted_by<C: Creds>(
    next: Writecap,
    trusting_creds: &C,
    issued_to: Principal,
    path_components: Vec<&str>,
) -> Writecap {
    let hour_hence = Epoch::now() + Duration::from_secs(3600);
    let mut writecap = Writecap {
        body: WritecapBody {
            issued_to,
            path: make_path_with_root(next.body.path.root().clone(), path_components),
            expires: hour_hence,
            signing_key: trusting_creds.public_sign().clone(),
        },
        signature: Signature::empty(Sign::RSA_PSS_3072_SHA_256),
        next: Some(Box::from(next)),
    };
    trusting_creds
        .sign_writecap(&mut writecap)
        .map_err(convert_err)
        .expect("failed to sign writecap");
    writecap
}

pub(crate) fn make_key_pair() -> impl Creds {
    ROOT_CREDS.clone()
}

pub(crate) fn make_self_signed_writecap() -> (Writecap, impl Creds) {
    let key = make_key_pair();
    (make_self_signed_writecap_with(&key), key)
}

pub(crate) fn make_self_signed_writecap_with<C: Creds>(key: &C) -> Writecap {
    let root_principal = key.principal();
    let hour_hence = Epoch::now() + Duration::from_secs(3600);
    let mut writecap = Writecap {
        body: WritecapBody {
            issued_to: root_principal.clone(),
            path: make_path_with_root(root_principal, vec![]),
            expires: hour_hence,
            signing_key: key.public_sign().clone(),
        },
        signature: Signature::empty(Sign::RSA_PSS_3072_SHA_256),
        next: None,
    };
    key.sign_writecap(&mut writecap)
        .map_err(convert_err)
        .expect("failed to sign writecap");
    writecap
}

pub(crate) fn make_readcap_for<C: Encrypter + Principaled>(creds: &C) -> Readcap {
    Readcap {
        issued_to: creds.principal(),
        key: creds
            .ser_encrypt(&BLOCK_KEY)
            .expect("failed to encrypt block key"),
    }
}

pub(crate) fn make_block_with(readcap: Readcap) -> Box<dyn Block> {
    let mut readcaps = BTreeMap::new();
    readcaps.insert(readcap.issued_to, readcap.key);
    // Notice that the writecap path contains the block path. If this were not the case, the block
    // would be invalid.
    let (writecap, creds) = make_writecap_and_creds(vec!["apps"]);
    let root_writecap = writecap.next.as_ref().unwrap();
    let header = BlockMetaBody {
        path: make_path_with_root(root_writecap.body.issued_to.clone(), vec!["apps", "verse"]),
        inherit: Some(INHERIT.clone()),
        readcaps,
        writecap: Some(writecap),
        integrity: Some(Hash::Sha2_256([0u8; HashKind::Sha2_256.len()])),
        signing_key: creds.public_sign().to_owned(),
    };
    let sig = Signature::copy_from(Sign::RSA_PSS_3072_SHA_256, &SIGNATURE);
    let mut stream =
        BlockStream::new(BtCursor::new(Vec::new()), creds).expect("create block stream failed");
    stream.meta.body = header;
    stream.meta.sig = sig;
    let block_key = stream.block_key().expect("get block key failed");
    let stream = MerkleStream::new(stream).expect("create merkle stream failed");
    let stream = SecretStream::new(block_key)
        .try_compose(stream)
        .expect("create secret stream failed");
    let stream = SectoredBuf::new()
        .try_compose(stream)
        .expect("create sectored buf failed");
    Box::new(stream)
}

/// This function can be run as a test to write a new RSA key pair, as two Rust arrays,
/// out to a file.
#[allow(dead_code)]
fn write_test_keys() -> Result<()> {
    write_rsa_keys_to_file("test_keys.rs")
}

struct NamedSlice<'a> {
    name: &'a str,
    slice: &'a [u8],
}

impl<'a> NamedSlice<'a> {
    fn new(name: &'a str, slice: &'a [u8]) -> NamedSlice<'a> {
        NamedSlice { name, slice }
    }
}

fn write_rsa_keys_to_file(path: &str) -> Result<()> {
    use openssl::rsa::Rsa;
    let rsa = Rsa::generate(3072).map_err(convert_err)?;
    let public_der = rsa.public_key_to_der().map_err(convert_err)?;
    let private_der = rsa.private_key_to_der().map_err(convert_err)?;
    let slices = [
        NamedSlice::new("PUBLIC", public_der.as_slice()),
        NamedSlice::new("PRIVATE", private_der.as_slice()),
    ];
    write_to_file(path, slices.into_iter())
}

fn write_to_file<'a, I: Iterator<Item = NamedSlice<'a>>>(path: &str, slices: I) -> Result<()> {
    let path = std::path::Path::new(path);
    let mut file = File::create(path).map_err(Error::Io)?;
    for NamedSlice { name, slice } in slices {
        write_slice(&mut file, name, slice)?;
    }
    Ok(())
}

fn write_slice<W: Write>(output: &mut W, name: &str, slice: &[u8]) -> Result<()> {
    const LINE_LEN: usize = 100;
    let mut line = String::with_capacity(LINE_LEN);
    write!(line, "    ").map_err(Error::Format)?;
    writeln!(output, "pub const {}: [u8; {}] = [", name, slice.len()).map_err(Error::Io)?;
    for byte in slice {
        if line.len() + 6 > LINE_LEN {
            writeln!(output, "{}", line).map_err(Error::Io)?;
            line.clear();
            write!(line, "    ").map_err(Error::Format)?;
        }
        write!(line, "0x{:02X?}, ", byte).map_err(Error::Format)?;
    }
    writeln!(output, "{}", line).map_err(Error::Io)?;
    writeln!(output, "];").map_err(Error::Io)?;
    writeln!(output).map_err(Error::Io)?;
    Ok(())
}

/// A naive randomizer implementation that is intended only for testing.
pub struct Randomizer {
    state: [u8; Self::HASH.len()],
    buf: [u8; Self::HASH.len()],
}

impl Randomizer {
    pub const HASH: HashKind = HashKind::Sha2_256;

    pub fn new(seed: [u8; Self::HASH.len()]) -> Randomizer {
        Randomizer {
            state: seed,
            buf: [0u8; Self::HASH.len()],
        }
    }
}

impl Iterator for Randomizer {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        const BYTES: usize = usize::BITS as usize / 8;
        Self::HASH
            .digest(&mut self.buf, std::iter::once(self.state.as_slice()))
            .expect("digest failed");
        self.state.copy_from_slice(&self.buf);
        let int_bytes = self.buf.as_slice()[..BYTES]
            .try_into()
            .expect("failed to convert array");
        Some(usize::from_ne_bytes(int_bytes))
    }
}

pub fn make_sector(mut buf: &mut [u8], sect_index: usize) {
    let data = (sect_index + 1).to_ne_bytes();
    for chunk in std::iter::repeat(data).take(buf.len() / data.len()) {
        (&mut buf[..chunk.len()]).copy_from_slice(chunk.as_slice());
        buf = &mut buf[chunk.len()..];
    }
}

pub fn write_fill<W: Write>(mut write: W, sect_sz: usize, sect_ct: usize) {
    let mut buf = vec![0u8; sect_sz];
    for sect_index in 0..sect_ct {
        make_sector(&mut buf, sect_index);
        write.write_all(&mut buf).expect("write failed");
    }
    write.flush().expect("flush failed");
}

pub fn read_check<R: Read>(mut read: R, sect_sz: usize, sect_ct: usize) {
    let mut actual = vec![0u8; sect_sz];
    let mut expected = vec![0u8; sect_sz];
    for sect_index in 0..sect_ct {
        make_sector(&mut expected, sect_index);
        read.read_exact(&mut actual).expect("read failed");
        assert_eq!(expected, actual);
    }
}

/// Trait for types which can be referenced as slices and which support conversion from `Vec<u8>`.
pub trait FromVec: AsRef<[u8]> {
    fn from_vec(vec: Vec<u8>) -> Self;
}

impl FromVec for Vec<u8> {
    fn from_vec(vec: Vec<u8>) -> Self {
        vec
    }
}

impl<const N: usize> FromVec for [u8; N] {
    fn from_vec(vec: Vec<u8>) -> Self {
        assert_eq!(N, vec.len());
        let mut buf = [0u8; N];
        buf.copy_from_slice(&vec);
        buf
    }
}

/// Module containing functions for serializing and deserializing buffers in `SectoredCursor<T>`.
mod serde_cursor {
    use super::FromVec;
    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use std::{
        cell::RefCell,
        io::{Cursor, Read},
        result::Result,
    };

    pub fn serialize<T: AsRef<[u8]>, S: Serializer>(
        cursor: &RefCell<Cursor<T>>,
        ser: S,
    ) -> Result<S::Ok, S::Error> {
        let mut cursor = cursor.borrow_mut();
        let pos = cursor.position();
        cursor.set_position(0);
        let mut data = Vec::new();
        cursor
            .read_to_end(&mut data)
            .expect("reading from cursor failed");
        cursor.set_position(pos);
        data.serialize(ser)
    }

    pub fn deserialize<'de, T: FromVec, D: Deserializer<'de>>(
        de: D,
    ) -> Result<RefCell<Cursor<T>>, D::Error> {
        let data = Vec::<u8>::deserialize(de)?;
        Ok(RefCell::new(Cursor::new(T::from_vec(data))))
    }
}

/// A wrapper for `Cursor<T>` which implements additional traits.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct BtCursor<T: FromVec> {
    #[serde(with = "serde_cursor")]
    cursor: RefCell<Cursor<T>>,
}

impl<T: FromVec> BtCursor<T> {
    pub(crate) fn new(inner: T) -> BtCursor<T> {
        BtCursor {
            cursor: RefCell::new(Cursor::new(inner)),
        }
    }
}

impl Write for BtCursor<Vec<u8>> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.cursor.get_mut().write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.cursor.get_mut().flush()
    }
}

impl<const N: usize> Write for BtCursor<[u8; N]> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.cursor.get_mut().write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.cursor.get_mut().flush()
    }
}

impl<T: FromVec> Read for BtCursor<T> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.cursor.get_mut().read(buf)
    }
}

impl<T: FromVec> Seek for BtCursor<T> {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        self.cursor.get_mut().seek(pos)
    }
}

impl WriteInteg for BtCursor<Vec<u8>> {
    fn flush_integ(&mut self, _: &[u8]) -> io::Result<()> {
        Ok(())
    }
}

impl<const N: usize> WriteInteg for BtCursor<[u8; N]> {
    fn flush_integ(&mut self, _: &[u8]) -> io::Result<()> {
        Ok(())
    }
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct SectoredCursor<T: FromVec> {
    cursor: BtCursor<T>,
    sect_sz: usize,
}

impl<T: FromVec> SectoredCursor<T> {
    pub fn new(inner: T, sect_sz: usize) -> SectoredCursor<T> {
        SectoredCursor {
            cursor: BtCursor::new(inner),
            sect_sz,
        }
    }
}

impl<T: FromVec> Sectored for SectoredCursor<T> {
    fn sector_sz(&self) -> usize {
        self.sect_sz
    }
}

impl Write for SectoredCursor<Vec<u8>> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.assert_sector_sz(buf.len())?;
        self.cursor.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.cursor.flush()
    }
}

impl<const N: usize> Write for SectoredCursor<[u8; N]> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.assert_sector_sz(buf.len())?;
        self.cursor.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.cursor.flush()
    }
}

impl<T: FromVec> Read for SectoredCursor<T> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.assert_sector_sz(buf.len())?;
        self.cursor.read(buf)
    }
}

impl<T: FromVec> Seek for SectoredCursor<T> {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        self.cursor.seek(pos)
    }
}

trait ExitStatusExt {
    fn success_or_err(&self) -> Result<()>;
}

impl ExitStatusExt for ExitStatus {
    fn success_or_err(&self) -> btserde::Result<()> {
        match self.code() {
            Some(0) => Ok(()),
            Some(code) => Err(btserde::Error::Message(format!(
                "ExitCode was non-zero: {}",
                code
            ))),
            None => Err(btserde::Error::Message("ExitCode was None.".to_string())),
        }
    }
}

/// A DBus message which is sent when the ownership of a name changes.
struct NameOwnerChanged {
    name: String,
    old_owner: String,
    new_owner: String,
}

impl dbus::arg::AppendAll for NameOwnerChanged {
    fn append(&self, iter: &mut dbus::arg::IterAppend) {
        dbus::arg::RefArg::append(&self.name, iter);
        dbus::arg::RefArg::append(&self.old_owner, iter);
        dbus::arg::RefArg::append(&self.new_owner, iter);
    }
}

impl dbus::arg::ReadAll for NameOwnerChanged {
    fn read(iter: &mut dbus::arg::Iter) -> std::result::Result<Self, dbus::arg::TypeMismatchError> {
        Ok(NameOwnerChanged {
            name: iter.read()?,
            old_owner: iter.read()?,
            new_owner: iter.read()?,
        })
    }
}

impl dbus::message::SignalArgs for NameOwnerChanged {
    const NAME: &'static str = "NameOwnerChanged";
    const INTERFACE: &'static str = "org.freedesktop.DBus";
}

/// A struct used to block until a specific name appears on DBus.
struct DbusBlocker {
    receiver: Receiver<()>,
    conn: dbus::blocking::Connection,
}

impl DbusBlocker {
    fn new_session(name: String) -> Result<DbusBlocker> {
        use dbus::{blocking::Connection, Message};
        const DEST: &str = "org.freedesktop.DBus";

        let (sender, receiver) = channel();
        let conn = Connection::new_session().map_err(Error::custom)?;
        let proxy = conn.with_proxy(DEST, "/org/freedesktop/DBus", Duration::from_secs(1));
        let _ = proxy.match_signal(move |h: NameOwnerChanged, _: &Connection, _: &Message| {
            let name_appeared = h.name == name;
            if name_appeared {
                if let Err(err) = sender.send(()) {
                    error!("failed to send unblocking signal: {err}");
                }
            }
            let remove_match = !name_appeared;
            remove_match
        });
        Ok(DbusBlocker { receiver, conn })
    }

    fn block(&mut self, timeout: Duration) -> Result<()> {
        let time_limit = SystemTime::now() + timeout;
        loop {
            self.conn
                .process(Duration::from_millis(100))
                .map_err(Error::custom)?;
            match self.receiver.try_recv() {
                Ok(_) => break,
                Err(err) => match err {
                    TryRecvError::Empty => (),
                    _ => return Err(Error::custom(err)),
                },
            }
            if SystemTime::now() > time_limit {
                return Err(Error::custom("timed out"));
            }
        }
        Ok(())
    }
}

pub(crate) struct SwtpmHarness {
    dir: TempDir,
    port: u16,
    state_path: PathBuf,
    pid_path: PathBuf,
    tabrmd: Child,
}

impl SwtpmHarness {
    const HOST: &'static str = "127.0.0.1";

    fn dbus_name(port: u16) -> String {
        let port_str: String = port
            .to_string()
            .chars()
            // Shifting each code point by 17 makes the digits into capital letters.
            .map(|e| ((e as u8) + 17) as char)
            .collect();
        format!("com.intel.tss2.Tabrmd.{port_str}")
    }

    pub(crate) fn new() -> crypto::Result<SwtpmHarness> {
        static PORT: AtomicU16 = AtomicU16::new(21901);
        let port = PORT.fetch_add(2, Ordering::SeqCst);
        let ctrl_port = port + 1;
        let dir = TempDir::new(format!("swtpm_harness.{port}").as_str())?;
        let dir_path = dir.path();
        let dir_path_display = dir_path.display();
        let conf_path = dir_path.join("swtpm_setup.conf");
        let state_path = dir_path.join("state.bt");
        let pid_path = dir_path.join("swtpm.pid");
        let dbus_name = Self::dbus_name(port);
        let addr = Self::HOST;
        std::fs::write(
            &conf_path,
            r#"# Program invoked for creating certificates
#create_certs_tool= /usr/bin/swtpm_localca
# Comma-separated list (no spaces) of PCR banks to activate by default
active_pcr_banks = sha256
"#,
        )?;
        Command::new("swtpm_setup")
            .stdout(Stdio::null())
            .args([
                "--tpm2",
                "--config",
                conf_path.to_str().unwrap(),
                "--tpm-state",
                format!("{dir_path_display}").as_str(),
            ])
            .status()?
            .success_or_err()?;
        Command::new("swtpm")
            .args([
                "socket",
                "--daemon",
                "--tpm2",
                "--server",
                format!("type=tcp,port={port},bindaddr={addr}").as_str(),
                "--ctrl",
                format!("type=tcp,port={ctrl_port},bindaddr={addr}").as_str(),
                "--log",
                format!("file={dir_path_display}/log.txt,level=5").as_str(),
                "--flags",
                "not-need-init,startup-clear",
                "--tpmstate",
                format!("dir={dir_path_display}").as_str(),
                "--pid",
                format!("file={}", pid_path.display()).as_str(),
            ])
            .status()?
            .success_or_err()?;
        let mut blocker = DbusBlocker::new_session(dbus_name.clone())?;
        let tabrmd = Command::new("tpm2-abrmd")
            .args([
                format!("--tcti=swtpm:host=127.0.0.1,port={port}").as_str(),
                "--dbus-name",
                dbus_name.as_str(),
                "--session",
            ])
            .spawn()?;
        blocker.block(Duration::from_secs(5))?;
        Ok(SwtpmHarness {
            dir,
            port,
            state_path,
            pid_path,
            tabrmd,
        })
    }

    pub(crate) fn context(&self) -> crypto::Result<Context> {
        let config_string = format!("bus_name={},bus_type=session", Self::dbus_name(self.port));
        let config = TabrmdConfig::from_str(config_string.as_str())?;
        Ok(Context::new(TctiNameConf::Tabrmd(config))?)
    }

    pub(crate) fn dir_path(&self) -> &std::path::Path {
        self.dir.path()
    }

    pub(crate) fn state_path(&self) -> &std::path::Path {
        &self.state_path
    }
}

impl Drop for SwtpmHarness {
    fn drop(&mut self) {
        if let Err(err) = self.tabrmd.kill() {
            error!("failed to kill tpm2-abrmd: {err}");
        }
        let pid_str = std::fs::read_to_string(&self.pid_path).unwrap();
        let pid_int = pid_str.parse::<i32>().unwrap();
        let pid = Pid::from_raw(pid_int);
        signal::kill(pid, Signal::SIGKILL).unwrap();
    }
}