// Copyright 2020-2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
#![deny(missing_docs)]
#![deny(warnings)]

/// The module which parses command parameters from command-line arguments.
pub mod commands_parser;
/// The module which provides mappings between NitroCliErrors and their corresponding code.
pub mod document_errors;
/// The module which provides JSON-ready information structures.
pub mod json_output;
/// The module which provides the per-process logger.
pub mod logger;
/// The module which provides signal handling.
pub mod signal_handler;

use chrono::offset::Utc;
use log::error;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use std::env;
use std::io::{Read, Write};

#[cfg(test)]
use std::os::raw::c_char;

use std::os::unix::net::UnixStream;
use std::path::{Path, PathBuf};

use document_errors::ERROR_CODES;
use logger::get_log_file_base_path;

/// The most common result type provided by Nitro CLI operations.
pub type NitroCliResult<T> = Result<T, NitroCliFailure>;

/// The CID for the vsock device of the parent VM.
pub const VMADDR_CID_PARENT: u32 = 3;

/// The vsock port used to confirm that the enclave has booted.
pub const ENCLAVE_READY_VSOCK_PORT: u32 = 9000;

/// The amount of time in milliseconds an enclave process will wait for certain operations.
pub const ENCLAVE_PROC_WAIT_TIMEOUT_MSEC: isize = 3000;

/// The confirmation code sent by an enclave process to a requesting CLI instance
/// in order to signal that it is alive.
pub const MSG_ENCLAVE_CONFIRM: u64 = 0xEEC0;

/// The environment variable which holds the path to the Unix sockets directory.
pub const SOCKETS_DIR_PATH_ENV_VAR: &str = "NITRO_CLI_SOCKETS_PATH";

/// The default path to the Unix sockets directory.
const SOCKETS_DIR_PATH: &str = "/run/nitro_enclaves";

/// Constant used for identifying the backtrace environment variable.
const BACKTRACE_VAR: &str = "BACKTRACE";

/// All possible errors which may occur.
#[derive(Debug, Clone, Copy, Hash, PartialEq)]
pub enum NitroCliErrorEnum {
    /// Unspecified error (should avoid using it thoughout the code).
    UnspecifiedError = 0,
    /// Error for handling missing arguments.
    MissingArgument,
    /// Error for handling conflicting arguments.
    ConflictingArgument,
    /// Invalid type argument.
    InvalidArgument,
    /// Failed to create socket pair.
    SocketPairCreationFailure,
    /// Failed to spawn a child process.
    ProcessSpawnFailure,
    /// Failed to daemonize current process.
    DaemonizeProcessFailure,
    /// Failed to read requested content from disk.
    ReadFromDiskFailure,
    /// Unusable connection error.
    UnusableConnectionError,
    /// Socket close error.
    SocketCloseError,
    /// Socket connect timeout error.
    SocketConnectTimeoutError,
    /// General error for handling socket-related errors.
    SocketError,
    /// General error for handling epoll-related errors.
    EpollError,
    /// General error for handling inotify-related errors.
    InotifyError,
    /// Invalid command format.
    InvalidCommand,
    /// Lock acquire failure.
    LockAcquireFailure,
    /// Thread join failure.
    ThreadJoinFailure,
    /// General error for handling serde-related errors.
    SerdeError,
    /// File permissions error.
    FilePermissionsError,
    /// File operation failure.
    FileOperationFailure,
    /// Invalid CPU list configuration.
    InvalidCpuConfiguration,
    /// Requested CPU not available in the pool.
    NoSuchCpuAvailableInPool,
    /// Not enough CPUs available in the pool.
    InsufficientCpus,
    /// Malformed CPU ID error.
    MalformedCpuId,
    /// General error to catch all other CPU-related errors.
    CpuError,
    /// No such hugepage map flag.
    NoSuchHugepageFlag,
    /// Insufficient memory requested.
    InsufficientMemoryRequested,
    /// Insufficient memory available.
    InsufficientMemoryAvailable,
    /// Invalid enclave file descriptor.
    InvalidEnclaveFd,
    /// General ioctl failure.
    IoctlFailure,
    /// Image load info ioctl failure.
    IoctlImageLoadInfoFailure,
    /// Enclave set memory region ioctl failure.
    IoctlSetMemoryRegionFailure,
    /// VCPU add ioctl failure.
    IoctlAddVcpuFailure,
    /// Enclave start ioctl failure.
    IoctlEnclaveStartFailure,
    /// Memory overflow.
    MemoryOverflow,
    /// General EIF parsing related error.
    EifParsingError,
    /// Error specific to enclave booting issues.
    EnclaveBootFailure,
    /// Enclave event wait error.
    EnclaveEventWaitError,
    /// Enclave process command was not executed.
    EnclaveProcessCommandNotExecuted,
    /// Could not connect to an enclave process.
    EnclaveProcessConnectionFailure,
    /// Socket path not found.
    SocketPathNotFound,
    /// Enclave process failed to send back reply.
    EnclaveProcessSendReplyFailure,
    /// Error when trying to allocate enclave memory regions.
    EnclaveMmapError,
    /// Error when trying to release enclave memory regions.
    EnclaveMunmapError,
    /// Enclave connection to console failed.
    EnclaveConsoleConnectionFailure,
    /// Error when reading from the console.
    EnclaveConsoleReadError,
    /// Error when writing console output to stream.
    EnclaveConsoleWriteOutputError,
    /// Integer parsing error.
    IntegerParsingError,
    /// Could not build EIF file.
    EifBuildingError,
    /// Could not build Docker image.
    DockerImageBuildError,
    /// Could not pull Docker image.
    DockerImagePullError,
    /// Artifacts path environment variable not set.
    ArtifactsPathNotSet,
    /// Blobs path environment variable not set.
    BlobsPathNotSet,
    /// Clock skew error.
    ClockSkewError,
    /// Signal masking error.
    SignalMaskingError,
    /// Signal unmasking error.
    SignalUnmaskingError,
    /// General error for handling logger-related errors.
    LoggerError,
    /// Hasher operation error
    HasherError,
    /// Enclave naming error
    EnclaveNamingError,
    /// Signature checker error
    EIFSignatureCheckerError,
}

impl Default for NitroCliErrorEnum {
    fn default() -> NitroCliErrorEnum {
        NitroCliErrorEnum::UnspecifiedError
    }
}

impl Eq for NitroCliErrorEnum {}

/// The type of commands that can be sent to an enclave process.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum EnclaveProcessCommandType {
    /// Launch (run) an enclave (sent by the CLI).
    Run = 0,
    /// Terminate an enclave (sent by the CLI).
    Terminate,
    /// Notify that the enclave has terminated (sent by the enclave process to itself).
    TerminateComplete,
    /// Describe an enclave (broadcast by the CLI).
    Describe,
    /// Request an enclave's CID (sent by the CLI).
    GetEnclaveCID,
    /// Request an enclave's flags (sent by the CLI).
    GetEnclaveFlags,
    /// Request an enclave's name (sent by the CLI).
    GetEnclaveName,
    /// Request the ID of an enclave only if the name matches (sent by the CLI).
    GetIDbyName,
    /// Notify the socket connection listener to shut down (sent by the enclave process to itself).
    ConnectionListenerStop,
    /// Do not execute a command due to insufficient privileges (sent by the CLI, modified by the enclave process).
    NotPermitted,
}

/// The type of replies that an enclave process can send to a CLI instance.
#[derive(Debug, Serialize, Deserialize)]
pub enum EnclaveProcessReply {
    /// A message which must be printed to the CLI's standard output.
    StdOutMessage(String),
    /// A messge which must be printed to the CLI's standard error.
    StdErrMessage(String),
    /// The status of the operation that the enclave process has performed.
    Status(i32),
}

/// Struct that is passed along the backtrace and accumulates error messages.
#[derive(Debug, Default)]
pub struct NitroCliFailure {
    /// Main action which was attempted and failed.
    pub action: String,
    /// (Possibly) more subactions which lead to the root cause of the failure.
    pub subactions: Vec<String>,
    /// Computer-readable error code.
    pub error_code: NitroCliErrorEnum,
    /// File in which the root error occurred.
    pub file: String,
    /// Line at which the root error occurred.
    pub line: u32,
    /// Additional info regarding the error, passed as individual components (for easier parsing).
    pub additional_info: Vec<String>,
}

impl NitroCliFailure {
    /// Returns an empty `NitroCliFailure` object.
    pub fn new() -> Self {
        NitroCliFailure {
            action: String::new(),
            subactions: vec![],
            error_code: NitroCliErrorEnum::default(),
            file: String::new(),
            line: 0,
            additional_info: vec![],
        }
    }

    /// Sets the main action which failed (i.e. RUN_ENCLAVE).
    pub fn set_action(mut self, action: String) -> Self {
        self.action = action;
        self
    }

    /// Adds a new layer into the backtrace, corresponding to a failing subaction (i.e. NOT_ENOUGH_MEM).
    pub fn add_subaction(mut self, subaction: String) -> Self {
        self.subactions.push(subaction);
        self
    }

    /// Sets the error code.
    pub fn set_error_code(mut self, error_code: NitroCliErrorEnum) -> Self {
        self.error_code = error_code;
        self
    }

    /// Sets the name of the file the error occurred in.
    pub fn set_file(mut self, file: &str) -> Self {
        self.file = file.to_string();
        self
    }

    /// Sets the number of the line the error occurred on.
    pub fn set_line(mut self, line: u32) -> Self {
        self.line = line;
        self
    }

    /// Sets both error file and error line.
    pub fn set_file_and_line(mut self, file: &str, line: u32) -> Self {
        self.file = file.to_string();
        self.line = line;
        self
    }

    /// Include additional error information.
    pub fn add_info(mut self, info: Vec<&str>) -> Self {
        for info_ in info {
            self.additional_info.push(info_.to_string());
        }
        self
    }
}

