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

/// The module which provides top-level enclave commands.
pub mod commands;
/// The module which provides a connection to the enclave process.
pub mod connection;
/// The module which provides an enclave socket monitor that listens for incoming connections.
pub mod connection_listener;
/// The module which provides CPU information utilities.
pub mod cpu_info;
/// The module which provides the enclave manager and its utilities.
pub mod resource_manager;
/// The module which provides the managed Unix socket needed to communicate with the enclave process.
pub mod socket;
/// The module which provides additional enclave process utilities.
pub mod utils;

use log::{info, warn};
use nix::sys::epoll::EpollFlags;
use nix::sys::signal::{Signal, SIGHUP};
use nix::unistd::{daemon, getpid, getppid};
use std::os::unix::io::{FromRawFd, IntoRawFd, RawFd};
use std::os::unix::net::UnixStream;
use std::process;
use std::thread::{self, JoinHandle};

use super::common::MSG_ENCLAVE_CONFIRM;
use super::common::{construct_error_message, enclave_proc_command_send_single, notify_error};
use super::common::{
    EnclaveProcessCommandType, ExitGracefully, NitroCliErrorEnum, NitroCliFailure, NitroCliResult,
};
use crate::common::commands_parser::{DescribeEnclavesArgs, EmptyArgs, RunEnclavesArgs};
use crate::common::logger::EnclaveProcLogWriter;
use crate::common::signal_handler::SignalHandler;
use crate::enclave_proc::connection::safe_conn_println;
use crate::new_nitro_cli_failure;

use commands::{describe_enclaves, run_enclaves, terminate_enclaves, DescribeThread};
use connection::Connection;
use connection_listener::ConnectionListener;
use resource_manager::EnclaveManager;

/// The type of enclave event that has been handled.
enum HandledEnclaveEvent {
    /// A hang-up event.
    HangUp,
    /// An unexpected but non-critical event.
    Unexpected,
    /// There was no event that needed handling.
    None,
}

/// Obtain the logger ID from the full enclave ID.
fn get_logger_id(enclave_id: &str) -> String {
    // 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();
    format!("enc-{}:{}", tokens[0], std::process::id())
}

/// Get the action associated with `cmd` as a String.
fn get_command_action(cmd: EnclaveProcessCommandType) -> String {
    match cmd {
        EnclaveProcessCommandType::Run => "Run Enclave".to_string(),
        EnclaveProcessCommandType::Terminate | EnclaveProcessCommandType::TerminateComplete => {
            "Terminate Enclave".to_string()
        }
        EnclaveProcessCommandType::Describe => "Describe Enclaves".to_string(),
        EnclaveProcessCommandType::GetEnclaveCID => "Get Enclave CID".to_string(),
        EnclaveProcessCommandType::GetEnclaveFlags => "Get Enclave Flags".to_string(),
        EnclaveProcessCommandType::ConnectionListenerStop => "Stop Connection Listener".to_string(),
        _ => "Unknown Command".to_string(),
    }
}

/// Send the given command, then close the channel that was used for sending it.
fn send_command_and_close(cmd: EnclaveProcessCommandType, stream: &mut UnixStream) {
    let action_str = &get_command_action(cmd);

    enclave_proc_command_send_single::<EmptyArgs>(cmd, None, stream)
        .ok_or_exit_with_errno(Some("Failed to send command"));
    stream
        .shutdown(std::net::Shutdown::Both)
        .map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to close stream after sending command: {:?}", e),
                NitroCliErrorEnum::SocketCloseError
            )
            .set_action(action_str.to_string())
        })
        .ok_or_exit_with_errno(Some("Failed to shut down stream"));
}

/// Notify that an error has occurred, also forwarding the error message to a connection.
fn notify_error_with_conn(err_msg: &str, conn: &Connection, action: EnclaveProcessCommandType) {
    let action_str = &get_command_action(action);

    notify_error(err_msg);
    conn.eprintln(err_msg)
        .map_err(|e| e.set_action(action_str.to_string()))
        .ok_or_exit_with_errno(Some("Failed to forward error message to connection"));
}

/// Perform enclave termination.
fn run_terminate(
    connection: Connection,
    mut thread_stream: UnixStream,
    mut enclave_manager: EnclaveManager,
) {
    terminate_enclaves(&mut enclave_manager, Some(&connection)).unwrap_or_else(|e| {
        notify_error_with_conn(
            construct_error_message(&e).as_str(),
            &connection,
            EnclaveProcessCommandType::Terminate,
        );
    });

    // Notify the main thread that enclave termination has completed.
    send_command_and_close(
        EnclaveProcessCommandType::TerminateComplete,
        &mut thread_stream,
    );
}

/// Start enclave termination.
fn notify_terminate(
    connection: Connection,
    conn_listener: &ConnectionListener,
    enclave_manager: EnclaveManager,
) -> NitroCliResult<JoinHandle<()>> {
    let (local_stream, thread_stream) = UnixStream::pair().map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Could not create stream pair: {:?}", e),
            NitroCliErrorEnum::SocketPairCreationFailure
        )
    })?;

    conn_listener.add_stream_to_epoll(local_stream)?;
    Ok(thread::spawn(move || {
        run_terminate(connection, thread_stream, enclave_manager)
    }))
}

/// Launch the POSIX signal handler on a dedicated thread and ensure its events are accessible.
fn enclave_proc_configure_signal_handler(conn_listener: &ConnectionListener) -> NitroCliResult<()> {
    let mut signal_handler = SignalHandler::new_with_defaults()
        .mask_all()
        .map_err(|e| e.add_subaction("Failed to configure signal handler".to_string()))?;
    let (local_stream, thread_stream) = UnixStream::pair()
        .map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to create stream pair: {:?}", e),
                NitroCliErrorEnum::SocketPairCreationFailure
            )
            .set_action("Run Enclave".to_string())
        })
        .ok_or_exit_with_errno(Some("Failed to create stream pair"));

    conn_listener
        .add_stream_to_epoll(local_stream)
        .map_err(|e| {
            e.add_subaction(
                "Failed to add stream to epoll when configuring signal handler".to_string(),
            )
        })?;
    signal_handler.start_handler(thread_stream.into_raw_fd(), enclave_proc_handle_signals);

    Ok(())
}

/// The default POSIX signal handling function, which notifies the enclave process to shut down gracefully.
fn enclave_proc_handle_signals(comm_fd: RawFd, signal: Signal) -> bool {
    let mut stream = unsafe { UnixStream::from_raw_fd(comm_fd) };

    warn!(
        "Received signal {:?}. The enclave process will now close.",
        signal
    );
    send_command_and_close(
        EnclaveProcessCommandType::ConnectionListenerStop,
        &mut stream,
    );

    true
}

/// Handle an event coming from an enclave.
fn try_handle_enclave_event(connection: &Connection) -> NitroCliResult<HandledEnclaveEvent> {
    // Check if this is an enclave connection.
    if let Some(mut enc_events) = connection
        .get_enclave_event_flags()
        .map_err(|e| e.add_subaction("Failed to get enclave events flag".to_string()))?
    {
        let enc_hup = enc_events.contains(EpollFlags::EPOLLHUP);

        // Check if non-hang-up events have occurred.
        enc_events.remove(EpollFlags::EPOLLHUP);
        if !enc_events.is_empty() {
            warn!("Received unexpected enclave event(s): {:?}", enc_events);
        }

        // If we received the hang-up event we need to terminate cleanly.
        if enc_hup {
            warn!("Received hang-up event from the enclave. Enclave process will shut down.");
            return Ok(HandledEnclaveEvent::HangUp);
        }

        // Non-hang-up enclave events are not fatal.
        return Ok(HandledEnclaveEvent::Unexpected);
    }

    Ok(HandledEnclaveEvent::None)
}

/// Fetch result of describe thread which was started during start of enclave
/// After result is fetched and stored to enclave manager thread is set to None
/// Thus, actual fetching happens only the first time after enclave is started
fn fetch_describe_result(
    describe_thread: &mut DescribeThread,
    enclave_manager: &mut EnclaveManager,
) -> NitroCliResult<()> {
    if let Some(thread) = describe_thread.take() {
        let result = thread
            .join()
            .map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Termination thread join failed: {:?}", e),
                    NitroCliErrorEnum::ThreadJoinFailure
                )
            })?
            .map_err(|e| e.add_subaction("Failed to save PCR values".to_string()))?;

        enclave_manager
            .set_measurements(result.measurements)
            .map_err(|e| {
                e.add_subaction("Failed to set measurements inside enclave handle.".to_string())
            })?;

        if let Some(metadata) = result.metadata {
            enclave_manager.set_metadata(metadata).map_err(|e| {
                e.add_subaction("Failed to set metadata inside enclave handle.".to_string())
            })?;
        }
    }

    Ok(())
}

