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

extern crate libc;

use std::io::{Error as IoError, Read, Result as IoResult, Write};
use std::mem::size_of;
use std::os::raw::c_uint;
use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
use std::os::unix::net::{UnixListener, UnixStream};
use std::time::Duration;

/// Generic stream trait - must support reading, writing, and setting a timeout
/// for each of those two.
pub trait Stream: Read + Write {
    fn set_read_timeout(&self, timeout: Option<Duration>) -> IoResult<()>;
    fn set_write_timeout(&self, timeout: Option<Duration>) -> IoResult<()>;
}

/// Generic listener trait, that can `accept()` connected streams of the associated type
/// `Self::Stream`.
pub trait Listener {
    /// The connected stream type, implementing `Read + Write`.
    type Stream: Stream;

    /// Accept a new connected stream.
    fn accept(&self) -> IoResult<Self::Stream>;
}

/// A vsock address spec.
#[derive(Clone, Copy, Debug)]
pub struct VsockAddr {
    /// vsock CID.
    pub cid: c_uint,
    /// vsock port.
    pub port: c_uint,
}

impl VsockAddr {
    /// Create a VsockAddr with the special VMADDR_CID_ANY CID and the provided port.
    /// This can be used to bind to the local vsock device for listening.
    pub fn any_cid_with_port(port: c_uint) -> Self {
        Self {
            cid: libc::VMADDR_CID_ANY,
            port,
        }
    }
}

/// An AF_VSOCK stream (client)
#[derive(Debug, Clone)]
pub struct VsockStream {
    fd: RawFd,
}

impl VsockStream {
    /// Create and return a VsockStream connected to `addr`.
    pub fn connect(addr: VsockAddr) -> IoResult<Self> {
        let fd = unsafe { libc::socket(libc::AF_VSOCK, libc::SOCK_STREAM, 0) };
        if fd < 0 {
            return Err(IoError::last_os_error());
        }

        let sa = libc::sockaddr_vm {
            svm_family: libc::AF_VSOCK as libc::sa_family_t,
            svm_cid: addr.cid,
            svm_port: addr.port,
            svm_reserved1: 0,
            svm_zero: [0u8; 4],
        };
        let rc = unsafe {
            libc::connect(
                fd,
                &sa as *const _ as *const libc::sockaddr,
                size_of::<libc::sockaddr_vm>() as u32,
            )
        };
        if rc < 0 {
            let err = IoError::last_os_error();
            unsafe { libc::close(fd) };
            return Err(err);
        }

        Ok(unsafe { VsockStream::from_raw_fd(fd) })
    }
}

impl FromRawFd for VsockStream {
    unsafe fn from_raw_fd(fd: RawFd) -> Self {
        Self { fd }
    }
}

impl AsRawFd for VsockStream {
    fn as_raw_fd(&self) -> RawFd {
        self.fd
    }
}

/// VsockStream Drop
impl Drop for VsockStream {
    fn drop(&mut self) {
        unsafe { libc::close(self.fd) };
    }
}

/// VsockStream Reader
impl std::io::Read for VsockStream {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        let rc = unsafe { libc::read(self.fd, buf.as_mut_ptr() as *mut libc::c_void, buf.len()) };
        if rc < 0 {
            Err(std::io::Error::last_os_error())
        } else {
            Ok(rc as usize)
        }
    }
}

/// VsockStream Writer
impl std::io::Write for VsockStream {
    fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
        let rc = unsafe { libc::write(self.fd, buf.as_ptr() as *mut libc::c_void, buf.len()) };
        if rc < 0 {
            Err(std::io::Error::last_os_error())
        } else {
            Ok(rc as usize)
        }
    }

    fn flush(&mut self) -> IoResult<()> {
        Ok(())
    }
}

impl Stream for VsockStream {
    fn set_read_timeout(&self, timeout: Option<Duration>) -> IoResult<()> {
        let tv = match timeout {
            Some(dur) => libc::timeval {
                tv_sec: dur.as_secs() as i64,
                tv_usec: dur.subsec_micros() as i64,
            },
            None => libc::timeval {
                tv_sec: 0,
                tv_usec: 0,
            },
        };
        let rc = unsafe {
            libc::setsockopt(
                self.as_raw_fd(),
                libc::SOL_SOCKET,
                libc::SO_RCVTIMEO,
                &tv as *const _ as *const libc::c_void,
                std::mem::size_of::<libc::timeval>() as u32,
            )
        };
        if rc != 0 {
            return Err(IoError::last_os_error());
        }
        Ok(())
    }

    fn set_write_timeout(&self, timeout: Option<Duration>) -> IoResult<()> {
        let tv = match timeout {
            Some(dur) => libc::timeval {
                tv_sec: dur.as_secs() as i64,
                tv_usec: dur.subsec_micros() as i64,
            },
            None => libc::timeval {
                tv_sec: 0,
                tv_usec: 0,
            },
        };
        let rc = unsafe {
            libc::setsockopt(
                self.as_raw_fd(),
                libc::SOL_SOCKET,
                libc::SO_SNDTIMEO,
                &tv as *const _ as *const libc::c_void,
                std::mem::size_of::<libc::timeval>() as u32,
            )
        };
        if rc != 0 {
            return Err(IoError::last_os_error());
        }
        Ok(())
    }
}

impl Stream for UnixStream {
    fn set_read_timeout(&self, timeout: Option<Duration>) -> IoResult<()> {
        UnixStream::set_read_timeout(self, timeout)
    }
    fn set_write_timeout(&self, timeout: Option<Duration>) -> IoResult<()> {
        UnixStream::set_write_timeout(self, timeout)
    }
}

/// An AF_VSOCK listener (server)
#[derive(Debug, Clone)]
pub struct VsockListener {
    fd: RawFd,
}

impl VsockListener {
    /// Create and return a VsockListener that is bound to `addr` and ready to accept
    /// client connections.
    pub fn bind(addr: VsockAddr, backlog: std::os::raw::c_int) -> IoResult<Self> {
        let fd = unsafe { libc::socket(libc::AF_VSOCK, libc::SOCK_STREAM, 0) };
        if fd < 0 {
            return Err(IoError::last_os_error());
        }

        let mut sa = libc::sockaddr_vm {
            svm_family: libc::AF_VSOCK as libc::sa_family_t,
            svm_cid: addr.cid,
            svm_port: addr.port,
            svm_reserved1: 0,
            svm_zero: [0u8; 4],
        };
        let mut rc = unsafe {
            libc::bind(
                fd,
                &mut sa as *mut _ as *mut libc::sockaddr,
                size_of::<libc::sockaddr_vm>() as u32,
            )
        };
        if rc < 0 {
            let err = IoError::last_os_error();
            unsafe { libc::close(fd) };
            return Err(err);
        }
        rc = unsafe { libc::listen(fd, backlog) };
        if rc < 0 {
            let err = IoError::last_os_error();
            unsafe { libc::close(fd) };
            return Err(err);
        }
        Ok(Self { fd })
    }
}

impl Listener for VsockListener {
    type Stream = VsockStream;

    /// Accept a client connection and return a client stream
    fn accept(&self) -> IoResult<VsockStream> {
        let mut addr: libc::sockaddr_vm = unsafe { std::mem::zeroed() };
        let mut addr_len = size_of::<libc::sockaddr_vm>() as libc::socklen_t;
        let cl_fd = unsafe {
            libc::accept(
                self.fd,
                &mut addr as *mut _ as *mut libc::sockaddr,
                &mut addr_len,
            )
        };
        if cl_fd < 0 {
            return Err(IoError::last_os_error());
        }

        Ok(unsafe { VsockStream::from_raw_fd(cl_fd) })
    }
}

/// VsockListener Drop
impl Drop for VsockListener {
    fn drop(&mut self) {
        unsafe { libc::close(self.fd) };
    }
}

impl Listener for UnixListener {
    type Stream = UnixStream;

    fn accept(&self) -> IoResult<UnixStream> {
        UnixListener::accept(self).map(|(s, _)| s)
    }
}


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

use std::io::{BufRead, BufReader, Read, Write};

use super::api::schema::{ApiRequest, ApiResponse};

#[derive(Debug)]
pub enum Error {
    BadUrl,
    IoError(std::io::Error),
    MsgLen,
    ParseError,
    SerdeError(serde_json::error::Error),
}

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

/// RPC transport trait. Implementors will have to provide the methods for send / receiving
/// requests and responses. The implementation will likely wrap an underlying stream
/// (i.e. Read + Write), and handle message / datagram traffic.
pub trait Transport {
    /// Receive an RPC request.
    fn recv_request(&mut self) -> Result<ApiRequest>;
    /// Send an RPC request.
    fn send_request(&mut self, req: &ApiRequest) -> Result<()>;
    /// Receive an RPC response.
    fn recv_response(&mut self) -> Result<ApiResponse>;
    /// Send an RPC response.
    fn send_response(&mut self, resp: &ApiResponse) -> Result<()>;
}

/// RPC transport implementation via a super-simple subset of HTTP.
///
/// Requests and responses are serialized as HTTP messages, with JSON bodies.
/// They always travel the underlying stream sequantially (i.e. once a request is sent,
/// the caller must always read its corresponding response before submitting a new request).
///
/// All API-related information is exchanged via the message (JSON) body. I.e. HTTP handling
/// stops at the transport layer, so methods, URLs, status codes, do not change across RPC
/// messages. I.e:
/// - the method is always POST;
/// - the URL is always `self.url`;
/// - the response status is always "200 OK".
///
///
/// Example:
///
/// ----- REQUEST -----
/// POST /rpc/v1.0 HTTP/1.1
/// Content-Type: application/json
/// Content-Length: 35
///
/// {"Hello": {"sender": "TestClient"}}
/// ----- REQUEST -----
///
/// ----- RESPONSE -----
/// 200 OK
/// Content-Type: application/json
/// Content-Length: 39
///
/// {"Hello": {"Ok": "Hello, TestClient!"}}
/// ----- RESPONSE -----
pub struct HttpTransport<S: Read + Write> {
    /// The (connected) data stream between caller and callee.
    stream: S,
    /// The HTTP URL that the RPC server responds to. This can be used to versionize the API
    /// (e.g. /rpc/v1).
    url: &'static str,
}

/// The HTTP headers our RPC transport is interested in.
struct HttpHeaders {
    content_length: usize,
}

impl<S: Read + Write> HttpTransport<S> {
    /// Maximum size (in bytes) of an HTTP message headers section.
    const MAX_HDR_LEN: usize = 1 * 1024;
    /// Maximum size (in bytes) of an HTTP message body.
    const MAX_BODY_LEN: usize = 64 * 1024;

    /// Create a new HTTP transport object from a connected stream.
    /// Args:
    ///   - stream: the connected `Read + Write` stream;
    ///   - url: the HTTP URL (can be used to provide API versioning);
    ///          Note: all requests use the same URL, since the actual API endpoint call is
    ///          part of the JSON body.
    pub fn new(stream: S, url: &'static str) -> Self {
        Self { stream, url }
    }

    /// Read / parse the HTTP headers (that we recognize) from the provided buffered reader.
    fn read_headers<R: BufRead>(reader: &mut R) -> Result<HttpHeaders> {
        let mut content_length = None;

        loop {
            let mut ln = String::new();
            reader.read_line(&mut ln).map_err(Error::IoError)?;
            let ln = ln.as_str().trim();

            if ln.len() == 0 {
                break;
            }

            let mut iter = ln.split(": ");
            match (iter.next(), iter.next()) {
                (Some("Content-Length"), Some(len)) => {
                    let len = len.parse::<usize>().map_err(|_| Error::ParseError)?;
                    content_length = Some(len);
                }
                (Some("Content-Type"), Some("application/json")) => (),
                _ => return Err(Error::ParseError),
            }
        }

        match content_length {
            Some(content_length) => {
                if content_length > Self::MAX_BODY_LEN {
                    return Err(Error::MsgLen);
                }
                Ok(HttpHeaders { content_length })
            }
            _ => Err(Error::ParseError),
        }
    }
}