/// Macro used for constructing a NitroCliFailure in a more convenient manner.
#[macro_export]
macro_rules! new_nitro_cli_failure {
    ($subaction:expr, $error_code:expr) => {
        NitroCliFailure::new()
            .add_subaction(($subaction).to_string())
            .set_error_code($error_code)
            .set_file_and_line(file!(), line!())
    };
}

/// Logs the given backtrace string to a separate, backtrace-specific file.
/// Returns a string denoting the path to the corresponding log file.
fn log_backtrace(backtrace: String) -> Result<String, &'static str> {
    let log_path_base = get_log_file_base_path();

    // Check if backtrace logs location exists and create it if necessary.
    if !Path::new(&log_path_base).exists() {
        let create_logs_dir = std::fs::create_dir_all(&log_path_base);
        if create_logs_dir.is_err() {
            return Err("Could not create backtrace logs directory");
        }
    }

    let utc_time_now = Utc::now().to_rfc3339();
    let log_path_str = format!("{}/err{}.log", &log_path_base, utc_time_now);
    let log_path = Path::new(&log_path_str);
    let log_file = std::fs::File::create(log_path);
    if log_file.is_err() {
        return Err("Could not create backtrace log file");
    }

    let write_result = log_file.unwrap().write_all(backtrace.as_bytes());
    if write_result.is_err() {
        return Err("Could not write to backtrace log file");
    }

    match log_path.to_str() {
        Some(log_path) => Ok(log_path.to_string()),
        None => Err("Could not return log file path"),
    }
}

/// Assembles the error message which gets displayed to the user.
pub fn construct_error_message(failure: &NitroCliFailure) -> String {
    // Suggestive error description comes first.
    let error_info: String = document_errors::get_detailed_info(
        (*ERROR_CODES.get(&failure.error_code).unwrap_or(&"E00")).to_string(),
        &failure.additional_info,
    );

    // Include a link to the documentation page.
    let help_link: String = document_errors::construct_help_link(
        (*ERROR_CODES.get(&failure.error_code).unwrap_or(&"E00")).to_string(),
    );
    let backtrace: String = document_errors::construct_backtrace(failure);

    // Write backtrace to a log file.
    let log_path = log_backtrace(backtrace.clone());

    // Return final output, depending on whether the user requested the backtrace or not.
    match std::env::var(BACKTRACE_VAR) {
        Ok(display_backtrace) => match display_backtrace.as_str() {
            "1" => {
                if let Ok(log_path) = log_path {
                    format!(
                        "{}\n\nFor more details, please visit {}\n\nBacktrace:\n{}\n\nIf you open a support ticket, please provide the error log found at \"{}\"",
                        error_info, help_link, backtrace, log_path
                    )
                } else {
                    format!(
                        "{}\n\nFor more details, please visit {}\n\nBacktrace:\n{}",
                        error_info, help_link, backtrace
                    )
                }
            }
            _ => {
                if let Ok(log_path) = log_path {
                    format!(
                        "{}\n\nFor more details, please visit {}\n\nIf you open a support ticket, please provide the error log found at \"{}\"",
                        error_info, help_link, log_path
                    )
                } else {
                    format!(
                        "{}\n\nFor more details, please visit {}",
                        error_info, help_link
                    )
                }
            }
        },
        _ => {
            if let Ok(log_path) = log_path {
                format!(
                    "{}\n\nFor more details, please visit {}\n\nIf you open a support ticket, please provide the error log found at \"{}\"",
                    error_info, help_link, log_path
                )
            } else {
                format!(
                    "{}\n\nFor more details, please visit {}",
                    error_info, help_link
                )
            }
        }
    }
}

/// A trait which allows a more graceful program exit instead of the standard `panic`.
/// Provides a custom exit code.
pub trait ExitGracefully<T> {
    /// Provide the inner value of a `Result` or exit gracefully with a message and custom errno.
    fn ok_or_exit_with_errno(self, additional_info: Option<&str>) -> T;
}

impl<T> ExitGracefully<T> for NitroCliResult<T> {
    /// Provide the inner value of a `Result` or exit gracefully with a message and custom errno.
    fn ok_or_exit_with_errno(self, additional_info: Option<&str>) -> T {
        match self {
            Ok(val) => val,
            Err(err) => {
                let err_str = construct_error_message(&err);
                if let Some(additional_info_str) = additional_info {
                    notify_error(&format!("{} | {}", additional_info_str, err_str));
                } else {
                    notify_error(&err_str);
                }
                std::process::exit(err.error_code as i32);
            }
        }
    }
}

/// Notify both the user and the logger of an error.
pub fn notify_error(err_msg: &str) {
    eprintln!("{}", err_msg);
    error!("{}", err_msg);
}

/// Read a LE-encoded 64-bit unsigned value from a socket.
pub fn read_u64_le(socket: &mut dyn Read) -> NitroCliResult<u64> {
    let mut bytes = [0u8; std::mem::size_of::<u64>()];
    socket.read_exact(&mut bytes).map_err(|e| {
        new_nitro_cli_failure!(
            &format!(
                "Failed to read {} bytes from the given socket: {:?}",
                std::mem::size_of::<u64>(),
                e
            ),
            NitroCliErrorEnum::SocketError
        )
    })?;

    Ok(u64::from_le_bytes(bytes))
}

/// Write a LE-encoded 64-bit unsigned value to a socket.
pub fn write_u64_le(socket: &mut dyn Write, value: u64) -> NitroCliResult<()> {
    let bytes = value.to_le_bytes();
    socket.write_all(&bytes).map_err(|e| {
        new_nitro_cli_failure!(
            &format!(
                "Failed to write {} bytes to the given socket: {:?}",
                std::mem::size_of::<u64>(),
                e
            ),
            NitroCliErrorEnum::SocketError
        )
    })
}

/// Send a command to a single socket.
pub fn enclave_proc_command_send_single<T>(
    cmd: EnclaveProcessCommandType,
    args: Option<&T>,
    mut socket: &mut UnixStream,
) -> NitroCliResult<()>
where
    T: Serialize,
{
    // Serialize the command type.
    let cmd_bytes = serde_cbor::to_vec(&cmd).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Invalid command format: {:?}", e),
            NitroCliErrorEnum::InvalidCommand
        )
    })?;

    // The command is written twice. The first read is done by the connection listener to check if this is
    // a shut-down command. The second read is done by the enclave process for all non-shut-down commands.
    for _ in 0..2 {
        write_u64_le(&mut socket, cmd_bytes.len() as u64)
            .map_err(|e| e.add_subaction("Failed to send single command size".to_string()))?;
        socket.write_all(&cmd_bytes[..]).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to send single command: {:?}", e),
                NitroCliErrorEnum::SocketError
            )
        })?;
    }

    // Serialize the command arguments.
    if let Some(args) = args {
        let arg_bytes = serde_cbor::to_vec(args).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Invalid single command arguments: {:?}", e),
                NitroCliErrorEnum::InvalidCommand
            )
        })?;

        // Write the serialized command arguments.
        write_u64_le(&mut socket, arg_bytes.len() as u64)
            .map_err(|e| e.add_subaction("Failed to send arguments size".to_string()))?;
        socket.write_all(&arg_bytes).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to send arguments: {:?}", e),
                NitroCliErrorEnum::SocketError
            )
        })?;
    }

    Ok(())
}

/// Receive an object of a specified type from an input stream.
pub fn receive_from_stream<T>(input_stream: &mut dyn Read) -> NitroCliResult<T>
where
    T: DeserializeOwned,
{
    let size = read_u64_le(input_stream)
        .map_err(|e| e.add_subaction("Failed to receive data size".to_string()))?
        as usize;
    let mut raw_data: Vec<u8> = vec![0; size];
    input_stream.read_exact(&mut raw_data[..]).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to receive data: {:?}", e),
            NitroCliErrorEnum::SocketError
        )
    })?;
    let data: T = serde_cbor::from_slice(&raw_data[..]).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to decode received data: {:?}", e),
            NitroCliErrorEnum::SerdeError
        )
    })?;
    Ok(data)
}

/// Get the path to the directory containing the Unix sockets owned by all enclave processes.
pub fn get_sockets_dir_path() -> PathBuf {
    let log_path = match env::var(SOCKETS_DIR_PATH_ENV_VAR) {
        Ok(env_path) => env_path,
        Err(_) => SOCKETS_DIR_PATH.to_string(),
    };
    Path::new(&log_path).to_path_buf()
}

/// Get the path to the Unix socket owned by an enclave process which also owns the enclave with the given ID.
pub fn get_socket_path(enclave_id: &str) -> NitroCliResult<PathBuf> {
    // The full enclave ID is "i-(...)-enc<enc_id>" and we want to extract only <enc_id>.
    let tokens: Vec<_> = enclave_id.rsplit("-enc").collect();
    let sockets_path = get_sockets_dir_path();
    Ok(sockets_path.join(tokens[0]).with_extension("sock"))
}

#[cfg(test)]
mod tests {
    #[allow(unused_imports)]
    use super::*;

    use crate::common::commands_parser::EmptyArgs;

    const TMP_DIR_STR: &str = "./tmp_sock_dir";

    fn unset_envvar(varname: &String) {
        let _ = unsafe {
            libc::unsetenv(varname.as_ptr() as *const c_char);
        };
    }

    /// Tests that a value wrote by `write_u64_le()` is read
    /// correctly by `read_u64_le()`.
    #[test]
    fn test_read_write_u64() {
        let (mut sock0, mut sock1) = UnixStream::pair().unwrap();

        let _ = write_u64_le(&mut sock0, 127);
        let result = read_u64_le(&mut sock1);

        if let Ok(result) = result {
            assert_eq!(result, 127);
        }
    }