/// Handle a single command, returning whenever an error occurs.
fn handle_command(
    cmd: EnclaveProcessCommandType,
    logger: &EnclaveProcLogWriter,
    connection: &Connection,
    conn_listener: &mut ConnectionListener,
    enclave_manager: &mut EnclaveManager,
    terminate_thread: &mut Option<std::thread::JoinHandle<()>>,
    describe_thread: &mut DescribeThread,
) -> NitroCliResult<(i32, bool)> {
    Ok(match cmd {
        EnclaveProcessCommandType::Run => {
            // We should never receive a Run command if we are already running.
            if !enclave_manager.enclave_id.is_empty() {
                (libc::EEXIST, false)
            } else {
                let run_args = connection.read::<RunEnclavesArgs>().map_err(|e| {
                    e.add_subaction("Failed to get run arguments".to_string())
                        .set_action("Run Enclave".to_string())
                })?;
                info!("Run args = {:?}", run_args);

                let run_result = run_enclaves(&run_args, Some(connection)).map_err(|e| {
                    e.add_subaction("Failed to trigger enclave run".to_string())
                        .set_action("Run Enclave".to_string())
                })?;
                *enclave_manager = run_result.enclave_manager;
                *describe_thread = run_result.describe_thread;

                info!("Enclave ID = {}", enclave_manager.enclave_id);
                logger
                    .update_logger_id(&get_logger_id(&enclave_manager.enclave_id))
                    .map_err(|e| e.set_action("Failed to update logger ID".to_string()))?;
                conn_listener
                    .start(&enclave_manager.enclave_id)
                    .map_err(|e| {
                        e.set_action("Failed to start connection listener thread".to_string())
                    })?;

                // Add the enclave descriptor to epoll to listen for enclave events.
                let enc_fd = enclave_manager
                    .get_enclave_descriptor()
                    .map_err(|e| e.set_action("Failed to get enclave descriptor".to_string()))?;
                conn_listener
                    .register_enclave_descriptor(enc_fd)
                    .map_err(|e| {
                        e.set_action("Failed to register enclave descriptor".to_string())
                    })?;
                (0, false)
            }
        }

        EnclaveProcessCommandType::Terminate => {
            *terminate_thread = Some(
                notify_terminate(connection.clone(), conn_listener, enclave_manager.clone())
                    .map_err(|e| {
                        e.set_action("Failed to send enclave termination request".to_string())
                    })?,
            );
            (0, false)
        }

        EnclaveProcessCommandType::TerminateComplete => {
            info!("Enclave has completed termination.");
            (0, true)
        }

        EnclaveProcessCommandType::GetEnclaveCID => {
            let enclave_cid = enclave_manager
                .get_console_resources_enclave_cid()
                .map_err(|e| {
                    e.set_action("Failed to get console resources (enclave CID)".to_string())
                })?;
            connection.write_u64(enclave_cid).map_err(|e| {
                e.add_subaction("Failed to write enclave CID to connection".to_string())
                    .set_action("Get Enclave CID".to_string())
            })?;
            (0, false)
        }

        EnclaveProcessCommandType::GetEnclaveFlags => {
            let enclave_flags = enclave_manager
                .get_console_resources_enclave_flags()
                .map_err(|e| {
                    e.set_action("Failed to get console resources (enclave flags)".to_string())
                })?;
            connection.write_u64(enclave_flags).map_err(|e| {
                e.add_subaction("Failed to write enclave flags to connection".to_string())
                    .set_action("Get Enclave Flags".to_string())
            })?;
            (0, false)
        }

        EnclaveProcessCommandType::GetEnclaveName => {
            connection.write_u64(MSG_ENCLAVE_CONFIRM).map_err(|e| {
                e.add_subaction("Failed to write confirmation".to_string())
                    .set_action("Get Enclave Name".to_string())
            })?;
            safe_conn_println(
                Some(connection),
                serde_json::to_string_pretty(&enclave_manager.enclave_name)
                    .map_err(|err| {
                        new_nitro_cli_failure!(
                            &format!("Failed to write enclave name to connection: {:?}", err),
                            NitroCliErrorEnum::SerdeError
                        )
                    })?
                    .as_str(),
            )?;
            (0, false)
        }

        EnclaveProcessCommandType::GetIDbyName => {
            connection.write_u64(MSG_ENCLAVE_CONFIRM).map_err(|e| {
                e.add_subaction("Failed to write confirmation".to_string())
                    .set_action("Name to ID".to_string())
            })?;
            let name = connection.read::<String>().map_err(|e| {
                e.add_subaction("Failed to get enclave name".to_string())
                    .set_action("Name to ID".to_string())
            })?;

            // Respond only if the current enclave name matches
            if enclave_manager.enclave_name == name {
                safe_conn_println(
                    Some(connection),
                    serde_json::to_string_pretty(&enclave_manager.enclave_id)
                        .map_err(|err| {
                            new_nitro_cli_failure!(
                                &format!("Failed to display RunEnclaves data: {:?}", err),
                                NitroCliErrorEnum::SerdeError
                            )
                        })?
                        .as_str(),
                )?;
            }
            (0, false)
        }

        EnclaveProcessCommandType::Describe => {
            let describe_args = connection.read::<DescribeEnclavesArgs>().map_err(|e| {
                e.add_subaction("Failed to get describe arguments".to_string())
                    .set_action("Describe Enclave".to_string())
            })?;
            connection.write_u64(MSG_ENCLAVE_CONFIRM).map_err(|e| {
                e.add_subaction("Failed to write confirmation".to_string())
                    .set_action("Describe Enclaves".to_string())
            })?;

            // Evaluate describe thread result if needed
            fetch_describe_result(describe_thread, enclave_manager)?;

            describe_enclaves(enclave_manager, connection, describe_args.metadata).map_err(
                |e| {
                    e.add_subaction("Failed to describe enclave".to_string())
                        .set_action("Describe Enclaves".to_string())
                },
            )?;
            (0, false)
        }

        EnclaveProcessCommandType::ConnectionListenerStop => (0, true),

        EnclaveProcessCommandType::NotPermitted => (libc::EACCES, false),
    })
}

/// The main event loop of the enclave process.
fn process_event_loop(
    comm_stream: UnixStream,
    logger: &EnclaveProcLogWriter,
) -> NitroCliResult<()> {
    let mut conn_listener = ConnectionListener::new()?;
    let mut enclave_manager = EnclaveManager::default();
    let mut terminate_thread: Option<std::thread::JoinHandle<()>> = None;
    let mut describe_thread: DescribeThread = None;
    let mut done = false;
    let mut ret_value = Ok(());

    // Start the signal handler before spawning any other threads. This is done since the
    // handler will mask all relevant signals from the current thread and this setting will
    // be automatically inherited by all threads spawned from this point on; we want this
    // because only the dedicated thread spawned by the handler should listen for signals.
    enclave_proc_configure_signal_handler(&conn_listener)
        .map_err(|e| e.add_subaction("Failed to configure signal handler".to_string()))?;

    // Add the CLI communication channel to epoll.
    conn_listener
        .handle_new_connection(comm_stream)
        .map_err(|e| {
            e.add_subaction("Failed to add CLI communication channel to epoll".to_string())
        })?;

    while !done {
        // We can get connections to CLI instances, to the enclave or to ourselves.
        let connection =
            conn_listener.get_next_connection(enclave_manager.get_enclave_descriptor().ok())?;

        // If this is an enclave event, handle it.
        match try_handle_enclave_event(&connection) {
            Ok(HandledEnclaveEvent::HangUp) => break,
            Ok(HandledEnclaveEvent::Unexpected) => continue,
            Ok(HandledEnclaveEvent::None) => (),
            Err(error_info) => {
                ret_value = Err(error_info
                    .add_subaction("Error while trying to handle enclave event".to_string()));
                break;
            }
        }

        // At this point we have a connection that is not coming from an enclave.
        // Read the command that should be executed.
        let cmd = match connection.read_command() {
            Ok(value) => value,
            Err(mut error_info) => {
                error_info = error_info
                    .add_subaction("Failed to read command".to_string())
                    .set_action("Run Enclave".to_string());
                notify_error_with_conn(
                    &construct_error_message(&error_info),
                    &connection,
                    EnclaveProcessCommandType::NotPermitted,
                );
                break;
            }
        };

        info!("Received command: {:?}", cmd);
        let status = handle_command(
            cmd,
            logger,
            &connection,
            &mut conn_listener,
            &mut enclave_manager,
            &mut terminate_thread,
            &mut describe_thread,
        );

        // Obtain the status code and whether the event loop must be exited.
        let (status_code, do_break) = match status {
            Ok(value) => value,
            Err(mut error_info) => {
                // Any encountered error is both logged and send to the other side of the connection.
                error_info = error_info
                    .add_subaction(format!("Failed to execute command `{:?}`", cmd))
                    .set_action("Run Enclave".to_string());
                notify_error_with_conn(&construct_error_message(&error_info), &connection, cmd);
                (libc::EINVAL, true)
            }
        };

        done = do_break;

        // Perform clean-up and stop the connection listener before returning the status to the CLI.
        // This is done to avoid race conditions where the enclave process has not yet removed the
        // socket and another CLI issues a command on that very-soon-to-be-removed socket.
        if done {
            // Stop the connection listener.
            conn_listener.stop()?;

            // Wait for the termination thread, if any.
            if terminate_thread.is_some() {
                terminate_thread.take().unwrap().join().map_err(|e| {
                    new_nitro_cli_failure!(
                        &format!("Termination thread join failed: {:?}", e),
                        NitroCliErrorEnum::ThreadJoinFailure
                    )
                })?;
            };
        }

        // Only the commands coming from the CLI must be replied to with the status code.
        match cmd {
            EnclaveProcessCommandType::Run
            | EnclaveProcessCommandType::Terminate
            | EnclaveProcessCommandType::Describe
            | EnclaveProcessCommandType::GetEnclaveName
            | EnclaveProcessCommandType::GetIDbyName => {
                connection.write_status(status_code).map_err(|_| {
                    new_nitro_cli_failure!(
                        "Process event loop failed",
                        NitroCliErrorEnum::EnclaveProcessSendReplyFailure
                    )
                })?
            }
            _ => (),
        }
    }

    info!("Enclave process {} exited event loop.", process::id());

    ret_value
}

/// Create the enclave process.
fn create_enclave_process(logger: &EnclaveProcLogWriter) -> NitroCliResult<()> {
    // To get a detached process, we first:
    // (1) Temporarily ignore specific signals (SIGHUP).
    // (2) Daemonize the current process.
    // (3) Wait until the detached process is orphaned.
    // (4) Restore signal handlers.
    let signal_handler = SignalHandler::new(&[SIGHUP])
        .mask_all()
        .map_err(|e| e.add_subaction("Failed to mask signals".to_string()))?;
    let ppid = getpid();

    // Daemonize the current process. The working directory remains
    // unchanged and the standard descriptors are routed to '/dev/null'.
    daemon(true, false).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to daemonize enclave process: {:?}", e),
            NitroCliErrorEnum::DaemonizeProcessFailure
        )
    })?;

    // This is our detached process.
    logger
        .update_logger_id(format!("enc-xxxxxxx:{}", std::process::id()).as_str())
        .map_err(|e| e.add_subaction("Failed to update logger id".to_string()))?;
    info!("Enclave process PID: {}", process::id());

    // We must wait until we're 100% orphaned. That is, our parent must
    // no longer be the pre-fork process.
    while getppid() == ppid {
        thread::sleep(std::time::Duration::from_millis(10));
    }

    // Restore signal handlers.
    signal_handler
        .unmask_all()
        .map_err(|e| e.add_subaction("Failed to restore signal handlers".to_string()))?;

    Ok(())
}