impl<S: Read + Write> Transport for HttpTransport<S> {
    fn recv_request(&mut self) -> Result<ApiRequest> {
        // This is a nifty trick to get line-by-line input, while also limitting the amount
        // of data we can read from the wire:
        // Read::take(N) will return a Take struct (itself and implementor of Read), that will
        // be limited at reading at most N bytes. We can then pass this Take to BufReader, and
        // get a limited buffered reader.
        // However both Take and BufReader need to own the underlying Read implementor. To get
        // around ownership issues, and since our BufReader only needs to survive reading
        // exactly one message, we can rely on the fact that if R: Read, then &mut R: Read,
        // and use a Take(&mut self.stream).
        let mut reader = BufReader::new((&mut self.stream).take(Self::MAX_HDR_LEN as u64));
        let mut ln = String::new();
        reader.read_line(&mut ln).map_err(Error::IoError)?;
        let mut iter = ln.as_str().trim().split_whitespace();
        match (iter.next(), iter.next(), iter.next()) {
            (Some("POST"), Some(url), Some("HTTP/1.1")) => {
                if url != self.url {
                    return Err(Error::BadUrl);
                }
            }
            _ => return Err(Error::ParseError),
        }

        Self::read_headers(&mut reader).and_then(|h| {
            reader.get_mut().set_limit(h.content_length as u64);
            let mut buf = vec![0u8; h.content_length];
            reader
                .read_exact(buf.as_mut_slice())
                .map_err(Error::IoError)?;
            serde_json::from_slice(buf.as_slice()).map_err(Error::SerdeError)
        })
    }

    fn send_request(&mut self, request: &ApiRequest) -> Result<()> {
        let body = serde_json::to_vec(&request).map_err(Error::SerdeError)?;
        self.stream
            .write(
                format!(
                    "POST {} HTTP/1.1\r\n\
                    Content-Type: application/json\r\n\
                    Content-Length: {}\r\n\
                    \r\n",
                    self.url,
                    body.len()
                )
                .as_bytes(),
            )
            .map_err(Error::IoError)?;
        self.stream.write(body.as_slice()).map_err(Error::IoError)?;
        Ok(())
    }

    fn recv_response(&mut self) -> Result<ApiResponse> {
        let mut reader = BufReader::new((&mut self.stream).take(Self::MAX_HDR_LEN as u64));
        let mut ln = String::new();
        reader.read_line(&mut ln).map_err(Error::IoError)?;
        let mut iter = ln.as_str().trim().split_whitespace();
        match (iter.next(), iter.next()) {
            (Some("200"), Some("OK")) => (),
            _ => return Err(Error::ParseError),
        }
        Self::read_headers(&mut reader).and_then(|h| {
            reader.get_mut().set_limit(h.content_length as u64);
            let mut buf = vec![0u8; h.content_length];
            reader
                .read_exact(buf.as_mut_slice())
                .map_err(Error::IoError)?;
            serde_json::from_slice(buf.as_slice()).map_err(Error::SerdeError)
        })
    }

    fn send_response(&mut self, response: &ApiResponse) -> Result<()> {
        let body = serde_json::to_vec(&response).map_err(Error::SerdeError)?;
        self.stream
            .write(
                format!(
                    "200 OK\r\n\
                    Content-Type: application/json\r\n\
                    Content-Length: {}\r\n\
                    \r\n",
                    body.len()
                )
                .as_bytes(),
            )
            .map_err(Error::IoError)?;
        self.stream.write(body.as_slice()).map_err(Error::IoError)?;
        Ok(())
    }
}


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

fn main() {
    println!("cargo:rustc-link-lib=dylib=aws-c-common");
    println!("cargo:rustc-link-lib=dylib=aws-nitro-enclaves-sdk-c");
}


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

extern crate libc;
extern crate vtok_common;
extern crate vtok_rpc;

mod aws_ne;
mod worker;

use std::fmt;
use std::os::unix::net::UnixListener;
use std::time::Duration;

use vtok_common::{config, defs};
use vtok_rpc::api::schema;
use vtok_rpc::proto::Stream;
use vtok_rpc::HttpTransport;
use vtok_rpc::{Listener, VsockAddr, VsockListener};
use worker::Worker;

const USAGE: &str = r#"p11ne RPC server
    Usage:
        p11ne-srv vsock <port>
        p11ne-srv unix <path>
"#;

#[derive(Debug)]
enum Error {
    ConfigError(config::Error),
    IoError(std::io::Error),
    UsageError,
    WorkerError(worker::Error),
}

impl From<Error> for i32 {
    fn from(_other: Error) -> i32 {
        // NOTE: we could discriminate between errors here to provide a more specific
        // exit code.
        1
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // TODO: pretty-format error messages
        match self {
            Self::ConfigError(e) => write!(f, "{:?}", e),
            Self::IoError(e) => write!(f, "{:?}", e),
            Self::UsageError => write!(f, "{}", USAGE),
            Self::WorkerError(e) => write!(f, "{:?}", e),
        }
    }
}

fn handle_client<S: Stream>(stream: S) -> Result<(), Error> {
    // TODO: worker management (e.g. concurrency, jailing, etc).
    stream
        .set_read_timeout(Some(Duration::from_millis(defs::RPC_STREAM_TIMEOUT_MS)))
        .map_err(Error::IoError)?;
    stream
        .set_write_timeout(Some(Duration::from_millis(defs::RPC_STREAM_TIMEOUT_MS)))
        .map_err(Error::IoError)?;

    let xport = HttpTransport::new(stream, schema::API_URL);
    let mut worker = Worker::new(xport);
    worker.run().map_err(Error::WorkerError)
}

fn run_server<L: Listener>(listener: L) -> Result<(), Error> {
    config::Config::init_new().map_err(Error::ConfigError)?;
    println!("[vToken] Provisioning server is now running");

    loop {
        let client_result = listener
            .accept()
            .map_err(Error::IoError)
            .and_then(|stream| handle_client(stream));
        if let Err(e) = client_result {
            eprintln!("Error handling client connection: {:?}", e);
        }
    }
}

fn rusty_main() -> Result<(), Error> {
    let mut args = std::env::args();

    args.next();

    match (
        args.next().as_ref().map(|s| s.as_str()),
        args.next().as_ref().map(|s| s.as_str()),
    ) {
        (Some("vsock"), Some(port)) => {
            let port = port
                .parse::<std::os::raw::c_uint>()
                .map_err(|_| Error::UsageError)?;
            VsockListener::bind(VsockAddr::any_cid_with_port(port), 5)
                .map_err(Error::IoError)
                .and_then(|l| run_server(l))?;
        }
        (Some("unix"), Some(path)) => {
            UnixListener::bind(path)
                .map_err(Error::IoError)
                .and_then(|l| run_server(l))?;
        }
        (_, _) => return Err(Error::UsageError),
    };
    Ok(())
}

fn main() {
    match rusty_main() {
        Ok(()) => std::process::exit(0),
        Err(e) => {
            eprintln!("[vToken] {}", e);
            std::process::exit(i32::from(e))
        }
    }
}


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

use vtok_common::{config, defs, util};
use vtok_rpc::api::schema;
use vtok_rpc::api::schema::{ApiError, ApiOk, ApiRequest, ApiResponse};
use vtok_rpc::{Transport, TransportError};

use super::aws_ne;
use base64;

#[derive(Debug)]
pub enum Error {
    TransportError(TransportError),
}

pub struct Worker<T: Transport> {
    transport: T,
}