    /// Tests that a command sent though a socket by `enclave_proc_command_send_single()`
    /// is received correctly at the other end, by `receive_command_type()`.
    #[test]
    fn test_enclave_proc_command_send_single() {
        let (mut sock0, mut sock1) = UnixStream::pair().unwrap();
        let cmd = EnclaveProcessCommandType::Describe;
        let args: std::option::Option<&EmptyArgs> = None;

        let result0 = enclave_proc_command_send_single::<EmptyArgs>(cmd, args, &mut sock0);
        assert!(result0.is_ok());

        let result1 = receive_from_stream::<EnclaveProcessCommandType>(&mut sock1);
        assert!(result1.is_ok());
        assert_eq!(result1.unwrap(), EnclaveProcessCommandType::Describe);
    }

    /// Tests that the returned sockets_dir_path matches the expected path,
    /// as retrieved from the corresponding environment variable.
    #[test]
    fn test_get_sockets_dir_path_default() {
        let sockets_dir = env::var(SOCKETS_DIR_PATH_ENV_VAR);
        let sockets_dir_path_f = get_sockets_dir_path();

        if let Ok(sockets_dir) = sockets_dir {
            assert_eq!(sockets_dir, sockets_dir_path_f.as_path().to_str().unwrap());
        } else {
            assert_eq!(
                SOCKETS_DIR_PATH,
                sockets_dir_path_f.as_path().to_str().unwrap()
            );
        }
    }

    /// Tests that altering the content of the sockets_dir_path environment variable
    /// changes the sockets_dir_path string returned by `get_sockets_dir_path()`.
    #[test]
    fn test_get_sockets_dir_path_custom_envvar() {
        let old_sockets_dir = env::var(SOCKETS_DIR_PATH_ENV_VAR);
        env::set_var(SOCKETS_DIR_PATH_ENV_VAR, TMP_DIR_STR);

        let sockets_dir_path_f = get_sockets_dir_path();

        assert_eq!(TMP_DIR_STR, sockets_dir_path_f.as_path().to_str().unwrap());

        // Restore previous environment variable value
        if let Ok(old_sockets_dir) = old_sockets_dir {
            env::set_var(SOCKETS_DIR_PATH_ENV_VAR, old_sockets_dir);
        } else {
            env::set_var(SOCKETS_DIR_PATH_ENV_VAR, "");
            unset_envvar(&String::from(SOCKETS_DIR_PATH_ENV_VAR));
        }
    }

    /// Tests that `get_socket_path()` returns the expected socket path,
    /// given a specific enclave id.
    #[test]
    fn test_get_socket_path_valid_id() {
        let enclave_id = "i-0000000000000000-enc0123456789012345";
        let tokens: Vec<_> = enclave_id.rsplit("-enc").collect();
        let sockets_path = get_sockets_dir_path();
        let result = get_socket_path(enclave_id);

        assert!(result.is_ok());
        assert_eq!(
            result.unwrap().as_path().to_str().unwrap(),
            format!(
                "{}/{}.sock",
                sockets_path.as_path().to_str().unwrap(),
                tokens[0]
            )
        );
    }

    /// Tests that `get_socket_path()` returns an invalid socket path,
    /// given a malformed enclave id.
    #[test]
    fn test_get_socket_path_invalid_id() {
        let enclave_id = "i-0000000000000000_enc0123456789012345";
        let sockets_path = get_sockets_dir_path();
        let result = get_socket_path(enclave_id);

        assert!(result.is_ok());
        assert_eq!(
            result.unwrap().as_path().to_str().unwrap(),
            format!(
                "{}/{}.sock",
                sockets_path.as_path().to_str().unwrap(),
                enclave_id
            )
        );
    }
}


// Copyright 2020-2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
#![deny(missing_docs)]
#![deny(warnings)]

use log::warn;
use nix::sys::signal::{SigSet, Signal};
use nix::sys::signal::{SIGHUP, SIGINT, SIGQUIT, SIGTERM};
use std::os::unix::io::RawFd;
use std::thread;

use crate::common::{NitroCliErrorEnum, NitroCliFailure, NitroCliResult};
use crate::new_nitro_cli_failure;

/// The custom handler of POSIX signals.
pub struct SignalHandler {
    sig_set: Option<SigSet>,
}

impl SignalHandler {
    /// Create a new `SignalHandler` instance from the given list of signals.
    pub fn new(signals: &[Signal]) -> Self {
        let mut sig_set = SigSet::empty();
        for signal in signals.iter() {
            sig_set.add(*signal);
        }

        SignalHandler {
            sig_set: Some(sig_set),
        }
    }

    /// Create a new `SignalHandler` instance from a default list of signals.
    pub fn new_with_defaults() -> Self {
        SignalHandler::new(&[SIGINT, SIGQUIT, SIGTERM, SIGHUP])
    }

    /// Mask (block) all signals covered by the handler.
    pub fn mask_all(self) -> NitroCliResult<Self> {
        if let Some(set) = self.sig_set {
            set.thread_block().map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Masking signals covered by handler failed: {:?}", e),
                    NitroCliErrorEnum::SignalMaskingError
                )
            })?;
        }

        Ok(self)
    }

    /// Unmask (unblock) all signals covered by the handler.
    pub fn unmask_all(self) -> NitroCliResult<Self> {
        if let Some(set) = self.sig_set {
            set.thread_unblock().map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Unmasking signals covered by handler failed: {:?}", e),
                    NitroCliErrorEnum::SignalUnmaskingError
                )
            })?;
        }

        Ok(self)
    }

    /// Start listening for events on a dedicated thread and handle them using the provided function.
    pub fn start_handler(&mut self, fd: RawFd, handler: fn(RawFd, Signal) -> bool) {
        if self.sig_set.is_none() {
            return;
        }

        let thread_sig_set = self.sig_set.take().unwrap();
        thread::spawn(move || {
            let mut stop = false;
            while !stop {
                stop = match thread_sig_set.wait() {
                    Ok(signal) => handler(fd, signal),
                    Err(e) => {
                        warn!("Error listening for signals: {}", e);
                        true
                    }
                };
            }
        });
    }
}


// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
#![deny(missing_docs)]
#![deny(warnings)]

use aws_nitro_enclaves_image_format::defs::EifIdentityInfo;
use aws_nitro_enclaves_image_format::utils::eif_reader::EifReader;
use aws_nitro_enclaves_image_format::utils::{get_pcrs, PcrSignatureChecker};
use log::debug;
use sha2::{Digest, Sha384};
use std::collections::BTreeMap;
use std::fs::File;
use std::thread::JoinHandle;

use crate::common::commands_parser::RunEnclavesArgs;
use crate::common::construct_error_message;
use crate::common::json_output::EnclaveTerminateInfo;
use crate::common::{NitroCliErrorEnum, NitroCliFailure, NitroCliResult};
use crate::enclave_proc::connection::Connection;
use crate::enclave_proc::connection::{safe_conn_eprintln, safe_conn_println};
use crate::enclave_proc::cpu_info::CpuInfo;
use crate::enclave_proc::resource_manager::{EnclaveManager, EnclaveState};
use crate::enclave_proc::utils::get_enclave_describe_info;
use crate::new_nitro_cli_failure;

/// Result of thread to calculate PCRs and fetch EIF metadata
pub struct DescribeThreadResult {
    /// PCR measurements
    pub measurements: BTreeMap<String, String>,
    /// EIF metadata
    pub metadata: Option<EifIdentityInfo>,
}

/// Thread handle from parallel computing of PCRs and fetching of metadata
pub type DescribeThread = Option<JoinHandle<NitroCliResult<DescribeThreadResult>>>;

/// Information returned by run_enclave function.
pub struct RunEnclaveResult {
    /// Manager structure describing the enclave.
    pub enclave_manager: EnclaveManager,
    /// Handle of the thread that computes PCRs.
    pub describe_thread: DescribeThread,
}

/// Launch an enclave with the specified arguments and provide the launch status through the given connection.
pub fn run_enclaves(
    args: &RunEnclavesArgs,
    connection: Option<&Connection>,
) -> NitroCliResult<RunEnclaveResult> {
    debug!("run_enclaves");

    let eif_file = File::open(&args.eif_path).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to open the EIF file: {:?}", e),
            NitroCliErrorEnum::FileOperationFailure
        )
        .add_info(vec![&args.eif_path, "Open"])
    })?;

    let cpu_ids = CpuInfo::new()
        .map_err(|e| e.add_subaction("Failed to construct CPU information".to_string()))?
        .get_cpu_config(args)
        .map_err(|e| e.add_subaction("Failed to get CPU configuration".to_string()))?;

    let enclave_name = match &args.enclave_name {
        Some(enclave_name) => enclave_name,
        None => {
            return Err(new_nitro_cli_failure!(
                "Failed to set enclave name in the manager".to_string(),
                NitroCliErrorEnum::FileOperationFailure
            ))
        }
    };

    let mut enclave_manager = EnclaveManager::new(
        args.enclave_cid,
        args.memory_mib,
        cpu_ids,
        eif_file,
        args.debug_mode,
        enclave_name.to_string(),
    )
    .map_err(|e| {
        e.add_subaction("Failed to construct EnclaveManager with given arguments".to_string())
    })?;

    let mut signature_checker = PcrSignatureChecker::from_eif(&args.eif_path).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to create EIF signature checker: {:?}", e),
            NitroCliErrorEnum::EIFSignatureCheckerError
        )
    })?;

    // Verify the certificate only if signature section exists
    if !signature_checker.is_empty() {
        signature_checker.verify().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Invalid signing certificate: {:?}", e),
                NitroCliErrorEnum::EIFSignatureCheckerError
            )
        })?;
    }

    // Launch parallel computing of PCRs
    let path = args.eif_path.clone();
    let handle = std::thread::spawn(move || {
        let mut eif_reader = EifReader::from_eif(path).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed initialize EIF reader: {:?}", e),
                NitroCliErrorEnum::EifParsingError
            )
        })?;
        let measurements = get_pcrs(
            &mut eif_reader.image_hasher,
            &mut eif_reader.bootstrap_hasher,
            &mut eif_reader.app_hasher,
            &mut eif_reader.cert_hasher,
            Sha384::new(),
            eif_reader.signature_section.is_some(),
        )
        .map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to calculate PCRs: {:?}", e),
                NitroCliErrorEnum::EifParsingError
            )
        })?;
        Ok(DescribeThreadResult {
            measurements,
            metadata: eif_reader.get_metadata(),
        })
    });
    enclave_manager
        .run_enclave(connection)
        .map_err(|e| e.add_subaction("Failed to run enclave".to_string()))?;
    enclave_manager
        .update_state(EnclaveState::Running)
        .map_err(|e| e.add_subaction("Failed to update enclave state".to_string()))?;

    Ok(RunEnclaveResult {
        enclave_manager,
        describe_thread: Some(handle),
    })
}