/// Launch the enclave process.
///
/// * `comm_fd` - A descriptor used for initial communication with the parent Nitro CLI instance.
/// * `logger` - The current log writer, whose ID gets updated when an enclave is launched.
pub fn enclave_process_run(comm_stream: UnixStream, logger: &EnclaveProcLogWriter) {
    create_enclave_process(logger)
        .map_err(|e| e.set_action("Run Enclave".to_string()))
        .ok_or_exit_with_errno(None);
    let res = process_event_loop(comm_stream, logger);
    if let Err(mut error_info) = res {
        error_info = error_info.set_action("Run Enclave".to_string());
        notify_error(construct_error_message(&error_info).as_str());
        process::exit(error_info.error_code as i32);
    }
    process::exit(0);
}


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

use aws_nitro_enclaves_image_format::defs::EifIdentityInfo;
use driver_bindings::*;
use eif_loader::{enclave_ready, TIMEOUT_MINUTE_MS};
use libc::c_int;
use log::{debug, info};
use std::collections::BTreeMap;
use std::convert::{From, Into};
use std::fs::{File, OpenOptions};
use std::io::prelude::*;
use std::io::Error;
use std::mem::size_of;
use std::os::unix::io::{AsRawFd, RawFd};
use std::str;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use vsock::{VsockAddr, VsockListener};

use crate::common::json_output::EnclaveBuildInfo;
use crate::common::{construct_error_message, notify_error};
use crate::common::{
    ExitGracefully, NitroCliErrorEnum, NitroCliFailure, NitroCliResult, ENCLAVE_READY_VSOCK_PORT,
    VMADDR_CID_PARENT,
};
use crate::enclave_proc::connection::Connection;
use crate::enclave_proc::connection::{safe_conn_eprintln, safe_conn_println};
use crate::enclave_proc::cpu_info::EnclaveCpuConfig;
use crate::enclave_proc::utils::get_run_enclaves_info;
use crate::enclave_proc::utils::{GiB, MiB};
use crate::new_nitro_cli_failure;
use crate::utils::ceil_div;

/// CamelCase alias for the bindgen generated driver struct (ne_enclave_start_info).
pub type EnclaveStartInfo = ne_enclave_start_info;

/// CamelCase alias for the bindgen generated driver struct (ne_user_memory_region).
pub type UserMemoryRegion = ne_user_memory_region;

/// CamelCase alias for the bindgen generate struct (ne_image_load_info).
pub type ImageLoadInfo = ne_image_load_info;

/// The internal data type needed for describing an enclave.
type UnpackedHandle = (u64, u64, u64, Vec<u32>, u64, u64, EnclaveState);

/// The bit indicating if an enclave has been launched in debug mode.
pub const NE_ENCLAVE_DEBUG_MODE: u64 = 0x1;

/// Constant number used for computing the lower memory limit.
const ENCLAVE_MEMORY_EIF_SIZE_RATIO: u64 = 4;

/// Enclave Image Format (EIF) flag.
const NE_EIF_IMAGE: u64 = 0x01;

/// Flag indicating a memory region for enclave general usage.
const NE_DEFAULT_MEMORY_REGION: u64 = 0;

/// Magic number for Nitro Enclave IOCTL codes.
const NE_MAGIC: u64 = 0xAE;

/// Path corresponding to the Nitro Enclaves device file.
const NE_DEV_FILEPATH: &str = "/dev/nitro_enclaves";

/// IOCTL code for `NE_CREATE_VM`.
pub const NE_CREATE_VM: u64 = nix::request_code_read!(NE_MAGIC, 0x20, size_of::<u64>()) as _;

/// IOCTL code for `NE_ADD_VCPU`.
pub const NE_ADD_VCPU: u64 = nix::request_code_readwrite!(NE_MAGIC, 0x21, size_of::<u32>()) as _;

/// IOCTL code for `NE_GET_IMAGE_LOAD_INFO`.
pub const NE_GET_IMAGE_LOAD_INFO: u64 =
    nix::request_code_readwrite!(NE_MAGIC, 0x22, size_of::<ImageLoadInfo>()) as _;

/// IOCTL code for `NE_SET_USER_MEMORY_REGION`.
pub const NE_SET_USER_MEMORY_REGION: u64 =
    nix::request_code_write!(NE_MAGIC, 0x23, size_of::<MemoryRegion>()) as _;

/// IOCTL code for `NE_START_ENCLAVE`.
pub const NE_START_ENCLAVE: u64 =
    nix::request_code_readwrite!(NE_MAGIC, 0x24, size_of::<EnclaveStartInfo>()) as _;

/// Mapping between hugepage size and allocation flag, in descending order of size.
const HUGE_PAGE_MAP: [(libc::c_int, u64); 9] = [
    (libc::MAP_HUGE_16GB, 16 * GiB),
    (libc::MAP_HUGE_2GB, 2 * GiB),
    (libc::MAP_HUGE_1GB, GiB),
    (libc::MAP_HUGE_512MB, 512 * MiB),
    (libc::MAP_HUGE_256MB, 256 * MiB),
    (libc::MAP_HUGE_32MB, 32 * MiB),
    (libc::MAP_HUGE_16MB, 16 * MiB),
    (libc::MAP_HUGE_8MB, 8 * MiB),
    (libc::MAP_HUGE_2MB, 2 * MiB),
];

/// A memory region used by the enclave memory allocator.
#[derive(Clone, Debug)]
pub struct MemoryRegion {
    /// Flags to determine the usage for the memory region.
    flags: u64,
    /// The region's size in bytes.
    mem_size: u64,
    /// The region's virtual address.
    mem_addr: u64,
}

/// The state an enclave may be in.
#[derive(Clone)]
pub enum EnclaveState {
    /// The enclave is not running (it's either not started or has been terminated).
    Empty,
    /// The enclave is running.
    Running,
    /// The enclave is in the process of terminating.
    Terminating,
}

/// Helper structure to allocate memory resources needed by an enclave.
#[derive(Clone, Default)]
struct ResourceAllocator {
    /// The requested memory size in bytes.
    requested_mem: u64,
    /// The memory regions that have actually been allocated.
    mem_regions: Vec<MemoryRegion>,
}

/// Helper structure for managing an enclave's resources.
#[derive(Default)]
struct EnclaveHandle {
    /// The CPU configuration as requested by the user.
    #[allow(dead_code)]
    cpu_config: EnclaveCpuConfig,
    /// List of CPU IDs provided to the enclave.
    cpu_ids: Vec<u32>,
    /// Amount of memory allocated for the enclave, in MB.
    allocated_memory_mib: u64,
    /// The enclave slot ID.
    slot_uid: u64,
    /// The enclave CID.
    enclave_cid: Option<u64>,
    /// Enclave flags (including the enclave debug mode flag).
    flags: u64,
    /// The driver-provided enclave descriptor.
    enc_fd: RawFd,
    /// The allocator used to manage enclave memory.
    resource_allocator: ResourceAllocator,
    /// The enclave image file.
    eif_file: Option<File>,
    /// The current state the enclave is in.
    state: EnclaveState,
    /// PCR values.
    build_info: EnclaveBuildInfo,
    /// EIF metadata
    metadata: Option<EifIdentityInfo>,
}

/// The structure which manages an enclave in a thread-safe manner.
#[derive(Clone, Default)]
pub struct EnclaveManager {
    /// The full ID of the managed enclave.
    pub enclave_id: String,
    /// Name of the managed enclave.
    pub enclave_name: String,
    /// A thread-safe handle to the enclave's resources.
    enclave_handle: Arc<Mutex<EnclaveHandle>>,
}

impl ToString for EnclaveState {
    fn to_string(&self) -> String {
        match self {
            EnclaveState::Empty => "EMPTY",
            EnclaveState::Running => "RUNNING",
            EnclaveState::Terminating => "TERMINATING",
        }
        .to_string()
    }
}

impl Default for EnclaveState {
    fn default() -> Self {
        EnclaveState::Empty
    }
}

impl Default for EnclaveBuildInfo {
    fn default() -> Self {
        EnclaveBuildInfo::new(BTreeMap::new())
    }
}

/// Construct a UserMemoryRegion object from a MemoryRegion instance.
/// Implementing the `From` trait automatically gives access to an
/// implementation of `Into` which can be used for a MemoryRegion instance.
impl From<&MemoryRegion> for UserMemoryRegion {
    fn from(mem_reg: &MemoryRegion) -> UserMemoryRegion {
        UserMemoryRegion {
            flags: mem_reg.flags,
            memory_size: mem_reg.mem_size,
            userspace_addr: mem_reg.mem_addr,
        }
    }
}

impl MemoryRegion {
    /// Create a new `MemoryRegion` instance with the specified size (in bytes).
    pub fn new(hugepage_flag: libc::c_int) -> NitroCliResult<Self> {
        let region_index = HUGE_PAGE_MAP
            .iter()
            .position(|&page_info| page_info.0 == hugepage_flag)
            .ok_or_else(|| {
                new_nitro_cli_failure!(
                    &format!(
                        "Failed to find huge page entry for flag {:X?}",
                        hugepage_flag
                    ),
                    NitroCliErrorEnum::NoSuchHugepageFlag
                )
            })?;
        let region_size = HUGE_PAGE_MAP[region_index].1;

        let addr = unsafe {
            libc::mmap(
                std::ptr::null_mut(),
                region_size as usize,
                libc::PROT_READ | libc::PROT_WRITE,
                libc::MAP_PRIVATE | libc::MAP_ANONYMOUS | libc::MAP_HUGETLB | hugepage_flag,
                -1,
                0,
            )
        };

        if addr == libc::MAP_FAILED {
            return Err(new_nitro_cli_failure!(
                "Failed to map memory",
                NitroCliErrorEnum::EnclaveMmapError
            ));
        }

        // Record the allocated region.
        Ok(MemoryRegion {
            flags: NE_DEFAULT_MEMORY_REGION,
            mem_size: region_size,
            mem_addr: addr as u64,
        })
    }