impl<T> Worker<T>
where
    T: Transport,
{
    pub fn new(transport: T) -> Self {
        Self { transport }
    }

    pub fn run(&mut self) -> Result<(), Error> {
        let request = self
            .transport
            .recv_request()
            .map_err(Error::TransportError)?;

        let response = request
            .validate_args()
            .map_err(ApiError::InvalidArgs)
            .and_then(|_| match request {
                ApiRequest::AddToken { token } => Self::add_token(token),
                ApiRequest::DescribeDevice => Self::describe_device(),
                ApiRequest::DescribeToken { label, pin } => Self::describe_token(label, pin),
                ApiRequest::RefreshToken {
                    label,
                    pin,
                    envelope_key,
                } => Self::refresh_token(label, pin, envelope_key),
                ApiRequest::RemoveToken { label, pin } => Self::remove_token(label, pin),
                ApiRequest::UpdateToken { label, pin, token } => {
                    Self::update_token(label, pin, token)
                }
            });

        self.transport
            .send_response(&response)
            .map_err(Error::TransportError)
    }

    fn add_token(token: schema::Token) -> ApiResponse {
        let mut config = config::Config::load_rw().map_err(|_| ApiError::InternalError)?;

        // Check if the token label is already in use by another token.
        let dup = config
            .slots()
            .iter()
            .filter(|s| match s {
                None => false,
                Some(tok) => tok.label == token.label,
            })
            .count();
        if dup > 0 {
            return Err(ApiError::TokenLabelInUse);
        }

        let private_keys = Self::decrypt_token_keys(token.keys, token.envelope_key)?;

        // Find the first free slot.
        let free_slot = config
            .slots_mut()
            .iter_mut()
            .find(|s| s.is_none())
            .ok_or(ApiError::TooManyTokens)?;

        free_slot.replace(config::Token {
            label: token.label,
            pin: token.pin,
            private_keys,
            expiry_ts: util::time::monotonic_secs() + defs::TOKEN_EXPIRY_SECS,
        });

        config.save().map_err(|_| ApiError::InternalError)?;

        Ok(ApiOk::None)
    }

    fn describe_device() -> ApiResponse {
        let config = config::Config::load_ro().map_err(|_| ApiError::InternalError)?;
        let tokens: Vec<schema::TokenDescription> = config
            .slots()
            .iter()
            .enumerate()
            .filter_map(|(slot_id, slot)| {
                slot.as_ref().and_then(|tok| {
                    Some(schema::TokenDescription {
                        label: tok.label.clone(),
                        slot_id,
                        ttl_secs: tok.expiry_ts.checked_sub(util::time::monotonic_secs())?,
                        keys: None,
                    })
                })
            })
            .collect();
        let free_slot_count = defs::DEVICE_MAX_SLOTS - tokens.len();

        Ok(ApiOk::DeviceDescription(schema::DeviceDescription {
            free_slot_count,
            tokens,
        }))
    }

    fn describe_token(label: String, pin: String) -> ApiResponse {
        let config = config::Config::load_ro().map_err(|_| ApiError::InternalError)?;
        let (slot_id, token) = config
            .slots()
            .iter()
            .enumerate()
            .filter_map(|(idx, slot)| {
                slot.as_ref()
                    .filter(|tok| tok.label == label)
                    .map(|tok| (idx, tok))
            })
            .next()
            .ok_or(ApiError::TokenNotFound)?;
        if token.pin != pin {
            return Err(ApiError::AccessDenied);
        }

        let token_desc = schema::TokenDescription {
            label: token.label.clone(),
            slot_id,
            ttl_secs: token
                .expiry_ts
                .checked_sub(util::time::monotonic_secs())
                .unwrap_or(0),
            keys: Some(
                token
                    .private_keys
                    .iter()
                    .map(|key| schema::PrivateKeyDescription {
                        label: key.label.clone(),
                        id: key.id,
                        uri: format!(
                            "pkcs11:model={};manufacturer={};serial=EVT{:02X};token={};id=%{:02x};object={};type=private",
                            defs::TOKEN_MODEL,
                            defs::MANUFACTURER,
                            slot_id,
                            token.label.as_str(),
                            key.id,
                            key.label.as_str(),
                        )
                    })
                    .collect(),
            ),
        };

        Ok(ApiOk::TokenDescription(token_desc))
    }

    fn refresh_token(label: String, pin: String, envelope_key: schema::EnvelopeKey) -> ApiResponse {
        let mut config = config::Config::load_rw().map_err(|_| ApiError::InternalError)?;
        let slot = config
            .slots_mut()
            .iter_mut()
            .find(|s| match s {
                None => false,
                Some(tok) => tok.label == label,
            })
            .ok_or(ApiError::TokenNotFound)?;

        // It's safe to unwrap here, since the above find() ensures slot != None
        let mut token = slot.as_mut().unwrap();
        if token.pin != pin {
            return Err(ApiError::AccessDenied);
        }

        // Since all the keys in one token are encrypted via the same envelope, we only need to
        // successfully decrypt one key, in order to achieve attestation.
        token
            .private_keys
            .iter()
            .next()
            .ok_or(ApiError::EmptyToken)
            .and_then(|key| {
                match envelope_key {
                    schema::EnvelopeKey::Kms {
                        ref region,
                        ref access_key_id,
                        ref secret_access_key,
                        ref session_token,
                    } => {
                        aws_ne::kms_decrypt(
                            region.as_bytes(),
                            access_key_id.as_bytes(),
                            secret_access_key.as_bytes(),
                            session_token.as_bytes(),
                            &base64::decode(key.encrypted_pem_b64.as_str())
                                .map_err(|_| ApiError::TokenKeyDecodingFailed)?,
                        )
                        .map_err(|_| ApiError::KmsDecryptFailed)
                        .and_then(|v| {
                            String::from_utf8(v).map_err(|_| ApiError::TokenRefreshFailed)
                        })
                        .ok()
                        .filter(|pem| pem == &key.pem)
                        .ok_or(ApiError::TokenRefreshFailed)?;
                    }
                };
                Ok(())
            })?;
        token.expiry_ts = util::time::monotonic_secs() + defs::TOKEN_EXPIRY_SECS;

        config.save().map_err(|_| ApiError::InternalError)?;

        Ok(ApiOk::None)
    }

    fn remove_token(label: String, pin: String) -> ApiResponse {
        let mut config = config::Config::load_rw().map_err(|_| ApiError::InternalError)?;
        let slot = config
            .slots_mut()
            .iter_mut()
            .find(|s| match s {
                None => false,
                Some(tok) => tok.label == label,
            })
            .ok_or(ApiError::TokenNotFound)?;
        if let Some(tok) = slot {
            if tok.pin != pin {
                return Err(ApiError::AccessDenied);
            }
        }
        slot.take();

        config.save().map_err(|_| ApiError::InternalError)?;

        Ok(ApiOk::None)
    }

    fn update_token(label: String, pin: String, token: schema::Token) -> ApiResponse {
        let mut config = config::Config::load_rw().map_err(|_| ApiError::InternalError)?;

        // If this update is changing the token label, we need to make sure that the
        // new label is unique.
        if label != token.label {
            let dup = config
                .slots()
                .iter()
                .find(|s| s.as_ref().filter(|t| t.label == token.label).is_some())
                .is_some();
            if dup {
                return Err(ApiError::TokenLabelInUse);
            }
        }

        config
            .slots_mut()
            .iter_mut()
            .find(|s| match s {
                None => false,
                Some(tok) => tok.label == label,
            })
            .ok_or(ApiError::TokenNotFound)
            .and_then(|slot| {
                // It's safe to unwrap here, since the above find() ensures slot != None
                if slot.as_ref().unwrap().pin != pin {
                    return Err(ApiError::AccessDenied);
                }
                slot.replace(config::Token {
                    label,
                    pin,
                    expiry_ts: util::time::monotonic_secs() + defs::TOKEN_EXPIRY_SECS,
                    private_keys: Self::decrypt_token_keys(token.keys, token.envelope_key)?,
                });
                Ok(())
            })?;

        config.save().map_err(|_| ApiError::InternalError)?;

        Ok(ApiOk::None)
    }

    fn decrypt_token_keys(
        encrypted_keys: Vec<schema::PrivateKey>,
        envelope_key: schema::EnvelopeKey,
    ) -> Result<Vec<config::PrivateKey>, schema::ApiError> {
        let mut private_keys = Vec::new();
        for key in encrypted_keys {
            private_keys.push(config::PrivateKey {
                pem: match envelope_key {
                    schema::EnvelopeKey::Kms {
                        ref region,
                        ref access_key_id,
                        ref secret_access_key,
                        ref session_token,
                    } => aws_ne::kms_decrypt(
                        region.as_bytes(),
                        access_key_id.as_bytes(),
                        secret_access_key.as_bytes(),
                        session_token.as_bytes(),
                        &base64::decode(key.encrypted_pem_b64.as_str())
                            .map_err(|_| ApiError::TokenKeyDecodingFailed)?,
                    )
                    .map_err(|_| ApiError::KmsDecryptFailed)
                    .and_then(|v| {
                        String::from_utf8(v).map_err(|_| ApiError::TokenProvisioningFailed)
                    })?,
                },
                encrypted_pem_b64: key.encrypted_pem_b64,
                id: key.id,
                label: key.label,
                cert_pem: key.cert_pem,
            })
        }
        Ok(private_keys)
    }
}


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

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

pub use std::os::raw::{c_int, c_ulong, c_void};
pub type c_size_t = usize;

/// Must be kept in sync with the endpoint proxy
pub const AWS_NE_VSOCK_PROXY_ADDR: [u8; 2] = [0x33u8, 0x00];
pub const AWS_NE_VSOCK_PROXY_PORT: u16 = 8000;
pub const AWS_SOCKET_VSOCK_DOMAIN: c_int = 3;
pub const AWS_ADDRESS_MAX_LEN: usize = 108;
pub const KMS_MAX_DECRYPT_LEN: usize = 4096;

#[repr(C)]
pub struct aws_byte_buf {
    pub len: c_size_t,
    pub buffer: *mut u8,
    capacity: c_size_t,
    allocator: *mut aws_allocator,
}
#[repr(C)]
pub struct aws_string {
    _ph: [u8; 0],
}
#[repr(C)]
pub struct aws_allocator {
    _ph: [u8; 0],
}
#[repr(C)]
pub struct aws_nitro_enclaves_kms_client_configuration {
    _ph: [u8; 0],
}
#[repr(C)]
pub struct aws_nitro_enclaves_kms_client {
    _ph: [u8; 0],
}

#[repr(C)]
pub struct aws_socket_endpoint {
    pub address: [u8; AWS_ADDRESS_MAX_LEN],
    pub port: u16,
}

/// Imported SDK APIs
extern "C" {
    /// Aws-C-Common
    pub fn aws_byte_buf_from_array(bytes: *mut c_void, len: c_size_t) -> aws_byte_buf;
    pub fn aws_string_new_from_array(
        allocator: *mut aws_allocator,
        bytes: *const u8,
        len: c_size_t,
    ) -> *mut aws_string;
    pub fn aws_string_destroy_secure(string: *mut aws_string);
    pub fn aws_byte_buf_clean_up_secure(buf: *mut aws_byte_buf);

    /// AWS Nitro Enclaves SDK
    pub fn aws_nitro_enclaves_library_init(allocator: *mut aws_allocator);
    pub fn aws_nitro_enclaves_library_clean_up();
    pub fn aws_nitro_enclaves_get_allocator() -> *mut aws_allocator;

    pub fn aws_nitro_enclaves_kms_client_config_default(
        region: *mut aws_string,
        endpoint: *mut aws_socket_endpoint,
        domain: c_int,
        access_key_id: *mut aws_string,
        secret_access_key: *mut aws_string,
        session_token: *mut aws_string,
    ) -> *mut aws_nitro_enclaves_kms_client_configuration;

    pub fn aws_nitro_enclaves_kms_client_config_destroy(
        config: *mut aws_nitro_enclaves_kms_client_configuration,
    );

    pub fn aws_nitro_enclaves_kms_client_new(
        config: *mut aws_nitro_enclaves_kms_client_configuration,
    ) -> *mut aws_nitro_enclaves_kms_client;

    pub fn aws_nitro_enclaves_kms_client_destroy(client: *mut aws_nitro_enclaves_kms_client);

    pub fn aws_kms_decrypt_blocking(
        client: *mut aws_nitro_enclaves_kms_client,
        key_id: *mut aws_string,
        encryption_algorithm: *mut aws_string,
        ciphertext: *const aws_byte_buf,
        plaintext: *mut aws_byte_buf,
    ) -> c_int;
}


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

mod ffi;

pub enum Error {
    SdkInitError,
    SdkGenericError,
    SdkKmsConfigError,
    SdkKmsClientError,
    SdkKmsDecryptError,
}

