//! Nitro-Enclave-specific material for Veracruz
//!
//! ## Authors
//!
//! The Veracruz Development Team.
//!
//! ## Licensing and copyright notice
//!
//! See the `LICENSE.markdown` file in the Veracruz root directory for
//! information on licensing and copyright.

use err_derive::Error;
use nix::{
    sys::socket::{
        accept, bind, listen, setsockopt, shutdown, socket,
        sockopt::{ReuseAddr, ReusePort},
        AddressFamily, Shutdown, SockAddr, SockFlag, SockType,
    },
    unistd::close,
};
use serde_json::Value;
use std::{os::unix::io::AsRawFd, process::Command, thread::JoinHandle};

/// Errors generated by Nitro enclave components of Veracruz
#[derive(Debug, Error)]
pub enum NitroError {
    /// An error occurred while serializing or deserializing
    #[error(display = "Nitro: Serde Error")]
    SerdeError,
    /// A JSON-specific error occured while serializing or deserializing
    #[error(display = "nitro: Serde JSON Error:{:?}", _0)]
    SerdeJsonError(#[error(source)] serde_json::error::Error),
    /// An error was generated from the nix crate
    #[error(display = "Nitrno: Nix Error:{:?}", _0)]
    NixError(#[error(source)] nix::Error),
    /// an IO error occcurred
    #[error(display = "Nitro: IO Error:{:?}", _0)]
    IOError(#[error(source)] std::io::Error),
    /// An error was generated by the AWS Nitro CLI
    #[error(display = "Nitro: CLI error")]
    CLIError,
    /// a Veracruz-specific socket error occurred
    #[error(display = "Nitro: Veracruz Socket Error:{:?}", _0)]
    VeracruzSocketError(#[error(source)] crate::io::error::SocketError),
    /// An error occured while processing UTF8 string data
    #[error(display = "Nitro: Utf8Error:{:?}", _0)]
    Utf8Error(#[error(source)] std::str::Utf8Error),
    /// an error occured while working with an EC2 instance
    #[error(display = "Nitro: EC2 Error")]
    EC2Error,
    /// A mutex error occurred
    #[error(display = "nitro: Mutex Error")]
    MutexError,
}

/// a struct for holding all of the information about a nitro enclave
pub struct NitroEnclave {
    /// The enclave ID, as generated from the Nitro CLI tool when the enclave
    /// is created - it's the EC2-instance ID appended with an enclave-specific
    /// value
    enclave_id: String,
    /// A convenience struct for handling VSOCK connections to the enclave
    vsocksocket: crate::io::vsocket::VsockSocket,
    /// The thread that will handle ocall operations from the enclave, if enabled
    /// (hence it's an `Option`
    ocall_thread: Option<JoinHandle<()>>,
    /// A communication channel to tell the ocall thread to terminate (if enabled)
    ocall_terminate_sender: Option<std::sync::Mutex<std::sync::mpsc::Sender<bool>>>,
    /// the path to the Nictro CLI function. Not all AMI images have it in the
    /// same place in the file system, so we need to keep track of it
    nitro_cli_path: String,
}

/// The port that is used to communicate with the enclave
const VERACRUZ_PORT: u32 = 5005;
/// A convenience value to be use when you want to listen to any CID for connections
const VMADDR_CID_ANY: u32 = 0xFFFFFFFF;
/// The port that the ocall thread will be waiting for socket connections on
const OCALL_PORT: u32 = 5006;
/// the maximum back log of events that will be maintained while listening to
/// a socket
const BACKLOG: usize = 128;

/// A type declaration for the function the ocall thread will call when it receives
/// data from the enclave
pub type OCallHandler = fn(Vec<u8>) -> Result<Vec<u8>, NitroError>;

impl NitroEnclave {
    /// create a new Nitro enclave, started with the file in eif_path
    /// If `ocall_handler` is not `None`, an ocall listening thread will be started
    /// and will call the function in `ocall_handler` when it recieves data from
    /// the enclave
    pub fn new(
        nitro_sbin: bool,
        eif_path: &str,
        debug: bool,
        ocall_handler: Option<OCallHandler>,
    ) -> Result<Self, NitroError> {
        let mut args = vec![
            "run-enclave",
            "--eif-path",
            eif_path,
            "--cpu-count",
            "2",
            "--memory",
            "256",
        ];
        if debug {
            args.push("--debug-mode=true");
        }
        let nitro_cli_path = {
            match nitro_sbin {
                true => "/usr/sbin/nitro-cli",
                false => "/usr/bin/nitro-cli",
            }
        };
        let stdout = loop {
            let enclave_result = Command::new(nitro_cli_path)
                .args(&args)
                .output()
                .map_err(|err| err);
            match enclave_result {
                Err(err) => {
                    println!("NitroEnclave::new failed to start enclave:{:?}", err);
                    println!("sleeping before trying again");
                    std::thread::sleep(std::time::Duration::from_millis(1000));
                    continue;
                }
                Ok(result) => {
                    if !result.status.success() {
                        let enclave_result_stderr = std::str::from_utf8(&result.stderr)?;
                        println!("NitroEnclave::new CLI error:{:?}", enclave_result_stderr);
                        println!("sleeping before trying again");
                        std::thread::sleep(std::time::Duration::from_millis(1000));
                        continue;
                    } else {
                        break result.stdout;
                    }
                }
            }
        };

        let enclave_result_stdout = std::str::from_utf8(&stdout)?;
        println!("enclave_result_stdout:{:?}", enclave_result_stdout);

        let enclave_data: Value = serde_json::from_str(enclave_result_stdout)?;
        let cid: u32 = if !enclave_data["EnclaveCID"].is_number() {
            return Err(NitroError::SerdeError);
        } else {
            serde_json::from_value(enclave_data["EnclaveCID"].clone()).unwrap()
        };

        let (ocall_thread_opt, sender) = match ocall_handler {
            None => (None, None), // Do nothing, we don't need to support ocalls
            Some(handler) => {
                let (tx, rx): (
                    std::sync::mpsc::Sender<bool>,
                    std::sync::mpsc::Receiver<bool>,
                ) = std::sync::mpsc::channel();
                let ocall_thread =
                    std::thread::spawn(move || NitroEnclave::ocall_loop(handler, rx));
                (Some(ocall_thread), Some(std::sync::Mutex::new(tx)))
            }
        };

        let enclave: Self = NitroEnclave {
            enclave_id: enclave_data["EnclaveID"]
                .to_string()
                .trim_matches('"')
                .to_string(),
            vsocksocket: crate::io::vsocket::VsockSocket::connect(cid, VERACRUZ_PORT)?,
            ocall_thread: ocall_thread_opt,
            ocall_terminate_sender: sender,
            nitro_cli_path: nitro_cli_path.to_string(),
        };
        return Ok(enclave);
    }

    /// The function started in it's own thread when an `ocall_handler` is specified
    /// This function cannot return, so `panic!`s are a necessary evil in here
    fn ocall_loop(handler: OCallHandler, terminate_rx: std::sync::mpsc::Receiver<bool>) {
        let socket_fd = socket(
            AddressFamily::Vsock,
            SockType::Stream,
            SockFlag::SOCK_NONBLOCK,
            None,
        )
        .expect("NitroEnclave::ocall_loop failed to create a socket");

        if let Err(err) = setsockopt(socket_fd, ReuseAddr, &true) {
            println!(
                "NitroEnclave::ocall_loop setsockopt failed for ReuseAddr({:?}). Terminating loop.",
                err
            );
            return;
        }
        if let Err(err) = setsockopt(socket_fd, ReusePort, &true) {
            println!(
                "NitroEnclave::ocall_loop setsockopt failed for ReusePort({:?}). Terminating loop.",
                err
            );
            return;
        }

        let sockaddr = SockAddr::new_vsock(VMADDR_CID_ANY, OCALL_PORT);

        let mut im_done: bool = false;

        while let Err(err) = bind(socket_fd, &sockaddr) {
            println!(
                "NitroEnclave::ocall_loop failed to bind to socket:{:?}",
                err
            );
            if err == nix::Error::Sys(nix::errno::Errno::EADDRINUSE) {
                // before we continue, check to see if we should terminate
                if let Ok(terminate) = terminate_rx.try_recv() {
                    if terminate {
                        im_done = true;
                        break;
                    }
                }
                println!("sleeping before trying again");
                std::thread::sleep(std::time::Duration::from_millis(1000));
            } else {
                panic!(
                    "NitroEnclave::ocall_loop received an unhandled error from bind:{:?}",
                    err
                );
            }
        }

        if !im_done {
            listen(socket_fd, BACKLOG)
                .map_err(|err| NitroError::NixError(err))
                .expect("NitroEnclave::ocall_loop listen failed");
            loop {
                match accept(socket_fd) {
                    Ok(fd) => {
                        let received_buffer = crate::io::raw_fd::receive_buffer(fd)
                            .expect("NitroEnclave::ocall_loop failed to receive buffer");
                        // call the handler
                        let return_buffer = handler(received_buffer)
                            .expect("NitroEnclave::ocall_loop handler failed");
                        crate::io::raw_fd::send_buffer(fd, &return_buffer)
                            .expect("NitroEnclave::ocall_loop failed to send buffer");
                    }
                    Err(err) => match err {
                        nix::Error::Sys(_) => {
                            if let Ok(terminate) = terminate_rx.try_recv() {
                                if terminate {
                                    break;
                                }
                            }
                        }
                        _ => println!("NitroEnclave::ocall_loop received error:{:?}", err),
                    },
                }
            }
        }

        // TODO: Use veracruz-utils::VsockSocket, which should handle shutdown, close itself in it's Drop function
        if let Err(err) = shutdown(socket_fd, Shutdown::Both) {
            match err {
                nix::Error::Sys(err) => {
                    match err {
                        nix::errno::Errno::ENOTCONN => (), // this is normal
                        _ =>  println!("NitroEnclave::ocall_loop failed to shutdown socket ({:?}. This might cause you problems in the future.", err),
                    }
                },
                _ => println!("NitroEnclave::ocall_loop failed to shutdown socket ({:?}. This might cause you problems in the future.", err),
            }
        }
        if let Err(err) = close(socket_fd) {
            println!("NitroEnclave::ocall_loop failed to close socket file handle({:?}). This might cause you problems in the future", err);
        }
    }

    /// send a buffer of data to the enclave
    pub fn send_buffer(&self, buffer: &Vec<u8>) -> Result<(), NitroError> {
        crate::io::raw_fd::send_buffer(self.vsocksocket.as_raw_fd(), buffer)
            .map_err(|err| NitroError::VeracruzSocketError(err))
    }

    /// receive a buffer of data from the enclave
    pub fn receive_buffer(&self) -> Result<Vec<u8>, NitroError> {
        crate::io::raw_fd::receive_buffer(self.vsocksocket.as_raw_fd())
            .map_err(|err| NitroError::VeracruzSocketError(err))
    }
}

impl Drop for NitroEnclave {
    /// Drop the enclave. In ideal conditions, this means that the enclave will
    /// be terminated.
    fn drop(&mut self) {
        // first, tell the ocall loop to terminate
        if let Some(tx_mutex) = &self.ocall_terminate_sender {
            let sender_guard = tx_mutex.lock().unwrap();
            if let Err(err) = sender_guard.send(true) {
                println!("NitroEnclave::drop failed to send terminate message to ocall_thread. You can't do anything about this, since we are in the drop method, but I thought you would want to know:{:?}", err);
            }
        }

        // second, wait for the ocall loop to terminate
        // This is referred to as the "Option dance" - https://users.rust-lang.org/t/spawn-threads-and-join-in-destructor/1613
        // we can only do the "take" because we are effectively a destructor
        if let Some(thread_handle) = self.ocall_thread.take() {
            if let Err(err) = thread_handle.join() {
                println!("NitroEnclave::drop failed to join to the ocall_thread. You can't do anything about this, since we are in the drop method, but I thought you would want to know:{:?}", err);
            }
        }

        // now, shutdown the enclave
        loop {
            let enclave_result = Command::new(&self.nitro_cli_path)
                .args(&["terminate-enclave", "--enclave-id", &self.enclave_id])
                .output();
            match enclave_result {
                Err(err) => {
                    println!("NitroEnclave::drop Command::new returned err:{:?}, sleeping and trying again", err);
                    std::thread::sleep(std::time::Duration::from_millis(1000));
                    continue;
                }
                Ok(result) => {
                    if !result.status.success() {
                        println!("NitroEnclave::drop failed to terminate the enclave (exit_status:{:?}. You will need to terminate it yourself.", result.status);
                        let result_stderr = std::str::from_utf8(&result.stderr).unwrap();
                        println!("NitroEnclave::drop CLI error:{:?}", result_stderr);
                    }
                    break;
                }
            }
        }
    }
}