/// Terminate an enclave and provide the termination status through the given connection.
pub fn terminate_enclaves(
    enclave_manager: &mut EnclaveManager,
    connection: Option<&Connection>,
) -> NitroCliResult<()> {
    let enclave_id = enclave_manager.enclave_id.clone();
    let enclave_name = Some(enclave_manager.enclave_name.clone());

    debug!("terminate_enclaves");
    enclave_manager
        .update_state(EnclaveState::Terminating)
        .map_err(|e| e.add_subaction("Failed to update enclave state".to_string()))?;
    if let Err(error_info) = enclave_manager.terminate_enclave() {
        safe_conn_eprintln(
            connection,
            format!(
                "Warning: Failed to stop enclave {}\nError message: {:?}",
                enclave_manager.enclave_id,
                construct_error_message(&error_info).as_str()
            )
            .as_str(),
        )?;
        return Err(error_info);
    }

    enclave_manager.update_state(EnclaveState::Empty)?;
    safe_conn_eprintln(
        connection,
        format!(
            "Successfully terminated enclave {}.",
            enclave_manager.enclave_id
        )
        .as_str(),
    )?;

    // We notify the CLI of the termination's status.
    safe_conn_println(
        connection,
        serde_json::to_string_pretty(&EnclaveTerminateInfo::new(enclave_name, enclave_id, true))
            .map_err(|err| {
                new_nitro_cli_failure!(
                    &format!("Failed to display enclave termination data: {:?}", err),
                    NitroCliErrorEnum::SerdeError
                )
            })?
            .as_str(),
    )
}

/// Obtain an enclave's description and provide it through the given connection.
pub fn describe_enclaves(
    enclave_manager: &EnclaveManager,
    connection: &Connection,
    with_metadata: bool,
) -> NitroCliResult<()> {
    debug!("describe_enclaves");

    let info = get_enclave_describe_info(enclave_manager, with_metadata)
        .map_err(|e| e.add_subaction(String::from("Execute Describe Enclave command")))?;

    connection.println(
        serde_json::to_string_pretty(&info)
            .map_err(|err| {
                new_nitro_cli_failure!(
                    &format!("Failed to display enclave describe data: {:?}", err),
                    NitroCliErrorEnum::SerdeError
                )
            })?
            .as_str(),
    )
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
#![deny(missing_docs)]
#![deny(warnings)]

use log::{debug, warn};
use nix::sys::epoll::EpollFlags;
use nix::sys::socket::sockopt::PeerCredentials;
use nix::sys::socket::UnixCredentials;
use serde::de::DeserializeOwned;
use std::collections::HashMap;
use std::io::Write;
use std::os::unix::io::AsRawFd;
use std::os::unix::net::UnixStream;
use std::string::ToString;
use std::sync::{Arc, Mutex};

use crate::common::{receive_from_stream, write_u64_le};
use crate::common::{
    EnclaveProcessCommandType, EnclaveProcessReply, ExitGracefully, NitroCliErrorEnum,
    NitroCliFailure, NitroCliResult,
};
use crate::new_nitro_cli_failure;

/// The types of requesters which may send commands to the enclave process.
#[derive(PartialEq, Eq, Hash)]
enum CommandRequesterType {
    /// The requester is the user with the given UID.
    User(libc::uid_t),
    /// The requester is the group with the given GID.
    Group(libc::gid_t),
    /// The requester is any other user.
    Others,
}

/// The policy used to filter received commands based on the requester's type.
struct CommandRequesterPolicy {
    /// A mapping between a requester's type and all of its allowed commands.
    policy: HashMap<CommandRequesterType, Vec<EnclaveProcessCommandType>>,
}

/// Data held by a connection.
struct ConnectionData {
    /// Flags received from `epoll` if this was an event-triggered connection.
    epoll_flags: EpollFlags,
    /// A communication stream with the peer, if this was a socket-triggered connection.
    input_stream: Option<UnixStream>,
}

/// An enclave process connection to a CLI instance, an enclave or itself.
#[derive(Clone)]
pub struct Connection {
    /// The thread-safe data used internally by the connection.
    data: Arc<Mutex<ConnectionData>>,
}

impl Drop for ConnectionData {
    fn drop(&mut self) {
        if let Some(input_stream) = &self.input_stream {
            // Close the stream.
            input_stream
                .shutdown(std::net::Shutdown::Both)
                .map_err(|e| {
                    new_nitro_cli_failure!(
                        &format!("Stream shutdown error: {:?}", e),
                        NitroCliErrorEnum::SocketCloseError
                    )
                })
                .ok_or_exit_with_errno(Some("Failed to shut down"));
        }
    }
}

impl CommandRequesterPolicy {
    /// Create a new `CommandRequesterPolicy` with the default rules. These rules allow
    /// the user which spawned the enclave, together with `root`, to make any request,
    /// whereas all other users are only allowed to make read-only requests (namely,
    /// to describe an enclave or to read its CID).
    fn new_with_defaults() -> Self {
        let cmds_read_write = vec![
            EnclaveProcessCommandType::Run,
            EnclaveProcessCommandType::Terminate,
            EnclaveProcessCommandType::TerminateComplete,
            EnclaveProcessCommandType::Describe,
            EnclaveProcessCommandType::GetEnclaveCID,
            EnclaveProcessCommandType::GetEnclaveFlags,
            EnclaveProcessCommandType::GetEnclaveName,
            EnclaveProcessCommandType::GetIDbyName,
            EnclaveProcessCommandType::ConnectionListenerStop,
        ];
        let cmds_read_only = vec![
            EnclaveProcessCommandType::Describe,
            EnclaveProcessCommandType::GetEnclaveCID,
            EnclaveProcessCommandType::GetEnclaveFlags,
            EnclaveProcessCommandType::GetEnclaveName,
            EnclaveProcessCommandType::GetIDbyName,
        ];
        let mut policy = HashMap::new();

        // The user which owns this enclave process may issue any command.
        policy.insert(
            CommandRequesterType::User(unsafe { libc::getuid() }),
            cmds_read_write.clone(),
        );

        // The root user may issue any command.
        policy.insert(CommandRequesterType::User(0_u32), cmds_read_write);

        // All other users may only issue read-only commands.
        policy.insert(CommandRequesterType::Others, cmds_read_only);

        CommandRequesterPolicy { policy }
    }

    /// Find the policy rule which applies to the given requester and command.
    fn find_policy_rule(
        &self,
        cmd: EnclaveProcessCommandType,
        requester: &CommandRequesterType,
    ) -> bool {
        match self.policy.get(requester) {
            None => false,
            Some(allowed_cmds) => allowed_cmds.contains(&cmd),
        }
    }

    /// Check if the user with the specified credentials has permission to run the specified command.
    fn can_execute_command(&self, cmd: EnclaveProcessCommandType, creds: &UnixCredentials) -> bool {
        // Search for a policy rule on the provided user ID.
        if self.find_policy_rule(cmd, &CommandRequesterType::User(creds.uid())) {
            return true;
        }

        // Search for a policy rule on the provided group ID.
        if self.find_policy_rule(cmd, &CommandRequesterType::Group(creds.gid())) {
            return true;
        }

        // Search for a policy rule on all other users.
        if self.find_policy_rule(cmd, &CommandRequesterType::Others) {
            return true;
        }

        // If we haven't found any applicable policy rule we can't allow the command to be executed.
        false
    }
}

impl Connection {
    /// Create a new connection instance.
    pub fn new(epoll_flags: EpollFlags, input_stream: Option<UnixStream>) -> Self {
        let conn_data = ConnectionData {
            epoll_flags,
            input_stream,
        };

        Connection {
            data: Arc::new(Mutex::new(conn_data)),
        }
    }

    /// Read a command and its corresponding credentials.
    pub fn read_command(&self) -> NitroCliResult<EnclaveProcessCommandType> {
        let mut lock = self.data.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        if lock.input_stream.is_none() {
            return Err(new_nitro_cli_failure!(
                "Cannot read a command from this connection",
                NitroCliErrorEnum::UnusableConnectionError
            ));
        }

        // First, read the incoming command.
        let mut cmd =
            receive_from_stream::<EnclaveProcessCommandType>(lock.input_stream.as_mut().unwrap())?;

        // Next, read the credentials of the command requester.
        let conn_fd = lock.input_stream.as_ref().unwrap().as_raw_fd();
        let socket_creds = nix::sys::socket::getsockopt(conn_fd, PeerCredentials);

        // If the credentials cannot be read, the command will be skipped.
        let user_creds = match socket_creds {
            Ok(creds) => creds,
            Err(e) => {
                warn!("Failed to get user credentials: {}", e);
                return Ok(EnclaveProcessCommandType::NotPermitted);
            }
        };

        // Apply the default command access policy based on the user's credentials.
        let policy = CommandRequesterPolicy::new_with_defaults();
        if !policy.can_execute_command(cmd, &user_creds) {
            // Log the failed execution attempt.
            warn!(
                "The requester with credentials ({:?}) is not allowed to perform '{:?}'.",
                user_creds, cmd
            );

            // Force the command to be skipped by the main event loop.
            cmd = EnclaveProcessCommandType::NotPermitted;
        } else {
            // Log the successful execution attempt.
            debug!(
                "The requester with credentials ({:?}) is allowed to perform '{:?}'.",
                user_creds, cmd
            );
        }

        Ok(cmd)
    }

    /// Read an object of the specified type from this connection.
    pub fn read<T>(&self) -> NitroCliResult<T>
    where
        T: DeserializeOwned,
    {
        let mut lock = self.data.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        if lock.input_stream.is_none() {
            return Err(new_nitro_cli_failure!(
                "Cannot read from this connection",
                NitroCliErrorEnum::SocketError
            ));
        }

        receive_from_stream::<T>(lock.input_stream.as_mut().unwrap())
    }

    /// Write a 64-bit unsigned value on this connection.
    pub fn write_u64(&self, value: u64) -> NitroCliResult<()> {
        let mut lock = self.data.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        if lock.input_stream.is_none() {
            return Err(new_nitro_cli_failure!(
                "Cannot write a 64-bit value to this connection",
                NitroCliErrorEnum::SocketError
            ));
        }

        write_u64_le(lock.input_stream.as_mut().unwrap(), value)
    }

    /// Write a message to the standard output of the connection's other end.
    pub fn println(&self, msg: &str) -> NitroCliResult<()> {
        let mut msg_str = msg.to_string();

        // Append a new-line at the end of the string.
        msg_str.push('\n');

        let reply = EnclaveProcessReply::StdOutMessage(msg_str);
        self.write_reply(&reply)
    }

    /// Write a message to the standard error of the connection's other end.
    pub fn eprintln(&self, msg: &str) -> NitroCliResult<()> {
        let mut msg_str = msg.to_string();

        // Append a new-line at the end of the string.
        msg_str.push('\n');

        let reply = EnclaveProcessReply::StdErrMessage(msg_str);
        self.write_reply(&reply)
    }

    /// Write an operation's status to the connection's other end.
    pub fn write_status(&self, status: i32) -> NitroCliResult<()> {
        let reply = EnclaveProcessReply::Status(status);
        self.write_reply(&reply)
    }

    /// Get the enclave event flags.
    pub fn get_enclave_event_flags(&self) -> NitroCliResult<Option<EpollFlags>> {
        let lock = self.data.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire connection lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        match lock.input_stream {
            None => Ok(Some(lock.epoll_flags)),
            _ => Ok(None),
        }
    }

    /// Write a string and its corresponding destination to a socket.
    fn write_reply(&self, reply: &EnclaveProcessReply) -> NitroCliResult<()> {
        let mut lock = self.data.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        if lock.input_stream.is_none() {
            return Err(new_nitro_cli_failure!(
                "Cannot write message to connection",
                NitroCliErrorEnum::SocketError
            ));
        }

        let mut stream = lock.input_stream.as_mut().unwrap();
        let reply_bytes = serde_cbor::to_vec(reply).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to serialize reply: {:?}", e),
                NitroCliErrorEnum::SerdeError
            )
        })?;

        write_u64_le(&mut stream, reply_bytes.len() as u64)
            .map_err(|e| e.add_subaction("Write reply".to_string()))?;
        stream.write_all(&reply_bytes).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to write to stream: {:?}", e),
                NitroCliErrorEnum::SocketError
            )
        })
    }
}