/// KMS decrypt FFI wrapper
/// TODO: Add a trait for Drop wrappers for SDK resources
pub fn kms_decrypt(
    aws_region: &[u8],
    aws_key_id: &[u8],
    aws_secret_key: &[u8],
    aws_session_token: &[u8],
    ciphertext: &[u8],
) -> Result<Vec<u8>, Error> {
    // Initialize the SDK
    unsafe {
        ffi::aws_nitro_enclaves_library_init(std::ptr::null_mut());
    };

    // Fetch allocator
    let allocator = unsafe { ffi::aws_nitro_enclaves_get_allocator() };
    if allocator.is_null() {
        unsafe {
            ffi::aws_nitro_enclaves_library_clean_up();
        }
        return Err(Error::SdkInitError);
    }
    // REGION
    let region = unsafe {
        let reg = ffi::aws_string_new_from_array(allocator, aws_region.as_ptr(), aws_region.len());
        if reg.is_null() {
            ffi::aws_nitro_enclaves_library_clean_up();
            return Err(Error::SdkGenericError);
        }
        reg
    };
    // ENDPOINT
    let mut endpoint = {
        let mut ep = ffi::aws_socket_endpoint {
            address: [0; ffi::AWS_ADDRESS_MAX_LEN],
            port: ffi::AWS_NE_VSOCK_PROXY_PORT,
        };
        ep.address[..ffi::AWS_NE_VSOCK_PROXY_ADDR.len()]
            .copy_from_slice(&ffi::AWS_NE_VSOCK_PROXY_ADDR);
        ep
    };
    // AWS_ACCESS_KEY_ID
    let key_id = unsafe {
        let kid = ffi::aws_string_new_from_array(allocator, aws_key_id.as_ptr(), aws_key_id.len());
        if kid.is_null() {
            ffi::aws_string_destroy_secure(region);
            ffi::aws_nitro_enclaves_library_clean_up();
            return Err(Error::SdkGenericError);
        }
        kid
    };
    // AWS_SECRET_ACCESS_KEY
    let secret_key = unsafe {
        let skey = ffi::aws_string_new_from_array(
            allocator,
            aws_secret_key.as_ptr(),
            aws_secret_key.len(),
        );
        if skey.is_null() {
            ffi::aws_string_destroy_secure(key_id);
            ffi::aws_string_destroy_secure(region);
            ffi::aws_nitro_enclaves_library_clean_up();
            return Err(Error::SdkGenericError);
        }
        skey
    };
    // AWS_SESSION_TOKEN
    let session_token = unsafe {
        let sess_token = ffi::aws_string_new_from_array(
            allocator,
            aws_session_token.as_ptr(),
            aws_session_token.len(),
        );
        if sess_token.is_null() {
            ffi::aws_string_destroy_secure(secret_key);
            ffi::aws_string_destroy_secure(key_id);
            ffi::aws_string_destroy_secure(region);
            ffi::aws_nitro_enclaves_library_clean_up();
            return Err(Error::SdkGenericError);
        }
        sess_token
    };
    // Construct KMS client configuration
    let kms_client_cfg = unsafe {
        // Configure
        let cfg = ffi::aws_nitro_enclaves_kms_client_config_default(
            region,
            &mut endpoint,
            ffi::AWS_SOCKET_VSOCK_DOMAIN,
            key_id,
            secret_key,
            session_token,
        );

        if cfg.is_null() {
            ffi::aws_string_destroy_secure(key_id);
            ffi::aws_string_destroy_secure(secret_key);
            ffi::aws_string_destroy_secure(session_token);
            ffi::aws_string_destroy_secure(region);
            ffi::aws_nitro_enclaves_library_clean_up();
            return Err(Error::SdkKmsConfigError);
        }
        cfg
    };
    // Construct KMS Client
    let kms_client = unsafe { ffi::aws_nitro_enclaves_kms_client_new(kms_client_cfg) };
    if kms_client.is_null() {
        unsafe {
            ffi::aws_string_destroy_secure(key_id);
            ffi::aws_string_destroy_secure(secret_key);
            ffi::aws_string_destroy_secure(session_token);
            ffi::aws_string_destroy_secure(region);
            ffi::aws_nitro_enclaves_kms_client_config_destroy(kms_client_cfg);
            ffi::aws_nitro_enclaves_library_clean_up();
        }
        return Err(Error::SdkKmsClientError);
    }
    // Ciphertext
    let ciphertext_buf = unsafe {
        ffi::aws_byte_buf_from_array(ciphertext.as_ptr() as *mut ffi::c_void, ciphertext.len())
    };

    // Decrypt
    let mut plaintext_buf: ffi::aws_byte_buf = unsafe { std::mem::zeroed() };
    let rc = unsafe {
        ffi::aws_kms_decrypt_blocking(
            kms_client,
            std::ptr::null_mut(),
            std::ptr::null_mut(),
            &ciphertext_buf,
            &mut plaintext_buf,
        )
    };
    if rc != 0 {
        unsafe {
            ffi::aws_string_destroy_secure(key_id);
            ffi::aws_string_destroy_secure(secret_key);
            ffi::aws_string_destroy_secure(session_token);
            ffi::aws_string_destroy_secure(region);
            ffi::aws_nitro_enclaves_kms_client_config_destroy(kms_client_cfg);
            ffi::aws_nitro_enclaves_kms_client_destroy(kms_client);
            ffi::aws_nitro_enclaves_library_clean_up();
        }
        return Err(Error::SdkKmsDecryptError);
    }

    // Cleanup
    unsafe {
        ffi::aws_string_destroy_secure(key_id);
        ffi::aws_string_destroy_secure(secret_key);
        ffi::aws_string_destroy_secure(session_token);
        ffi::aws_string_destroy_secure(region);
        ffi::aws_nitro_enclaves_kms_client_config_destroy(kms_client_cfg);
        ffi::aws_nitro_enclaves_kms_client_destroy(kms_client);
        ffi::aws_nitro_enclaves_library_clean_up();
    }

    // Plaintext
    let plaintext = unsafe {
        std::slice::from_raw_parts(plaintext_buf.buffer, plaintext_buf.len as usize).to_vec()
    };
    unsafe { ffi::aws_byte_buf_clean_up_secure(&mut plaintext_buf) };

    Ok(plaintext)
}


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

extern crate serde;
extern crate serde_json;
extern crate vtok_rpc;

use std::collections::HashMap;
use std::fmt;
use std::os::unix::net::UnixStream;

use vtok_rpc::api::schema;
use vtok_rpc::{HttpTransport, Transport};
use vtok_rpc::{VsockAddr, VsockStream};

const USAGE: &str = r#"p11ne client
    Usage:
        p11ne-client <command> [<global options>] [<command options>]

    Global options:
        --server <address>
            [REQUIRED]
            Address used to connect to the p11ne server, using the format
            [vsock:CID:PORT | unix:PATH].
            vsock example (say CID=3 and port=5252): --server vsock:3:5252
            Unix sockets example: --server unix:/tmp/p11ne-rpc.sock

        --help
            Show this usage message.

    Commands:

        raw-rpc
            Perform a raw remote procedure call. This provides a low-level interface
            to the p11ne RPC server.
            A raw (json-serialized) API request is read from STDIN, sent to the p11ne
            RPC server, and the response is written to STDOUT, also in its raw,
            json-serialized form.

        help
            Show this usage message.
"#;

enum Error {
    IoError(std::io::Error),
    TransportError(vtok_rpc::transport::Error),
    UsageError(String),
    SerdeError(serde_json::Error),
}

enum ServerAddr {
    Vsock(VsockAddr),
    Unix(String),
}

enum CliOption {
    Server(ServerAddr),
}

impl From<Error> for i32 {
    fn from(_other: Error) -> i32 {
        // NOTE: we could discriminate between errors here to provide a more specific
        // exit code.
        1
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // TODO: pretty-format error messages
        match self {
            Self::IoError(e) => write!(f, "{:?}", e),
            Self::SerdeError(_) => write!(f, "internal error"),
            Self::UsageError(s) => write!(f, "Error: {}.\nUse --help for help.", s),
            Self::TransportError(e) => write!(f, "{:?}", e),
        }
    }
}

fn parse_server_addr(addr_str: &str) -> Result<ServerAddr, Error> {
    let mut iter = addr_str.split(":");
    match (iter.next(), iter.next(), iter.next()) {
        (Some("unix"), Some(path), None) => Ok(ServerAddr::Unix(path.to_string())),
        (Some("vsock"), Some(cid_str), Some(port_str)) => {
            match (
                cid_str.parse::<std::os::raw::c_uint>(),
                port_str.parse::<std::os::raw::c_uint>(),
            ) {
                (Ok(cid), Ok(port)) => Ok(ServerAddr::Vsock(VsockAddr { cid, port })),
                _ => Err(Error::UsageError(format!(
                    "invalid server addr: {}",
                    addr_str
                ))),
            }
        }
        _ => Err(Error::UsageError(format!(
            "invalid server addr: {}",
            addr_str
        ))),
    }
}

fn cmd_raw_rpc<I: Iterator<Item = String>>(mut arg_iter: I) -> Result<(), Error> {
    let mut cli_opts = HashMap::new();

    while let Some(word) = arg_iter.next() {
        match word.as_str() {
            "--server" => {
                let addr = arg_iter
                    .next()
                    .ok_or(Error::UsageError(format!("invalid server addr")))
                    .and_then(|addr| parse_server_addr(addr.as_str()))?;
                cli_opts.insert(word, CliOption::Server(addr));
            }
            _ => return Err(Error::UsageError(format!("unexpected argument: {}", word))),
        }
    }

    let server_addr = match cli_opts.get("--server") {
        Some(CliOption::Server(addr)) => addr,
        _ => return Err(Error::UsageError(format!("missing server address"))),
    };

    fn do_raw_rpc<T: Transport>(mut transport: T) -> Result<(), Error> {
        let request = serde_json::from_reader(std::io::stdin())
            .map_err(|_| Error::UsageError(format!("Invalid RPC request")))?;
        transport
            .send_request(&request)
            .map_err(Error::TransportError)?;
        let response = transport.recv_response().map_err(Error::TransportError)?;

        serde_json::to_writer(std::io::stdout(), &response).map_err(Error::SerdeError)?;

        Ok(())
    }

    match server_addr {
        ServerAddr::Unix(path) => UnixStream::connect(path)
            .map_err(Error::IoError)
            .map(|stream| HttpTransport::new(stream, schema::API_URL))
            .and_then(|xport| do_raw_rpc(xport)),
        ServerAddr::Vsock(addr) => VsockStream::connect(*addr)
            .map_err(Error::IoError)
            .map(|stream| HttpTransport::new(stream, schema::API_URL))
            .and_then(|xport| do_raw_rpc(xport)),
    }
}

/// Usage:
/// AF_VSOCK: p11ne-client "vsock" "10000"
/// AF_UNIX:  p11ne-client "unix" "some/path"
fn rusty_main() -> Result<(), Error> {
    let mut args = std::env::args();

    // Skip executable name.
    args.next();

    match args.next().as_ref().map(|s| s.as_str()) {
        Some("raw-rpc") => cmd_raw_rpc(args),
        Some("help") | Some("--help") | Some("-h") => {
            println!("{}", USAGE);
            Ok(())
        }
        Some(cmd) => Err(Error::UsageError(format!("invalid command: {}", cmd))),
        None => Err(Error::UsageError("missing command".to_string())),
    }
}

fn main() {
    match rusty_main() {
        Ok(()) => std::process::exit(0),
        Err(e) => {
            eprintln!("{}", e);
            std::process::exit(i32::from(e))
        }
    }
}


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

use std::env;
use std::fs::OpenOptions;
use std::io::prelude::*;
use std::io::SeekFrom;
use std::io::{Read, Write};
use std::num::Wrapping;

fn usage(binname: &String) {
    eprintln!(
        "Usage: {} <option>
Options:
\tbitflip <filename> <bit>
\tprng_file <outputpath> <filesize> <seed>
\tprng_int <seed> <startvalue> <endvalue>",
        binname
    );
}

struct Prng {
    seed: Wrapping<u32>,
}

impl Prng {
    fn new(seed: u32) -> Self {
        Self {
            seed: Wrapping(seed),
        }
    }

    fn random_u8(&mut self) -> u8 {
        let rand: u8 = (self.seed.0 & 0xff) as u8;
        self.seed = Wrapping(1103515245) * self.seed + Wrapping(1234567);
        rand
    }

    fn random_fill(&mut self, v: &mut [u8]) {
        for x in v.iter_mut() {
            *x = self.random_u8()
        }
    }

    fn random_u32(&mut self) -> u32 {
        let mut vec = [0u8; 4];
        self.random_fill(&mut vec);
        u32::from_be_bytes(vec)
    }
}

fn prng_file(filename: String, filesize: u64, seed: u32) -> Result<(), String> {
    let mut prng = Prng::new(seed);
    let mut buf = [0u8; 4096];
    let mut total_written: u64 = 0;
    let mut file = OpenOptions::new()
        .write(true)
        .create_new(true)
        .open(filename)
        .map_err(|e| format!("Error creating file: {}", e))?;

    while total_written < filesize {
        let wsize = std::cmp::min((filesize - total_written) as usize, buf.len());
        prng.random_fill(&mut buf[..wsize]);
        file.write_all(&buf[..wsize])
            .map_err(|e| format!("Error writing file: {}", e))?;
        total_written += wsize as u64;
    }

    Ok(())
}

/// Renerate random value in range [start, end) from seed
fn prng_int(seed: u32, start: u32, end: u32) -> Result<(), String> {
    let mut prng = Prng::new(seed);
    if end <= start {
        return Err("invalid interval".to_string());
    }
    println!("{}\n", (prng.random_u32() % (end - start)) + start);

    Ok(())
}

fn bitflip(filename: String, bit: u64) -> Result<(), String> {
    let mut file = OpenOptions::new()
        .read(true)
        .write(true)
        .open(&filename)
        .map_err(|e| format!("Could not open file {}: {}", filename, e))?;
    file.seek(SeekFrom::Start(bit / 8))
        .map_err(|e| format!("Could not seek: {}", e))?;
    let mut buf = [0u8; 1];
    file.read_exact(&mut buf)
        .map_err(|e| format!("Could not read: {}", e))?;
    buf[0] ^= 1 << (bit % 8);
    file.seek(SeekFrom::Start(bit / 8))
        .map_err(|e| format!("Could not seek: {}", e))?;
    file.write_all(&mut buf)
        .map_err(|e| format!("Could not read: {}", e))?;

    Ok(())
}