    /// Create a new `MemoryRegion` instance with the specified values.
    pub fn new_with(flags: u64, mem_addr: u64, mem_size: u64) -> Self {
        MemoryRegion {
            flags,
            mem_size,
            mem_addr,
        }
    }

    /// Free the memory region, if it has been allocated earlier.
    fn free(&mut self) -> NitroCliResult<()> {
        // Do nothing if the region has already been freed.
        if self.mem_addr == 0 {
            return Ok(());
        }

        let rc =
            unsafe { libc::munmap(self.mem_addr as *mut libc::c_void, self.mem_size as usize) };

        if rc < 0 {
            return Err(new_nitro_cli_failure!(
                "Failed to unmap memory",
                NitroCliErrorEnum::EnclaveMunmapError
            ));
        }

        // Set the address and length to 0 to avoid double-freeing.
        self.mem_addr = 0;
        self.mem_size = 0;

        Ok(())
    }

    /// Write the content from a file into memory at a given offset.
    fn fill_from_file(
        &self,
        file: &mut File,
        region_offset: usize,
        size: usize,
    ) -> NitroCliResult<()> {
        let offset_plus_size = region_offset.checked_add(size).ok_or_else(|| {
            new_nitro_cli_failure!(
                "Memory overflow when writing EIF file to region",
                NitroCliErrorEnum::MemoryOverflow
            )
        })?;

        if offset_plus_size > self.mem_size as usize {
            return Err(new_nitro_cli_failure!(
                "Out of region",
                NitroCliErrorEnum::MemoryOverflow
            ));
        }

        let bytes = unsafe {
            std::slice::from_raw_parts_mut(self.mem_addr as *mut u8, self.mem_size as usize)
        };

        file.read_exact(&mut bytes[region_offset..region_offset + size])
            .map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Error while reading from enclave image: {:?}", e),
                    NitroCliErrorEnum::EifParsingError
                )
            })?;

        Ok(())
    }

    /// Get the virtual address of the memory region.
    pub fn mem_addr(&self) -> u64 {
        self.mem_addr
    }

    /// Get the size in bytes of the memory region.
    pub fn mem_size(&self) -> u64 {
        self.mem_size
    }
}

impl Drop for MemoryRegion {
    fn drop(&mut self) {
        self.free()
            .ok_or_exit_with_errno(Some("Failed to drop memory region"));
    }
}

impl ResourceAllocator {
    /// Create a new `ResourceAllocator` instance which must cover at least the requested amount of memory (in bytes).
    fn new(requested_mem: u64) -> NitroCliResult<Self> {
        if requested_mem == 0 {
            return Err(new_nitro_cli_failure!(
                "Cannot start an enclave with no memory",
                NitroCliErrorEnum::InsufficientMemoryRequested
            )
            .add_info(vec!["memory", &(requested_mem >> 20).to_string()]));
        }

        Ok(ResourceAllocator {
            requested_mem,
            mem_regions: Vec::new(),
        })
    }

    /// Allocate and provide a list of memory regions. This function creates a list of
    /// memory regions which contain at least `self.requested_mem` bytes. Each region
    /// is equivalent to a huge-page and is allocated using memory mapping.
    fn allocate(&mut self) -> NitroCliResult<&Vec<MemoryRegion>> {
        let mut allocated_pages = BTreeMap::<u64, u32>::new();
        let mut needed_mem = self.requested_mem as i64;
        let mut split_index = 0;

        info!(
            "Allocating memory regions to hold {} bytes.",
            self.requested_mem
        );

        // Always allocate larger pages first, to reduce fragmentation and page count.
        // Once an allocation of a given page size fails, proceed to the next smaller
        // page size and retry.
        for (_, page_info) in HUGE_PAGE_MAP.iter().enumerate() {
            while needed_mem >= page_info.1 as i64 {
                match MemoryRegion::new(page_info.0) {
                    Ok(value) => {
                        needed_mem -= value.mem_size as i64;
                        self.mem_regions.push(value);
                    }
                    Err(_) => break,
                }
            }
        }

        // If the user requested exactly the amount of memory that was reserved earlier,
        // we should be left with no more memory that needs allocation. But if the user
        // requests a smaller amount, we must then aim to reduce waster memory from
        // larger-page allocations (Ex: if we have 1 x 1 GB page and 1 x 2 MB page, but
        // we want to allocate only 512 MB, the above algorithm will have allocated only
        // the 2 MB page, since the 1 GB page was too large for what was needed; we now
        // need to allocate in increasing order of page size in order to reduce wastage).

        if needed_mem > 0 {
            for (_, page_info) in HUGE_PAGE_MAP.iter().rev().enumerate() {
                while needed_mem > 0 {
                    match MemoryRegion::new(page_info.0) {
                        Ok(value) => {
                            needed_mem -= value.mem_size as i64;
                            self.mem_regions.push(value);
                        }
                        Err(_) => break,
                    }
                }
            }
        }

        // If we still have memory to allocate, it means we have insufficient resources.
        if needed_mem > 0 {
            return Err(new_nitro_cli_failure!(
                &format!(
                    "Failed to allocate entire memory ({} MB remained)",
                    needed_mem >> 20
                ),
                NitroCliErrorEnum::InsufficientMemoryAvailable
            )
            .add_info(vec!["memory", &(self.requested_mem >> 20).to_string()]));
        }

        // At this point, we may have allocated more than we need, so we release all
        // regions we no longer need, starting with the smallest ones.
        self.mem_regions
            .sort_by(|reg1, reg2| reg2.mem_size.cmp(&reg1.mem_size));

        needed_mem = self.requested_mem as i64;
        for (_, region) in self.mem_regions.iter().enumerate() {
            if needed_mem <= 0 {
                break;
            }

            needed_mem -= region.mem_size as i64;
            split_index += 1
        }

        // The regions that we no longer need are freed automatically on draining, since
        // MemRegion implements Drop.
        self.mem_regions.drain(split_index..);

        // Generate a summary of the allocated memory.
        for (_, region) in self.mem_regions.iter().enumerate() {
            if let Some(page_count) = allocated_pages.get_mut(&region.mem_size) {
                *page_count += 1;
            } else {
                allocated_pages.insert(region.mem_size, 1);
            }
        }

        info!(
            "Allocated {} region(s): {}",
            self.mem_regions.len(),
            allocated_pages
                .iter()
                .map(|(size, count)| format!("{} page(s) of {} MB", count, size >> 20))
                .collect::<Vec<String>>()
                .join(", ")
        );

        Ok(&self.mem_regions)
    }

    /// Free all previously-allocated memory regions.
    fn free(&mut self) -> NitroCliResult<()> {
        for region in self.mem_regions.iter_mut() {
            region
                .free()
                .map_err(|e| e.add_subaction("Failed to free enclave memory region".to_string()))?;
        }

        self.mem_regions.clear();
        Ok(())
    }
}

impl Drop for ResourceAllocator {
    fn drop(&mut self) {
        self.free()
            .ok_or_exit_with_errno(Some("Failed to drop resource allocator"));
    }
}