/// Print a message to a connection's standard output, if the connection is available.
pub fn safe_conn_println(conn: Option<&Connection>, msg: &str) -> NitroCliResult<()> {
    if conn.is_none() {
        return Ok(());
    }

    conn.unwrap().println(msg)
}

/// Print a message to a connection's standard error, if the connection is available.
pub fn safe_conn_eprintln(conn: Option<&Connection>, msg: &str) -> NitroCliResult<()> {
    if conn.is_none() {
        return Ok(());
    }

    conn.unwrap().eprintln(msg)
}


// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
#![deny(missing_docs)]
#![deny(warnings)]

use log::{debug, info, warn};
use nix::sys::epoll::{self, EpollEvent, EpollFlags, EpollOp};
use std::fs::set_permissions;
use std::fs::Permissions;
use std::io;

#[cfg(test)]
use std::os::raw::c_char;

use std::os::unix::fs::PermissionsExt;
use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
use std::os::unix::net::{UnixListener, UnixStream};
use std::thread::{self, JoinHandle};

use super::connection::Connection;
use super::socket::EnclaveProcSock;
use crate::common::commands_parser::EmptyArgs;
use crate::common::{enclave_proc_command_send_single, receive_from_stream};
use crate::common::{
    EnclaveProcessCommandType, ExitGracefully, NitroCliErrorEnum, NitroCliFailure, NitroCliResult,
};
use crate::new_nitro_cli_failure;

/// A listener which waits for incoming connections on the enclave process socket.
#[derive(Default)]
pub struct ConnectionListener {
    /// The epoll descriptor used to register new connections.
    epoll_fd: RawFd,
    /// A dedicated thread that listens for new connections.
    listener_thread: Option<JoinHandle<()>>,
    /// The Unix socket that the listener binds to.
    socket: EnclaveProcSock,
}

/// The listener must be cloned when launching the listening thread.
impl Clone for ConnectionListener {
    fn clone(&self) -> Self {
        // Actually clone only what's relevant for the listening thread.
        ConnectionListener {
            epoll_fd: self.epoll_fd,
            listener_thread: None,
            socket: self.socket.clone(),
        }
    }
}

