// SPDX-License-Identifier: Apache-2.0

mod error;
mod linux;
mod types;

pub use error::*;
pub use types::*;

use crate::device::Device;
use linux::*;
use std::os::fd::{AsRawFd, RawFd};
use vsock::VMADDR_CID_HOST;

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

/// Facilitates the execution of the nitro enclaves launch process.
#[derive(Default)]
pub struct Launcher {
    vm_fd: RawFd,
    slot_uid: u64,
    cpu_ids: Vec<u32>,
}

impl Launcher {
    /// Begin the nitro enclaves launch process by creating a new enclave VM.
    pub fn new(dev: &Device) -> Result<Self> {
        let mut slot_uid: u64 = 0;
        let vm_fd = unsafe { ne_create_vm(dev.as_raw_fd(), &mut slot_uid) }?;

        if vm_fd < 0 || slot_uid == 0 {
            return Err(LaunchError::ioctl_err_from_errno());
        }

        Ok(Self {
            vm_fd,
            slot_uid,
            cpu_ids: Vec::new(),
        })
    }

    /// Get the enclave's file descriptor.
    pub fn vm_fd(&self) -> RawFd {
        self.vm_fd
    }

    /// Get the enclave's slot UID.
    pub fn slot_uid(&self) -> u64 {
        self.slot_uid
    }

    /// Allocate enclave memory and populate it with the enclave image.
    pub fn set_memory(&mut self, mem: MemoryInfo) -> Result<()> {
        // Load the VM's enclave image type and fetch the offset in enclave memory of where to
        // start placing the enclave image.
        let mut load_info = ImageLoadInfo::from(&mem.image_type);

        // Get the image offset.
        unsafe { ne_get_image_load_info(self.vm_fd.as_raw_fd(), &mut load_info) }?;

        // Allocate the memory regions from the requested size.
        let mut regions = UserMemoryRegions::new(mem.size_mib).map_err(LaunchError::MemInit)?;

        // Populate the memory regions with the contents of the enclave image.
        regions
            .image_fill(load_info.memory_offset as usize, mem.image_type)
            .map_err(LaunchError::MemInit)?;

        // Add each memory region.
        for r in regions.inner_ref() {
            unsafe { ne_set_user_memory_region(self.vm_fd, r) }?;
        }

        Ok(())
    }

    /// Set a vCPU for an enclave. The vCPU can be auto-chosen from the NE CPU pool or it can be
    /// set by the caller.
    ///
    /// If set by the caller, the CPU needs to be available in the NE CPU pool.
    pub fn add_vcpu(&mut self, id: Option<u32>) -> Result<()> {
        let mut id = id.unwrap_or(0);

        unsafe { ne_add_vcpu(self.vm_fd, &mut id) }?;

        self.cpu_ids.push(id);

        Ok(())
    }

    /// Start running an enclave. Supply start flags and optional enclave CID. If successful, will
    /// return the actual enclave's CID (which may be different than the supplied CID).
    pub fn start(&self, flags: StartFlags, cid: Option<u32>) -> Result<u64> {
        let cid = match cid {
            Some(cid) => {
                // Ensure that the provided CID is valid.
                if cid <= VMADDR_CID_HOST || cid == u32::MAX {
                    return Err(LaunchError::CidInvalid);
                }

                cid
            }
            None => 0,
        };

        // Start the enclave VM.
        let mut start_info = StartInfo::new(flags, cid.into());

        unsafe { ne_start_enclave(self.vm_fd, &mut start_info) }?;

        Ok(start_info.cid)
    }
}