impl EnclaveHandle {
    /// Create a new enclave handle instance.
    fn new(
        enclave_cid: Option<u64>,
        memory_mib: u64,
        cpu_config: EnclaveCpuConfig,
        eif_file: File,
        debug_mode: bool,
    ) -> NitroCliResult<Self> {
        let requested_mem = memory_mib << 20;
        let eif_size = eif_file
            .metadata()
            .map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to get enclave image file metadata: {:?}", e),
                    NitroCliErrorEnum::FileOperationFailure
                )
            })?
            .len();
        if ENCLAVE_MEMORY_EIF_SIZE_RATIO * eif_size > requested_mem {
            return Err(new_nitro_cli_failure!(
                &format!(
                    "At least {} MB must be allocated (which is {} times the EIF file size)",
                    ceil_div(ceil_div(eif_size, 1024), 1024) * ENCLAVE_MEMORY_EIF_SIZE_RATIO,
                    ENCLAVE_MEMORY_EIF_SIZE_RATIO
                ),
                NitroCliErrorEnum::InsufficientMemoryRequested
            )
            .add_info(vec![
                "memory",
                &memory_mib.to_string(),
                &(ceil_div(ceil_div(eif_size, 1024), 1024) * ENCLAVE_MEMORY_EIF_SIZE_RATIO)
                    .to_string(),
            ]));
        }

        // Open the device file.
        let dev_file = OpenOptions::new()
            .read(true)
            .write(true)
            .open(NE_DEV_FILEPATH)
            .map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to open device file: {:?}", e),
                    NitroCliErrorEnum::FileOperationFailure
                )
                .add_info(vec![NE_DEV_FILEPATH, "Open"])
            })?;

        let mut slot_uid: u64 = 0;
        let enc_fd = EnclaveHandle::do_ioctl(dev_file.as_raw_fd(), NE_CREATE_VM, &mut slot_uid)
            .map_err(|e| e.add_subaction("Create VM ioctl failed".to_string()))?;
        let flags: u64 = if debug_mode { NE_ENCLAVE_DEBUG_MODE } else { 0 };

        if enc_fd < 0 {
            return Err(new_nitro_cli_failure!(
                &format!("Invalid enclave file descriptor ({})", enc_fd),
                NitroCliErrorEnum::InvalidEnclaveFd
            ));
        }

        Ok(EnclaveHandle {
            cpu_config,
            cpu_ids: vec![],
            allocated_memory_mib: 0,
            slot_uid,
            enclave_cid,
            flags,
            enc_fd,
            resource_allocator: ResourceAllocator::new(requested_mem)
                .map_err(|e| e.add_subaction("Create resource allocator".to_string()))?,
            eif_file: Some(eif_file),
            state: EnclaveState::default(),
            build_info: EnclaveBuildInfo::new(BTreeMap::new()),
            metadata: None,
        })
    }

    /// Initialize the enclave environment and start the enclave.
    fn create_enclave(
        &mut self,
        enclave_name: String,
        connection: Option<&Connection>,
    ) -> NitroCliResult<String> {
        self.init_memory(connection)
            .map_err(|e| e.add_subaction("Memory initialization issue".to_string()))?;
        self.init_cpus()
            .map_err(|e| e.add_subaction("vCPUs initialization issue".to_string()))?;

        let sockaddr = VsockAddr::new(VMADDR_CID_PARENT, ENCLAVE_READY_VSOCK_PORT);
        let listener = VsockListener::bind(&sockaddr).map_err(|_| {
            new_nitro_cli_failure!(
                "Enclave boot heartbeat vsock connection - vsock bind error",
                NitroCliErrorEnum::EnclaveBootFailure
            )
        })?;

        let enclave_start = self
            .start(connection)
            .map_err(|e| e.add_subaction("Enclave start issue".to_string()))?;

        // Get eif size to feed it to calculate_necessary_timeout helper function
        let eif_size = self
            .eif_file
            .as_ref()
            .ok_or_else(|| {
                new_nitro_cli_failure!(
                    "Failed to get EIF file",
                    NitroCliErrorEnum::FileOperationFailure
                )
            })?
            .metadata()
            .map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to get enclave image file metadata: {:?}", e),
                    NitroCliErrorEnum::FileOperationFailure
                )
            })?
            .len();

        // Update the poll timeout based on the eif size or allocated memory
        let poll_timeout = calculate_necessary_timeout(eif_size);

        enclave_ready(listener, poll_timeout).map_err(|err| {
            let err_msg = format!("Waiting on enclave to boot failed with error {:?}", err);
            self.terminate_enclave_error(&err_msg);
            new_nitro_cli_failure!(&err_msg, NitroCliErrorEnum::EnclaveBootFailure)
        })?;

        self.enclave_cid = Some(enclave_start.enclave_cid);

        let info = get_run_enclaves_info(
            enclave_name,
            enclave_start.enclave_cid,
            self.slot_uid,
            self.cpu_ids.clone(),
            self.allocated_memory_mib,
        )
        .map_err(|e| e.add_subaction("Get RunEnclaves information issue".to_string()))?;

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

        Ok(info.enclave_id)
    }

    /// Allocate memory and provide it to the enclave.
    fn init_memory(&mut self, connection: Option<&Connection>) -> NitroCliResult<()> {
        // Allocate the memory regions needed by the enclave.
        safe_conn_eprintln(connection, "Start allocating memory...")?;

        let requested_mem_mib = self.resource_allocator.requested_mem >> 20;
        let regions = self
            .resource_allocator
            .allocate()
            .map_err(|e| e.add_subaction("Failed to allocate enclave memory".to_string()))?;

        self.allocated_memory_mib = regions.iter().fold(0, |mut acc, val| {
            acc += val.mem_size;
            acc
        }) >> 20;

        if self.allocated_memory_mib < requested_mem_mib {
            return Err(new_nitro_cli_failure!(
                &format!(
                    "Failed to allocate sufficient memory (requested {} MB, but got {} MB)",
                    requested_mem_mib, self.allocated_memory_mib
                ),
                NitroCliErrorEnum::InsufficientMemoryAvailable
            )
            .add_info(vec!["memory", &requested_mem_mib.to_string()]));
        }

        let eif_file = self.eif_file.as_mut().ok_or_else(|| {
            new_nitro_cli_failure!(
                "Failed to get mutable reference to EIF file",
                NitroCliErrorEnum::FileOperationFailure
            )
        })?;

        let mut image_load_info = ImageLoadInfo {
            flags: NE_EIF_IMAGE,
            memory_offset: 0,
        };
        EnclaveHandle::do_ioctl(self.enc_fd, NE_GET_IMAGE_LOAD_INFO, &mut image_load_info)
            .map_err(|e| e.add_subaction("Get image load info ioctl failed".to_string()))?;

        debug!("Memory load information: {:?}", image_load_info);
        write_eif_to_regions(eif_file, regions, image_load_info.memory_offset as usize)
            .map_err(|e| e.add_subaction("Write EIF to enclave memory regions".to_string()))?;

        // Provide the regions to the driver for ownership change.
        for region in regions {
            let mut user_mem_region: UserMemoryRegion = region.into();
            EnclaveHandle::do_ioctl(self.enc_fd, NE_SET_USER_MEMORY_REGION, &mut user_mem_region)
                .map_err(|e| e.add_subaction("Set user memory region ioctl failed".to_string()))?;
        }

        info!("Finished initializing memory.");

        Ok(())
    }

    /// Initialize a single vCPU from a given ID.
    fn init_single_cpu(&mut self, mut cpu_id: u32) -> NitroCliResult<()> {
        EnclaveHandle::do_ioctl(self.enc_fd, NE_ADD_VCPU, &mut cpu_id)
            .map_err(|e| e.add_subaction("Add vCPU ioctl failed".to_string()))?;

        self.cpu_ids.push(cpu_id);
        debug!("Added CPU with ID {}.", cpu_id);

        Ok(())
    }

    /// Provide CPUs from the parent instance to the enclave.
    fn init_cpus(&mut self) -> NitroCliResult<()> {
        let cpu_config = self.cpu_config.clone();

        match cpu_config {
            EnclaveCpuConfig::List(cpu_ids) => {
                for cpu_id in cpu_ids {
                    self.init_single_cpu(cpu_id).map_err(|e| {
                        e.add_subaction(format!("Failed to add CPU with ID {}", cpu_id))
                    })?;
                }
            }
            EnclaveCpuConfig::Count(cpu_count) => {
                for _ in 0..cpu_count {
                    self.init_single_cpu(0)?;
                }
            }
        }

        Ok(())
    }

    /// Start an enclave after providing it with its necessary resources.
    fn start(&mut self, connection: Option<&Connection>) -> NitroCliResult<EnclaveStartInfo> {
        let mut start = EnclaveStartInfo {
            flags: self.flags,
            enclave_cid: self.enclave_cid.unwrap_or(0),
        };

        EnclaveHandle::do_ioctl(self.enc_fd, NE_START_ENCLAVE, &mut start)
            .map_err(|e| e.add_subaction("Start enclave ioctl failed".to_string()))?;

        safe_conn_eprintln(
            connection,
            format!(
                "Started enclave with enclave-cid: {}, memory: {} MiB, cpu-ids: {:?}",
                { start.enclave_cid },
                self.allocated_memory_mib,
                self.cpu_ids
            )
            .as_str(),
        )?;

        Ok(start)
    }

    /// Terminate an enclave.
    fn terminate_enclave(&mut self) -> NitroCliResult<()> {
        if self.enclave_cid.unwrap_or(0) != 0 {
            release_enclave_descriptor(self.enc_fd)
                .map_err(|e| e.add_subaction("Failed to release enclave descriptor".to_string()))?;

            // Release used memory.
            self.resource_allocator
                .free()
                .map_err(|e| e.add_subaction("Failed to release used memory".to_string()))?;
            info!("Enclave terminated.");

            // Mark enclave as terminated.
            self.clear();
        }

        Ok(())
    }

    /// Terminate an enclave and notify in case of errors.
    fn terminate_enclave_and_notify(&mut self) {
        // Attempt to terminate the enclave we are holding.
        if let Err(error_info) = self.terminate_enclave() {
            let mut err_msg = format!(
                "Terminating enclave '{:X}' failed with error: {:?}",
                self.slot_uid,
                construct_error_message(&error_info).as_str()
            );
            err_msg.push_str(
                "!!! The instance could be in an inconsistent state, please reboot it !!!",
            );

            // The error message should reach both the user and the logger.
            notify_error(&err_msg);
        }
    }

    /// Clear handle resources after terminating an enclave.
    fn clear(&mut self) {
        self.cpu_ids.clear();
        self.allocated_memory_mib = 0;
        self.enclave_cid = Some(0);
        self.enc_fd = -1;
        self.slot_uid = 0;
    }

    /// Terminate the enclave if `run-enclave` failed.
    fn terminate_enclave_error(&mut self, err: &str) {
        let err_msg = format!("{}. Terminating the enclave...", err);

        // Notify the user and the logger of the error, then terminate the enclave.
        notify_error(&err_msg);
        self.terminate_enclave_and_notify();
    }

    /// Wrapper over an `ioctl()` operation
    fn do_ioctl<T>(fd: RawFd, ioctl_code: u64, arg: &mut T) -> NitroCliResult<i32> {
        let rc = unsafe { libc::ioctl(fd, ioctl_code as _, arg) };
        if rc >= 0 {
            return Ok(rc);
        }

        let err_msg = match Error::last_os_error().raw_os_error().unwrap_or(0) as u32 {
            NE_ERR_VCPU_ALREADY_USED => "The provided vCPU is already used".to_string(),
            NE_ERR_VCPU_NOT_IN_CPU_POOL => {
                "The provided vCPU is not available in the CPU pool".to_string()
            }
            NE_ERR_VCPU_INVALID_CPU_CORE => {
                "The vCPU core ID is invalid for the CPU pool".to_string()
            }
            NE_ERR_INVALID_MEM_REGION_SIZE => {
                "The memory region's size is not a multiple of 2 MiB".to_string()
            }
            NE_ERR_INVALID_MEM_REGION_ADDR => "The memory region's address is invalid".to_string(),
            NE_ERR_UNALIGNED_MEM_REGION_ADDR => {
                "The memory region's address is not aligned".to_string()
            }
            NE_ERR_MEM_REGION_ALREADY_USED => "The memory region is already used".to_string(),
            NE_ERR_MEM_NOT_HUGE_PAGE => {
                "The memory region is not backed by contiguous physical huge page(s)".to_string()
            }
            NE_ERR_MEM_DIFFERENT_NUMA_NODE => {
                "The memory region's pages and the CPUs belong to different NUMA nodes".to_string()
            }
            NE_ERR_MEM_MAX_REGIONS => {
                "The maximum number of memory regions per enclave has been reached".to_string()
            }
            NE_ERR_NO_MEM_REGIONS_ADDED => {
                "The enclave cannot start because no memory regions have been added".to_string()
            }
            NE_ERR_NO_VCPUS_ADDED => {
                "The enclave cannot start because no vCPUs have been added".to_string()
            }
            NE_ERR_ENCLAVE_MEM_MIN_SIZE => {
                "The enclave's memory size is lower than the minimum supported".to_string()
            }
            NE_ERR_FULL_CORES_NOT_USED => {
                "The enclave cannot start because full CPU cores have not been set".to_string()
            }
            NE_ERR_NOT_IN_INIT_STATE => {
                "The enclave is in an incorrect state to set resources or start".to_string()
            }
            NE_ERR_INVALID_VCPU => {
                "The provided vCPU is out of range of the available CPUs".to_string()
            }
            NE_ERR_NO_CPUS_AVAIL_IN_POOL => {
                "The enclave cannot be created because no CPUs are available in the pool"
                    .to_string()
            }
            NE_ERR_INVALID_PAGE_SIZE => {
                "The memory region is not backed by page(s) multiple of 2 MiB".to_string()
            }
            NE_ERR_INVALID_FLAG_VALUE => {
                "The provided flags value in the ioctl arg data structure is invalid".to_string()
            }
            NE_ERR_INVALID_ENCLAVE_CID => {
                "The provided enclave CID is invalid, being a well-known CID or the parent VM CID"
                    .to_string()
            }
            e => format!("An error has occurred: {} (rc: {})", e, rc),
        };

        Err(new_nitro_cli_failure!(
            &err_msg,
            NitroCliErrorEnum::IoctlFailure
        ))
    }
}