impl ConnectionListener {
    /// Create a new `ConnectionListener` instance.
    pub fn new() -> NitroCliResult<Self> {
        Ok(ConnectionListener {
            epoll_fd: epoll::epoll_create().map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to initialize epoll: {:?}", e),
                    NitroCliErrorEnum::EpollError
                )
            })?,
            listener_thread: None,
            socket: EnclaveProcSock::default(),
        })
    }

    /// Expose the `epoll` descriptor.
    pub fn get_epoll_fd(&self) -> RawFd {
        self.epoll_fd
    }

    /// Initialize the connection listener from a specified enclave ID.
    pub fn start(&mut self, enclave_id: &str) -> NitroCliResult<()> {
        // Obtain the socket to listen on.
        self.socket = EnclaveProcSock::new(enclave_id)
            .map_err(|e| e.add_subaction("Failed to create enclave process socket".to_string()))?;

        // Bind the listener to the socket and spawn the listener thread.
        let listener = UnixListener::bind(self.socket.get_path()).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to bind connection listener: {:?}", e),
                NitroCliErrorEnum::SocketError
            )
        })?;
        self.enable_credentials_passing(&listener);
        self.socket
            .start_monitoring(true)
            .map_err(|e| e.add_subaction("Failed to start monitoring socket".to_string()))?;
        debug!(
            "Connection listener started on socket {:?}.",
            self.socket.get_path()
        );

        let self_clone = self.clone();
        self.listener_thread = Some(thread::spawn(move || {
            self_clone
                .connection_listener_run(listener)
                .map_err(|e| {
                    e.add_subaction("Failed to start the listener thread".to_string())
                        .set_action("Run Enclave".to_string())
                })
                .ok_or_exit_with_errno(None);
        }));

        Ok(())
    }

    /// Add a stream to `epoll`.
    pub fn add_stream_to_epoll(&self, stream: UnixStream) -> NitroCliResult<()> {
        let stream_fd = stream.as_raw_fd();
        let mut cli_evt = EpollEvent::new(EpollFlags::EPOLLIN, stream.into_raw_fd() as u64);
        epoll::epoll_ctl(self.epoll_fd, EpollOp::EpollCtlAdd, stream_fd, &mut cli_evt).map_err(
            |e| {
                new_nitro_cli_failure!(
                    &format!("Failed to add stream to epoll: {:?}", e),
                    NitroCliErrorEnum::EpollError
                )
            },
        )?;

        Ok(())
    }

    /// Add the enclave descriptor to `epoll`.
    pub fn register_enclave_descriptor(&mut self, enc_fd: RawFd) -> NitroCliResult<()> {
        let mut enc_event = EpollEvent::new(
            EpollFlags::EPOLLIN | EpollFlags::EPOLLERR | EpollFlags::EPOLLHUP,
            enc_fd as u64,
        );
        epoll::epoll_ctl(self.epoll_fd, EpollOp::EpollCtlAdd, enc_fd, &mut enc_event).map_err(
            |e| {
                new_nitro_cli_failure!(
                    &format!("Failed to add enclave descriptor to epoll: {:?}", e),
                    NitroCliErrorEnum::EpollError
                )
            },
        )?;

        Ok(())
    }

    /// Handle an incoming connection.
    pub fn handle_new_connection(
        &self,
        mut stream: UnixStream,
    ) -> NitroCliResult<EnclaveProcessCommandType> {
        let cmd_type =
            receive_from_stream::<EnclaveProcessCommandType>(&mut stream).map_err(|e| {
                e.add_subaction("Failed to receive command type from stream".to_string())
            })?;

        // All connections must be registered with epoll, with the exception of the shutdown one.
        if cmd_type != EnclaveProcessCommandType::ConnectionListenerStop {
            self.add_stream_to_epoll(stream)
                .map_err(|e| e.add_subaction("Failed to add stream to epoll".to_string()))?;
        }

        Ok(cmd_type)
    }

    /// Listen for incoming connections and handle them as they appear.
    fn connection_listener_run(self, listener: UnixListener) -> NitroCliResult<()> {
        // Accept connections and process them (this is a blocking call).
        for stream in listener.incoming() {
            match stream {
                Ok(stream) => {
                    // Received a new connection. Shut down if required.
                    let cmd = self.handle_new_connection(stream);
                    if let Ok(cmd) = cmd {
                        if cmd == EnclaveProcessCommandType::ConnectionListenerStop {
                            break;
                        }
                    }
                }
                Err(err) => {
                    // Connection failed.
                    warn!("Connection error: {:?}", err);
                    break;
                }
            }
        }

        // Remove the listener's socket.
        self.socket
            .close()
            .map_err(|e| e.add_subaction("Failed to close socket".to_string()))?;
        debug!("Connection listener has finished.");
        Ok(())
    }

    /// Terminate the connection listener.
    pub fn stop(&mut self) -> NitroCliResult<()> {
        // Nothing to do if the connection listener thread has not been started.
        if self.listener_thread.is_none() {
            return Ok(());
        }

        // Send termination notification to the listener thread.
        let mut self_conn = UnixStream::connect(self.socket.get_path()).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to connect to listener thread: {:?}", e),
                NitroCliErrorEnum::SocketError
            )
        })?;
        enclave_proc_command_send_single::<EmptyArgs>(
            EnclaveProcessCommandType::ConnectionListenerStop,
            None,
            &mut self_conn,
        )
        .map_err(|e| e.add_subaction("Failed to notify listener thread of shutdown".to_string()))?;

        // Shut the connection down.
        self_conn.shutdown(std::net::Shutdown::Both).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to close connection: {:?}", e),
                NitroCliErrorEnum::SocketCloseError
            )
        })?;

        // Ensure that the listener thread has terminated.
        self.listener_thread.take().unwrap().join().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to join listener thread: {:?}", e),
                NitroCliErrorEnum::ThreadJoinFailure
            )
        })?;
        info!("The connection listener has been stopped.");

        Ok(())
    }

    /// Fetch the next available connection.
    pub fn get_next_connection(&self, enc_fd: Option<RawFd>) -> NitroCliResult<Connection> {
        // Wait on epoll until a valid event is received.
        let mut events = [EpollEvent::empty(); 1];
        loop {
            match epoll::epoll_wait(self.epoll_fd, &mut events, -1) {
                Ok(_) => break,
                Err(nix::errno::Errno::EINTR) => continue,
                Err(e) => {
                    return Err(new_nitro_cli_failure!(
                        &format!("Failed to wait on epoll: {:?}", e),
                        NitroCliErrorEnum::EpollError
                    ))
                }
            }
        }

        let fd = events[0].data() as RawFd;
        let input_stream = match enc_fd {
            // This is a connection to an enclave.
            Some(enc_fd) if enc_fd == fd => None,
            // This is a connection to a CLI instance or to ourselves.
            _ => Some(unsafe { UnixStream::from_raw_fd(fd) }),
        };

        // Remove the fetched descriptor from epoll. We are doing this here since
        // otherwise the Connection would have to do it when dropped and we prefer
        // the Connection not touch epoll directly.
        epoll::epoll_ctl(self.epoll_fd, EpollOp::EpollCtlDel, fd, None).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to remove descriptor from epoll: {:?}", e),
                NitroCliErrorEnum::EpollError
            )
        })?;

        Ok(Connection::new(events[0].events(), input_stream))
    }

    /// Enable the sending of credentials from incoming connections.
    fn enable_credentials_passing(&self, listener: &UnixListener) {
        let val: libc::c_int = 1;
        let rc = unsafe {
            libc::setsockopt(
                listener.as_raw_fd(),
                libc::SOL_SOCKET,
                libc::SO_PASSCRED,
                &val as *const libc::c_int as *const libc::c_void,
                std::mem::size_of::<libc::c_int>() as libc::socklen_t,
            )
        };

        if rc < 0 {
            warn!(
                "Failed to enable credentials passing on socket listener: {}",
                io::Error::last_os_error()
            );
        }

        // Since access policy is handled within the enclave process explicitly, we
        // allow full access to the socket itself (otherwise other users will not
        // be allowed to connect to the socket in the first place).
        if let Ok(sock_addr) = listener.local_addr() {
            if let Some(sock_path) = sock_addr.as_pathname() {
                let perms = Permissions::from_mode(0o766);
                if let Err(e) = set_permissions(sock_path, perms) {
                    warn!("Failed to update socket permissions: {}", e);
                }
            } else {
                warn!("Failed to get the listener's socket path.");
            }
        } else {
            warn!("Failed to get the socket listener's local address.")
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::common::{get_sockets_dir_path, SOCKETS_DIR_PATH_ENV_VAR};

    use std::env;
    use std::fs;
    use std::path::PathBuf;
    use std::process::Command;
    use std::sync::{Arc, Condvar, Mutex};

    const THREADS_STR: &str = "Threads:";
    const TMP_DIR: &str = "./npe";

    fn unset_envvar(varname: &String) {
        unsafe { libc::unsetenv(varname.as_ptr() as *const c_char) };
    }

    /// Inspects the content of /proc/<PID>/status in order to
    /// retrieve the number of threads running in the context of
    /// process <PID>.
    fn get_num_threads_from_status_output(status_str: String) -> u32 {
        let start_idx = status_str.find(THREADS_STR);
        let mut iter = status_str.chars();
        iter.by_ref().nth(start_idx.unwrap() + THREADS_STR.len()); // skip "Threads:\t"
        let slice = iter.as_str();

        let new_str = slice.to_string();
        let end_idx = new_str.find('\n'); // skip after the first '\n'
        let substr = &slice[..end_idx.unwrap()];

        substr.parse().unwrap()
    }

    /// Tests that get_epoll_fd() returns the expected epoll_fd.
    #[test]
    fn test_get_epoll_fd() {
        let connection_listener = ConnectionListener::new().unwrap();
        let epoll_fd = connection_listener.epoll_fd;

        assert_eq!(epoll_fd, connection_listener.get_epoll_fd());
    }

    /// Tests that new connections are monitored and that a command
    /// sent through the connection is received correctly.
    #[test]
    fn test_handle_new_connection() {
        let (mut sock0, sock1) = UnixStream::pair().unwrap();

        let connection_listener = ConnectionListener::new().unwrap();

        let cmd = EnclaveProcessCommandType::Describe;
        let _ = enclave_proc_command_send_single::<EmptyArgs>(cmd, None, &mut sock0);

        let result = connection_listener.handle_new_connection(sock1);

        assert!(result.is_ok());
        assert_eq!(result.unwrap(), EnclaveProcessCommandType::Describe);
    }

    /// Test that add_stream_to_epoll registers a sockfd and that next subsequent
    /// attempts to register the same sockfd fail (since the sockfd is already registered
    /// once).
    #[test]
    fn test_add_stream_to_epoll() {
        let (_, sock1) = UnixStream::pair().unwrap();

        let connection_listener = ConnectionListener::new().unwrap();
        let copy_sock1 = sock1.try_clone();

        if let Ok(copy_sock1) = copy_sock1 {
            let mut cli_evt = EpollEvent::new(EpollFlags::EPOLLIN, copy_sock1.into_raw_fd() as u64);
            let _ = epoll::epoll_ctl(
                connection_listener.epoll_fd,
                EpollOp::EpollCtlAdd,
                sock1.as_raw_fd(),
                &mut cli_evt,
            );
            // Second add should return Err(Sys(EEXIST)), as sock1 is already registed
            // with connection_listener.epoll_fd
            let result = epoll::epoll_ctl(
                connection_listener.epoll_fd,
                EpollOp::EpollCtlAdd,
                sock1.as_raw_fd(),
                &mut cli_evt,
            );

            assert_eq!(result.is_err(), true);
        }
    }

    /// Test that connection_listener_run closes a previously-spawned thread when
    /// processing a ConnectionListenerStop command.
    #[test]
    fn test_connection_listener_run_connection_stop() {
        let old_log_path = env::var(SOCKETS_DIR_PATH_ENV_VAR);

        env::set_var(SOCKETS_DIR_PATH_ENV_VAR, TMP_DIR);

        let resources_dir = get_sockets_dir_path();

        let path_existed = resources_dir.as_path().exists();
        let _ = fs::create_dir(resources_dir.as_path());
        let dummy_sock_name = "run_connection_stop.sock";
        let dummy_sock_path = format!(
            "{}/{}",
            resources_dir.as_path().to_str().unwrap(),
            dummy_sock_name
        );

        // Remove pre-existing socket file
        let _ = std::fs::remove_file(&dummy_sock_path);

        let mut connection_listener = ConnectionListener::new().unwrap();
        connection_listener
            .socket
            .set_path(PathBuf::from(&dummy_sock_path));

        // Get number of running threads before spawning the listener thread
        let out_cmd0 = Command::new("cat")
            .arg(format!("/proc/{}/status", std::process::id()))
            .output()
            .expect("Failed to run cat");
        let out0 = std::str::from_utf8(&out_cmd0.stdout).unwrap();
        let crt_num_threads0 = get_num_threads_from_status_output(out0.to_string());

        let pair = Arc::new((Mutex::new(false), Condvar::new()));
        let pair2 = pair.clone();

        let listener_thread = thread::spawn(move || {
            {
                let (lock, cvar) = &*pair2;
                let mut started = lock.lock().unwrap();
                *started = true;
                cvar.notify_one();
            }

            // Bind the listener to the socket and spawn the listener thread.
            let listener = UnixListener::bind(connection_listener.socket.get_path())
                .map_err(|e| format!("Failed to bind connection listener: {:?}", e))
                .unwrap();
            connection_listener.enable_credentials_passing(&listener);
            connection_listener
                .socket
                .start_monitoring(true)
                .map_err(|e| format!("Failed to start socket monitoring: {:?}", e))
                .unwrap();

            let res = connection_listener.connection_listener_run(listener);
            assert!(res.is_ok());
        });

        // Allow thread to finish spawning
        let (lock, cvar) = &*pair;
        let mut started = lock.lock().unwrap();
        while !*started {
            started = cvar.wait(started).unwrap();
        }

        // Check that the listener thread is running
        let out_cmd1 = Command::new("cat")
            .arg(format!("/proc/{}/status", std::process::id()))
            .output()
            .expect("Failed to run cat");
        let out1 = std::str::from_utf8(&out_cmd1.stdout).unwrap();
        let crt_num_threads1 = get_num_threads_from_status_output(out1.to_string());
        assert!(crt_num_threads0 < crt_num_threads1);

        let my_stream = UnixStream::connect(&dummy_sock_path);

        if let Ok(mut my_stream) = my_stream {
            // Close the listener thread
            let cmd = EnclaveProcessCommandType::ConnectionListenerStop;
            let _ = enclave_proc_command_send_single::<EmptyArgs>(cmd, None, &mut my_stream);
        }

        // Wait for thread to join after exiting
        listener_thread
            .join()
            .expect("Failed to join on the associated thread");

        // Check number of threads after closing the listener thread
        let out_cmd2 = Command::new("cat")
            .arg(format!("/proc/{}/status", std::process::id()))
            .output()
            .expect("Failed to run cat");
        let out2 = std::str::from_utf8(&out_cmd2.stdout).unwrap();
        let crt_num_threads2 = get_num_threads_from_status_output(out2.to_string());
        assert_eq!(crt_num_threads0, crt_num_threads2);
        assert!(crt_num_threads2 < crt_num_threads1);

        if !path_existed {
            // Remove whole resources_dir
            let _ = fs::remove_dir_all(resources_dir.as_path().to_str().unwrap());
        } else {
            // Only remove the socket file
            let _ = fs::remove_file(&dummy_sock_path);
        }

        // Restore previous environment variable value
        if let Ok(old_log_path) = old_log_path {
            env::set_var(SOCKETS_DIR_PATH_ENV_VAR, old_log_path);
        } else {
            env::set_var(SOCKETS_DIR_PATH_ENV_VAR, "");
            unset_envvar(&String::from(SOCKETS_DIR_PATH_ENV_VAR));
        }
    }

    /// Test that connection_listener_run closes a previously-spawned thread when
    /// processing a ConnectionListenerStop command.
    #[test]
    fn test_connection_listener_run_describe() {
        let old_log_path = env::var(SOCKETS_DIR_PATH_ENV_VAR);

        env::set_var(SOCKETS_DIR_PATH_ENV_VAR, TMP_DIR);

        let resources_dir = get_sockets_dir_path();
        let path_existed = resources_dir.as_path().exists();

        let _ = fs::create_dir(resources_dir.as_path());

        let dummy_sock_name = "run_describe.sock";
        let dummy_sock_path = format!(
            "{}/{}",
            resources_dir.as_path().to_str().unwrap(),
            dummy_sock_name
        );

        // Remove pre-existing socket file
        let _ = std::fs::remove_file(&dummy_sock_path);

        let mut connection_listener = ConnectionListener::new().unwrap();
        connection_listener
            .socket
            .set_path(PathBuf::from(&dummy_sock_path));

        // Get number of running threads before spawning the listener thread
        let out_cmd0 = Command::new("cat")
            .arg(format!("/proc/{}/status", std::process::id()))
            .output()
            .expect("Failed to run cat");
        let out0 = std::str::from_utf8(&out_cmd0.stdout).unwrap();
        let crt_num_threads0 = get_num_threads_from_status_output(out0.to_string());

        let pair = Arc::new((Mutex::new(false), Condvar::new()));
        let pair2 = pair.clone();

        let conn_clone = connection_listener.clone();
        let listener_thread = thread::spawn(move || {
            {
                let (lock, cvar) = &*pair2;
                let mut started = lock.lock().unwrap();
                *started = true;
                cvar.notify_one();
            }

            // Bind the listener to the socket and spawn the listener thread.
            let listener = UnixListener::bind(connection_listener.socket.get_path())
                .map_err(|e| format!("Failed to bind connection listener: {:?}", e))
                .unwrap();
            connection_listener.enable_credentials_passing(&listener);
            connection_listener
                .socket
                .start_monitoring(true)
                .map_err(|e| format!("Failed to start socket monitoring: {:?}", e))
                .unwrap();

            conn_clone.connection_listener_run(listener).unwrap();
        });

        // Allow thread to finish spawning
        let (lock, cvar) = &*pair;
        let mut started = lock.lock().unwrap();
        while !*started {
            started = cvar.wait(started).unwrap();
        }

        // Check that the listener thread is running
        let out_cmd1 = Command::new("cat")
            .arg(format!("/proc/{}/status", std::process::id()))
            .output()
            .expect("Failed to run cat");
        let out1 = std::str::from_utf8(&out_cmd1.stdout).unwrap();
        let crt_num_threads1 = get_num_threads_from_status_output(out1.to_string());
        assert!(crt_num_threads0 < crt_num_threads1);

        let my_stream = UnixStream::connect(&dummy_sock_path);

        if let Ok(mut my_stream) = my_stream {
            // Run a command other than ConnectionListenerStop
            let cmd = EnclaveProcessCommandType::Describe;
            let _ = enclave_proc_command_send_single::<EmptyArgs>(cmd, None, &mut my_stream);
        }

        // Check that the listener thread is still running
        let out_cmd2 = Command::new("cat")
            .arg(format!("/proc/{}/status", std::process::id()))
            .output()
            .expect("Failed to run cat");
        let out2 = std::str::from_utf8(&out_cmd2.stdout).unwrap();
        let crt_num_threads2 = get_num_threads_from_status_output(out2.to_string());
        assert!(crt_num_threads0 < crt_num_threads2);

        let my_stream = UnixStream::connect(&dummy_sock_path);

        if let Ok(mut my_stream) = my_stream {
            // Close the listener thread
            let cmd = EnclaveProcessCommandType::ConnectionListenerStop;
            let _ = enclave_proc_command_send_single::<EmptyArgs>(cmd, None, &mut my_stream);

            // Wait for the thread to join after exiting
            listener_thread
                .join()
                .expect("Failed to join on the associated thread");
        }

        // Check number of threads after closing the listener thread
        let out_cmd3 = Command::new("cat")
            .arg(format!("/proc/{}/status", std::process::id()))
            .output()
            .expect("Failed to run cat");
        let out3 = std::str::from_utf8(&out_cmd3.stdout).unwrap();
        let crt_num_threads3 = get_num_threads_from_status_output(out3.to_string());
        assert_eq!(crt_num_threads0, crt_num_threads3);
        assert!(crt_num_threads3 < crt_num_threads1);

        if !path_existed {
            // Remove whole resources_dir
            let _ = fs::remove_dir_all(resources_dir.as_path().to_str().unwrap());
        } else {
            // Only remove the socket file
            let _ = fs::remove_file(&dummy_sock_path);
        }

        // Restore previous enviornment variable value
        if let Ok(old_log_path) = old_log_path {
            env::set_var(SOCKETS_DIR_PATH_ENV_VAR, old_log_path);
        } else {
            env::set_var(SOCKETS_DIR_PATH_ENV_VAR, "");
            unset_envvar(&String::from(SOCKETS_DIR_PATH_ENV_VAR));
        }
    }
}


// Copyright 2019-2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
#![deny(missing_docs)]
#![deny(warnings)]

use std::collections::BTreeSet;
use std::fs::File;
use std::io::{BufRead, BufReader};

use crate::common::commands_parser::RunEnclavesArgs;
use crate::common::{NitroCliErrorEnum, NitroCliFailure, NitroCliResult};
use crate::new_nitro_cli_failure;

/// Path corresponding to the NE CPU pool.
const POOL_FILENAME: &str = "/sys/module/nitro_enclaves/parameters/ne_cpus";

/// The CPU configuration requested by the user.
#[derive(Clone, PartialEq, Eq)]
pub enum EnclaveCpuConfig {
    /// A list with the desired CPU IDs.
    List(Vec<u32>),
    /// The number of desired CPU IDs.
    Count(u32),
}

/// Aggregate CPU information for multiple CPUs.
#[derive(Debug)]
pub struct CpuInfo {
    /// The list with the CPUs available for enclaves.
    cpu_ids: Vec<u32>,
}

impl Default for EnclaveCpuConfig {
    fn default() -> Self {
        EnclaveCpuConfig::Count(0)
    }
}

impl CpuInfo {
    /// Create a new `CpuInfo` instance from the current system configuration.
    pub fn new() -> NitroCliResult<Self> {
        let pool_file = File::open(POOL_FILENAME).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to open CPU pool file: {}", e),
                NitroCliErrorEnum::FileOperationFailure
            )
            .add_info(vec![POOL_FILENAME, "Open"])
        })?;
        let file_reader = BufReader::new(pool_file);

        CpuInfo::from_reader(file_reader)
    }

    /// Create a new `CpuInfo` instance from a buffered reader
    pub(crate) fn from_reader<B: BufRead>(reader: B) -> NitroCliResult<Self> {
        Ok(CpuInfo {
            cpu_ids: CpuInfo::get_cpu_info(reader)?,
        })
    }

    /// Get the CPU configuration from the command-line arguments.
    pub fn get_cpu_config(&self, args: &RunEnclavesArgs) -> NitroCliResult<EnclaveCpuConfig> {
        if let Some(cpu_ids) = args.cpu_ids.clone() {
            self.check_cpu_ids(&cpu_ids).map_err(|e| {
                e.add_subaction("Failed to check whether CPU list is valid".to_string())
            })?;
            Ok(EnclaveCpuConfig::List(cpu_ids))
        } else if let Some(cpu_count) = args.cpu_count {
            if self.cpu_ids.len() < cpu_count as usize {
                return Err(new_nitro_cli_failure!(
                    &format!(
                        "Insufficient CPUs available (requested {}, but maximum is {})",
                        cpu_count,
                        self.cpu_ids.len()
                    ),
                    NitroCliErrorEnum::InsufficientCpus
                )
                .add_info(vec!["cpu-count", &cpu_count.to_string()]));
            }
            Ok(EnclaveCpuConfig::Count(cpu_count))
        } else {
            // Should not happen.
            Err(new_nitro_cli_failure!(
                "Invalid CPU configuration argument",
                NitroCliErrorEnum::InvalidArgument
            ))
        }
    }

    /// Verify that a provided list of CPU IDs is valid.
    pub fn check_cpu_ids(&self, cpu_ids: &[u32]) -> NitroCliResult<()> {
        // Ensure there are no duplicate IDs.
        let mut unique_ids = BTreeSet::new();

        for cpu_id in cpu_ids {
            unique_ids.insert(cpu_id);
        }

        if unique_ids.len() < cpu_ids.len() {
            let duplicate_cpus = CpuInfo::get_duplicate_cpus(&unique_ids, cpu_ids);

            return Err(new_nitro_cli_failure!(
                &format!(
                    "CPU IDs list contains {} duplicate(s)",
                    cpu_ids.len() - unique_ids.len()
                ),
                NitroCliErrorEnum::InvalidCpuConfiguration
            )
            .add_info(vec!["cpu-ids", duplicate_cpus.as_str()]));
        }

        // Ensure the requested CPUs are available in the CPU pool.
        for cpu_id in unique_ids {
            if !self.cpu_ids.contains(cpu_id) {
                return Err(new_nitro_cli_failure!(
                    &format!(
                        "The CPU with ID {} is not available in the NE CPU pool",
                        cpu_id
                    ),
                    NitroCliErrorEnum::NoSuchCpuAvailableInPool
                )
                .add_info(vec!["cpu-ids", &cpu_id.to_string()]));
            }
        }

        // At this point, all requested CPU IDs are part of the enclave CPU pool.
        Ok(())
    }

    /// Get a list of all available CPU IDs.
    pub fn get_cpu_candidates(&self) -> Vec<u32> {
        self.cpu_ids.clone()
    }

    /// Parse a `lscpu` line to obtain a numeric value.
    pub fn get_value(line: &str) -> NitroCliResult<u32> {
        let mut line_str = line.to_string();
        line_str.retain(|c| !c.is_whitespace());
        line_str.parse::<u32>().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to parse CPU ID: {}", e),
                NitroCliErrorEnum::MalformedCpuId
            )
        })
    }

    fn parse_cpu_pool_line(line_str: &str) -> NitroCliResult<Vec<u32>> {
        let mut result: Vec<u32> = Vec::new();

        // The CPU pool format is: "id1-id2,id3-id4,..."
        for interval in line_str.split(',') {
            let bounds: Vec<&str> = interval.split('-').collect();
            match bounds.len() {
                1 => result.push(CpuInfo::get_value(bounds[0])?),
                2 => {
                    let start_id = CpuInfo::get_value(bounds[0])?;
                    let end_id = CpuInfo::get_value(bounds[1])?;

                    for cpu_id in start_id..=end_id {
                        result.push(cpu_id);
                    }
                }
                _ => {
                    return Err(new_nitro_cli_failure!(
                        &format!("Invalid CPU ID interval ({})", interval),
                        NitroCliErrorEnum::CpuError
                    ))
                }
            }
        }

        Ok(result)
    }

    /// Parse the CPU pool and build the list of off-line CPUs.
    fn get_cpu_info<B: BufRead>(reader: B) -> NitroCliResult<Vec<u32>> {
        let mut result: Vec<u32> = Vec::new();

        for line in reader.lines() {
            let line_str = line.map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to read line from CPU pool file: {}", e),
                    NitroCliErrorEnum::FileOperationFailure
                )
                .add_info(vec![POOL_FILENAME, "Read"])
            })?;
            if line_str.trim().is_empty() {
                continue;
            }

            result.append(&mut CpuInfo::parse_cpu_pool_line(&line_str)?);
        }

        Ok(result)
    }

    /// Get a list of duplicate CPUs.
    fn get_duplicate_cpus(uniques: &BTreeSet<&u32>, cpu_ids: &[u32]) -> String {
        let mut result = String::new();
        for unique_cpu_id in uniques {
            if cpu_ids.iter().filter(|x| x == unique_cpu_id).count() > 1 {
                result.push_str(&(unique_cpu_id.to_string()));
            }
        }

        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::common::construct_error_message;

    #[test]
    fn test_parse_cpu_pool_line() {
        let result0 = CpuInfo::parse_cpu_pool_line("1-3,4-6,7-9");
        assert!(result0.is_ok());
        assert_eq!(result0.unwrap(), vec![1, 2, 3, 4, 5, 6, 7, 8, 9]);

        let result1 = CpuInfo::parse_cpu_pool_line("1-4,7");
        assert!(result1.is_ok());
        assert_eq!(result1.unwrap(), vec![1, 2, 3, 4, 7]);

        let result2 = CpuInfo::parse_cpu_pool_line("3,5,7,9");
        assert!(result2.is_ok());
        assert_eq!(result2.unwrap(), vec![3, 5, 7, 9]);

        let result3 = CpuInfo::parse_cpu_pool_line("3+5,7-10");
        assert!(result3.is_err());

        let result4 = CpuInfo::parse_cpu_pool_line("3-a,7-b");
        assert!(result4.is_err());
    }

    #[test]
    fn test_get_value_correct_format() {
        let result0 = CpuInfo::get_value("\t3");
        assert!(result0.is_ok());
        assert_eq!(result0.unwrap(), 3);

        let result1 = CpuInfo::get_value("   \t  4");
        assert!(result1.is_ok());
        assert_eq!(result1.unwrap(), 4);

        let result2 = CpuInfo::get_value("  \n 12");
        assert!(result2.is_ok());
        assert_eq!(result2.unwrap(), 12);
    }

    #[test]
    fn test_get_value_incorrect_format() {
        let result0 = CpuInfo::get_value("\t-2");
        assert!(result0.is_err());
        if let Err(err_info) = result0 {
            let err_str = construct_error_message(&err_info);
            assert!(err_str.contains("Malformed CPU ID error"));
        }

        let result1 = CpuInfo::get_value("\n\n0x06");
        assert!(result1.is_err());
        if let Err(err_info) = result1 {
            let err_str = construct_error_message(&err_info);
            assert!(err_str.contains("Malformed CPU ID error"));
        }

        let result2 = CpuInfo::get_value("     processor");
        assert!(result2.is_err());
        if let Err(err_info) = result2 {
            let err_str = construct_error_message(&err_info);
            assert!(err_str.contains("Malformed CPU ID error"));
        }
    }

    #[test]
    fn test_get_cpu_config_invalid_input() {
        let cpu_info = CpuInfo::from_reader("1".as_bytes()).unwrap();
        let mut run_args = RunEnclavesArgs {
            eif_path: String::new(),
            enclave_cid: None,
            memory_mib: 0,
            debug_mode: false,
            attach_console: false,
            cpu_ids: None,
            cpu_count: Some(343),
            enclave_name: Some("testName".to_string()),
        };

        let mut result = cpu_info.get_cpu_config(&run_args);
        assert!(result.is_err());

        if let Err(err_info) = result {
            let err_str = construct_error_message(&err_info);
            assert!(err_str.contains("Insufficient CPUs available"));
        }

        run_args.cpu_count = None;
        run_args.cpu_ids = Some(vec![1, 2, 3, 4, 5, 6, 7]);
        result = cpu_info.get_cpu_config(&run_args);
        assert!(result.is_err());

        if let Err(err_info) = result {
            let err_str = construct_error_message(&err_info);
            assert!(err_str.contains("No such CPU available in the pool"));
        }
    }

    #[test]
    fn test_get_cpu_config_valid_input() {
        let cpu_info = CpuInfo::from_reader("1,3".as_bytes()).unwrap();
        let mut run_args = RunEnclavesArgs {
            eif_path: String::new(),
            enclave_cid: None,
            memory_mib: 0,
            debug_mode: false,
            attach_console: false,
            cpu_ids: None,
            cpu_count: Some(2),
            enclave_name: Some("testName".to_string()),
        };

        let mut result = cpu_info.get_cpu_config(&run_args);
        assert!(result.is_ok());
        assert!(result.unwrap() == EnclaveCpuConfig::Count(2));

        run_args.cpu_count = None;
        run_args.cpu_ids = Some(vec![1, 3]);

        result = cpu_info.get_cpu_config(&run_args);
        assert!(result.is_ok());
        assert!(result.unwrap() == EnclaveCpuConfig::List(vec![1, 3]));
    }

    #[test]
    fn test_get_cpu_candidates() {
        let cpu_info = CpuInfo::from_reader("1".as_bytes()).unwrap();
        let candidate_cpus = cpu_info.get_cpu_candidates();

        assert!(!candidate_cpus.is_empty());
    }

    #[test]
    fn test_check_cpu_ids() {
        let cpu_info = CpuInfo::from_reader("1,3".as_bytes()).unwrap();
        let mut cpu_ids: Vec<u32> = vec![1];

        let mut result = cpu_info.check_cpu_ids(&cpu_ids);
        assert!(result.is_ok());

        cpu_ids = vec![1, 1];
        result = cpu_info.check_cpu_ids(&cpu_ids);
        assert!(result.is_err());
        if let Err(err_info) = result {
            let err_str = construct_error_message(&err_info);
            assert!(err_str.contains("Invalid CPU configuration"));
        }

        cpu_ids = vec![1, 3];
        result = cpu_info.check_cpu_ids(&cpu_ids);
        assert!(result.is_ok());

        cpu_ids = vec![1, 3, 5];
        result = cpu_info.check_cpu_ids(&cpu_ids);
        assert!(result.is_err());
        if let Err(err_info) = result {
            let err_str = construct_error_message(&err_info);
            assert!(err_str.contains("No such CPU available in the pool"));
        }
    }
}