fn main() -> Result<(), String> {
    let mut args = env::args();

    let binname = args.next().ok_or("Missing binary name".to_string())?;
    let command = args.next().ok_or("Missing option".to_string())?;

    match command.as_str() {
        "bitflip" => {
            let filename = args.next().ok_or("Missing argument".to_string())?;
            let bit: u64 = args
                .next()
                .ok_or("Missing argument".to_string())?
                .parse::<u64>()
                .map_err(|e| format!("Invalid number {}", e))?;
            bitflip(filename, bit)
        }
        "prng_file" => {
            let filename = args.next().ok_or("Missing argument".to_string())?;
            let filesize: u64 = args
                .next()
                .ok_or("Missing argument".to_string())?
                .parse::<u64>()
                .map_err(|e| format!("Invalid number {}", e))?;
            let seed: u32 = args
                .next()
                .ok_or("Missing argument".to_string())?
                .parse::<u32>()
                .map_err(|e| format!("Invalid number {}", e))?;
            prng_file(filename, filesize, seed)
        }
        "prng_int" => {
            let seed: u32 = args
                .next()
                .ok_or("Missing argument".to_string())?
                .parse::<u32>()
                .map_err(|e| format!("Invalid number {}", e))?;
            let start: u32 = args
                .next()
                .ok_or("Missing argument".to_string())?
                .parse::<u32>()
                .map_err(|e| format!("Invalid number {}", e))?;
            let end: u32 = args
                .next()
                .ok_or("Missing argument".to_string())?
                .parse::<u32>()
                .map_err(|e| format!("Invalid number {}", e))?;
            prng_int(seed, start, end)
        }
        _ => {
            usage(&binname);
            Err("Unexpected option".to_string())
        }
    }
}


/* automatically generated by rust-bindgen 0.62.0 */

pub const NE_ERR_VCPU_ALREADY_USED: u32 = 256;
pub const NE_ERR_VCPU_NOT_IN_CPU_POOL: u32 = 257;
pub const NE_ERR_VCPU_INVALID_CPU_CORE: u32 = 258;
pub const NE_ERR_INVALID_MEM_REGION_SIZE: u32 = 259;
pub const NE_ERR_INVALID_MEM_REGION_ADDR: u32 = 260;
pub const NE_ERR_UNALIGNED_MEM_REGION_ADDR: u32 = 261;
pub const NE_ERR_MEM_REGION_ALREADY_USED: u32 = 262;
pub const NE_ERR_MEM_NOT_HUGE_PAGE: u32 = 263;
pub const NE_ERR_MEM_DIFFERENT_NUMA_NODE: u32 = 264;
pub const NE_ERR_MEM_MAX_REGIONS: u32 = 265;
pub const NE_ERR_NO_MEM_REGIONS_ADDED: u32 = 266;
pub const NE_ERR_NO_VCPUS_ADDED: u32 = 267;
pub const NE_ERR_ENCLAVE_MEM_MIN_SIZE: u32 = 268;
pub const NE_ERR_FULL_CORES_NOT_USED: u32 = 269;
pub const NE_ERR_NOT_IN_INIT_STATE: u32 = 270;
pub const NE_ERR_INVALID_VCPU: u32 = 271;
pub const NE_ERR_NO_CPUS_AVAIL_IN_POOL: u32 = 272;
pub const NE_ERR_INVALID_PAGE_SIZE: u32 = 273;
pub const NE_ERR_INVALID_FLAG_VALUE: u32 = 274;
pub const NE_ERR_INVALID_ENCLAVE_CID: u32 = 275;
pub type __u64 = ::std::os::raw::c_ulonglong;
#[doc = " struct ne_image_load_info - Info necessary for in-memory enclave image"]
#[doc = "\t\t\t       loading (in / out)."]
#[doc = " @flags:\t\tFlags to determine the enclave image type"]
#[doc = "\t\t\t(e.g. Enclave Image Format - EIF) (in)."]
#[doc = " @memory_offset:\tOffset in enclave memory where to start placing the"]
#[doc = "\t\t\tenclave image (out)."]
#[repr(C)]
#[derive(Debug, Default, Copy, Clone)]
pub struct ne_image_load_info {
    pub flags: __u64,
    pub memory_offset: __u64,
}
#[test]
fn bindgen_test_layout_ne_image_load_info() {
    const UNINIT: ::std::mem::MaybeUninit<ne_image_load_info> = ::std::mem::MaybeUninit::uninit();
    let ptr = UNINIT.as_ptr();
    assert_eq!(
        ::std::mem::size_of::<ne_image_load_info>(),
        16usize,
        concat!("Size of: ", stringify!(ne_image_load_info))
    );
    assert_eq!(
        ::std::mem::align_of::<ne_image_load_info>(),
        8usize,
        concat!("Alignment of ", stringify!(ne_image_load_info))
    );
    assert_eq!(
        unsafe { ::std::ptr::addr_of!((*ptr).flags) as usize - ptr as usize },
        0usize,
        concat!(
            "Offset of field: ",
            stringify!(ne_image_load_info),
            "::",
            stringify!(flags)
        )
    );
    assert_eq!(
        unsafe { ::std::ptr::addr_of!((*ptr).memory_offset) as usize - ptr as usize },
        8usize,
        concat!(
            "Offset of field: ",
            stringify!(ne_image_load_info),
            "::",
            stringify!(memory_offset)
        )
    );
}
#[doc = " struct ne_user_memory_region - Memory region to be set for an enclave (in)."]
#[doc = " @flags:\t\tFlags to determine the usage for the memory region (in)."]
#[doc = " @memory_size:\tThe size, in bytes, of the memory region to be set for"]
#[doc = "\t\t\tan enclave (in)."]
#[doc = " @userspace_addr:\tThe start address of the userspace allocated memory of"]
#[doc = "\t\t\tthe memory region to set for an enclave (in)."]
#[repr(C)]
#[derive(Debug, Default, Copy, Clone)]
pub struct ne_user_memory_region {
    pub flags: __u64,
    pub memory_size: __u64,
    pub userspace_addr: __u64,
}
#[test]
fn bindgen_test_layout_ne_user_memory_region() {
    const UNINIT: ::std::mem::MaybeUninit<ne_user_memory_region> =
        ::std::mem::MaybeUninit::uninit();
    let ptr = UNINIT.as_ptr();
    assert_eq!(
        ::std::mem::size_of::<ne_user_memory_region>(),
        24usize,
        concat!("Size of: ", stringify!(ne_user_memory_region))
    );
    assert_eq!(
        ::std::mem::align_of::<ne_user_memory_region>(),
        8usize,
        concat!("Alignment of ", stringify!(ne_user_memory_region))
    );
    assert_eq!(
        unsafe { ::std::ptr::addr_of!((*ptr).flags) as usize - ptr as usize },
        0usize,
        concat!(
            "Offset of field: ",
            stringify!(ne_user_memory_region),
            "::",
            stringify!(flags)
        )
    );
    assert_eq!(
        unsafe { ::std::ptr::addr_of!((*ptr).memory_size) as usize - ptr as usize },
        8usize,
        concat!(
            "Offset of field: ",
            stringify!(ne_user_memory_region),
            "::",
            stringify!(memory_size)
        )
    );
    assert_eq!(
        unsafe { ::std::ptr::addr_of!((*ptr).userspace_addr) as usize - ptr as usize },
        16usize,
        concat!(
            "Offset of field: ",
            stringify!(ne_user_memory_region),
            "::",
            stringify!(userspace_addr)
        )
    );
}
#[doc = " struct ne_enclave_start_info - Setup info necessary for enclave start (in / out)."]
#[doc = " @flags:\t\tFlags for the enclave to start with (e.g. debug mode) (in)."]
#[doc = " @enclave_cid:\tContext ID (CID) for the enclave vsock device. If 0 as"]
#[doc = "\t\t\tinput, the CID is autogenerated by the hypervisor and"]
#[doc = "\t\t\treturned back as output by the driver (in / out)."]
#[repr(C)]
#[derive(Debug, Default, Copy, Clone)]
pub struct ne_enclave_start_info {
    pub flags: __u64,
    pub enclave_cid: __u64,
}
#[test]
fn bindgen_test_layout_ne_enclave_start_info() {
    const UNINIT: ::std::mem::MaybeUninit<ne_enclave_start_info> =
        ::std::mem::MaybeUninit::uninit();
    let ptr = UNINIT.as_ptr();
    assert_eq!(
        ::std::mem::size_of::<ne_enclave_start_info>(),
        16usize,
        concat!("Size of: ", stringify!(ne_enclave_start_info))
    );
    assert_eq!(
        ::std::mem::align_of::<ne_enclave_start_info>(),
        8usize,
        concat!("Alignment of ", stringify!(ne_enclave_start_info))
    );
    assert_eq!(
        unsafe { ::std::ptr::addr_of!((*ptr).flags) as usize - ptr as usize },
        0usize,
        concat!(
            "Offset of field: ",
            stringify!(ne_enclave_start_info),
            "::",
            stringify!(flags)
        )
    );
    assert_eq!(
        unsafe { ::std::ptr::addr_of!((*ptr).enclave_cid) as usize - ptr as usize },
        8usize,
        concat!(
            "Offset of field: ",
            stringify!(ne_enclave_start_info),
            "::",
            stringify!(enclave_cid)
        )
    );
}


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

//! Rust FFI bindings to Linux Nitro Enclaves driver, generated using
//! [bindgen](https://crates.io/crates/bindgen).

#![allow(missing_docs)]
#![allow(non_camel_case_types)]

mod bindings;
pub use self::bindings::*;


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

use libc::c_int;
use nix::poll::poll;
use nix::poll::{PollFd, PollFlags};
use std::io::Read;
use std::io::Write;
use std::os::unix::io::AsRawFd;
use vsock::VsockListener;

/// Timeout of 1 second in milliseconds
pub const TIMEOUT_SECOND_MS: i32 = 1000;

/// Timeout of 1 minute 30 seconds in milliseconds
pub const TIMEOUT_MINUTE_MS: i32 = 90 * TIMEOUT_SECOND_MS;

const HEART_BEAT: u8 = 0xB7;

#[derive(Debug, PartialEq, Eq)]
/// Internal errors while sending an Eif file
pub enum EifLoaderError {
    SocketPollingError,
    VsockAcceptingError,
    VsockBindingError,
    VsockReceivingError,
    VsockTimeoutError,
}

pub fn enclave_ready(
    listener: VsockListener,
    poll_timeout_ms: c_int,
) -> Result<(), EifLoaderError> {
    let mut poll_fds = [PollFd::new(listener.as_raw_fd(), PollFlags::POLLIN)];
    let result = poll(&mut poll_fds, poll_timeout_ms);
    if result == Ok(0) {
        return Err(EifLoaderError::VsockTimeoutError);
    } else if result != Ok(1) {
        return Err(EifLoaderError::SocketPollingError);
    }

    let mut stream = listener
        .accept()
        .map_err(|_err| EifLoaderError::VsockAcceptingError)?;

    // Wait until the other end is closed
    let mut buf = [0u8];
    let bytes = stream
        .0
        .read(&mut buf)
        .map_err(|_err| EifLoaderError::VsockReceivingError)?;

    if bytes != 1 || buf[0] != HEART_BEAT {
        return Err(EifLoaderError::VsockReceivingError);
    }

    stream
        .0
        .write_all(&buf)
        .map_err(|_err| EifLoaderError::VsockReceivingError)?;

    Ok(())
}


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

use crate::docker::DockerError::CredentialsError;
use base64::{engine::general_purpose, Engine as _};
use futures::stream::StreamExt;
use log::{debug, error, info};
use serde_json::{json, Value};
use shiplift::RegistryAuth;
use shiplift::{BuildOptions, Docker, PullOptions};
use std::fs::File;
use std::io::Write;
use std::path::Path;
use tempfile::NamedTempFile;
use tokio::runtime::Runtime;
use url::Url;

