// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::util::copy_data_to_ck_out_slice;
use crate::backend::Mechanism;
use crate::crypto::OpCtxState;
use crate::pkcs11;
use crate::util::ckraw::CkRawMechanism;
use crate::Error;

/// See PKCS#11 v2.40 Section 5.9 Decryption functions

pub extern "C" fn C_DecryptInit(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pMechanism: pkcs11::CK_MECHANISM_PTR,
    hKey: pkcs11::CK_OBJECT_HANDLE,
) -> pkcs11::CK_RV {
    trace!("C_DecryptInit() called");

    if pMechanism.is_null() {
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    let raw_mech = unsafe { CkRawMechanism::from_raw_ptr_unchecked(pMechanism) };
    let mech = match Mechanism::from_ckraw_mech(&raw_mech) {
        Ok(mech) => mech,
        Err(_) => return pkcs11::CKR_MECHANISM_INVALID,
    };

    lock_session_mut!(hSession, session, _sarc);
    session
        .decrypt_init(&mech, hKey.into())
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_Decrypt(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pEncryptedData: pkcs11::CK_BYTE_PTR,
    ulEncryptedDataLen: pkcs11::CK_ULONG,
    pData: pkcs11::CK_BYTE_PTR,
    pulDataLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_Decrypt() called");

    lock_session_mut!(hSession, session, _sarc);

    let out_len = match session
        .decrypt_ctx()
        .as_mut()
        .ok_or(Error::OperationNotInitialized)
        .and_then(|ctx| {
            ctx.enter_state(OpCtxState::SinglepartActive)
                .map_err(Error::CryptoError)?;
            Ok(ctx.len())
        }) {
        Ok(l) => l,
        Err(e) => {
            session.decrypt_ctx().take();
            return e.into();
        }
    };

    let out_slice = ck_out_buf_to_mut_slice!(pData, pulDataLen, out_len, {
        session.decrypt_ctx().take();
    });
    let in_slice = ck_in_buf_to_slice!(pEncryptedData, ulEncryptedDataLen, {
        session.decrypt_ctx().take();
    });

    session
        .decrypt_ctx()
        .take()
        // It's safe to unwrap here since we've already checked for an existing DecryptCtx
        .unwrap()
        .decrypt(in_slice)
        .map(|v| copy_data_to_ck_out_slice(v.as_slice(), out_slice, pulDataLen))
        .map_err(Error::CryptoError)
        .unwrap_or_else(|e| e.into())
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::util::copy_data_to_ck_out_slice;
use crate::crypto::OpCtxState;
use crate::pkcs11;
use crate::Error;

/// See PKCS#11 v2.40 Section 5.10 Message digesting functions

pub extern "C" fn C_DigestInit(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pMechanism: *mut pkcs11::CK_MECHANISM,
) -> pkcs11::CK_RV {
    trace!("C_DigestInit() called");

    if pMechanism.is_null() {
        error!("C_DigestInit() called with NULL mech ptr");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }
    let mech_type = unsafe { (*pMechanism).mechanism };

    lock_session_mut!(hSession, session, _sarc);
    session
        .digest_init(mech_type)
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_Digest(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pData: *mut pkcs11::CK_BYTE,
    ulDataLen: pkcs11::CK_ULONG,
    pDigest: *mut pkcs11::CK_BYTE,
    pulDigestLen: *mut pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_Digest() called");

    lock_session_mut!(hSession, session, _sarc);

    let out_len = match session
        .digest_ctx()
        .as_mut()
        .ok_or(Error::OperationNotInitialized)
        .and_then(|ctx| {
            ctx.enter_state(OpCtxState::SinglepartActive)
                .map_err(Error::CryptoError)?;
            Ok(ctx.len())
        }) {
        Ok(l) => l,
        Err(e) => {
            session.digest_ctx().take();
            return e.into();
        }
    };

    let out_slice = ck_out_buf_to_mut_slice!(pDigest, pulDigestLen, out_len, {
        session.digest_ctx().take()
    });
    let in_slice = ck_in_buf_to_slice!(pData, ulDataLen, {
        session.digest_ctx().take();
    });

    session
        .digest_ctx()
        .take()
        // It's safe to unwrap here since we've already checked for a digest ctx.
        .unwrap()
        .digest(in_slice)
        .map(|v| copy_data_to_ck_out_slice(v.as_slice(), out_slice, pulDigestLen))
        .map_err(Error::CryptoError)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_DigestUpdate(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pPart: *mut pkcs11::CK_BYTE,
    ulPartLen: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_DigestUpdate() called");

    lock_session_mut!(hSession, session, _sarc);
    let in_slice = ck_in_buf_to_slice!(pPart, ulPartLen, {
        session.digest_ctx().take();
    });
    session
        .digest_ctx()
        .as_mut()
        .ok_or(Error::OperationNotInitialized)
        .and_then(|ctx| ctx.update(in_slice).map_err(Error::CryptoError))
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_DigestFinal(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pDigest: *mut pkcs11::CK_BYTE,
    pulDigestLen: *mut pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_DigestFinal() called");

    lock_session_mut!(hSession, session, _sarc);

    let out_len = match session
        .digest_ctx()
        .as_mut()
        .ok_or(Error::OperationNotInitialized)
        .and_then(|ctx| {
            ctx.enter_state(OpCtxState::MultipartReady)
                .map_err(Error::CryptoError)?;
            Ok(ctx.len())
        }) {
        Ok(l) => l,
        Err(e) => {
            session.digest_ctx().take();
            return e.into();
        }
    };

    let out_slice = ck_out_buf_to_mut_slice!(pDigest, pulDigestLen, out_len, {
        session.digest_ctx().take()
    });

    session
        .digest_ctx()
        .take()
        // It's safe to unwrap here since we've already checked for a digest ctx.
        .unwrap()
        .finalize()
        .map(|v| copy_data_to_ck_out_slice(v.as_slice(), out_slice, pulDigestLen))
        .map_err(Error::CryptoError)
        .unwrap_or_else(|e| e.into())
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::util::copy_data_to_ck_out_slice;
use crate::backend::Mechanism;
use crate::crypto::OpCtxState;
use crate::pkcs11;
use crate::util::ckraw::CkRawMechanism;
use crate::Error;

/// See PKCS#11 v2.40 Section 5.8 Encryption functions

pub extern "C" fn C_EncryptInit(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pMechanism: pkcs11::CK_MECHANISM_PTR,
    hKey: pkcs11::CK_OBJECT_HANDLE,
) -> pkcs11::CK_RV {
    if pMechanism.is_null() {
        error!("C_EncryptInit() called with NULL mech");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    lock_session_mut!(hSession, session, _sarc);

    let raw_mech = unsafe { CkRawMechanism::from_raw_ptr_unchecked(pMechanism) };
    let mech = match Mechanism::from_ckraw_mech(&raw_mech) {
        Ok(mech) => mech,
        Err(_) => return pkcs11::CKR_MECHANISM_INVALID,
    };

    session
        .encrypt_init(&mech, hKey.into())
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_Encrypt(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pData: pkcs11::CK_BYTE_PTR,
    ulDataLen: pkcs11::CK_ULONG,
    pEncryptedData: pkcs11::CK_BYTE_PTR,
    pulEncryptedDataLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_Encrypt() called");

    lock_session_mut!(hSession, session, _sarc);

    let out_len = match session
        .encrypt_ctx()
        .as_mut()
        .ok_or(Error::OperationNotInitialized)
        .and_then(|ctx| {
            ctx.enter_state(OpCtxState::SinglepartActive)
                .map_err(Error::CryptoError)?;
            Ok(ctx.len())
        }) {
        Ok(l) => l,
        Err(e) => {
            session.encrypt_ctx().take();
            return e.into();
        }
    };

    let out_slice = ck_out_buf_to_mut_slice!(pEncryptedData, pulEncryptedDataLen, out_len, {
        session.sign_ctx().take();
    });
    let in_slice = ck_in_buf_to_slice!(pData, ulDataLen, {
        session.sign_ctx().take();
    });

    session
        .encrypt_ctx()
        .take()
        // It's safe to unwrap here since we've already checked for an existing EncryptCtx
        .unwrap()
        .encrypt(in_slice)
        .map(|v| copy_data_to_ck_out_slice(v.as_slice(), out_slice, pulEncryptedDataLen))
        .map_err(Error::CryptoError)
        .unwrap_or_else(|e| e.into())
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

#![allow(non_snake_case)]

/// Utility macro for locking the vToken device
/// It creates two variables in the calling scope, one holding the lock guard
/// and the other a reference to the device.
macro_rules! lock_device {
    ($guard:ident, $device:ident) => {
        let $guard = crate::data::DEVICE.lock().unwrap();
        if $guard.is_none() {
            return crate::pkcs11::CKR_CRYPTOKI_NOT_INITIALIZED;
        }
        let $device = $guard.as_ref().unwrap();
    };
}

/// Utility macro for locking the vToken device
/// It creates two variables in the calling scope, one holding the lock guard
/// and the other a mutable reference to the device.
macro_rules! lock_device_mut {
    ($guard:ident, $device:ident) => {
        let mut $guard = crate::data::DEVICE.lock().unwrap();
        if $guard.is_none() {
            return crate::pkcs11::CKR_CRYPTOKI_NOT_INITIALIZED;
        }
        let $device = $guard.as_mut().unwrap();
    };
}

/// Utility macro for locking a vToken session
/// It creates two variables in the calling scope, one holding the lock guard
/// and the other a reference to the session.
macro_rules! lock_session {
    ($handle:expr, $session:ident, $sess_arc:ident) => {
        let $sess_arc;
        let $session;
        {
            lock_device!(guard, device);
            $sess_arc = device.session($handle);
            $session = match $sess_arc {
                Some(ref sess) => sess.lock().unwrap(),
                None => return crate::pkcs11::CKR_SESSION_HANDLE_INVALID,
            };
        }
    };
}

/// Utility macro for locking a vToken session
/// It creates two variables in the calling scope, one holding the lock guard
/// and the other a mutable reference to the session.
macro_rules! lock_session_mut {
    ($handle:expr, $session:ident, $sess_arc:ident) => {
        let $sess_arc;
        let mut $session;
        {
            lock_device!(guard, device);
            $sess_arc = device.session($handle);
            $session = match $sess_arc {
                Some(ref sess) => sess.lock().unwrap(),
                None => return crate::pkcs11::CKR_SESSION_HANDLE_INVALID,
            };
        }
    };
}

/// Utility macro implementing the PKCS#11 Section 5.2 on producing output
/// for serveral API functions
macro_rules! ck_out_buf_to_mut_slice {
    ($buf_ptr:ident, $buf_len_ptr:ident, $min_len:expr, $on_err:block) => {{
        if $buf_len_ptr.is_null() {
            $on_err;
            return crate::pkcs11::CKR_ARGUMENTS_BAD;
        }
        let user_len = unsafe { std::ptr::read($buf_len_ptr) };
        unsafe {
            std::ptr::write($buf_len_ptr, $min_len as pkcs11::CK_ULONG);
        }
        if $buf_ptr.is_null() {
            return crate::pkcs11::CKR_OK;
        }
        if user_len < $min_len as pkcs11::CK_ULONG {
            return crate::pkcs11::CKR_BUFFER_TOO_SMALL;
        }
        unsafe { std::slice::from_raw_parts_mut($buf_ptr, user_len as usize) }
    }};
}

/// Utility macro for passing raw foreign pointers to slices passed
/// to the module backends
macro_rules! ck_in_buf_to_slice {
    ($buf_ptr:ident, $buf_len_ptr:ident, $on_err:block) => {{
        if $buf_ptr.is_null() {
            $on_err;
            return pkcs11::CKR_ARGUMENTS_BAD;
        }
        unsafe { std::slice::from_raw_parts($buf_ptr, $buf_len_ptr as usize) }
    }};
}

pub mod util {
    use crate::pkcs11;
    use std::ptr;

    pub fn copy_data_to_ck_out_slice(
        src: &[u8],
        dst: &mut [u8],
        dst_len_ptr: *mut pkcs11::CK_ULONG,
    ) -> pkcs11::CK_RV {
        if dst.len() < src.len() {
            // This should never happen, as the output slice should've already been constucted
            // from a validated buffer. If it does, something must be terribly wrong.
            pkcs11::CKR_DEVICE_ERROR
        } else {
            dst[..src.len()].copy_from_slice(src);
            unsafe {
                ptr::write(dst_len_ptr, src.len() as pkcs11::CK_ULONG);
            }
            pkcs11::CKR_OK
        }
    }
}

pub mod decrypt;
pub mod digest;
pub mod encrypt;
pub mod nyi;
pub mod object;
pub mod session;
pub mod sign;
pub mod token;
pub mod verify;

use crate::backend::device::Device;
use crate::data;
use crate::defs;
use crate::pkcs11;
use crate::util::logger::Logger;

/// See PKCS#11 v2.40 Section 5.4 General-purpose functions
#[no_mangle]
pub extern "C" fn C_GetFunctionList(
    pp_fn_list: *mut *const pkcs11::CK_FUNCTION_LIST,
) -> pkcs11::CK_RV {
    if pp_fn_list.is_null() {
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    unsafe {
        std::ptr::write(pp_fn_list, &data::FN_LIST);
    }
    pkcs11::CKR_OK
}

pub extern "C" fn C_Initialize(pInitArgs: pkcs11::CK_VOID_PTR) -> pkcs11::CK_RV {
    // TODO: implement a proper logger, so we
    // don't have to hack this piggybacking on top of C_Initialize.
    Logger::init();

    trace!("C_Initialize() called");

    if defs::CRYPTOKI_VERSION == ck_version!(2, 40) && !pInitArgs.is_null() {
        let args = pInitArgs as pkcs11::CK_C_INITIALIZE_ARGS_PTR;
        unsafe {
            if !(*args).pReserved.is_null() {
                return pkcs11::CKR_ARGUMENTS_BAD;
            }
        }
    }
    let mut maybe_device = data::DEVICE.lock().unwrap();
    if maybe_device.is_some() {
        return pkcs11::CKR_CRYPTOKI_ALREADY_INITIALIZED;
    }

    Device::new()
        .map(|device| maybe_device.replace(device))
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_Finalize(pReserved: pkcs11::CK_VOID_PTR) -> pkcs11::CK_RV {
    trace!("C_Finalize() called");

    if defs::CRYPTOKI_VERSION == ck_version!(2, 40) && !pReserved.is_null() {
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    let mut maybe_device = data::DEVICE.lock().unwrap();
    if maybe_device.is_none() {
        pkcs11::CKR_CRYPTOKI_NOT_INITIALIZED
    } else {
        maybe_device.take();
        pkcs11::CKR_OK
    }
}

pub extern "C" fn C_GetInfo(info: pkcs11::CK_INFO_PTR) -> pkcs11::CK_RV {
    trace!("C_GetInfo() called");

    if info.is_null() {
        error!("C_GetInfo() received NULL pointer.");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    lock_device!(guard, device);
    unsafe {
        std::ptr::write(info, device.ck_info());
    }
    pkcs11::CKR_OK
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

#![allow(dead_code)]
#![allow(unused_variables)]
#![allow(non_snake_case)]
#![allow(non_camel_case_types)]

use crate::pkcs11;

/// Not-Yet-Implemented Cryptoki functions
/// See PKCS#11 v2.40 for information on these functions

pub extern "C" fn C_InitPIN(
    hSession: pkcs11::CK_SESSION_HANDLE,
    uPin: pkcs11::CK_UTF8CHAR_PTR,
    ulPinLen: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_InitPIN() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_SetPIN(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pOldPin: pkcs11::CK_UTF8CHAR_PTR,
    ulOldLen: pkcs11::CK_ULONG,
    pNewPin: pkcs11::CK_UTF8CHAR_PTR,
    ulNewLen: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_SetPIN() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_GetOperationState(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pOperationState: pkcs11::CK_BYTE_PTR,
    pulOperationStateLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_GetOperationState() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_SetOperationState(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pOperationState: pkcs11::CK_BYTE_PTR,
    ulOperationStateLen: pkcs11::CK_ULONG,
    hEncryptionKey: pkcs11::CK_OBJECT_HANDLE,
    hAuthenticationKey: pkcs11::CK_OBJECT_HANDLE,
) -> pkcs11::CK_RV {
    trace!("C_SetOperationState() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_CreateObject(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pTemplate: pkcs11::CK_ATTRIBUTE_PTR,
    ulCount: pkcs11::CK_ULONG,
    phObject: pkcs11::CK_OBJECT_HANDLE_PTR,
) -> pkcs11::CK_RV {
    trace!("C_CreateObject() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_CopyObject(
    hSession: pkcs11::CK_SESSION_HANDLE,
    hObject: pkcs11::CK_OBJECT_HANDLE,
    pTemplate: pkcs11::CK_ATTRIBUTE_PTR,
    ulCount: pkcs11::CK_ULONG,
    phNewObject: pkcs11::CK_OBJECT_HANDLE_PTR,
) -> pkcs11::CK_RV {
    trace!("C_CopyObject() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_DestroyObject(
    hSession: pkcs11::CK_SESSION_HANDLE,
    hObject: pkcs11::CK_OBJECT_HANDLE,
) -> pkcs11::CK_RV {
    trace!("C_DestroyObject() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_SetAttributeValue(
    hSession: pkcs11::CK_SESSION_HANDLE,
    hObject: pkcs11::CK_OBJECT_HANDLE,
    pTemplate: pkcs11::CK_ATTRIBUTE_PTR,
    ulCount: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_SetAttributeValue() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_EncryptUpdate(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pPart: pkcs11::CK_BYTE_PTR,
    ulPartLen: pkcs11::CK_ULONG,
    pEncryptedPart: pkcs11::CK_BYTE_PTR,
    pulEncryptedPartLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_EncryptUpdate() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_EncryptFinal(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pLastEncryptedPart: pkcs11::CK_BYTE_PTR,
    pulLastEncryptedPartLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_EncryptFinal() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_DecryptUpdate(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pEncryptedPart: pkcs11::CK_BYTE_PTR,
    ulEncryptedPartLen: pkcs11::CK_ULONG,
    pPart: pkcs11::CK_BYTE_PTR,
    pulPartLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_DecryptUpdate() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_DecryptFinal(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pLastPart: pkcs11::CK_BYTE_PTR,
    pulLastPartLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_DecryptFinal() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_DigestKey(
    hSession: pkcs11::CK_SESSION_HANDLE,
    hKey: pkcs11::CK_OBJECT_HANDLE,
) -> pkcs11::CK_RV {
    trace!("C_DigestKey() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_SignRecoverInit(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pMechanism: pkcs11::CK_MECHANISM_PTR,
    hKey: pkcs11::CK_OBJECT_HANDLE,
) -> pkcs11::CK_RV {
    trace!("C_SignRecoverInit() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_SignRecover(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pData: pkcs11::CK_BYTE_PTR,
    ulDataLen: pkcs11::CK_ULONG,
    pSignature: pkcs11::CK_BYTE_PTR,
    pulSignatureLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_SignRecover() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_VerifyRecoverInit(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pMechanism: pkcs11::CK_MECHANISM_PTR,
    hKey: pkcs11::CK_OBJECT_HANDLE,
) -> pkcs11::CK_RV {
    trace!("C_VerifyRecoverInit() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_VerifyRecover(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pSignature: pkcs11::CK_BYTE_PTR,
    ulSignatureLen: pkcs11::CK_ULONG,
    pData: pkcs11::CK_BYTE_PTR,
    pulDataLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_VerifyRecover() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_DigestEncryptUpdate(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pPart: pkcs11::CK_BYTE_PTR,
    ulPartLen: pkcs11::CK_ULONG,
    pEncryptedPart: pkcs11::CK_BYTE_PTR,
    pulEncryptedPartLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_DigestEncryptUpdate() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_DecryptDigestUpdate(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pEncryptedPart: pkcs11::CK_BYTE_PTR,
    ulEncryptedPartLen: pkcs11::CK_ULONG,
    pPart: pkcs11::CK_BYTE_PTR,
    pulPartLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_DecryptDigestUpdate() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_SignEncryptUpdate(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pPart: pkcs11::CK_BYTE_PTR,
    ulPartLen: pkcs11::CK_ULONG,
    pEncryptedPart: pkcs11::CK_BYTE_PTR,
    pulEncryptedPartLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_SignEncryptUpdate() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_DecryptVerifyUpdate(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pEncryptedPart: pkcs11::CK_BYTE_PTR,
    ulEncryptedPartLen: pkcs11::CK_ULONG,
    pPart: pkcs11::CK_BYTE_PTR,
    pulPartLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_DecryptVerifyUpdate() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_GenerateKey(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pMechanism: pkcs11::CK_MECHANISM_PTR,
    pTemplate: pkcs11::CK_ATTRIBUTE_PTR,
    ulCount: pkcs11::CK_ULONG,
    phKey: pkcs11::CK_OBJECT_HANDLE_PTR,
) -> pkcs11::CK_RV {
    trace!("C_GenerateKey() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_GenerateKeyPair(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pMechanism: pkcs11::CK_MECHANISM_PTR,
    pPublicKeyTemplate: pkcs11::CK_ATTRIBUTE_PTR,
    ulPublicKeyAttributeCount: pkcs11::CK_ULONG,
    pPrivateKeyTemplate: pkcs11::CK_ATTRIBUTE_PTR,
    ulPrivateKeyAttributeCount: pkcs11::CK_ULONG,
    phPublicKey: pkcs11::CK_OBJECT_HANDLE_PTR,
    phPrivateKey: pkcs11::CK_OBJECT_HANDLE_PTR,
) -> pkcs11::CK_RV {
    trace!("C_GenerateKeyPair() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_WrapKey(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pMechanism: pkcs11::CK_MECHANISM_PTR,
    hWrappingKey: pkcs11::CK_OBJECT_HANDLE,
    hKey: pkcs11::CK_OBJECT_HANDLE,
    pWrappedKey: pkcs11::CK_BYTE_PTR,
    pulWrappedKeyLen: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_WrapKey() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_UnwrapKey(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pMechanism: pkcs11::CK_MECHANISM_PTR,
    hUnwrappingKey: pkcs11::CK_OBJECT_HANDLE,
    pWrappedKey: pkcs11::CK_BYTE_PTR,
    ulWrappedKeyLen: pkcs11::CK_ULONG,
    pTemplate: pkcs11::CK_ATTRIBUTE_PTR,
    ulAttributeCount: pkcs11::CK_ULONG,
    phKey: pkcs11::CK_OBJECT_HANDLE_PTR,
) -> pkcs11::CK_RV {
    trace!("C_UnwrapKey() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_DeriveKey(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pMechanism: pkcs11::CK_MECHANISM_PTR,
    hBaseKey: pkcs11::CK_OBJECT_HANDLE,
    pTemplate: pkcs11::CK_ATTRIBUTE_PTR,
    ulAttributeCount: pkcs11::CK_ULONG,
    phKey: pkcs11::CK_OBJECT_HANDLE_PTR,
) -> pkcs11::CK_RV {
    trace!("C_DeriveKey() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_SeedRandom(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pSeed: pkcs11::CK_BYTE_PTR,
    ulSeedLen: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_SeedRandom() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_GenerateRandom(
    hSession: pkcs11::CK_SESSION_HANDLE,
    RandomData: pkcs11::CK_BYTE_PTR,
    ulRandomLen: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_GenerateRandom() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_WaitForSlotEvent(
    flags: pkcs11::CK_FLAGS,
    pSlot: pkcs11::CK_SLOT_ID_PTR,
    pReserved: pkcs11::CK_VOID_PTR,
) -> pkcs11::CK_RV {
    trace!("C_WaitForSlotEvent() called - NYI");
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::pkcs11;
use crate::util::CkRawAttrTemplate;
use std::ptr;

/// See PKCS#11 v2.40 Section 5.7 Object management functions

pub extern "C" fn C_FindObjectsInit(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pTemplate: pkcs11::CK_ATTRIBUTE_PTR,
    ulCount: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_FindObjectsInit() called");

    if ulCount > 0 && pTemplate.is_null() {
        error!("C_FindObjectsInit() called with non-zero count and NULL template");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    lock_session_mut!(hSession, session, _sess_arc);

    if session.enum_active() {
        return pkcs11::CKR_OPERATION_ACTIVE;
    }

    let attr_template = if !pTemplate.is_null() {
        Some(unsafe { CkRawAttrTemplate::from_raw_ptr_unchecked(pTemplate, ulCount as usize) })
    } else {
        None
    };
    session.enum_init(attr_template);

    pkcs11::CKR_OK
}

pub extern "C" fn C_FindObjects(
    hSession: pkcs11::CK_SESSION_HANDLE,
    phObject: pkcs11::CK_OBJECT_HANDLE_PTR,
    ulMaxObjectCount: pkcs11::CK_ULONG,
    pulObjectCount: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_FindObjects() called");

    if phObject.is_null() || pulObjectCount.is_null() {
        error!("C_FindObects(): NULL output pointer provided.");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }
    lock_session_mut!(hSession, session, _sess_arc);

    if !session.enum_active() {
        return pkcs11::CKR_OPERATION_NOT_INITIALIZED;
    }

    match session.enum_next_chunk(ulMaxObjectCount as usize) {
        Some(chunk) => {
            info!("C_FindObjects returning count={}", chunk.len());
            for (i, _obj) in chunk.iter().enumerate() {
                unsafe {
                    ptr::write(phObject.add(i), chunk[i].into());
                };
            }
            unsafe {
                ptr::write(pulObjectCount, chunk.len() as pkcs11::CK_ULONG);
            };
        }
        None => {
            // Enumeration context consumed
            unsafe {
                ptr::write(pulObjectCount, 0 as pkcs11::CK_ULONG);
            }
        }
    }
    pkcs11::CKR_OK
}

pub extern "C" fn C_FindObjectsFinal(hSession: pkcs11::CK_SESSION_HANDLE) -> pkcs11::CK_RV {
    trace!("C_FindObjectsFinal() called");

    lock_session_mut!(hSession, session, _sess_arc);

    session
        .enum_finalize()
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_GetAttributeValue(
    hSession: pkcs11::CK_SESSION_HANDLE,
    hObject: pkcs11::CK_OBJECT_HANDLE,
    pTemplate: pkcs11::CK_ATTRIBUTE_PTR,
    ulCount: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_GetAttributeValue() called, count={}", ulCount);
    if pTemplate.is_null() {
        error!("C_GetAttributeValue() called with NULL output ptr");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    lock_session!(hSession, session, _sarc);
    let obj = match session.object(hObject.into()) {
        Some(obj) => obj,
        None => return pkcs11::CKR_OBJECT_HANDLE_INVALID,
    };

    // Already checked for null
    let mut attr_tmpl =
        unsafe { CkRawAttrTemplate::from_raw_ptr_unchecked(pTemplate, ulCount as usize) };

    obj.fill_attr_template(&mut attr_tmpl)
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_GetObjectSize(
    hSession: pkcs11::CK_SESSION_HANDLE,
    hObject: pkcs11::CK_OBJECT_HANDLE,
    pulSize: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_GetObjectSize() called");

    if pulSize.is_null() {
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    lock_session!(hSession, session, _src);

    if session.object(hObject.into()).is_some() {
        // The tokens do not export memory consumption info
        // thus neither token objects should. We just implement this
        // function for standard compliance and testing.
        unsafe {
            ptr::write(
                pulSize,
                pkcs11::CK_UNAVAILABLE_INFORMATION as pkcs11::CK_ULONG,
            );
        }
        return pkcs11::CKR_OK;
    }
    pkcs11::CKR_OBJECT_HANDLE_INVALID
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::pkcs11;

/// See PKCS#11 v2.40 Section 5.6 Session management functions

pub extern "C" fn C_OpenSession(
    slotID: pkcs11::CK_SLOT_ID,
    flags: pkcs11::CK_FLAGS,
    _pApplication: pkcs11::CK_VOID_PTR,
    _Notify: pkcs11::CK_NOTIFY,
    phSession: pkcs11::CK_SESSION_HANDLE_PTR,
) -> pkcs11::CK_RV {
    trace!("C_OpenSession() called");

    if phSession.is_null() {
        error!("C_OpenSession() called with NULL output ptr.");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }
    if flags & pkcs11::CKF_SERIAL_SESSION == 0 {
        error!("C_OpenSession() legacy serialization flag must be set.");
        return pkcs11::CKR_SESSION_PARALLEL_NOT_SUPPORTED;
    }
    if flags & pkcs11::CKF_RW_SESSION != 0 {
        error!("C_OpenSession() R/W sessions are not supported");
        return pkcs11::CKR_TOKEN_WRITE_PROTECTED;
    }

    lock_device_mut!(guard, device);

    device
        .open_session(slotID)
        .and_then(|handle| {
            unsafe {
                std::ptr::write(phSession, handle);
            }
            Ok(())
        })
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_CloseSession(hSession: pkcs11::CK_SESSION_HANDLE) -> pkcs11::CK_RV {
    trace!("C_CloseSession() called");
    lock_device_mut!(guard, device);

    device
        .close_session(hSession)
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_CloseAllSessions(slotID: pkcs11::CK_SLOT_ID) -> pkcs11::CK_RV {
    trace!("C_CloseAllSessions() called");
    lock_device_mut!(guard, device);

    device
        .close_all_slot_sessions(slotID)
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_GetSessionInfo(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pInfo: pkcs11::CK_SESSION_INFO_PTR,
) -> pkcs11::CK_RV {
    trace!("C_GetSessionInfo() called");

    if pInfo.is_null() {
        return pkcs11::CKR_ARGUMENTS_BAD;
    }
    lock_session!(hSession, session, _sess_arc);

    // Already checked for NULL ptr.
    let ck_info = session.ck_info();
    unsafe {
        std::ptr::write(pInfo, ck_info);
    }

    pkcs11::CKR_OK
}

pub extern "C" fn C_GetFunctionStatus(hSession: pkcs11::CK_SESSION_HANDLE) -> pkcs11::CK_RV {
    trace!("C_GetFunctionStatus() called");
    lock_session!(hSession, _session, _sess_arc);
    pkcs11::CKR_FUNCTION_NOT_PARALLEL
}

pub extern "C" fn C_CancelFunction(hSession: pkcs11::CK_SESSION_HANDLE) -> pkcs11::CK_RV {
    trace!("C_CancelFunction() called");
    lock_session!(hSession, _session, _sess_arc);
    pkcs11::CKR_FUNCTION_NOT_PARALLEL
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::util::copy_data_to_ck_out_slice;
use crate::backend::Mechanism;
use crate::crypto::OpCtxState;
use crate::pkcs11;
use crate::util::ckraw::CkRawMechanism;
use crate::Error;

/// See PKCS#11 v2.40 5.11 Signing and MACing functions

pub extern "C" fn C_SignInit(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pMechanism: *mut pkcs11::CK_MECHANISM,
    hKey: pkcs11::CK_OBJECT_HANDLE,
) -> pkcs11::CK_RV {
    trace!("C_SignInit() called");

    if pMechanism.is_null() {
        error!("C_SignInit() called with NULL mech");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }
    let raw_mech = unsafe { CkRawMechanism::from_raw_ptr_unchecked(pMechanism) };
    let mech = match Mechanism::from_ckraw_mech(&raw_mech) {
        Ok(mech) => mech,
        Err(_) => return pkcs11::CKR_MECHANISM_INVALID,
    };

    lock_session_mut!(hSession, session, _sarc);
    session
        .sign_init(&mech, hKey.into())
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_Sign(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pData: *mut pkcs11::CK_BYTE,
    ulDataLen: pkcs11::CK_ULONG,
    pSignature: *mut pkcs11::CK_BYTE,
    pulSignatureLen: *mut pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_Sign() called");

    lock_session_mut!(hSession, session, _sarc);

    let out_len = match session
        .sign_ctx()
        .as_mut()
        .ok_or(Error::OperationNotInitialized)
        .and_then(|ctx| {
            ctx.enter_state(OpCtxState::SinglepartActive)
                .map_err(Error::CryptoError)?;
            Ok(ctx.sig_len_ck())
        }) {
        Ok(l) => l,
        Err(e) => {
            session.sign_ctx().take();
            return e.into();
        }
    };

    let out_slice = ck_out_buf_to_mut_slice!(pSignature, pulSignatureLen, out_len, {
        session.sign_ctx().take();
    });
    let in_slice = ck_in_buf_to_slice!(pData, ulDataLen, {
        session.sign_ctx().take();
    });

    session
        .sign_ctx()
        .take()
        // It's safe to unwrap here since we've already checked for a sign ctx.
        .unwrap()
        .sign(in_slice)
        .map(|v| copy_data_to_ck_out_slice(v.as_slice(), out_slice, pulSignatureLen))
        .map_err(Error::CryptoError)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_SignUpdate(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pPart: *mut pkcs11::CK_BYTE,
    ulPartLen: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_SignUpdate() called");
    if pPart.is_null() {
        error!("C_SignUpdate() called with null data ptr.");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    lock_session_mut!(hSession, session, _sarc);
    let in_slice = ck_in_buf_to_slice!(pPart, ulPartLen, {
        session.sign_ctx().take();
    });
    session
        .sign_ctx()
        .take()
        .as_mut()
        .ok_or(Error::OperationNotInitialized)
        .and_then(|ctx| ctx.update(in_slice).map_err(Error::CryptoError))
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_SignFinal(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pSignature: *mut pkcs11::CK_BYTE,
    pulSignatureLen: *mut pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_SignFinal() called");

    lock_session_mut!(hSession, session, _sarc);

    let out_len = match session
        .sign_ctx()
        .as_mut()
        .ok_or(Error::OperationNotInitialized)
        .and_then(|ctx| {
            ctx.enter_state(OpCtxState::MultipartReady)
                .map_err(Error::CryptoError)?;
            Ok(ctx.sig_len_ck())
        }) {
        Ok(l) => l,
        Err(e) => {
            session.sign_ctx().take();
            return e.into();
        }
    };
    let out_slice = ck_out_buf_to_mut_slice!(pSignature, pulSignatureLen, out_len, {
        session.sign_ctx().take();
    });

    session
        .sign_ctx()
        .take()
        // It's safe to unwrap here since we've already checked for a sign ctx.
        .unwrap()
        .finalize()
        .map(|v| copy_data_to_ck_out_slice(v.as_slice(), out_slice, pulSignatureLen))
        .map_err(Error::CryptoError)
        .unwrap_or_else(|e| e.into())
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::pkcs11;
use crate::Error;

/// See PKCS#11 v2.40 Section 5.5 Slot and token management functions

pub extern "C" fn C_GetSlotList(
    tokenPresent: pkcs11::CK_BBOOL,
    pSlotList: pkcs11::CK_SLOT_ID_PTR,
    pulCount: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_GetSlotList() called.");

    lock_device!(guard, device);

    if pulCount.is_null() {
        error!("C_GetSlotList() called with NULL count");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    // Get the current slot count
    let slot_count = device.slot_count(tokenPresent != 0);

    if pSlotList.is_null() {
        unsafe {
            // Already checked for NULL.
            std::ptr::write(pulCount, slot_count as pkcs11::CK_ULONG);
        }
        return pkcs11::CKR_OK;
    } else {
        // Valid slot list - check if it has enough space
        unsafe {
            // Already checked for NULL.
            let in_count = std::ptr::read(pulCount);
            if in_count < (slot_count as pkcs11::CK_ULONG) {
                std::ptr::write(pulCount, slot_count as pkcs11::CK_ULONG);
                return pkcs11::CKR_BUFFER_TOO_SMALL;
            }
        }
    }

    let slot_id_vec = device.ck_slot_ids(tokenPresent != 0);

    // As safe as we can, since this buffer is managed by the lib user, and we've already
    // checked for NULL and avail size.
    unsafe {
        std::ptr::copy_nonoverlapping(slot_id_vec.as_ptr(), pSlotList, slot_id_vec.len());
        std::ptr::write(pulCount, slot_count as pkcs11::CK_ULONG);
    }

    pkcs11::CKR_OK
}

pub extern "C" fn C_GetSlotInfo(
    slotID: pkcs11::CK_SLOT_ID,
    pInfo: pkcs11::CK_SLOT_INFO_PTR,
) -> pkcs11::CK_RV {
    trace!("C_GetSlotInfo() called: slotID={}", slotID);

    lock_device!(guard, device);

    if pInfo.is_null() {
        error!("C_GetSlotInfo() called with NULL info");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    device
        .slot(slotID)
        .ok_or(Error::SlotIdInvalid)
        .and_then(|slot| {
            let ck_info = slot.ck_info();
            unsafe {
                std::ptr::write(pInfo, ck_info);
            }
            Ok(pkcs11::CKR_OK)
        })
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_GetTokenInfo(
    slotID: pkcs11::CK_SLOT_ID,
    pInfo: *mut pkcs11::CK_TOKEN_INFO,
) -> pkcs11::CK_RV {
    trace!("C_GetTokenInfo() called: slotID={}", slotID);

    lock_device!(guard, device);

    if pInfo.is_null() {
        error!("C_GetTokenInfo() called with NULL OUTPUT pointer.");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    device
        .token(slotID)
        .and_then(|token| {
            let ck_info = token.ck_info();
            unsafe {
                std::ptr::write(pInfo, ck_info);
            }
            Ok(pkcs11::CKR_OK)
        })
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_InitToken(
    _slotID: pkcs11::CK_SLOT_ID,
    _pPin: *mut pkcs11::CK_UTF8CHAR,
    _ulPinLen: pkcs11::CK_ULONG,
    _pLabel: *mut pkcs11::CK_UTF8CHAR,
) -> pkcs11::CK_RV {
    pkcs11::CKR_FUNCTION_NOT_SUPPORTED
}

pub extern "C" fn C_GetMechanismList(
    slotID: pkcs11::CK_SLOT_ID,
    pMechanismList: pkcs11::CK_MECHANISM_TYPE_PTR,
    pulCount: pkcs11::CK_ULONG_PTR,
) -> pkcs11::CK_RV {
    trace!("C_GetMechanismList() called");

    if pulCount.is_null() {
        error!("C_GetMechanismList called with NULL count");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    lock_device!(guard, device);

    let mech_count = match device
        .token(slotID)
        .and_then(|tok| tok.mech_count().map_err(Error::TokenError))
    {
        Ok(count) => count,
        Err(e) => return e.into(),
    };

    if pMechanismList.is_null() {
        // Already checked for NULL.
        unsafe {
            std::ptr::write(pulCount, mech_count as pkcs11::CK_ULONG);
        }
        return pkcs11::CKR_OK;
    }

    // Already checked for NULL.
    let in_count = unsafe { std::ptr::read(pulCount) };

    // *pulCount needs to be set to the correct value, even if we aren't able to fill
    // in pMechanismList.
    unsafe {
        std::ptr::write(pulCount, mech_count as pkcs11::CK_ULONG);
    }

    // Check if the user provided a large enough buffer.
    if in_count < (mech_count as pkcs11::CK_ULONG) {
        return pkcs11::CKR_BUFFER_TOO_SMALL;
    }

    device
        .token(slotID)
        .and_then(|tok| tok.mech_list().map_err(Error::TokenError))
        .and_then(|mechs| {
            for (i, mk) in mechs.iter().enumerate() {
                unsafe {
                    std::ptr::write(pMechanismList.add(i), mk.ck_type());
                }
            }
            Ok(pkcs11::CKR_OK)
        })
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_GetMechanismInfo(
    slotID: pkcs11::CK_SLOT_ID,
    type_: pkcs11::CK_MECHANISM_TYPE,
    pInfo: pkcs11::CK_MECHANISM_INFO_PTR,
) -> pkcs11::CK_RV {
    trace!("C_GetMechanismInfo() called");

    if pInfo.is_null() {
        error!("C_GetMechanismInfo() called with NULL info");
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    lock_device!(guard, device);

    device
        .token(slotID)
        .and_then(|tok| tok.mech(type_).map_err(Error::TokenError))
        .and_then(|mech| {
            let mech_info = mech.ck_info();
            unsafe {
                std::ptr::write(pInfo, mech_info);
            }
            Ok(pkcs11::CKR_OK)
        })
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_Login(
    hSession: pkcs11::CK_SESSION_HANDLE,
    userType: pkcs11::CK_USER_TYPE,
    pPin: pkcs11::CK_UTF8CHAR_PTR,
    ulPinLen: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_Login() called");

    if pPin.is_null() {
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    lock_device_mut!(guard, device);

    match userType {
        pkcs11::CKU_SO => pkcs11::CKR_SESSION_READ_ONLY_EXISTS,
        pkcs11::CKU_CONTEXT_SPECIFIC => pkcs11::CKR_USER_TYPE_INVALID,
        pkcs11::CKU_USER => {
            // Already checked pPin for NULL.
            let pin_slice = unsafe { std::slice::from_raw_parts(pPin, ulPinLen as usize) };
            let pin = match std::str::from_utf8(pin_slice) {
                Ok(pin) => pin,
                Err(_) => return pkcs11::CKR_PIN_INCORRECT,
            };
            device
                .login(hSession, pin)
                .map(|_| pkcs11::CKR_OK)
                .unwrap_or_else(|e| e.into())
        }
        _ => pkcs11::CKR_USER_TYPE_INVALID,
    }
}

pub extern "C" fn C_Logout(hSession: pkcs11::CK_SESSION_HANDLE) -> pkcs11::CK_RV {
    trace!("C_Logout() called");

    lock_device_mut!(guard, device);

    device
        .logout(hSession)
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::backend::Mechanism;
use crate::crypto::OpCtxState;
use crate::pkcs11;
use crate::util::ckraw::CkRawMechanism;
use crate::Error;

/// See PKCS#11 v2.40 5.11 Signing and MACing functions

pub extern "C" fn C_VerifyInit(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pMechanism: pkcs11::CK_MECHANISM_PTR,
    hKey: pkcs11::CK_OBJECT_HANDLE,
) -> pkcs11::CK_RV {
    trace!("C_VerifyInit() called");

    if pMechanism.is_null() {
        return pkcs11::CKR_ARGUMENTS_BAD;
    }

    let raw_mech = unsafe { CkRawMechanism::from_raw_ptr_unchecked(pMechanism) };
    let mech = match Mechanism::from_ckraw_mech(&raw_mech) {
        Ok(mech) => mech,
        Err(_) => return pkcs11::CKR_MECHANISM_INVALID,
    };

    lock_session_mut!(hSession, session, _sarc);

    session
        .verify_init(&mech, hKey.into())
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_Verify(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pData: pkcs11::CK_BYTE_PTR,
    ulDataLen: pkcs11::CK_ULONG,
    pSignature: pkcs11::CK_BYTE_PTR,
    ulSignatureLen: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_Verify() called");

    lock_session_mut!(hSession, session, _sarc);

    let ctx = match session.verify_ctx().take() {
        Some(mut ctx) => {
            if let Err(e) = ctx.enter_state(OpCtxState::SinglepartActive) {
                return Error::CryptoError(e).into();
            }
            ctx
        }
        None => return pkcs11::CKR_OPERATION_NOT_INITIALIZED,
    };
    let data_slice = ck_in_buf_to_slice!(pData, ulDataLen, {});
    let sig_slice = ck_in_buf_to_slice!(pSignature, ulSignatureLen, {});

    if !ctx.verify_sig_len_ck(sig_slice.len()) {
        return pkcs11::CKR_SIGNATURE_LEN_RANGE;
    }

    ctx.verify(data_slice, sig_slice)
        .map(|_| pkcs11::CKR_OK)
        .map_err(Error::CryptoError)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_VerifyUpdate(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pPart: pkcs11::CK_BYTE_PTR,
    ulPartLen: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_VerifyUpdate() called");

    lock_session_mut!(hSession, session, _sarc);

    let data_slice = ck_in_buf_to_slice!(pPart, ulPartLen, {});
    session
        .verify_ctx()
        .as_mut()
        .ok_or(Error::OperationNotInitialized)
        .and_then(|ctx| ctx.update(data_slice).map_err(Error::CryptoError))
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}

pub extern "C" fn C_VerifyFinal(
    hSession: pkcs11::CK_SESSION_HANDLE,
    pSignature: pkcs11::CK_BYTE_PTR,
    ulSignatureLen: pkcs11::CK_ULONG,
) -> pkcs11::CK_RV {
    trace!("C_VerifyFinal() called");

    lock_session_mut!(hSession, session, _sarc);

    let ctx = match session.verify_ctx().take() {
        Some(mut ctx) => {
            if let Err(e) = ctx.enter_state(OpCtxState::MultipartReady) {
                return Error::CryptoError(e).into();
            }
            ctx
        }
        None => return pkcs11::CKR_OPERATION_NOT_INITIALIZED,
    };
    let sig_slice = ck_in_buf_to_slice!(pSignature, ulSignatureLen, {});

    if !ctx.verify_sig_len_ck(sig_slice.len()) {
        return pkcs11::CKR_SIGNATURE_LEN_RANGE;
    }

    ctx.finalize(sig_slice)
        .map_err(Error::CryptoError)
        .map(|_| pkcs11::CKR_OK)
        .unwrap_or_else(|e| e.into())
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use vtok_common::config::Config;

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use super::{Session, Slot, Token};
use crate::defs;
use crate::pkcs11;
use crate::{Error, Result};

/// The root device container. Holds a known number of device slots.
/// A successful provisioning operation will attach a token to the respective
/// slot. Newly created sessions are attached to their respective slots and
/// the user authenticates against the token PIN from those slots respectively.
pub struct Device {
    slots: Vec<Slot>,
    session_slot_map: HashMap<pkcs11::CK_SESSION_HANDLE, pkcs11::CK_SLOT_ID>,
    next_session_handle: pkcs11::CK_SESSION_HANDLE,
}

impl Device {
    pub fn new() -> Result<Self> {
        let config = Config::load_ro().map_err(|_| Error::GeneralError)?;

        let mut slots = Vec::with_capacity(defs::MAX_SLOTS);
        for (slot_id, slot_config) in config.slots().iter().enumerate() {
            slots.push(match slot_config {
                None => Slot::new(slot_id as pkcs11::CK_SLOT_ID),
                Some(token_config) => Slot::new_with_token(
                    slot_id as pkcs11::CK_SLOT_ID,
                    Token::from_config(slot_id as pkcs11::CK_SLOT_ID, &token_config)
                        .map_err(Error::TokenError)?,
                ),
            });
        }

        Ok(Self {
            slots,
            session_slot_map: HashMap::new(),
            next_session_handle: 1,
        })
    }

    pub fn ck_info(&self) -> pkcs11::CK_INFO {
        pkcs11::CK_INFO {
            cryptokiVersion: defs::CRYPTOKI_VERSION,
            manufacturerID: ck_padded_str!(defs::DEVICE_MANUFACTURER, 32),
            flags: 0,
            libraryDescription: ck_padded_str!(defs::DEVICE_DESCRIPTION, 32),
            libraryVersion: defs::DEVICE_VERSION,
        }
    }

    pub fn ck_slot_ids(&self, must_have_token: bool) -> Vec<pkcs11::CK_SLOT_ID> {
        self.slots
            .iter()
            .enumerate()
            .filter_map(|(idx, slot)| {
                if must_have_token && !slot.has_token() {
                    None
                } else {
                    Some(idx as pkcs11::CK_SLOT_ID)
                }
            })
            .collect()
    }

    pub fn slot_count(&self, must_have_token: bool) -> usize {
        self.slots
            .iter()
            .filter(|slot| !must_have_token || slot.has_token())
            .count()
    }

    pub fn open_session(
        &mut self,
        slot_id: pkcs11::CK_SLOT_ID,
    ) -> Result<pkcs11::CK_SESSION_HANDLE> {
        let handle = self.next_session_handle;

        self.slot_mut(slot_id)
            .ok_or(Error::CkError(pkcs11::CKR_SLOT_ID_INVALID))?
            .token_mut()
            .ok_or(Error::CkError(pkcs11::CKR_TOKEN_NOT_PRESENT))?
            .open_session(handle)
            .map_err(Error::TokenError)?;

        self.session_slot_map.insert(handle, slot_id);
        self.next_session_handle += 1;

        Ok(handle)
    }

    pub fn close_session(&mut self, handle: pkcs11::CK_SESSION_HANDLE) -> Result<()> {
        let slot_id = *self
            .session_slot_map
            .get(&handle)
            .ok_or(Error::SessionHandleInvalid)?;
        self.token_mut(slot_id)?
            .close_session(handle)
            .map_err(Error::TokenError)?;
        self.session_slot_map.remove(&handle);
        Ok(())
    }

    pub fn close_all_slot_sessions(&mut self, slot_id: pkcs11::CK_SLOT_ID) -> Result<()> {
        self.token_mut(slot_id)
            .map(|tok| tok.close_all_sessions())?;
        self.session_slot_map.retain(|_, &mut sid| sid != slot_id);
        Ok(())
    }

    pub fn session(&self, handle: pkcs11::CK_SESSION_HANDLE) -> Option<Arc<Mutex<Session>>> {
        self.session_slot_map
            .get(&handle)
            .and_then(|&slot_id| self.token(slot_id).ok())
            .and_then(|token| token.session(handle))
    }

    pub fn login(&mut self, session_handle: pkcs11::CK_SESSION_HANDLE, pin: &str) -> Result<()> {
        let slot_id = *self
            .session_slot_map
            .get(&session_handle)
            .ok_or(Error::SessionHandleInvalid)?;
        self.token_mut(slot_id)?
            .login(session_handle, pin)
            .map_err(Error::TokenError)
    }

    pub fn logout(&mut self, session_handle: pkcs11::CK_SESSION_HANDLE) -> Result<()> {
        let slot_id = *self
            .session_slot_map
            .get(&session_handle)
            .ok_or(Error::SessionHandleInvalid)?;
        self.token_mut(slot_id)?
            .logout(session_handle)
            .map_err(Error::TokenError)
    }

    pub fn slot(&self, slot_id: pkcs11::CK_SLOT_ID) -> Option<&Slot> {
        if (slot_id as usize) >= self.slots.len() {
            return None;
        }
        Some(&self.slots[slot_id as usize])
    }

    pub fn slot_mut(&mut self, slot_id: pkcs11::CK_SLOT_ID) -> Option<&mut Slot> {
        if (slot_id as usize) >= self.slots.len() {
            return None;
        }
        Some(&mut self.slots[slot_id as usize])
    }

    pub fn token(&self, slot_id: pkcs11::CK_SLOT_ID) -> Result<&Token> {
        self.slot(slot_id)
            .ok_or(Error::SlotIdInvalid)
            .and_then(|slot| slot.token().ok_or(Error::TokenNotPresent))
    }

    pub fn token_mut(&mut self, slot_id: pkcs11::CK_SLOT_ID) -> Result<&mut Token> {
        self.slot_mut(slot_id)
            .ok_or(Error::SlotIdInvalid)
            .and_then(|slot| slot.token_mut().ok_or(Error::TokenNotPresent))
    }
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::pkcs11;
use crate::util::{CkRawMechanism, Error as CkRawError};

pub enum Error {
    CkRaw(CkRawError),
    DigestMechMismatch,
    UnknownMech,
}

#[derive(Clone, Copy, Debug)]
pub enum MechDigest {
    Sha1,
    Sha224,
    Sha256,
    Sha384,
    Sha512,
}

#[derive(Clone, Copy, Debug)]
pub enum Mechanism {
    Digest(MechDigest),
    RsaPkcs(Option<MechDigest>),
    RsaPkcsPss(Option<MechDigest>, Option<pkcs11::CK_RSA_PKCS_PSS_PARAMS>),
    RsaX509,
    Ecdsa(Option<MechDigest>),
}

/// The token supported mechanisms and their capabilities.
/// See PKCS#11 Mechanisms Specification Version 2.40 for details on how these
/// mechanisms should behave.
///
/// Mechanisms are split into single-part (i.e. sign()) and/or multi-part
/// operations (i.e. sign_update() + sign_final()) depending on type and capabilities.
impl Mechanism {
    const RSA_MIN_KEY_BITS: pkcs11::CK_ULONG = 1024;
    const RSA_MAX_KEY_BITS: pkcs11::CK_ULONG = 8192;
    const EC_MIN_KEY_BITS: pkcs11::CK_ULONG = 224;
    const EC_MAX_KEY_BITS: pkcs11::CK_ULONG = 521;

    pub fn from_ckraw_mech(raw_mech: &CkRawMechanism) -> Result<Self, Error> {
        let mech = match raw_mech.type_() {
            pkcs11::CKM_SHA_1 => Self::Digest(MechDigest::Sha1),
            pkcs11::CKM_SHA224 => Self::Digest(MechDigest::Sha224),
            pkcs11::CKM_SHA256 => Self::Digest(MechDigest::Sha256),
            pkcs11::CKM_SHA384 => Self::Digest(MechDigest::Sha384),
            pkcs11::CKM_SHA512 => Self::Digest(MechDigest::Sha512),
            pkcs11::CKM_RSA_PKCS => Self::RsaPkcs(None),
            pkcs11::CKM_SHA1_RSA_PKCS => Self::RsaPkcs(Some(MechDigest::Sha1)),
            pkcs11::CKM_SHA224_RSA_PKCS => Self::RsaPkcs(Some(MechDigest::Sha224)),
            pkcs11::CKM_SHA256_RSA_PKCS => Self::RsaPkcs(Some(MechDigest::Sha256)),
            pkcs11::CKM_SHA384_RSA_PKCS => Self::RsaPkcs(Some(MechDigest::Sha384)),
            pkcs11::CKM_SHA512_RSA_PKCS => Self::RsaPkcs(Some(MechDigest::Sha512)),
            pkcs11::CKM_RSA_PKCS_PSS => Self::RsaPkcsPss(
                None,
                // Safe because Self::RsaPkcsPss defines the correct params struct type
                // (i.e. CK_RSA_PKCS_PSS_PARAMS).
                unsafe { raw_mech.params() }.map_err(Error::CkRaw)?,
            ),
            pkcs11::CKM_SHA1_RSA_PKCS_PSS => Self::RsaPkcsPss(
                Some(MechDigest::Sha1),
                // Safe because Self::RsaPkcsPss defines the correct params struct type
                // (i.e. CK_RSA_PKCS_PSS_PARAMS).
                unsafe { raw_mech.params() }.map_err(Error::CkRaw)?,
            ),
            pkcs11::CKM_SHA224_RSA_PKCS_PSS => Self::RsaPkcsPss(
                Some(MechDigest::Sha224),
                // Safe because Self::RsaPkcsPss defines the correct params struct type
                // (i.e. CK_RSA_PKCS_PSS_PARAMS).
                unsafe { raw_mech.params() }.map_err(Error::CkRaw)?,
            ),
            pkcs11::CKM_SHA256_RSA_PKCS_PSS => Self::RsaPkcsPss(
                Some(MechDigest::Sha256),
                // Safe because Self::RsaPkcsPss defines the correct params struct type
                // (i.e. CK_RSA_PKCS_PSS_PARAMS).
                unsafe { raw_mech.params() }.map_err(Error::CkRaw)?,
            ),
            pkcs11::CKM_SHA384_RSA_PKCS_PSS => Self::RsaPkcsPss(
                Some(MechDigest::Sha384),
                // Safe because Self::RsaPkcsPss defines the correct params struct type
                // (i.e. CK_RSA_PKCS_PSS_PARAMS).
                unsafe { raw_mech.params() }.map_err(Error::CkRaw)?,
            ),
            pkcs11::CKM_SHA512_RSA_PKCS_PSS => Self::RsaPkcsPss(
                Some(MechDigest::Sha512),
                // Safe because Self::RsaPkcsPss defines the correct params struct type
                // (i.e. CK_RSA_PKCS_PSS_PARAMS).
                unsafe { raw_mech.params() }.map_err(Error::CkRaw)?,
            ),
            pkcs11::CKM_RSA_X_509 => Self::RsaX509,
            pkcs11::CKM_ECDSA => Self::Ecdsa(None),
            pkcs11::CKM_ECDSA_SHA1 => Self::Ecdsa(Some(MechDigest::Sha1)),
            pkcs11::CKM_ECDSA_SHA224 => Self::Ecdsa(Some(MechDigest::Sha224)),
            pkcs11::CKM_ECDSA_SHA256 => Self::Ecdsa(Some(MechDigest::Sha256)),
            pkcs11::CKM_ECDSA_SHA384 => Self::Ecdsa(Some(MechDigest::Sha384)),
            pkcs11::CKM_ECDSA_SHA512 => Self::Ecdsa(Some(MechDigest::Sha512)),

            _ => return Err(Error::UnknownMech),
        };

        // Sanity check: the mechanism-defined hash algorithm must match the algorithm defined
        // by its parameters (if any).
        if let Self::RsaPkcsPss(Some(digest), Some(params)) = mech {
            if params.hashAlg != Self::Digest(digest).ck_type() {
                return Err(Error::DigestMechMismatch);
            }
        }

        Ok(mech)
    }

    pub fn ck_type(&self) -> pkcs11::CK_MECHANISM_TYPE {
        match self {
            Self::Digest(digest) => match digest {
                MechDigest::Sha1 => pkcs11::CKM_SHA_1,
                MechDigest::Sha224 => pkcs11::CKM_SHA224,
                MechDigest::Sha256 => pkcs11::CKM_SHA256,
                MechDigest::Sha384 => pkcs11::CKM_SHA384,
                MechDigest::Sha512 => pkcs11::CKM_SHA512,
            },
            Self::RsaPkcs(digest) => match digest {
                None => pkcs11::CKM_RSA_PKCS,
                Some(MechDigest::Sha1) => pkcs11::CKM_SHA1_RSA_PKCS,
                Some(MechDigest::Sha224) => pkcs11::CKM_SHA224_RSA_PKCS,
                Some(MechDigest::Sha256) => pkcs11::CKM_SHA256_RSA_PKCS,
                Some(MechDigest::Sha384) => pkcs11::CKM_SHA384_RSA_PKCS,
                Some(MechDigest::Sha512) => pkcs11::CKM_SHA512_RSA_PKCS,
            },
            Self::RsaPkcsPss(digest, _) => match digest {
                None => pkcs11::CKM_RSA_PKCS_PSS,
                Some(MechDigest::Sha1) => pkcs11::CKM_SHA1_RSA_PKCS_PSS,
                Some(MechDigest::Sha224) => pkcs11::CKM_SHA224_RSA_PKCS_PSS,
                Some(MechDigest::Sha256) => pkcs11::CKM_SHA256_RSA_PKCS_PSS,
                Some(MechDigest::Sha384) => pkcs11::CKM_SHA384_RSA_PKCS_PSS,
                Some(MechDigest::Sha512) => pkcs11::CKM_SHA512_RSA_PKCS_PSS,
            },
            Self::RsaX509 => pkcs11::CKM_RSA_X_509,
            Self::Ecdsa(digest) => match digest {
                None => pkcs11::CKM_ECDSA,
                Some(MechDigest::Sha1) => pkcs11::CKM_ECDSA_SHA1,
                Some(MechDigest::Sha224) => pkcs11::CKM_ECDSA_SHA224,
                Some(MechDigest::Sha256) => pkcs11::CKM_ECDSA_SHA256,
                Some(MechDigest::Sha384) => pkcs11::CKM_ECDSA_SHA384,
                Some(MechDigest::Sha512) => pkcs11::CKM_ECDSA_SHA512,
            },
        }
    }

    pub fn ck_info(&self) -> pkcs11::CK_MECHANISM_INFO {
        let (min_bits, max_bits) = match self {
            Self::Digest(_) => (0, 0),
            Self::RsaPkcs(_) | Self::RsaPkcsPss(_, _) | Self::RsaX509 => {
                (Self::RSA_MIN_KEY_BITS, Self::RSA_MAX_KEY_BITS)
            }
            Self::Ecdsa(_) => (Self::EC_MIN_KEY_BITS, Self::EC_MAX_KEY_BITS),
        };
        pkcs11::CK_MECHANISM_INFO {
            ulMinKeySize: min_bits,
            ulMaxKeySize: max_bits,
            flags: self.ck_flags(),
        }
    }

    pub fn ck_flags(&self) -> pkcs11::CK_FLAGS {
        // NOTE: Though we have a soft-token, we stamp the pkcs11::CKF_HW flag since most unit
        // tests out there seem to check for it
        pkcs11::CKF_HW
            | match self {
                Self::Digest(_) => pkcs11::CKF_DIGEST,
                // Single-part CKM_RSA_PKCS also has encrypt/decrypt
                Self::RsaPkcs(None) => {
                    pkcs11::CKF_SIGN
                        | pkcs11::CKF_VERIFY
                        | pkcs11::CKF_DECRYPT
                        | pkcs11::CKF_ENCRYPT
                }
                // Multi-part CKM_RSA_PKCS has sign/verify only
                Self::RsaPkcs(Some(_)) => pkcs11::CKF_SIGN | pkcs11::CKF_VERIFY,
                Self::RsaPkcsPss(_, _) => pkcs11::CKF_SIGN | pkcs11::CKF_VERIFY,
                Self::RsaX509 => {
                    pkcs11::CKF_SIGN
                        | pkcs11::CKF_VERIFY
                        | pkcs11::CKF_DECRYPT
                        | pkcs11::CKF_ENCRYPT
                }
                Self::Ecdsa(_) => {
                    pkcs11::CKF_SIGN
                        | pkcs11::CKF_VERIFY
                        | pkcs11::CKF_EC_F_P
                        | pkcs11::CKF_EC_NAMEDCURVE
                        | pkcs11::CKF_EC_UNCOMPRESS
                }
            }
    }

    pub fn is_multipart(&self) -> bool {
        match self {
            Self::RsaPkcs(ref digest)
            | Self::RsaPkcsPss(ref digest, _)
            | Self::Ecdsa(ref digest) => digest.is_some(),
            _ => false,
        }
    }
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

pub mod db;
pub mod device;
pub mod mech;
pub mod session;
pub mod slot;
pub mod token;

pub use db::Db;
pub use device::Device;
pub use mech::Mechanism;
pub use session::Session;
pub use slot::Slot;
pub use token::Token;


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::cmp;

use super::db::{Db, Object, ObjectHandle, ObjectKind};
use super::Mechanism;
use crate::crypto;
use crate::crypto::{
    DecryptCtx, DigestCtx, DigestSignCtx, DigestVerifyCtx, DirectDecryptCtx, DirectEncryptCtx,
    DirectSignCtx, DirectVerifyCtx, EncryptCtx, SignCtx, VerifyCtx,
};
use crate::pkcs11;
use crate::util::CkRawAttrTemplate;
use crate::{Error, Result};

struct EnumCtx {
    handles: Vec<ObjectHandle>,
    index: usize,
}

impl EnumCtx {
    fn new(handles: Vec<ObjectHandle>) -> Self {
        Self { handles, index: 0 }
    }
    fn next_chunk(&mut self, count: usize) -> &[ObjectHandle] {
        let end = cmp::min(self.index + count, self.handles.len());
        let ret = &self.handles[self.index..end];
        self.index += ret.len();
        ret
    }
}

#[derive(Clone, Copy)]
pub enum SessionState {
    RoPublic,
    RoUser,
}

/// The active session states
/// New sessions start with R/O Public state and once
/// the user logins, the session shall enter R/O User Functions
/// state. Logout will cause the state to transition back to
/// the R/O Public state.
impl SessionState {
    fn to_ck_state(&self) -> pkcs11::CK_STATE {
        match self {
            Self::RoPublic => pkcs11::CKS_RO_PUBLIC_SESSION,
            Self::RoUser => pkcs11::CKS_RO_USER_FUNCTIONS,
        }
    }
}

/// The session container. Holds the provisioned object database
/// and also the operation context currently active (if any).
/// The enumeration context is used only when user requests to
/// find all objects attached. The other contexts are during actual
/// session cryptopgrahic operations. Due to the nature of the Cryptoki
/// PKCS#11 specification all operation contexts (digest, sign, verify,
/// encrypt, decrypt) are mutually exclusive.
pub struct Session {
    slot_id: pkcs11::CK_SLOT_ID,
    state: SessionState,
    db: Db,
    enum_ctx: Option<EnumCtx>,
    digest_ctx: Option<DigestCtx>,
    sign_ctx: Option<Box<dyn SignCtx>>,
    verify_ctx: Option<Box<dyn VerifyCtx>>,
    decrypt_ctx: Option<Box<dyn DecryptCtx>>,
    encrypt_ctx: Option<Box<dyn EncryptCtx>>,
}

impl Session {
    pub fn new(slot_id: pkcs11::CK_SLOT_ID, db: Db, state: SessionState) -> Self {
        Self {
            slot_id,
            state,
            db,
            enum_ctx: None,
            digest_ctx: None,
            sign_ctx: None,
            verify_ctx: None,
            decrypt_ctx: None,
            encrypt_ctx: None,
        }
    }

    pub fn ck_info(&self) -> pkcs11::CK_SESSION_INFO {
        pkcs11::CK_SESSION_INFO {
            slotID: self.slot_id,
            state: self.state.to_ck_state(),
            flags: pkcs11::CKF_SERIAL_SESSION,
            ulDeviceError: pkcs11::CKR_OK,
        }
    }

    pub fn state(&self) -> SessionState {
        self.state
    }

    pub fn set_state(&mut self, state: SessionState) {
        self.state = state
    }

    pub fn object(&self, handle: ObjectHandle) -> Option<&Object> {
        self.db.object(handle)
    }

    /// Enumerate session objects, possibly matching a given template.
    /// If no template is provided, all available objects are enumerated.
    pub fn enum_init(&mut self, template: Option<CkRawAttrTemplate>) {
        let enable_private = self.check_user_logged_in().is_ok();
        let handles: Vec<ObjectHandle> = self
            .db
            .enumerate()
            .filter_map(|(h, o)| {
                let pass = (enable_private || !o.is_private())
                    && template
                        .as_ref()
                        .map(|tpl| o.match_attr_template(tpl))
                        // Per the PKCS#11 2.40 spec, mechanism objects should not be listed
                        // unless specifically searched for by using a template with
                        // CKA_CLASS = CKO_MECHANISM.
                        .unwrap_or(!o.is_mechanism());
                if pass {
                    Some(h)
                } else {
                    None
                }
            })
            .collect();
        self.enum_ctx = Some(EnumCtx::new(handles));
    }

    pub fn enum_next_chunk(&mut self, count: usize) -> Option<&[ObjectHandle]> {
        self.enum_ctx.as_mut().map(|x| x.next_chunk(count))
    }

    /// Enumeration is already in progress
    pub fn enum_active(&self) -> bool {
        self.enum_ctx.is_some()
    }

    /// Finalize an active enumeration operation
    pub fn enum_finalize(&mut self) -> Result<()> {
        self.enum_ctx
            .take()
            .ok_or(Error::CkError(pkcs11::CKR_OPERATION_NOT_INITIALIZED))?;
        Ok(())
    }

    /// Initialize a digest context for a digest operation
    pub fn digest_init(&mut self, mech_type: pkcs11::CK_MECHANISM_TYPE) -> Result<()> {
        self.digest_ctx = Some(
            DigestCtx::new(mech_type).map_err(|_| Error::CkError(pkcs11::CKR_MECHANISM_INVALID))?,
        );
        Ok(())
    }

    pub fn digest_ctx(&mut self) -> &mut Option<crypto::DigestCtx> {
        &mut self.digest_ctx
    }

    /// Initialize a signing context for a signing operation
    pub fn sign_init(&mut self, mech: &Mechanism, key_handle: ObjectHandle) -> Result<()> {
        self.check_user_logged_in()?;
        let pkey = self.private_key_for_mech(mech, key_handle)?;
        self.sign_ctx = Some(if mech.is_multipart() {
            Box::new(DigestSignCtx::new(mech, pkey).map_err(Error::CryptoError)?)
        } else {
            Box::new(DirectSignCtx::new(mech, pkey).map_err(Error::CryptoError)?)
        });
        Ok(())
    }

    pub fn sign_ctx(&mut self) -> &mut Option<Box<dyn SignCtx>> {
        &mut self.sign_ctx
    }

    /// Initialize a verification context for a verification operation
    pub fn verify_init(&mut self, mech: &Mechanism, key_handle: ObjectHandle) -> Result<()> {
        self.check_user_logged_in()?;
        let pkey = self.public_key_for_mech(mech, key_handle)?;
        self.verify_ctx = Some(if mech.is_multipart() {
            Box::new(DigestVerifyCtx::new(mech, pkey).map_err(Error::CryptoError)?)
        } else {
            Box::new(DirectVerifyCtx::new(mech, pkey).map_err(Error::CryptoError)?)
        });
        Ok(())
    }

    pub fn verify_ctx(&mut self) -> &mut Option<Box<dyn VerifyCtx>> {
        &mut self.verify_ctx
    }

    /// Initialize an encryption context for an encryption operation
    pub fn encrypt_init(&mut self, mech: &Mechanism, key_handle: ObjectHandle) -> Result<()> {
        self.check_user_logged_in()?;
        let pkey = self.public_key_for_mech(mech, key_handle)?;
        self.encrypt_ctx = Some(if mech.is_multipart() {
            return Err(Error::MechanismInvalid);
        } else {
            Box::new(DirectEncryptCtx::new(mech, pkey).map_err(Error::CryptoError)?)
        });
        Ok(())
    }

    pub fn encrypt_ctx(&mut self) -> &mut Option<Box<dyn EncryptCtx>> {
        &mut self.encrypt_ctx
    }

    /// Initialize an decryption context for a decryption operation
    pub fn decrypt_init(&mut self, mech: &Mechanism, key_handle: ObjectHandle) -> Result<()> {
        self.check_user_logged_in()?;
        let pkey = self.private_key_for_mech(mech, key_handle)?;
        self.decrypt_ctx = Some(if mech.is_multipart() {
            return Err(Error::MechanismInvalid);
        } else {
            Box::new(DirectDecryptCtx::new(mech, pkey).map_err(Error::CryptoError)?)
        });
        Ok(())
    }

    pub fn decrypt_ctx(&mut self) -> &mut Option<Box<dyn DecryptCtx>> {
        &mut self.decrypt_ctx
    }

    /// Helper for checking if the user has been logged in. Some cryptograhic
    /// operations (i.e. signing with private keys) require login.
    fn check_user_logged_in(&self) -> Result<()> {
        match self.state {
            SessionState::RoUser => Ok(()),
            _ => Err(Error::UserNotLoggedIn),
        }
    }

    /// Constructs an EVP_PKEY wrapper object from a p11 object based on an input
    /// PEM formatted string while also keeping consistency with the requested mechanism.
    fn private_key_for_mech(
        &self,
        mech: &Mechanism,
        key_handle: ObjectHandle,
    ) -> Result<crypto::Pkey> {
        let key_obj = self.db.object(key_handle).ok_or(Error::KeyHandleInvalid)?;
        match mech {
            Mechanism::RsaX509 | Mechanism::RsaPkcs(..) | Mechanism::RsaPkcsPss(..) => {
                if let ObjectKind::RsaPrivateKey(pem) = key_obj.kind() {
                    crypto::Pkey::from_private_pem(pem.as_str()).map_err(Error::CryptoError)
                } else {
                    Err(Error::KeyTypeInconsistent)
                }
            }
            Mechanism::Ecdsa(..) => {
                if let ObjectKind::EcPrivateKey(pem) = key_obj.kind() {
                    crypto::Pkey::from_private_pem(pem.as_str()).map_err(Error::CryptoError)
                } else {
                    Err(Error::KeyTypeInconsistent)
                }
            }
            _ => Err(Error::MechanismInvalid),
        }
    }

    /// Constructs an internal public key object type from the provisioned PEM file
    fn public_key_for_mech(
        &self,
        mech: &Mechanism,
        key_handle: ObjectHandle,
    ) -> Result<crypto::Pkey> {
        let key_obj = self.db.object(key_handle).ok_or(Error::KeyHandleInvalid)?;
        match mech {
            Mechanism::RsaX509 | Mechanism::RsaPkcs(..) | Mechanism::RsaPkcsPss(..) => {
                if let ObjectKind::RsaPublicKey(pem) = key_obj.kind() {
                    crypto::Pkey::from_private_pem(pem.as_str()).map_err(Error::CryptoError)
                } else {
                    Err(Error::KeyTypeInconsistent)
                }
            }
            Mechanism::Ecdsa(..) => {
                if let ObjectKind::EcPublicKey(pem) = key_obj.kind() {
                    crypto::Pkey::from_private_pem(pem.as_str()).map_err(Error::CryptoError)
                } else {
                    Err(Error::KeyTypeInconsistent)
                }
            }
            _ => Err(Error::MechanismInvalid),
        }
    }
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::token::Token;
use crate::defs;
use crate::pkcs11;

/// Device slot container. A token shall be attached to it
/// (i.e. CKF_TOKEN_PRESENT bit set) once it is successfully
/// provisioned with a database.
pub struct Slot {
    _id: pkcs11::CK_SLOT_ID,
    token: Option<Token>,
}

impl Slot {
    pub fn new(id: pkcs11::CK_SLOT_ID) -> Self {
        Self {
            _id: id,
            token: None,
        }
    }

    pub fn new_with_token(id: pkcs11::CK_SLOT_ID, token: Token) -> Self {
        Self {
            _id: id,
            token: Some(token),
        }
    }

    pub fn ck_info(&self) -> pkcs11::CK_SLOT_INFO {
        let mut flags = 0;
        if self.has_token() {
            flags |= pkcs11::CKF_TOKEN_PRESENT;
        }
        pkcs11::CK_SLOT_INFO {
            slotDescription: ck_padded_str!(defs::SLOT_DESCRIPTION, 64),
            manufacturerID: ck_padded_str!(defs::SLOT_MANUFACTURER, 32),
            flags,
            hardwareVersion: defs::SLOT_HARDWARE_VERSION,
            firmwareVersion: defs::SLOT_FIRMWARE_VERSION,
        }
    }

    pub fn has_token(&self) -> bool {
        self.token
            .as_ref()
            .filter(|tok| !tok.has_expired())
            .is_some()
    }

    pub fn token(&self) -> Option<&Token> {
        self.token.as_ref().filter(|tok| !tok.has_expired())
    }

    pub fn token_mut(&mut self) -> Option<&mut Token> {
        self.token.as_mut().filter(|tok| !tok.has_expired())
    }
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use vtok_common::{config, util};

use crate::defs;
use crate::pkcs11;

use super::db;
use super::db::{Db, ObjectKind};
use super::session::{Session, SessionState};
use super::Mechanism;

#[derive(Clone, Copy, Debug)]
pub enum Error {
    DbLoad(db::Error),
    MechNotFound,
    PinIncorrect,
    SessionCount,
    SessionHandleInvalid,
    SessionLockPoisoned,
    UserAlreadyLoggedIn,
    UserNotLoggedIn,
    TokenUninit,
}
pub type Result<T> = std::result::Result<T, Error>;

/// Slot token container. Contains all active sessions created against
/// the parent slot of this token. The database is attached to the token
/// and each session receives a copy of the provisioned database for
/// working with the database objects. The user must login successfuly
/// in order to actually do cryptographic operations in the context of
/// that respective session.
pub struct Token {
    slot_id: pkcs11::CK_SLOT_ID,
    label: String,
    sessions: HashMap<pkcs11::CK_SESSION_HANDLE, Arc<Mutex<Session>>>,
    db: Option<Db>,
    user_login: bool,
    expiry_ts: u64,
}

impl Token {
    pub fn from_config(slot_id: pkcs11::CK_SLOT_ID, token_config: &config::Token) -> Result<Self> {
        Ok(Self {
            slot_id,
            label: token_config.label.clone(),
            sessions: HashMap::new(),
            db: Some(Db::from_token_config(token_config).map_err(Error::DbLoad)?),
            user_login: false,
            expiry_ts: token_config.expiry_ts,
        })
    }

    pub fn ck_info(&self) -> pkcs11::CK_TOKEN_INFO {
        let mut flags = match self.db.as_ref() {
            Some(db) => {
                // If provisioning was successful, then it means
                // the user initialized the token along with a CK_USER PIN
                pkcs11::CKF_TOKEN_INITIALIZED
                    | pkcs11::CKF_USER_PIN_INITIALIZED
                    | if db.enumerate().any(|(_, o)| o.is_private()) {
                        // If any private objects are present then
                        // the user must login before being able to
                        // do certain operations (i.e. signing).
                        pkcs11::CKF_LOGIN_REQUIRED
                    } else {
                        0
                    }
            }
            None => 0,
        };

        // Persistent token flags: only CK_USER R/O are sessions supported
        flags |= pkcs11::CKF_WRITE_PROTECTED | pkcs11::CKF_SO_PIN_LOCKED;

        // TODO: set serial properly
        let serial = format!("EVT{:02X}", self.slot_id);

        pkcs11::CK_TOKEN_INFO {
            label: ck_padded_str!(self.label.as_str(), 32),
            manufacturerID: ck_padded_str!(defs::TOKEN_MANUFACTURER, 32),
            model: ck_padded_str!(defs::TOKEN_MODEL, 16),
            serialNumber: ck_padded_str!(serial.as_str(), 16),
            flags,
            ulMaxSessionCount: defs::TOKEN_MAX_SESSIONS,
            ulSessionCount: self.sessions.len() as pkcs11::CK_ULONG,
            ulMaxRwSessionCount: defs::TOKEN_MAX_RW_SESSIONS,
            ulRwSessionCount: 0,
            ulMaxPinLen: defs::TOKEN_MAX_PIN_LEN,
            ulMinPinLen: defs::TOKEN_MIN_PIN_LEN,
            ulTotalPublicMemory: pkcs11::CK_UNAVAILABLE_INFORMATION,
            ulFreePublicMemory: pkcs11::CK_UNAVAILABLE_INFORMATION,
            ulTotalPrivateMemory: pkcs11::CK_UNAVAILABLE_INFORMATION,
            ulFreePrivateMemory: pkcs11::CK_UNAVAILABLE_INFORMATION,
            hardwareVersion: defs::TOKEN_HARDWARE_VERSION,
            firmwareVersion: defs::TOKEN_FIRMWARE_VERSION,
            utcTime: ck_padded_str!(defs::TOKEN_UTC_TIME, 16),
        }
    }

    pub fn has_expired(&self) -> bool {
        util::time::monotonic_secs() > self.expiry_ts
    }

    pub fn open_session(&mut self, handle: pkcs11::CK_SESSION_HANDLE) -> Result<()> {
        if self.sessions.len() >= defs::TOKEN_MAX_SESSIONS as usize {
            return Err(Error::SessionCount);
        }
        let db_clone = self.db.as_ref().cloned().ok_or(Error::TokenUninit)?;
        let state = match self.user_login {
            // If an active user login is present on the token, all future
            // sessions from the user shall enter RoUser state implicitly
            true => SessionState::RoUser,
            false => SessionState::RoPublic,
        };
        self.sessions.insert(
            handle,
            Arc::new(Mutex::new(Session::new(self.slot_id, db_clone, state))),
        );

        Ok(())
    }

    pub fn close_session(&mut self, handle: pkcs11::CK_SESSION_HANDLE) -> Result<()> {
        self.sessions
            .remove(&handle)
            .ok_or(Error::SessionHandleInvalid)
            .map(|_| ())
    }

    pub fn close_all_sessions(&mut self) {
        self.sessions.clear();
    }

    pub fn session(&self, handle: pkcs11::CK_SESSION_HANDLE) -> Option<Arc<Mutex<Session>>> {
        self.sessions.get(&handle).cloned()
    }

    pub fn login(&mut self, session_handle: pkcs11::CK_SESSION_HANDLE, pin: &str) -> Result<()> {
        if self.db.as_ref().ok_or(Error::TokenUninit)?.token_pin() != pin {
            return Err(Error::PinIncorrect);
        }
        let sarc = self
            .session(session_handle)
            .ok_or(Error::SessionHandleInvalid)?;
        let mut session = sarc.lock().map_err(|_| Error::SessionLockPoisoned)?;
        match session.state() {
            SessionState::RoUser => return Err(Error::UserAlreadyLoggedIn),
            SessionState::RoPublic => session.set_state(SessionState::RoUser),
        };

        for (handle, sarc) in self.sessions.iter() {
            if *handle != session_handle {
                sarc.lock()
                    .map_err(|_| Error::SessionLockPoisoned)?
                    .set_state(SessionState::RoUser);
            }
        }

        // Token now has an active user login
        self.user_login = true;

        Ok(())
    }

    pub fn logout(&mut self, session_handle: pkcs11::CK_SESSION_HANDLE) -> Result<()> {
        let sarc = self
            .session(session_handle)
            .ok_or(Error::SessionHandleInvalid)?;
        let mut session = sarc.lock().map_err(|_| Error::SessionLockPoisoned)?;
        match session.state() {
            SessionState::RoPublic => return Err(Error::UserNotLoggedIn),
            SessionState::RoUser => session.set_state(SessionState::RoPublic),
        };
        for (handle, sarc) in self.sessions.iter() {
            if *handle != session_handle {
                sarc.lock()
                    .map_err(|_| Error::SessionLockPoisoned)?
                    .set_state(SessionState::RoPublic);
            }
        }

        // Token now has an inactive user login
        self.user_login = false;

        Ok(())
    }

    pub fn mech_count(&self) -> Result<usize> {
        Ok(self
            .db
            .as_ref()
            .ok_or(Error::TokenUninit)?
            .enumerate()
            .filter_map(|(_, obj)| match obj.kind() {
                ObjectKind::Mechanism(_) => Some(true),
                _ => None,
            })
            .count())
    }

    pub fn mech_list(&self) -> Result<Vec<Mechanism>> {
        Ok(self
            .db
            .as_ref()
            .ok_or(Error::TokenUninit)?
            .enumerate()
            .filter_map(|(_, obj)| match obj.kind() {
                ObjectKind::Mechanism(mech) => Some(*mech),
                _ => None,
            })
            .collect())
    }

    pub fn mech(&self, mech_type: pkcs11::CK_MECHANISM_TYPE) -> Result<Mechanism> {
        self.db
            .as_ref()
            .ok_or(Error::TokenUninit)?
            .enumerate()
            .filter_map(|(_, obj)| match obj.kind() {
                ObjectKind::Mechanism(mech) if mech.ck_type() == mech_type => Some(*mech),
                _ => None,
            })
            .take(1)
            .next()
            .ok_or(Error::MechNotFound)
    }
}


// Copyright 2020-2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use vtok_common::config;

use crate::crypto;
use crate::defs;
use crate::pkcs11;

pub mod object;

pub use object::{Object, ObjectHandle, ObjectKind};

// NOTE: for now, we use these *Info structs to construct key objects. The source PEM is
// preserved, so that a crypto::Pkey (an EVP_PKEY wrapper) can be constructed whenever
// it is needed (e.g. at operation context initialization).
// If the PEM to EVP_PKEY conversion turns out to impact performance, we could construct
// the crypto::Pkey object at DB creation time, and replace the *Info structs with it,
// provided we also implement a proper cloning mechanism for crypto::Pkey. This is needed
// in order to make sure that each session gets its own copy of each key, and maintain
// thread safety.
// Cloning could be done via RSAPrivateKey_dup() and EC_KEY_dup(), together with a TryClone
// trait, since these operations can fail.
#[derive(Clone)]
pub struct RsaKeyInfo {
    pub priv_pem: String,
    pub id: pkcs11::CK_BYTE,
    pub label: String,
    pub num_bits: pkcs11::CK_ULONG,
    pub modulus: Vec<u8>,
    pub public_exponent: Vec<u8>,
}

#[derive(Clone)]
pub struct EcKeyInfo {
    pub priv_pem: String,
    pub id: pkcs11::CK_BYTE,
    pub label: String,
    pub params_x962: Vec<u8>,
    pub point_q_x962: Vec<u8>,
}

#[derive(Clone)]
pub struct CertInfo {
    pub categ: crypto::CertCategory,
    pub id: pkcs11::CK_BYTE,
    pub label: String,
    pub subject_der: Vec<u8>,
    pub issuer_der: Vec<u8>,
    pub serno_der: Vec<u8>,
    pub cert_der: Vec<u8>,
}

#[derive(Clone, Copy, Debug)]
pub enum Error {
    GeneralError,
    CryptoError(crypto::Error),
    PemError(crypto::Error),
}

type Result<T> = std::result::Result<T, Error>;

#[derive(Clone)]
pub struct Db {
    token_pin: String,
    objects: Vec<Object>,
}

impl Db {
    pub fn from_token_config(token_config: &config::Token) -> Result<Self> {
        let mut objects = Vec::new();

        for mech in defs::TOKEN_MECH_LIST.iter() {
            objects.push(Object::new_mechanism(*mech));
        }

        for key_config in token_config.private_keys.iter() {
            let pkey =
                crypto::Pkey::from_private_pem(key_config.pem.as_str()).map_err(Error::PemError)?;
            match pkey.algo().map_err(Error::CryptoError)? {
                crypto::KeyAlgo::Rsa => {
                    let info = RsaKeyInfo {
                        id: key_config.id,
                        label: key_config.label.clone(),
                        priv_pem: key_config.pem.to_string(),
                        num_bits: pkey.num_bits().map_err(Error::CryptoError)? as u64,
                        modulus: pkey.rsa_modulus().map_err(Error::CryptoError)?,
                        public_exponent: pkey.rsa_public_exponent().map_err(Error::CryptoError)?,
                    };
                    objects.push(Object::new_rsa_private_key(info.clone()));
                    objects.push(Object::new_rsa_public_key(info));
                }
                crypto::KeyAlgo::Ec => {
                    let info = EcKeyInfo {
                        id: key_config.id,
                        label: key_config.label.clone(),
                        priv_pem: key_config.pem.to_string(),
                        params_x962: pkey.ec_params_x962().map_err(Error::CryptoError)?,
                        point_q_x962: pkey.ec_point_q_x962().map_err(Error::CryptoError)?,
                    };
                    objects.push(Object::new_ec_private_key(info.clone()));
                    objects.push(Object::new_ec_public_key(info));
                }
            }

            let x509_vfy = || -> Result<()> {
                // If the provisioned private key has an optional X509 certificate
                // or certificate chain attached to it. Soft failure
                if let Some(x509_pem) = &key_config.cert_pem {
                    let x509_chain = crypto::X509Chain::chain_from_pem(&x509_pem.as_str())
                        .map_err(Error::PemError)?;

                    // If multiple certs, build and verify the chain, otherwise check if it
                    // is a self-signed certificate.
                    match x509_chain.multiple_certs() {
                        true => x509_chain.verify_chain().map_err(Error::CryptoError)?,
                        false => x509_chain.verify(pkey).map_err(Error::CryptoError)?,
                    }
                    for (pos, x509) in x509_chain.enumerate() {
                        // The first certificate shares the same ID with the private key
                        // and the rest of the certificate chain get incremental IDs and
                        // are marked as CAs.
                        let (categ, id) = {
                            if pos == 0 {
                                (crypto::CertCategory::Token, key_config.id)
                            } else {
                                (crypto::CertCategory::Authority, key_config.id + pos)
                            }
                        };
                        let x509_obj = CertInfo {
                            categ: categ,
                            id: id,
                            label: format!("acm-ne-cert-{}", pos).to_string(),
                            subject_der: x509.subject_name().map_err(Error::CryptoError)?,
                            issuer_der: x509.issuer().map_err(Error::CryptoError)?,
                            serno_der: x509.serial_no().map_err(Error::CryptoError)?,
                            cert_der: x509.to_der().map_err(Error::CryptoError)?,
                        };
                        objects.push(Object::new_x509_cert(x509_obj));
                    }
                }
                Ok(())
            };
            if let Err(_) = x509_vfy() {
                warn!("X509 certificate chain verification failed.");
            }
        }

        Ok(Self {
            token_pin: token_config.pin.clone(),
            objects,
        })
    }

    pub fn enumerate(&self) -> impl Iterator<Item = (ObjectHandle, &Object)> {
        self.objects
            .iter()
            .enumerate()
            .map(|(i, o)| (ObjectHandle::from(i), o))
    }

    pub fn object(&self, handle: ObjectHandle) -> Option<&Object> {
        if self.objects.len() <= usize::from(handle) {
            return None;
        }
        Some(&self.objects[usize::from(handle)])
    }

    pub fn token_pin(&self) -> &str {
        self.token_pin.as_str()
    }
}