impl Drop for EnclaveHandle {
    fn drop(&mut self) {
        // Check if we are (still) owning an enclave.
        if self.enclave_cid.unwrap_or(0) == 0 {
            debug!("Resource manager does not hold an enclave.");
            return;
        }

        // Terminate the enclave, notifying of any errors.
        self.terminate_enclave_and_notify();
    }
}

impl EnclaveManager {
    /// Create a new `EnclaveManager` instance.
    pub fn new(
        enclave_cid: Option<u64>,
        memory_mib: u64,
        cpu_ids: EnclaveCpuConfig,
        eif_file: File,
        debug_mode: bool,
        enclave_name: String,
    ) -> NitroCliResult<Self> {
        let enclave_handle =
            EnclaveHandle::new(enclave_cid, memory_mib, cpu_ids, eif_file, debug_mode)
                .map_err(|e| e.add_subaction("Failed to create enclave handle".to_string()))?;
        Ok(EnclaveManager {
            enclave_id: String::new(),
            enclave_name,
            enclave_handle: Arc::new(Mutex::new(enclave_handle)),
        })
    }

    /// Launch an enclave using the previously-set configuration.
    ///
    /// The enclave handle is locked throughout enclave creation. This is fine, since
    /// the socket for receiving commands is exposed only after creation has completed.
    pub fn run_enclave(&mut self, connection: Option<&Connection>) -> NitroCliResult<()> {
        self.enclave_id = self
            .enclave_handle
            .lock()
            .map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to acquire lock: {:?}", e),
                    NitroCliErrorEnum::LockAcquireFailure
                )
            })?
            .create_enclave(self.enclave_name.clone(), connection)
            .map_err(|e| e.add_subaction("Failed to create enclave".to_string()))?;
        Ok(())
    }

    /// Set measurements field inside EnclaveHandle
    pub fn set_measurements(
        &mut self,
        measurements: BTreeMap<String, String>,
    ) -> NitroCliResult<()> {
        self.enclave_handle
            .lock()
            .map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to acquire lock: {:?}", e),
                    NitroCliErrorEnum::LockAcquireFailure
                )
            })?
            .build_info = EnclaveBuildInfo::new(measurements);
        Ok(())
    }

    /// Set metadata field inside EnclaveHandle
    pub fn set_metadata(&mut self, metadata: EifIdentityInfo) -> NitroCliResult<()> {
        self.enclave_handle
            .lock()
            .map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to acquire lock: {:?}", e),
                    NitroCliErrorEnum::LockAcquireFailure
                )
            })?
            .metadata = Some(metadata);
        Ok(())
    }

    /// Get the resources needed for describing an enclave.
    ///
    /// The enclave handle is locked during this operation.
    pub fn get_description_resources(&self) -> NitroCliResult<UnpackedHandle> {
        let locked_handle = self.enclave_handle.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        Ok((
            locked_handle.slot_uid,
            locked_handle.enclave_cid.unwrap(),
            locked_handle.cpu_ids.len() as u64,
            locked_handle.cpu_ids.clone(),
            locked_handle.allocated_memory_mib,
            locked_handle.flags,
            locked_handle.state.clone(),
        ))
    }

    /// Get measurements from enclave handle
    pub fn get_measurements(&self) -> NitroCliResult<EnclaveBuildInfo> {
        let locked_handle = self.enclave_handle.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        Ok(locked_handle.build_info.clone())
    }

    /// Get metadata from enclave handle
    pub fn get_metadata(&self) -> NitroCliResult<Option<EifIdentityInfo>> {
        let locked_handle = self.enclave_handle.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        Ok(locked_handle.metadata.clone())
    }

    /// Get the resources (enclave CID) needed for connecting to the enclave console.
    ///
    /// The enclave handle is locked during this operation.
    pub fn get_console_resources_enclave_cid(&self) -> NitroCliResult<u64> {
        let locked_handle = self.enclave_handle.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        Ok(locked_handle.enclave_cid.unwrap())
    }

    /// Get the resources (enclave flags) needed for connecting to the enclave console.
    ///
    /// The enclave handle is locked during this operation.
    pub fn get_console_resources_enclave_flags(&self) -> NitroCliResult<u64> {
        let locked_handle = self.enclave_handle.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        Ok(locked_handle.flags)
    }

    /// Get the resources needed for enclave termination.
    ///
    /// The enclave handle is locked during this operation.
    fn get_termination_resources(&self) -> NitroCliResult<(RawFd, ResourceAllocator)> {
        let locked_handle = self.enclave_handle.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        Ok((
            locked_handle.enc_fd,
            locked_handle.resource_allocator.clone(),
        ))
    }

    /// Get the enclave descriptor.
    ///
    /// The enclave handle is locked during this operation.
    pub fn get_enclave_descriptor(&self) -> NitroCliResult<RawFd> {
        let locked_handle = self.enclave_handle.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        Ok(locked_handle.enc_fd)
    }

    /// Update the state the enclave is in.
    ///
    /// The enclave handle is locked during this operation.
    pub fn update_state(&mut self, state: EnclaveState) -> NitroCliResult<()> {
        let mut locked_handle = self.enclave_handle.lock().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to acquire lock: {:?}", e),
                NitroCliErrorEnum::LockAcquireFailure
            )
        })?;
        locked_handle.state = state;
        Ok(())
    }

    /// Terminate the owned enclave.
    ///
    /// The enclave handle is locked only when getting the resources needed for termination.
    /// This will allow the enclave process to execute other commands while termination
    /// is taking place.
    pub fn terminate_enclave(&mut self) -> NitroCliResult<()> {
        let (enc_fd, mut resource_allocator) = self.get_termination_resources().map_err(|e| {
            e.add_subaction("Enclave manager failed to get termination resources".to_string())
        })?;
        release_enclave_descriptor(enc_fd).map_err(|e| {
            e.add_subaction("Enclave manager failed to release enclave descriptor".to_string())
        })?;
        resource_allocator.free().map_err(|e| {
            e.add_subaction("Enclave manager failed to free enclave memory".to_string())
        })?;
        self.enclave_handle
            .lock()
            .map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to acquire lock: {:?}", e),
                    NitroCliErrorEnum::LockAcquireFailure
                )
            })?
            .clear();
        Ok(())
    }
}

/// Write an enclave image file to the specified list of memory regions.
fn write_eif_to_regions(
    eif_file: &mut File,
    regions: &[MemoryRegion],
    image_write_offset: usize,
) -> NitroCliResult<()> {
    let file_size = eif_file
        .metadata()
        .map_err(|_| {
            new_nitro_cli_failure!(
                "Failed to obtain EIF file metadata",
                NitroCliErrorEnum::FileOperationFailure
            )
        })?
        .len() as usize;

    eif_file.rewind().map_err(|_| {
        new_nitro_cli_failure!(
            "Failed to seek to the beginning of the EIF file",
            NitroCliErrorEnum::FileOperationFailure
        )
    })?;

    let mut total_written: usize = 0;

    for region in regions {
        let offset_plus_file_size = file_size.checked_add(image_write_offset).ok_or_else(|| {
            new_nitro_cli_failure!(
                "Memory overflow when trying to write EIF file",
                NitroCliErrorEnum::MemoryOverflow
            )
        })?;

        if total_written >= offset_plus_file_size {
            // All bytes have been written.
            break;
        }

        let written_plus_region_size = total_written
            .checked_add(region.mem_size as usize)
            .ok_or_else(|| {
                new_nitro_cli_failure!(
                    "Memory overflow when trying to write EIF file",
                    NitroCliErrorEnum::MemoryOverflow
                )
            })?;

        if written_plus_region_size <= image_write_offset {
            // All bytes need to be skipped to get to the image write offset.
        } else {
            let region_offset = image_write_offset.saturating_sub(total_written);
            let file_offset = total_written.saturating_sub(image_write_offset);
            let size = std::cmp::min(
                region.mem_size as usize - region_offset,
                file_size - file_offset,
            );
            region
                .fill_from_file(eif_file, region_offset, size)
                .map_err(|e| {
                    e.add_subaction("Failed to fill region with file content".to_string())
                })?;
        }
        total_written += region.mem_size as usize;
    }

    Ok(())
}

