// Copyright 2022 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use {
    fuchsia_async::TimeoutExt as _,
    futures::{TryFutureExt as _, TryStreamExt as _},
    std::ffi::{CStr, CString},
};

/// Hardware derived key is expected to be a 128-bit AES key.
const DERIVED_KEY_SIZE: usize = 16;
const KEY_INFO_SIZE: usize = 32;

#[derive(Copy, Clone, Debug)]
pub enum TaKeysafeCommand {
    GetUserDataStorageKey = 8,
    RotateHardwareDerivedKey = 9,
}

/// Error values generated by this library
#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("tee command {0:?} failed: {1}")]
    TeeCommand(TaKeysafeCommand, u32),
    #[error("tee command {0:?} failed: not supported")]
    TeeCommandNotSupported(TaKeysafeCommand),
    #[error(
        "tee command {:?} failed: buffer with hardcoded size {} bytes was too small",
        .0,
        DERIVED_KEY_SIZE,
    )]
    TeeCommandBufferTooSmall(TaKeysafeCommand),
    #[error("failed to open dev directory")]
    DevDirectoryOpen(#[from] fuchsia_fs::node::OpenError),
    #[error("timeout waiting for tee device")]
    TeeDeviceWaitTimeout,
    #[error("failure waiting for tee device")]
    TeeDeviceWaitFailure(#[from] anyhow::Error),
}

/// The info used to identify a key.
pub struct KeyInfo {
    info: [u8; KEY_INFO_SIZE],
}

impl KeyInfo {
    /// Creates a new key info buffer using the provided string as the identifier.
    pub fn new(info: impl ToString) -> Self {
        Self::new_bytes(info.to_string().as_bytes())
    }

    /// Creates a new key info buffer using "zxcrypt" as the identifier.
    pub fn new_zxcrypt() -> Self {
        Self::new_bytes("zxcrypt".as_bytes())
    }

    fn new_bytes(info_bytes: &[u8]) -> Self {
        let mut info = [0; KEY_INFO_SIZE];
        // This will panic if the provided info is longer than 32 bytes, which is fine because
        // that's a programming error.
        info[..info_bytes.len()].copy_from_slice(info_bytes);
        Self { info }
    }
}

fn call_command(
    device: Option<&CStr>,
    op: &mut tee::TeecOperation,
    id: TaKeysafeCommand,
) -> Result<(), Error> {
    match device {
        Some(dev) => tee::call_command_on_device(dev, op, id as u32),
        None => tee::call_command(op, id as u32),
    }
    .map_err(|e| match e {
        tee::TEEC_ERROR_NOT_SUPPORTED => Error::TeeCommandNotSupported(id),
        tee::TEEC_ERROR_SHORT_BUFFER => Error::TeeCommandBufferTooSmall(id),
        e => Error::TeeCommand(id, e),
    })
}

fn get_key_from_tee_device(device: Option<&CStr>, info: KeyInfo) -> Result<Vec<u8>, Error> {
    let mut key_buf = [0u8; DERIVED_KEY_SIZE];

    let mut op = tee::create_operation(
        tee::teec_param_types(
            tee::TEEC_MEMREF_TEMP_INPUT,
            tee::TEEC_NONE,
            tee::TEEC_NONE,
            tee::TEEC_MEMREF_TEMP_OUTPUT,
        ),
        [
            tee::get_memref_input_parameter(&info.info),
            tee::get_zero_parameter(),
            tee::get_zero_parameter(),
            tee::get_memref_output_parameter(&mut key_buf),
        ],
    );

    call_command(device, &mut op, TaKeysafeCommand::GetUserDataStorageKey)?;

    Ok(key_buf.to_vec())
}

fn rotate_key_from_tee_device(device: Option<&CStr>, info: KeyInfo) -> Result<(), Error> {
    let mut op = tee::create_operation(
        tee::teec_param_types(
            tee::TEEC_MEMREF_TEMP_INPUT,
            tee::TEEC_NONE,
            tee::TEEC_NONE,
            tee::TEEC_NONE,
        ),
        [
            tee::get_memref_input_parameter(&info.info),
            tee::get_zero_parameter(),
            tee::get_zero_parameter(),
            tee::get_zero_parameter(),
        ],
    );

    call_command(device, &mut op, TaKeysafeCommand::RotateHardwareDerivedKey)
}

/// Gets a hardware derived key using the first device found in /dev/class/tee.
/// This is useful in early boot when other services may not be up.
pub async fn get_hardware_derived_key(info: KeyInfo) -> Result<Vec<u8>, Error> {
    const DEV_CLASS_TEE: &str = "/dev/class/tee";

    let dir =
        fuchsia_fs::directory::open_in_namespace(DEV_CLASS_TEE, fuchsia_fs::OpenFlags::empty())?;
    let mut stream = device_watcher::watch_for_files(&dir).await?;
    let first = stream
        .try_next()
        .map_err(Error::from)
        .on_timeout(std::time::Duration::from_secs(5), || Err(Error::TeeDeviceWaitTimeout))
        .await?;
    let first = first.ok_or_else(|| {
        Error::TeeDeviceWaitFailure(anyhow::anyhow!(
            "'{DEV_CLASS_TEE}' watcher closed unexpectedly"
        ))
    })?;
    let first = first.to_str().expect("paths are utf-8");

    let dev = format!("{DEV_CLASS_TEE}/{first}");
    let dev = CString::new(dev).expect("paths do not contain nul bytes");
    get_key_from_tee_device(Some(&dev), info)
}

/// Gets a hardware derived key using the service fuchsia.tee.Application. This should be used from
/// components.
pub async fn get_hardware_derived_key_from_service(info: KeyInfo) -> Result<Vec<u8>, Error> {
    get_key_from_tee_device(None, info)
}

/// Rotates the hardware derived key from a tee device at the /dev/class/tee.
/// This is useful in early boot when other services may not be up.
pub async fn rotate_hardware_derived_key(info: KeyInfo) -> Result<(), Error> {
    const DEV_CLASS_TEE: &str = "/dev/class/tee";

    let dir =
        fuchsia_fs::directory::open_in_namespace(DEV_CLASS_TEE, fuchsia_fs::OpenFlags::empty())?;
    let mut stream = device_watcher::watch_for_files(&dir).await?;
    let first = stream
        .try_next()
        .map_err(Error::from)
        .on_timeout(std::time::Duration::from_secs(5), || Err(Error::TeeDeviceWaitTimeout))
        .await?;
    let first = first.ok_or_else(|| {
        Error::TeeDeviceWaitFailure(anyhow::anyhow!(
            "'{DEV_CLASS_TEE}' watcher closed unexpectedly"
        ))
    })?;
    let first = first.to_str().expect("paths are utf-8");

    let dev = format!("{DEV_CLASS_TEE}/{first}");
    let dev = CString::new(dev).expect("paths do not contain nul bytes");
    rotate_key_from_tee_device(Some(&dev), info)
}

/// Rotates an existing hardware derived key identified by [`info`] using the service
/// fuchsia.tee.Application. This should be used from components.
pub async fn rotate_hardware_derived_key_from_service(info: KeyInfo) -> Result<(), Error> {
    rotate_key_from_tee_device(None, info)
}