/// Docker inspect architecture constants
pub const DOCKER_ARCH_ARM64: &str = "arm64";
pub const DOCKER_ARCH_AMD64: &str = "amd64";

#[derive(Debug, PartialEq, Eq)]
pub enum DockerError {
    BuildError,
    InspectError,
    PullError,
    RuntimeError,
    TempfileError,
    CredentialsError(String),
    UnsupportedEntryPoint,
}

/// Struct exposing the Docker functionalities to the EIF builder
pub struct DockerUtil {
    docker: Docker,
    docker_image: String,
}

impl DockerUtil {
    /// Constructor that takes as argument a tag for the docker image to be used
    pub fn new(docker_image: String) -> Self {
        let mut docker_image = docker_image;

        if !docker_image.contains(':') {
            docker_image.push_str(":latest");
        }

        DockerUtil {
            // DOCKER_HOST environment variable is parsed inside
            // if docker daemon address needs to be substituted.
            // By default it tries to connect to 'unix:///var/run/docker.sock'
            docker: Docker::new(),
            docker_image,
        }
    }

    /// Returns the credentials by reading ${HOME}/.docker/config.json or ${DOCKER_CONFIG}
    ///
    /// config.json doesn't seem to have a schema that we could use to validate
    /// we are parsing it correctly, so the parsing mechanism had been infered by
    /// reading a config.json created by:
    //         Docker version 19.03.2
    fn get_credentials(&self) -> Result<RegistryAuth, DockerError> {
        let image = self.docker_image.clone();
        let host = if let Ok(uri) = Url::parse(&image) {
            uri.host().map(|s| s.to_string())
        } else {
            // Some Docker URIs don't have the protocol included, so just use
            // a dummy one to trick Url that it's a properly defined Uri.
            let uri = format!("dummy://{image}");
            if let Ok(uri) = Url::parse(&uri) {
                uri.host().map(|s| s.to_string())
            } else {
                None
            }
        };

        if let Some(registry_domain) = host {
            let config_file = self.get_config_file()?;

            let config_json: serde_json::Value = serde_json::from_reader(&config_file)
                .map_err(|err| CredentialsError(format!("JSON was not well-formatted: {err}")))?;

            let auths = config_json.get("auths").ok_or_else(|| {
                CredentialsError("Could not find auths key in config JSON".to_string())
            })?;

            if let Value::Object(auths) = auths {
                for (registry_name, registry_auths) in auths.iter() {
                    if !registry_name.to_string().contains(&registry_domain) {
                        continue;
                    }

                    let auth = registry_auths
                        .get("auth")
                        .ok_or_else(|| {
                            CredentialsError("Could not find auth key in config JSON".to_string())
                        })?
                        .to_string();

                    let auth = auth.replace('"', "");
                    let decoded = general_purpose::STANDARD.decode(auth).map_err(|err| {
                        CredentialsError(format!("Invalid Base64 encoding for auth: {err}"))
                    })?;
                    let decoded = std::str::from_utf8(&decoded).map_err(|err| {
                        CredentialsError(format!("Invalid utf8 encoding for auth: {err}"))
                    })?;

                    if let Some(index) = decoded.rfind(':') {
                        let (user, after_user) = decoded.split_at(index);
                        let (_, password) = after_user.split_at(1);
                        return Ok(RegistryAuth::builder()
                            .username(user)
                            .password(password)
                            .build());
                    }
                }
            }
        }

        Err(CredentialsError(
            "No credentials found for the current image".to_string(),
        ))
    }

    fn get_config_file(&self) -> Result<File, DockerError> {
        if let Ok(file) = std::env::var("DOCKER_CONFIG") {
            let config_file = File::open(file).map_err(|err| {
                DockerError::CredentialsError(format!(
                    "Could not open file pointed by env\
                     DOCKER_CONFIG: {err}"
                ))
            })?;
            Ok(config_file)
        } else {
            if let Ok(home_dir) = std::env::var("HOME") {
                let default_config_path = format!("{home_dir}/.docker/config.json");
                let config_path = Path::new(&default_config_path);
                if config_path.exists() {
                    let config_file = File::open(config_path).map_err(|err| {
                        DockerError::CredentialsError(format!(
                            "Could not open file {:?}: {}",
                            config_path.to_str(),
                            err
                        ))
                    })?;
                    return Ok(config_file);
                }
            }
            Err(DockerError::CredentialsError(
                "Config file not present, please set env \
                 DOCKER_CONFIG accordingly"
                    .to_string(),
            ))
        }
    }

    /// Pull the image, with the tag provided in constructor, from the Docker registry
    pub fn pull_image(&self) -> Result<(), DockerError> {
        let act = async {
            // Check if the Docker image is locally available.
            // If available, early exit.
            if self
                .docker
                .images()
                .get(&self.docker_image)
                .inspect()
                .await
                .is_ok()
            {
                eprintln!("Using the locally available Docker image...");
                return Ok(());
            }

            let mut pull_options_builder = PullOptions::builder();
            pull_options_builder.image(&self.docker_image);

            match self.get_credentials() {
                Ok(auth) => {
                    pull_options_builder.auth(auth);
                }
                // It is not mandatory to have the credentials set, but this is
                // the most likely reason for failure when pulling, so log the
                // error.
                Err(err) => {
                    debug!("WARNING!! Credential could not be set {:?}", err);
                }
            };

            let mut stream = self.docker.images().pull(&pull_options_builder.build());

            loop {
                if let Some(item) = stream.next().await {
                    match item {
                        Ok(output) => {
                            let msg = &output;

                            if let Some(err_msg) = msg.get("error") {
                                error!("{:?}", err_msg.clone());
                                break Err(DockerError::PullError);
                            } else {
                                info!("{}", msg);
                            }
                        }
                        Err(e) => {
                            error!("{:?}", e);
                            break Err(DockerError::PullError);
                        }
                    }
                } else {
                    break Ok(());
                }
            }
        };

        let runtime = Runtime::new().map_err(|_| DockerError::RuntimeError)?;

        runtime.block_on(act)
    }

    /// Build an image locally, with the tag provided in constructor, using a
    /// directory that contains a Dockerfile
    pub fn build_image(&self, dockerfile_dir: String) -> Result<(), DockerError> {
        let act = async {
            let mut stream = self.docker.images().build(
                &BuildOptions::builder(dockerfile_dir)
                    .tag(self.docker_image.clone())
                    .build(),
            );

            loop {
                if let Some(item) = stream.next().await {
                    match item {
                        Ok(output) => {
                            let msg = &output;

                            if let Some(err_msg) = msg.get("error") {
                                error!("{:?}", err_msg.clone());
                                break Err(DockerError::BuildError);
                            } else {
                                info!("{}", msg);
                            }
                        }
                        Err(e) => {
                            error!("{:?}", e);
                            break Err(DockerError::BuildError);
                        }
                    }
                } else {
                    break Ok(());
                }
            }
        };

        let runtime = Runtime::new().map_err(|_| DockerError::RuntimeError)?;

        runtime.block_on(act)
    }

    /// Inspect docker image and return its description as a json String
    pub fn inspect_image(&self) -> Result<serde_json::Value, DockerError> {
        let act = async {
            match self.docker.images().get(&self.docker_image).inspect().await {
                Ok(image) => Ok(json!(image)),
                Err(e) => {
                    error!("{:?}", e);
                    Err(DockerError::InspectError)
                }
            }
        };

        let runtime = Runtime::new().map_err(|_| DockerError::RuntimeError)?;
        runtime.block_on(act)
    }

    fn extract_image(&self) -> Result<(Vec<String>, Vec<String>), DockerError> {
        // First try to find CMD parameters (together with potential ENV bindings)
        let act_cmd = async {
            match self.docker.images().get(&self.docker_image).inspect().await {
                Ok(image) => image.config.cmd.ok_or(DockerError::UnsupportedEntryPoint),
                Err(e) => {
                    error!("{:?}", e);
                    Err(DockerError::InspectError)
                }
            }
        };
        let act_env = async {
            match self.docker.images().get(&self.docker_image).inspect().await {
                Ok(image) => image.config.env.ok_or(DockerError::UnsupportedEntryPoint),
                Err(e) => {
                    error!("{:?}", e);
                    Err(DockerError::InspectError)
                }
            }
        };

        let check_cmd_runtime = Runtime::new()
            .map_err(|_| DockerError::RuntimeError)?
            .block_on(act_cmd);
        let check_env_runtime = Runtime::new()
            .map_err(|_| DockerError::RuntimeError)?
            .block_on(act_env);

        // If no CMD instructions are found, try to locate an ENTRYPOINT command
        if check_cmd_runtime.is_err() || check_env_runtime.is_err() {
            let act_entrypoint = async {
                match self.docker.images().get(&self.docker_image).inspect().await {
                    Ok(image) => image
                        .config
                        .entrypoint
                        .ok_or(DockerError::UnsupportedEntryPoint),
                    Err(e) => {
                        error!("{:?}", e);
                        Err(DockerError::InspectError)
                    }
                }
            };

            let check_entrypoint_runtime = Runtime::new()
                .map_err(|_| DockerError::RuntimeError)?
                .block_on(act_entrypoint);

            if check_entrypoint_runtime.is_err() {
                return Err(DockerError::UnsupportedEntryPoint);
            }

            let act = async {
                match self.docker.images().get(&self.docker_image).inspect().await {
                    Ok(image) => Ok((
                        image.config.entrypoint.unwrap(),
                        image.config.env.ok_or_else(Vec::<String>::new).unwrap(),
                    )),
                    Err(e) => {
                        error!("{:?}", e);
                        Err(DockerError::InspectError)
                    }
                }
            };

            let runtime = Runtime::new().map_err(|_| DockerError::RuntimeError)?;

            return runtime.block_on(act);
        }

        let act = async {
            match self.docker.images().get(&self.docker_image).inspect().await {
                Ok(image) => Ok((image.config.cmd.unwrap(), image.config.env.unwrap())),
                Err(e) => {
                    error!("{:?}", e);
                    Err(DockerError::InspectError)
                }
            }
        };

        let runtime = Runtime::new().map_err(|_| DockerError::RuntimeError)?;

        runtime.block_on(act)
    }

    /// The main function of this struct. This needs to be called in order to
    /// extract the necessary configuration values from the docker image with
    /// the tag provided in the constructor
    pub fn load(&self) -> Result<(NamedTempFile, NamedTempFile), DockerError> {
        let (cmd, env) = self.extract_image()?;

        let cmd_file = write_config(cmd)?;
        let env_file = write_config(env)?;

        Ok((cmd_file, env_file))
    }

    /// Fetch architecture information from an image
    pub fn architecture(&self) -> Result<String, DockerError> {
        let arch = async {
            match self.docker.images().get(&self.docker_image).inspect().await {
                Ok(image) => Ok(image.architecture),
                Err(e) => {
                    error!("{:?}", e);
                    Err(DockerError::InspectError)
                }
            }
        };

        let runtime = Runtime::new().map_err(|_| DockerError::RuntimeError)?;

        runtime.block_on(arch)
    }
}

fn write_config(config: Vec<String>) -> Result<NamedTempFile, DockerError> {
    let mut file = NamedTempFile::new().map_err(|_| DockerError::TempfileError)?;

    for line in config {
        file.write_fmt(format_args!("{line}\n"))
            .map_err(|_| DockerError::TempfileError)?;
    }

    Ok(file)
}

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

    /// Test extracted configuration is as expected
    #[test]
    fn test_config() {
        let docker = DockerUtil::new(String::from("public.ecr.aws/aws-nitro-enclaves/hello:v1"));

        let (cmd_file, env_file) = docker.load().unwrap();
        let mut cmd_file = File::open(cmd_file.path()).unwrap();
        let mut env_file = File::open(env_file.path()).unwrap();

        let mut cmd = String::new();
        cmd_file.read_to_string(&mut cmd).unwrap();
        assert_eq!(cmd, "/bin/hello.sh\n");

        let mut env = String::new();
        env_file.read_to_string(&mut env).unwrap();
        assert_eq!(
            env,
            "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin\n\
             HELLO=Hello from the enclave side!\n"
        );
    }
}


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