/// Release the enclave descriptor.
fn release_enclave_descriptor(enc_fd: RawFd) -> NitroCliResult<()> {
    // Close enclave descriptor.
    let rc = unsafe { libc::close(enc_fd) };
    if rc < 0 {
        return Err(new_nitro_cli_failure!(
            "Failed to close enclave descriptor",
            NitroCliErrorEnum::FileOperationFailure
        ));
    }

    Ok(())
}

/// Check if the `NITRO_BETWEEN_PACKETS_MILLIS` environment variable is set.
/// If it is, return a `Duration` representing its value.
pub fn between_packets_delay() -> Option<Duration> {
    if let Ok(value) = std::env::var("NITRO_BETWEEN_PACKETS_MILLIS") {
        if let Ok(value) = value.parse::<u64>() {
            return Some(Duration::from_millis(value));
        }
    }

    None
}

/// Helper function which contains heuristic for enclave build timeout calculation
///
/// # Arguments
///
/// * `eif_size` - The EIF size in bytes
///
/// # Examples
///
/// ```
/// use nitro_cli::enclave_proc::resource_manager::calculate_necessary_timeout;
/// use nitro_cli::enclave_proc::utils::GiB;
/// // Returns the timeout based on the 8GiB EIF size
/// let timeout = calculate_necessary_timeout(8 * GiB);
/// ```
pub fn calculate_necessary_timeout(eif_size: u64) -> c_int {
    // in case we have a valid eif_size give TIMEOUT_MINUTE_MS ms for each 6GiB
    let poll_timeout: c_int =
        ((1 + (eif_size - 1) / (6 * GiB)) as i32).saturating_mul(TIMEOUT_MINUTE_MS);

    poll_timeout
}

#[cfg(test)]
mod tests {
    use super::calculate_necessary_timeout;
    use crate::enclave_proc::utils::GiB;
    use eif_loader::TIMEOUT_MINUTE_MS;

    #[test]
    fn test_timeout_calculation() {
        assert_eq!(calculate_necessary_timeout(2 * GiB), TIMEOUT_MINUTE_MS);
        assert_eq!(calculate_necessary_timeout(6 * GiB), TIMEOUT_MINUTE_MS);
        assert_eq!(calculate_necessary_timeout(10 * GiB), 2 * TIMEOUT_MINUTE_MS);
    }
}


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

use inotify::{EventMask, Inotify, WatchMask};
use log::{debug, warn};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::thread::{self, JoinHandle};

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

/// The structure which manages the Unix socket that an enclave process listens on for commands.
#[derive(Default)]
pub struct EnclaveProcSock {
    /// The socket's file-system path.
    socket_path: PathBuf,
    /// The thread which listens for external events which delete the socket from the file-system.
    remove_listener_thread: Option<JoinHandle<()>>,
    /// A flag indicating if socket removal was requested.
    requested_remove: Arc<AtomicBool>,
}

/// The enclave process socket must allow cloning, since that is needed by the socket-listening thread.
impl Clone for EnclaveProcSock {
    fn clone(&self) -> Self {
        // Actually clone only what's relevant for the listening thread.
        EnclaveProcSock {
            socket_path: self.socket_path.clone(),
            remove_listener_thread: None,
            requested_remove: self.requested_remove.clone(),
        }
    }
}

impl Drop for EnclaveProcSock {
    fn drop(&mut self) {
        self.close_mut()
            .ok_or_exit_with_errno(Some("Failed to drop socket"));
    }
}

impl EnclaveProcSock {
    /// Create a new `EnclaveProcSock` instance from a given enclave ID.
    pub fn new(enclave_id: &str) -> NitroCliResult<Self> {
        let socket_path = get_socket_path(enclave_id).map_err(|_| {
            new_nitro_cli_failure!(
                "Failed to create enclave process socket",
                NitroCliErrorEnum::SocketPathNotFound
            )
        })?;

        Ok(EnclaveProcSock {
            socket_path,
            remove_listener_thread: None,
            requested_remove: Arc::new(AtomicBool::new(false)),
        })
    }

    /// Get the path to the managed Unix socket.
    pub fn get_path(&self) -> &Path {
        self.socket_path.as_path()
    }

    /// Set the path of the managed Unix socket.
    pub fn set_path(&mut self, socket_path: PathBuf) {
        self.socket_path = socket_path;
    }

    /// Start monitoring the Unix socket's state using `inotify`.
    pub fn start_monitoring(&mut self, exit_on_delete: bool) -> NitroCliResult<()> {
        let path_clone = self.socket_path.clone();
        let requested_remove_clone = self.requested_remove.clone();
        let mut socket_inotify = Inotify::init().map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to initialize socket notifications: {:?}", e),
                NitroCliErrorEnum::InotifyError
            )
        })?;

        // Relevant events to listen for are:
        // - IN_DELETE_SELF: triggered when the socket file inode gets removed.
        // - IN_ATTRIB: triggered when the reference count of the file inode changes.
        socket_inotify
            .add_watch(
                self.socket_path.as_path(),
                WatchMask::ATTRIB | WatchMask::DELETE_SELF,
            )
            .map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Failed to add watch to inotify: {:?}", e),
                    NitroCliErrorEnum::InotifyError
                )
            })?;
        self.remove_listener_thread = Some(thread::spawn(move || {
            socket_removal_listener(
                path_clone,
                requested_remove_clone,
                socket_inotify,
                exit_on_delete,
            )
        }));
        Ok(())
    }

    /// Remove the managed Unix socket and clean up after it. This is called with a mutable self-reference.
    fn close_mut(&mut self) -> NitroCliResult<()> {
        // Delete the socket from the disk. Also mark that this operation is intended, so that the
        // socket file monitoring thread doesn't exit forcefully when notifying the deletion.
        self.requested_remove.store(true, Ordering::SeqCst);
        if self.socket_path.exists() {
            std::fs::remove_file(&self.socket_path).map_err(|e| {
                new_nitro_cli_failure!(
                    &format!(
                        "Failed to remove socket file {:?} from disk: {:?}",
                        self.socket_path, e
                    ),
                    NitroCliErrorEnum::FileOperationFailure
                )
                .add_info(vec![
                    self.socket_path
                        .to_str()
                        .unwrap_or("Invalid unicode socket file name"),
                    "Remove",
                ])
            })?;
        }

        // Since the socket file has been deleted, we also wait for the event listener thread to finish.
        if self.remove_listener_thread.is_some() {
            self.remove_listener_thread
                .take()
                .unwrap()
                .join()
                .map_err(|e| {
                    new_nitro_cli_failure!(
                        &format!("Failed to join socket notification thread: {:?}", e),
                        NitroCliErrorEnum::ThreadJoinFailure
                    )
                })?;
        }

        Ok(())
    }

    /// Remove the managed Unix socket and clean up after it.
    pub fn close(mut self) -> NitroCliResult<()> {
        self.close_mut()
            .map_err(|e| e.add_subaction("Close socket".to_string()))
    }
}

/// Listen for an `inotify` event when the socket gets deleted from the disk.
fn socket_removal_listener(
    socket_path: PathBuf,
    requested_remove: Arc<AtomicBool>,
    mut socket_inotify: Inotify,
    exit_on_delete: bool,
) {
    let mut buffer = [0u8; 4096];
    let mut done = false;

    debug!("Socket file event listener started for {:?}.", socket_path);

    while !done {
        // Read events.
        let events = socket_inotify
            .read_events_blocking(&mut buffer)
            .map_err(|e| {
                new_nitro_cli_failure!(
                    &format!("Socket removal listener error: {:?}", e),
                    NitroCliErrorEnum::InotifyError
                )
                .set_action("Run Enclave".to_string())
            })
            .ok_or_exit_with_errno(Some("Failed to read inotify events"));

        for event in events {
            // We monitor the DELETE_SELF event, which occurs when the inode is no longer referenced by anybody. We
            // also monitor the IN_ATTRIB event, which gets triggered whenever the inode reference count changes. To
            // make sure this is a deletion, we also verify if the socket file is still present in the file-system.
            if (event.mask.contains(EventMask::ATTRIB)
                || event.mask.contains(EventMask::DELETE_SELF))
                && !socket_path.exists()
            {
                if requested_remove.load(Ordering::SeqCst) {
                    // At this point, the socket is shutting itself down and has notified the
                    // monitoring thread, so we just exit the loop gracefully.
                    debug!("The enclave process socket has deleted itself.");
                    done = true;
                } else {
                    // At this point, the socket has been deleted by an external action, so
                    // we exit forcefully, since there is no longer any way for a CLI instance
                    // to tell the current enclave process to terminate.
                    warn!("The enclave process socket has been deleted!");
                    if exit_on_delete {
                        std::process::exit(1);
                    }
                    done = true;
                }
            }
        }
    }

    debug!("Enclave process socket monitoring is done.");
}

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

    use std::os::unix::net::UnixListener;
    use std::process::Command;

    const DUMMY_ENCLAVE_ID: &str = "i-0000000000000000-enc0123456789012345";
    const THREADS_STR: &str = "Threads:";
    const WAIT_REMOVE_MILLIS: u64 = 10;

    /// 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 the initial values of the EnclaveProcSock attributes match the
    /// expected ones.
    #[test]
    fn test_enclaveprocsock_init() {
        let socket = EnclaveProcSock::new(&DUMMY_ENCLAVE_ID.to_string());

        assert!(socket.is_ok());

        if let Ok(socket) = socket {
            assert!(socket
                .socket_path
                .as_path()
                .to_str()
                .unwrap()
                .contains("0123456789012345"));
            assert!(socket.remove_listener_thread.is_none());
            assert!(!socket.requested_remove.load(Ordering::SeqCst));
        }
    }

    /// Tests that removing the socket file by means other than `close()` does not
    /// trigger a `socket.requested_remove` change.
    #[test]
    fn test_start_monitoring() {
        let socket = EnclaveProcSock::new(&DUMMY_ENCLAVE_ID.to_string());

        assert!(socket.is_ok());

        if let Ok(mut socket) = socket {
            UnixListener::bind(socket.get_path())
                .map_err(|e| {
                    new_nitro_cli_failure!(
                        &format!("Failed to bind to socket: {:?}", e),
                        NitroCliErrorEnum::SocketError
                    )
                })
                .ok_or_exit_with_errno(Some("Error binding"));
            let result = socket.start_monitoring(false);

            assert!(result.is_ok());

            // Remove socket file and expect `socket.requested_remove` to remain False
            std::fs::remove_file(&socket.socket_path).unwrap();
            std::thread::sleep(std::time::Duration::from_millis(WAIT_REMOVE_MILLIS));

            assert!(!socket.requested_remove.load(Ordering::SeqCst));
        }
    }

    /// Test that calling `close()` changes `socket.requested_remove` to True and
    /// that the listener thread joins.
    #[test]
    fn test_close() {
        let socket = EnclaveProcSock::new(&DUMMY_ENCLAVE_ID.to_string());

        assert!(socket.is_ok());

        // Get number of running threads before spawning the socket removal 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());

        if let Ok(mut socket) = socket {
            let _ = UnixListener::bind(socket.get_path())
                .map_err(|e| {
                    new_nitro_cli_failure!(
                        &format!("Failed to bind to socket: {:?}", e),
                        NitroCliErrorEnum::SocketError
                    )
                })
                .ok_or_exit_with_errno(Some("Error binding"));
            let result = socket.start_monitoring(true);

            assert!(result.is_ok());

            // Call `close_mut()` and expect `socket.requested_remove` to change to True
            let result = socket.close_mut();

            assert!(result.is_ok());
            assert!(socket.requested_remove.load(Ordering::SeqCst));
        }

        // Get number of running threads after closing the socket removal listener thread
        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());

        // Check that the number of threads remains the same before and after running the test
        assert_eq!(crt_num_threads0, crt_num_threads1);
    }
}


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

use std::fs::metadata;
use std::fs::File;
use std::io::Read;

use crate::common::json_output::{EnclaveDescribeInfo, EnclaveRunInfo, MetadataDescribeInfo};
use crate::common::{NitroCliErrorEnum, NitroCliFailure, NitroCliResult};
use crate::enclave_proc::resource_manager::EnclaveManager;
use crate::enclave_proc::resource_manager::NE_ENCLAVE_DEBUG_MODE;
use crate::new_nitro_cli_failure;

/// Kibibytes.
#[allow(non_upper_case_globals)]
pub const KiB: u64 = 1024;

/// Mebibytes.
#[allow(non_upper_case_globals)]
pub const MiB: u64 = 1024 * KiB;

/// Gibibytes.
#[allow(non_upper_case_globals)]
pub const GiB: u64 = 1024 * MiB;

/// Get a string representation of the bit-mask which holds the enclave launch flags.
pub fn flags_to_string(flags: u64) -> String {
    if flags & NE_ENCLAVE_DEBUG_MODE == NE_ENCLAVE_DEBUG_MODE {
        "DEBUG_MODE"
    } else {
        "NONE"
    }
    .to_string()
}

/// Obtain the enclave information requested by the `describe-enclaves` command.
pub fn get_enclave_describe_info(
    enclave_manager: &EnclaveManager,
    with_metadata: bool,
) -> NitroCliResult<EnclaveDescribeInfo> {
    let (slot_uid, enclave_cid, cpu_count, cpu_ids, memory_mib, flags, state) =
        enclave_manager.get_description_resources()?;
    let mut describe_meta: Option<MetadataDescribeInfo> = None;
    let mut img_name: Option<String> = None;
    let mut img_version: Option<String> = None;

    if with_metadata {
        if let Some(meta) = enclave_manager.get_metadata()? {
            img_name = Some(meta.img_name.clone());
            img_version = Some(meta.img_version.clone());
            describe_meta = Some(MetadataDescribeInfo::new(meta));
        }
    }

    let info = EnclaveDescribeInfo {
        enclave_name: Some(enclave_manager.enclave_name.clone()),
        enclave_id: generate_enclave_id(slot_uid)?,
        process_id: std::process::id(),
        enclave_cid,
        cpu_count,
        cpu_ids,
        memory_mib,
        state: state.to_string(),
        flags: flags_to_string(flags),
        build_info: Some(enclave_manager.get_measurements()?),
        img_name,
        img_version,
        metadata: describe_meta,
    };
    Ok(info)
}

/// Obtain the enclave information requested by the `run-enclaves` command.
pub fn get_run_enclaves_info(
    enclave_name: String,
    enclave_cid: u64,
    slot_id: u64,
    cpu_ids: Vec<u32>,
    memory: u64,
) -> NitroCliResult<EnclaveRunInfo> {
    let info = EnclaveRunInfo::new(
        enclave_name,
        generate_enclave_id(slot_id)?,
        enclave_cid,
        cpu_ids.len(),
        cpu_ids,
        memory,
    );
    Ok(info)
}

/// Generate a unique ID for a new enclave with the specified slot ID.
pub fn generate_enclave_id(slot_id: u64) -> NitroCliResult<String> {
    let file_path = "/sys/devices/virtual/dmi/id/board_asset_tag";
    if metadata(file_path).is_ok() {
        let mut file = File::open(file_path).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to open file: {:?}", e),
                NitroCliErrorEnum::FileOperationFailure
            )
            .add_info(vec![file_path, "Open"])
        })?;
        let mut contents = String::new();
        file.read_to_string(&mut contents).map_err(|e| {
            new_nitro_cli_failure!(
                &format!("Failed to read from file: {:?}", e),
                NitroCliErrorEnum::FileOperationFailure
            )
            .add_info(vec![file_path, "Read"])
        })?;
        contents.retain(|c| !c.is_whitespace());
        return Ok(format!("{}-enc{:x}", contents, slot_id));
    }
    Ok(format!("i-0000000000000000-enc{:x}", slot_id))
}

/// Obtain an enclave's slot ID from its full ID.
pub fn get_slot_id(enclave_id: String) -> Result<u64, String> {
    let tokens: Vec<&str> = enclave_id.split("-enc").collect();

    match tokens.get(1) {
        Some(slot_id) => {
            u64::from_str_radix(slot_id, 16).map_err(|_err| "Invalid enclave id format".to_string())
        }
        None => Err("Invalid enclave_id.".to_string()),
    }
}

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

    #[test]
    fn test_generate_enclave_id() {
        let slot_id: u64 = 7;
        let enc_id = generate_enclave_id(slot_id);
        let file_path = "/sys/devices/virtual/dmi/id/board_asset_tag";

        if metadata(file_path).is_err() {
            assert!(enc_id
                .unwrap()
                .eq(&format!("i-0000000000000000-enc{:?}", slot_id)));
        } else {
            assert!(!enc_id
                .unwrap()
                .split('-')
                .collect::<Vec<&str>>()
                .get(1)
                .unwrap()
                .eq(&"0000000000000000"));
        }
    }

    #[test]
    fn test_get_slot_id_valid() {
        let slot_id: u64 = 8;
        let enc_id = generate_enclave_id(slot_id);

        if let Ok(enc_id) = enc_id {
            let result = get_slot_id(enc_id);
            assert!(result.is_ok());
            assert_eq!(slot_id, result.unwrap());
        }
    }

    #[test]
    fn test_get_slot_id_invalid() {
        let enclave_id = String::from("i-0000_enc1234");
        let result = get_slot_id(enclave_id);

        assert!(result.is_err());
        if let Err(err_str) = result {
            assert!(err_str.eq("Invalid enclave_id."));
        }
    }

    /// Tests that `flags_to_string()` returns the correct String representation
    /// when the NE_ENCLAVE_DEBUG_MODE is either set or unset.
    #[test]
    fn test_flags_to_string() {
        let mut flags: u64 = 0;

        flags |= NE_ENCLAVE_DEBUG_MODE;
        let mut result = flags_to_string(flags);

        assert!(result.eq("DEBUG_MODE"));

        flags = 0;
        result = flags_to_string(flags);

        assert!(result.eq("NONE"));
    }

    /// Asserts that `get_run_enclaves_info()` returns a result containing
    /// exactly the same values as the supplied arguments.
    #[test]
    fn test_get_run_enclaves_info() {
        let enclave_name = "testName".to_string();
        let enclave_cid: u64 = 0;
        let slot_id: u64 = 7;
        let cpu_ids: Vec<u32> = vec![1, 3];
        let memory: u64 = 64;

        let result =
            get_run_enclaves_info(enclave_name, enclave_cid, slot_id, cpu_ids.clone(), memory);

        assert!(result.is_ok());

        if let Ok(result) = result {
            assert_eq!(enclave_cid, result.enclave_cid);
            assert_eq!(cpu_ids.len(), result.cpu_ids.len());
            for (idx, cpu_id) in result.cpu_ids.iter().enumerate() {
                assert_eq!(cpu_ids[idx], *cpu_id);
            }
            assert_eq!(memory, result.memory_mib);
        }
    }

    /// Asserts that `get_enclave_id()` returns the expected enclave
    /// id, which is obtained through a call to `get_run_enclaves_info()`.
    #[test]
    fn test_get_enclave_id() {
        let enclave_name = "testName".to_string();
        let enclave_cid: u64 = 0;
        let slot_id: u64 = 8;
        let cpu_ids: Vec<u32> = vec![1, 3];
        let memory: u64 = 64;

        let result =
            get_run_enclaves_info(enclave_name, enclave_cid, slot_id, cpu_ids.clone(), memory);

        assert!(result.is_ok());

        if let Ok(result) = result {
            let this_enclave_id = &result.enclave_id;
            assert!(this_enclave_id.eq(&result.enclave_id));

            assert_eq!(enclave_cid, result.enclave_cid);
            assert_eq!(cpu_ids.len(), result.cpu_ids.len());
            for (idx, cpu_id) in result.cpu_ids.iter().enumerate() {
                assert_eq!(cpu_ids[idx], *cpu_id);
            }
            assert_eq!(memory, result.memory_mib);
        }
    }
}