use std::fs::File;
use std::path::Path;
use std::process::Command;

mod docker;
mod yaml_generator;

use aws_nitro_enclaves_image_format::defs::{EifBuildInfo, EifIdentityInfo, EIF_HDR_ARCH_ARM64};
use aws_nitro_enclaves_image_format::utils::identity::parse_custom_metadata;
use aws_nitro_enclaves_image_format::utils::{EifBuilder, SignEnclaveInfo};
use docker::DockerUtil;
use serde_json::json;
use sha2::Digest;
use std::collections::BTreeMap;
use yaml_generator::YamlGenerator;

pub const DEFAULT_TAG: &str = "1.0";

pub struct Docker2Eif<'a> {
    docker_image: String,
    docker: DockerUtil,
    init_path: String,
    nsm_path: String,
    kernel_img_path: String,
    cmdline: String,
    linuxkit_path: String,
    artifacts_prefix: String,
    output: &'a mut File,
    sign_info: Option<SignEnclaveInfo>,
    img_name: Option<String>,
    img_version: Option<String>,
    metadata_path: Option<String>,
    build_info: EifBuildInfo,
}

#[derive(Debug, PartialEq, Eq)]
pub enum Docker2EifError {
    DockerError,
    DockerfilePathError,
    ImagePullError,
    InitPathError,
    NsmPathError,
    KernelPathError,
    LinuxkitExecError,
    LinuxkitPathError,
    MetadataPathError,
    MetadataError(String),
    ArtifactsPrefixError,
    RamfsError,
    RemoveFileError,
    SignImageError(String),
    SignArgsError,
    UnsupportedArchError,
}

impl<'a> Docker2Eif<'a> {
    pub fn new(
        docker_image: String,
        init_path: String,
        nsm_path: String,
        kernel_img_path: String,
        cmdline: String,
        linuxkit_path: String,
        output: &'a mut File,
        artifacts_prefix: String,
        certificate_path: &Option<String>,
        key_path: &Option<String>,
        img_name: Option<String>,
        img_version: Option<String>,
        metadata_path: Option<String>,
        build_info: EifBuildInfo,
    ) -> Result<Self, Docker2EifError> {
        let docker = DockerUtil::new(docker_image.clone());

        if !Path::new(&init_path).is_file() {
            return Err(Docker2EifError::InitPathError);
        } else if !Path::new(&nsm_path).is_file() {
            return Err(Docker2EifError::NsmPathError);
        } else if !Path::new(&kernel_img_path).is_file() {
            return Err(Docker2EifError::KernelPathError);
        } else if !Path::new(&linuxkit_path).is_file() {
            return Err(Docker2EifError::LinuxkitPathError);
        } else if !Path::new(&artifacts_prefix).is_dir() {
            return Err(Docker2EifError::ArtifactsPrefixError);
        }

        if let Some(ref path) = metadata_path {
            if !Path::new(path).is_file() {
                return Err(Docker2EifError::MetadataPathError);
            }
        }

        let sign_info = match (certificate_path, key_path) {
            (None, None) => None,
            (Some(cert_path), Some(key_path)) => Some(
                SignEnclaveInfo::new(cert_path, key_path)
                    .map_err(|err| Docker2EifError::SignImageError(format!("{err:?}")))?,
            ),
            _ => return Err(Docker2EifError::SignArgsError),
        };

        Ok(Docker2Eif {
            docker_image,
            docker,
            init_path,
            nsm_path,
            kernel_img_path,
            cmdline,
            linuxkit_path,
            output,
            artifacts_prefix,
            sign_info,
            img_name,
            img_version,
            metadata_path,
            build_info,
        })
    }

    pub fn pull_docker_image(&self) -> Result<(), Docker2EifError> {
        self.docker.pull_image().map_err(|e| {
            eprintln!("Docker error: {e:?}");
            Docker2EifError::DockerError
        })?;

        Ok(())
    }

    pub fn build_docker_image(&self, dockerfile_dir: String) -> Result<(), Docker2EifError> {
        if !Path::new(&dockerfile_dir).is_dir() {
            return Err(Docker2EifError::DockerfilePathError);
        }
        self.docker.build_image(dockerfile_dir).map_err(|e| {
            eprintln!("Docker error: {e:?}");
            Docker2EifError::DockerError
        })?;

        Ok(())
    }

    fn generate_identity_info(&self) -> Result<EifIdentityInfo, Docker2EifError> {
        let docker_info = self
            .docker
            .inspect_image()
            .map_err(|e| Docker2EifError::MetadataError(format!("Docker inspect error: {e:?}")))?;

        let uri_split: Vec<&str> = self.docker_image.split(':').collect();
        if uri_split.is_empty() {
            return Err(Docker2EifError::MetadataError(
                "Wrong image name specified".to_string(),
            ));
        }

        // Image hash is used by default in case image version is not provided.
        // It's taken from JSON generated by `docker inspect` and a bit fragile.
        // May be later we should change it to fetching this data
        // from a specific struct and not JSON
        let img_hash = docker_info
            .get("Id")
            .and_then(|val| val.as_str())
            .and_then(|str| str.strip_prefix("sha256:"))
            .ok_or_else(|| {
                Docker2EifError::MetadataError(
                    "Image info must contain string Id field".to_string(),
                )
            })?;

        let img_name = self
            .img_name
            .clone()
            .unwrap_or_else(|| uri_split[0].to_string());
        let img_version = self
            .img_version
            .clone()
            .unwrap_or_else(|| img_hash.to_string());

        let mut custom_info = json!(null);
        if let Some(ref path) = self.metadata_path {
            custom_info = parse_custom_metadata(path).map_err(Docker2EifError::MetadataError)?
        }

        Ok(EifIdentityInfo {
            img_name,
            img_version,
            build_info: self.build_info.clone(),
            docker_info,
            custom_info,
        })
    }

    pub fn create(&mut self) -> Result<BTreeMap<String, String>, Docker2EifError> {
        let (cmd_file, env_file) = self.docker.load().map_err(|e| {
            eprintln!("Docker error: {e:?}");
            Docker2EifError::DockerError
        })?;

        let yaml_generator = YamlGenerator::new(
            self.docker_image.clone(),
            self.init_path.clone(),
            self.nsm_path.clone(),
            cmd_file.path().to_str().unwrap().to_string(),
            env_file.path().to_str().unwrap().to_string(),
        );

        let ramfs_config_file = yaml_generator.get_bootstrap_ramfs().map_err(|e| {
            eprintln!("Ramfs error: {e:?}");
            Docker2EifError::RamfsError
        })?;
        let ramfs_with_rootfs_config_file = yaml_generator.get_customer_ramfs().map_err(|e| {
            eprintln!("Ramfs error: {e:?}");
            Docker2EifError::RamfsError
        })?;

        let bootstrap_ramfs = format!("{}/bootstrap-initrd.img", self.artifacts_prefix);
        let customer_ramfs = format!("{}/customer-initrd.img", self.artifacts_prefix);

        let output = Command::new(&self.linuxkit_path)
            .args([
                "build",
                "-name",
                &bootstrap_ramfs,
                "-format",
                "kernel+initrd",
                ramfs_config_file.path().to_str().unwrap(),
            ])
            .output()
            .map_err(|_| Docker2EifError::LinuxkitExecError)?;
        if !output.status.success() {
            eprintln!(
                "Linuxkit reported an error while creating the bootstrap ramfs: {:?}",
                String::from_utf8_lossy(&output.stderr)
            );
            return Err(Docker2EifError::LinuxkitExecError);
        }

        // Prefix the docker image filesystem, as expected by init
        let output = Command::new(&self.linuxkit_path)
            .args([
                "build",
                "-docker",
                "-name",
                &customer_ramfs,
                "-format",
                "kernel+initrd",
                "-prefix",
                "rootfs/",
                ramfs_with_rootfs_config_file.path().to_str().unwrap(),
            ])
            .output()
            .map_err(|_| Docker2EifError::LinuxkitExecError)?;
        if !output.status.success() {
            eprintln!(
                "Linuxkit reported an error while creating the customer ramfs: {:?}",
                String::from_utf8_lossy(&output.stderr)
            );
            return Err(Docker2EifError::LinuxkitExecError);
        }

        let arch = self.docker.architecture().map_err(|e| {
            eprintln!("Docker error: {e:?}");
            Docker2EifError::DockerError
        })?;

        let flags = match arch.as_str() {
            docker::DOCKER_ARCH_ARM64 => EIF_HDR_ARCH_ARM64,
            docker::DOCKER_ARCH_AMD64 => 0,
            _ => return Err(Docker2EifError::UnsupportedArchError),
        };

        let mut build = EifBuilder::new(
            Path::new(&self.kernel_img_path),
            self.cmdline.clone(),
            self.sign_info.clone(),
            sha2::Sha384::new(),
            flags,
            self.generate_identity_info()?,
        );

        // Linuxkit adds -initrd.img sufix to the file names.
        let bootstrap_ramfs = format!("{bootstrap_ramfs}-initrd.img");
        let customer_ramfs = format!("{customer_ramfs}-initrd.img");

        build.add_ramdisk(Path::new(&bootstrap_ramfs));
        build.add_ramdisk(Path::new(&customer_ramfs));

        Ok(build.write_to(self.output))
    }
}


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

use clap::{App, AppSettings, Arg};
use std::fs::OpenOptions;

use aws_nitro_enclaves_image_format::generate_build_info;
use enclave_build::Docker2Eif;

fn main() {
    let matches = App::new("Docker2Eif builder")
        .about("Generate consistent EIF image from a Docker image")
        .setting(AppSettings::DisableVersion)
        .arg(
            Arg::with_name("docker_image")
                .short('t')
                .long("tag")
                .help("Docker image tag")
                .takes_value(true)
                .required(true),
        )
        .arg(
            Arg::with_name("init_path")
                .short('i')
                .long("init")
                .help("Path to a binary representing the init process for the enclave")
                .takes_value(true)
                .required(true),
        )
        .arg(
            Arg::with_name("nsm_path")
                .short('n')
                .long("nsm")
                .help("Path to the NitroSecureModule Kernel Driver")
                .takes_value(true)
                .required(true),
        )
        .arg(
            Arg::with_name("kernel_img_path")
                .short('k')
                .long("kernel")
                .help("Path to a bzImage/Image file for x86_64/aarch64 linux kernel")
                .takes_value(true)
                .required(true),
        )
        .arg(
            Arg::with_name("kernel_cfg_path")
                .long("kernel_config")
                .help("Path to a bzImage.config/Image.config file for x86_64/aarch64 linux kernel config")
                .takes_value(true)
                .required(true),
        )
        .arg(
            Arg::with_name("cmdline")
                .short('c')
                .long("cmdline")
                .help("Cmdline for kernel")
                .takes_value(true)
                .required(true),
        )
        .arg(
            Arg::with_name("linuxkit_path")
                .short('l')
                .long("linuxkit")
                .help("Linuxkit executable path")
                .takes_value(true)
                .required(true),
        )
        .arg(
            Arg::with_name("output")
                .short('o')
                .long("output")
                .help("Output file for EIF image")
                .takes_value(true)
                .required(true),
        )
        .arg(
            Arg::with_name("signing-certificate")
                .long("signing-certificate")
                .help("Specify the path to the signing certificate")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("private-key")
                .long("private-key")
                .help("Specify the path to the private-key")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("build")
                .short('b')
                .long("build")
                .help("Build image from Dockerfile")
                .takes_value(true)
                .required(false),
        )
        .arg(
            Arg::with_name("pull")
                .short('p')
                .long("pull")
                .help("Pull the Docker image before generating EIF")
                .required(false)
                .conflicts_with("build"),
        )
        .arg(
            Arg::with_name("image_name")
                .long("name")
                .help("Name for enclave image")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("image_version")
                .long("version")
                .help("Version of the enclave image")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("metadata")
                .long("metadata")
                .help("Path to JSON containing the custom metadata provided by the user.")
                .takes_value(true),
        )
        .get_matches();

    let docker_image = matches.value_of("docker_image").unwrap();
    let init_path = matches.value_of("init_path").unwrap();
    let nsm_path = matches.value_of("nsm_path").unwrap();
    let kernel_img_path = matches.value_of("kernel_img_path").unwrap();
    let kernel_cfg_path = matches.value_of("kernel_cfg_path").unwrap();
    let cmdline = matches.value_of("cmdline").unwrap();
    let linuxkit_path = matches.value_of("linuxkit_path").unwrap();
    let output = matches.value_of("output").unwrap();
    let signing_certificate = matches
        .value_of("signing_certificate")
        .map(|val| val.to_string());
    let private_key = matches
        .value_of("private_certificate")
        .map(|val| val.to_string());
    let img_name = matches.value_of("image_name").map(|val| val.to_string());
    let img_version = matches.value_of("image_version").map(|val| val.to_string());
    let metadata = matches.value_of("metadata").map(|val| val.to_string());

    let mut output = OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .open(output)
        .expect("Failed to create output file");

    let mut img = Docker2Eif::new(
        docker_image.to_string(),
        init_path.to_string(),
        nsm_path.to_string(),
        kernel_img_path.to_string(),
        cmdline.to_string(),
        linuxkit_path.to_string(),
        &mut output,
        ".".to_string(),
        &signing_certificate,
        &private_key,
        img_name,
        img_version,
        metadata,
        generate_build_info!(kernel_cfg_path).expect("Can not generate build info"),
    )
    .unwrap();

    if matches.is_present("build") {
        let dockerfile_dir = matches.value_of("build").unwrap();
        img.build_docker_image(dockerfile_dir.to_string()).unwrap();
    } else if matches.is_present("pull") {
        img.pull_docker_image().unwrap();
    }

    img.create().unwrap();
}


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

use serde::{Deserialize, Serialize};
use std::io::Write;
use tempfile::NamedTempFile;

#[derive(Debug, Serialize, Deserialize)]
struct BootstrapRamfsTemplate {
    files: (DirTemplate, FileTemplate, FileTemplate),
}

#[derive(Debug, Serialize, Deserialize)]
struct CustomerRamfsTemplate {
    init: Vec<String>,
    files: (
        DirTemplate,
        DirTemplate,
        DirTemplate,
        DirTemplate,
        DirTemplate,
        DirTemplate,
        FileTemplate,
        FileTemplate,
    ),
}

#[derive(Debug, Serialize, Deserialize)]
struct FileTemplate {
    path: String,
    source: String,
    mode: String,
}

#[derive(Debug, Serialize, Deserialize)]
struct DirTemplate {
    path: String,
    directory: bool,
    mode: String,
}

#[derive(Debug, PartialEq, Eq)]
pub enum YamlGeneratorError {
    TempfileError,
}

pub struct YamlGenerator {
    docker_image: String,
    init_path: String,
    nsm_path: String,
    cmd_path: String,
    env_path: String,
}

impl YamlGenerator {
    pub fn new(
        docker_image: String,
        init_path: String,
        nsm_path: String,
        cmd_path: String,
        env_path: String,
    ) -> Self {
        YamlGenerator {
            docker_image,
            init_path,
            nsm_path,
            cmd_path,
            env_path,
        }
    }

    pub fn get_bootstrap_ramfs(&self) -> Result<NamedTempFile, YamlGeneratorError> {
        let ramfs = BootstrapRamfsTemplate {
            files: (
                DirTemplate {
                    path: String::from("dev"),
                    directory: true,
                    mode: String::from("0755"),
                },
                FileTemplate {
                    path: String::from("init"),
                    source: self.init_path.clone(),
                    mode: String::from("0755"),
                },
                FileTemplate {
                    path: String::from("nsm.ko"),
                    source: self.nsm_path.clone(),
                    mode: String::from("0755"),
                },
            ),
        };

        let yaml = serde_yaml::to_string(&ramfs);

        let mut file = NamedTempFile::new().map_err(|_| YamlGeneratorError::TempfileError)?;

        file.write_all(yaml.unwrap().as_bytes())
            .map_err(|_| YamlGeneratorError::TempfileError)?;

        Ok(file)
    }

    pub fn get_customer_ramfs(&self) -> Result<NamedTempFile, YamlGeneratorError> {
        let ramfs = CustomerRamfsTemplate {
            init: vec![self.docker_image.clone()],
            // Each directory must stay under rootfs, as expected by init
            files: (
                DirTemplate {
                    path: String::from("rootfs/dev"),
                    directory: true,
                    mode: String::from("0755"),
                },
                DirTemplate {
                    path: String::from("rootfs/run"),
                    directory: true,
                    mode: String::from("0755"),
                },
                DirTemplate {
                    path: String::from("rootfs/sys"),
                    directory: true,
                    mode: String::from("0755"),
                },
                DirTemplate {
                    path: String::from("rootfs/var"),
                    directory: true,
                    mode: String::from("0755"),
                },
                DirTemplate {
                    path: String::from("rootfs/proc"),
                    directory: true,
                    mode: String::from("0755"),
                },
                DirTemplate {
                    path: String::from("rootfs/tmp"),
                    directory: true,
                    mode: String::from("0755"),
                },
                FileTemplate {
                    path: String::from("cmd"),
                    source: self.cmd_path.clone(),
                    mode: String::from("0644"),
                },
                FileTemplate {
                    path: String::from("env"),
                    source: self.env_path.clone(),
                    mode: String::from("0644"),
                },
            ),
        };

        let yaml = serde_yaml::to_string(&ramfs);

        let mut file = NamedTempFile::new().map_err(|_| YamlGeneratorError::TempfileError)?;

        file.write_all(yaml.unwrap().as_bytes())
            .map_err(|_| YamlGeneratorError::TempfileError)?;

        Ok(file)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs::File;
    use std::io::Read;

    /// Test YAML config files are the same as the ones written by hand
    #[test]
    fn test_ramfs() {
        let yaml_generator = YamlGenerator::new(
            String::from("docker_image"),
            String::from("path_to_init"),
            String::from("path_to_nsm"),
            String::from("path_to_cmd"),
            String::from("path_to_env"),
        );

        let mut bootstrap_data = String::new();
        let bootstrap_ramfs = yaml_generator.get_bootstrap_ramfs().unwrap();
        let mut bootstrap_ramfs = File::open(bootstrap_ramfs.path()).unwrap();
        bootstrap_ramfs.read_to_string(&mut bootstrap_data).unwrap();
        assert_eq!(
            bootstrap_data,
            "---\
             \nfiles:\
             \n  - path: dev\
             \n    directory: true\
             \n    mode: \"0755\"\
             \n  - path: init\
             \n    source: path_to_init\
             \n    mode: \"0755\"\
             \n  - path: nsm.ko\
             \n    source: path_to_nsm\
             \n    mode: \"0755\"\
             \n\
             "
        );

        let mut customer_data = String::new();
        let customer_ramfs = yaml_generator.get_customer_ramfs().unwrap();
        let mut customer_ramfs = File::open(customer_ramfs.path()).unwrap();
        customer_ramfs.read_to_string(&mut customer_data).unwrap();
        assert_eq!(
            customer_data,
            "---\
             \ninit:\
             \n  - docker_image\
             \nfiles:\
             \n  - path: rootfs/dev\
             \n    directory: true\
             \n    mode: \"0755\"\
             \n  - path: rootfs/run\
             \n    directory: true\
             \n    mode: \"0755\"\
             \n  - path: rootfs/sys\
             \n    directory: true\
             \n    mode: \"0755\"\
             \n  - path: rootfs/var\
             \n    directory: true\
             \n    mode: \"0755\"\
             \n  - path: rootfs/proc\
             \n    directory: true\
             \n    mode: \"0755\"\
             \n  - path: rootfs/tmp\
             \n    directory: true\
             \n    mode: \"0755\"\
             \n  - path: cmd\
             \n    source: path_to_cmd\
             \n    mode: \"0644\"\
             \n  - path: env\
             \n    source: path_to_env\
             \n    mode: \"0644\"\
             \n\
             "
        );
    }
}


use serde::{Deserialize, Serialize};
use std::process::Output;

use clap::ArgMatches;

#[derive(Debug, Clone)]
pub struct ListenArgs {
    pub port: u32,
}

impl ListenArgs {
    pub fn new_with(args: &ArgMatches) -> Result<Self, String> {
        Ok(ListenArgs {
            port: parse_port(args)?,
        })
    }
}

#[derive(Debug, Clone)]
pub struct RunArgs {
    pub cid: u32,
    pub port: u32,
    pub command: String,
    pub no_wait: bool,
}

impl RunArgs {
    pub fn new_with(args: &ArgMatches) -> Result<Self, String> {
        Ok(RunArgs {
            cid: parse_cid(args)?,
            port: parse_port(args)?,
            command: parse_command(args)?,
            no_wait: parse_no_wait(args),
        })
    }
}

#[derive(Debug, Clone)]
pub struct FileArgs {
    pub cid: u32,
    pub port: u32,
    pub localfile: String,
    pub remotefile: String,
}

impl FileArgs {
    pub fn new_with(args: &ArgMatches) -> Result<Self, String> {
        Ok(FileArgs {
            cid: parse_cid(args)?,
            port: parse_port(args)?,
            localfile: parse_localfile(args)?,
            remotefile: parse_remotefile(args)?,
        })
    }
}

#[derive(Serialize, Deserialize, Debug)]
pub struct CommandOutput {
    pub stdout: String,
    pub stderr: String,
    pub rc: Option<i32>,
}

impl CommandOutput {
    pub fn new(stdout: String, stderr: String, code: i32) -> Self {
        CommandOutput {
            stdout,
            stderr,
            rc: Some(code),
        }
    }

    pub fn new_from(output: Output) -> Result<Self, String> {
        Ok(CommandOutput {
            stdout: String::from_utf8(output.stdout).map_err(|err| format!("{:?}", err))?,
            stderr: String::from_utf8(output.stderr).map_err(|err| format!("{:?}", err))?,
            rc: output.status.code(),
        })
    }
}

fn parse_cid(args: &ArgMatches) -> Result<u32, String> {
    let port = args.value_of("cid").ok_or("Could not find cid argument")?;
    port.parse()
        .map_err(|_err| "cid is not a number".to_string())
}

fn parse_port(args: &ArgMatches) -> Result<u32, String> {
    let port = args
        .value_of("port")
        .ok_or("Could not find port argument")?;
    port.parse()
        .map_err(|_err| "port is not a number".to_string())
}

fn parse_command(args: &ArgMatches) -> Result<String, String> {
    let command = args
        .value_of("command")
        .ok_or("Could not find command argument")?;
    Ok(String::from(command))
}

fn parse_no_wait(args: &ArgMatches) -> bool {
    if args.is_present("no-wait") {
        true
    } else {
        false
    }
}

fn parse_localfile(args: &ArgMatches) -> Result<String, String> {
    let output = args
        .value_of("localpath")
        .ok_or("Could not find localpath")?;
    Ok(String::from(output))
}

fn parse_remotefile(args: &ArgMatches) -> Result<String, String> {
    let output = args
        .value_of("remotepath")
        .ok_or("Could not find remotepath")?;
    Ok(String::from(output))
}