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

#[allow(unused_imports)]
#[cfg(test)]
mod tests {
    use nitro_cli::common::commands_parser::{
        BuildEnclavesArgs, RunEnclavesArgs, TerminateEnclavesArgs,
    };
    use nitro_cli::common::json_output::EnclaveDescribeInfo;
    use nitro_cli::enclave_proc::commands::{describe_enclaves, run_enclaves, terminate_enclaves};
    use nitro_cli::enclave_proc::resource_manager::NE_ENCLAVE_DEBUG_MODE;
    use nitro_cli::enclave_proc::utils::{
        flags_to_string, generate_enclave_id, get_enclave_describe_info,
    };
    use nitro_cli::utils::{Console, PcrType};
    use nitro_cli::{
        build_enclaves, build_from_docker, describe_eif, enclave_console, get_file_pcr,
        new_enclave_name,
    };
    use nitro_cli::{CID_TO_CONSOLE_PORT_OFFSET, VMADDR_CID_HYPERVISOR};
    use serde_json::json;
    use std::convert::TryInto;
    use std::fs::{File, OpenOptions};
    use std::io::Write;
    use tempfile::{tempdir, TempDir};

    use openssl::asn1::Asn1Time;
    use openssl::ec::{EcGroup, EcKey};
    use openssl::hash::MessageDigest;
    use openssl::nid::Nid;
    use openssl::pkey::{PKey, Private};
    use openssl::x509::{X509Name, X509};

    // Remote Docker image
    const SAMPLE_DOCKER: &str = "public.ecr.aws/aws-nitro-enclaves/hello:v1";

    #[cfg(target_arch = "x86_64")]
    mod sample_docker_pcrs {
        /// PCR0
        pub const IMAGE_PCR: &str = "7f3287dd1c4dbc49513abfaabc7f6afe79ab8269743c0c4ee55bb9e92d4f0a36f0cae7c0356d0bfec78b59b4d20c689c";
        /// PCR1
        pub const KERNEL_PCR: &str = "bcdf05fefccaa8e55bf2c8d6dee9e79bbff31e34bf28a99aa19e6b29c37ee80b214a414b7607236edf26fcb78654e63f";
        /// PCR2
        pub const APP_PCR: &str = "dd61366a5424eea46f60c4e9d59e6c645a46420ccf962550ee1f3c109d230f88ec23667617aeaac425a1f50fe8e384d7";
    }

    #[cfg(target_arch = "aarch64")]
    mod sample_docker_pcrs {
        /// PCR0
        pub const IMAGE_PCR: &str = "b32a774b09fff4324a6405dacf3f5aa462a75e554e3a563ee64708abd585456bb480fdf70b2e2c2ab9ec205717bc690e";
        /// PCR1
        pub const KERNEL_PCR: &str = "5d3938eb05288e20a981038b1861062ff4174884968a39aee5982b312894e60561883576cc7381d1a7d05b809936bd16";
        /// PCR2
        pub const APP_PCR: &str = "9397173aa14e47fe087e8aeb63928a233db048e290830de6ce2041f4580f83b599c48432467601bed8a4883e9d94ff10";
    }

    // Local Docker image
    const COMMAND_EXECUTER_DOCKER: &str = "command_executer:eif";

    pub const MAX_BOOT_TIMEOUT_SEC: u64 = 9;

    use std::convert::TryFrom;
    use std::time::Duration;

    fn setup_env() {
        if std::env::var("NITRO_CLI_BLOBS").is_err() {
            std::env::set_var("NITRO_CLI_BLOBS", "/usr/share/nitro_enclaves/blobs");
        }
    }

    #[test]
    fn build_enclaves_invalid_uri() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");
        setup_env();
        let args = BuildEnclavesArgs {
            docker_uri: "667861386598.dkr.ecr.us-east-1.amazonaws.com/enclaves-devel".to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: None,
            img_version: None,
            metadata: None,
        };

        assert_eq!(build_enclaves(args).is_err(), true);
    }

    #[test]
    fn build_enclaves_simple_image() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");
        setup_env();
        let args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: None,
            img_version: None,
            metadata: None,
        };

        let measurements = build_from_docker(
            &args.docker_uri,
            &args.docker_dir,
            &args.output,
            &args.signing_certificate,
            &args.private_key,
            &args.img_name,
            &args.img_version,
            &args.metadata,
        )
        .expect("Docker build failed")
        .1;
        assert_eq!(
            measurements.get("PCR0").unwrap(),
            sample_docker_pcrs::IMAGE_PCR
        );
        assert_eq!(
            measurements.get("PCR1").unwrap(),
            sample_docker_pcrs::KERNEL_PCR
        );
        assert_eq!(
            measurements.get("PCR2").unwrap(),
            sample_docker_pcrs::APP_PCR
        );
    }

    #[test]
    fn build_hello_world() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");
        setup_env();
        let args = BuildEnclavesArgs {
            docker_uri: "hello-world:latest".to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &args.docker_uri,
            &args.docker_dir,
            &args.output,
            &args.signing_certificate,
            &args.private_key,
            &args.img_name,
            &args.img_version,
            &args.metadata,
        )
        .expect("Docker build failed");
    }

    #[test]
    fn build_enclaves_command_executer() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");
        setup_env();
        let args = BuildEnclavesArgs {
            docker_uri: COMMAND_EXECUTER_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &args.docker_uri,
            &args.docker_dir,
            &args.output,
            &args.signing_certificate,
            &args.private_key,
            &args.img_name,
            &args.img_version,
            &args.metadata,
        )
        .expect("Docker build failed");
    }

    fn generate_signing_cert_and_key(cert_path: &str, key_path: &str) {
        let ec_group = EcGroup::from_curve_name(Nid::SECP384R1).unwrap();
        let key = EcKey::generate(&ec_group).unwrap();
        let pkey = PKey::from_ec_key(key.clone()).unwrap();

        let mut name = X509Name::builder().unwrap();
        name.append_entry_by_nid(Nid::COMMONNAME, "aws.nitro-enclaves")
            .unwrap();
        let name = name.build();

        let before = Asn1Time::days_from_now(0).unwrap();
        let after = Asn1Time::days_from_now(365).unwrap();

        let mut builder = X509::builder().unwrap();
        builder.set_version(2).unwrap();
        builder.set_subject_name(&name).unwrap();
        builder.set_issuer_name(&name).unwrap();
        builder.set_pubkey(&pkey).unwrap();
        builder.set_not_before(&before).unwrap();
        builder.set_not_after(&after).unwrap();
        builder.sign(&pkey, MessageDigest::sha384()).unwrap();

        let cert = builder.build();

        let mut key_file = OpenOptions::new()
            .write(true)
            .create(true)
            .open(key_path)
            .unwrap();
        key_file
            .write_all(&key.private_key_to_pem().unwrap())
            .unwrap();

        let mut cert_file = OpenOptions::new()
            .write(true)
            .create(true)
            .open(cert_path)
            .unwrap();
        cert_file.write_all(&cert.to_pem().unwrap()).unwrap();
    }

    #[test]
    fn build_enclaves_signed_simple_image() {
        let dir = tempdir().unwrap();
        let dir_path = dir.path().to_str().unwrap();
        let eif_path = format!("{}/test.eif", dir_path);
        let cert_path = format!("{}/cert.pem", dir_path);
        let key_path = format!("{}/key.pem", dir_path);
        generate_signing_cert_and_key(&cert_path, &key_path);

        setup_env();
        let args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path,
            signing_certificate: Some(cert_path),
            private_key: Some(key_path),
            img_name: None,
            img_version: None,
            metadata: None,
        };

        let measurements = build_from_docker(
            &args.docker_uri,
            &args.docker_dir,
            &args.output,
            &args.signing_certificate,
            &args.private_key,
            &args.img_name,
            &args.img_version,
            &args.metadata,
        )
        .expect("Docker build failed")
        .1;

        assert_eq!(
            measurements.get("PCR0").unwrap(),
            sample_docker_pcrs::IMAGE_PCR
        );
        assert_eq!(
            measurements.get("PCR1").unwrap(),
            sample_docker_pcrs::KERNEL_PCR
        );
        assert_eq!(
            measurements.get("PCR2").unwrap(),
            sample_docker_pcrs::APP_PCR
        );
    }

    #[test]
    fn run_describe_terminate_simple_docker_image() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");
        setup_env();
        let build_args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &build_args.docker_uri,
            &build_args.docker_dir,
            &build_args.output,
            &build_args.signing_certificate,
            &build_args.private_key,
            &build_args.img_name,
            &build_args.img_version,
            &build_args.metadata,
        )
        .expect("Docker build failed");

        let args = RunEnclavesArgs {
            enclave_cid: None,
            eif_path: build_args.output,
            cpu_ids: None,
            cpu_count: Some(2),
            memory_mib: 128,
            debug_mode: true,
            attach_console: false,
            enclave_name: Some("testName".to_string()),
        };
        run_describe_terminate(args);
    }

    #[test]
    fn run_describe_terminate_signed_enclave_image() {
        let dir = tempdir().unwrap();
        let dir_path = dir.path().to_str().unwrap();
        let eif_path = format!("{}/test.eif", dir_path);
        let cert_path = format!("{}/cert.pem", dir_path);
        let key_path = format!("{}/key.pem", dir_path);
        generate_signing_cert_and_key(&cert_path, &key_path);

        setup_env();
        let build_args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path,
            signing_certificate: Some(cert_path),
            private_key: Some(key_path),
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &build_args.docker_uri,
            &build_args.docker_dir,
            &build_args.output,
            &build_args.signing_certificate,
            &build_args.private_key,
            &build_args.img_name,
            &build_args.img_version,
            &build_args.metadata,
        )
        .expect("Docker build failed");

        let args = RunEnclavesArgs {
            enclave_cid: None,
            eif_path: build_args.output,
            cpu_ids: None,
            cpu_count: Some(2),
            memory_mib: 256,
            debug_mode: true,
            attach_console: false,
            enclave_name: Some("testName".to_string()),
        };
        run_describe_terminate(args);
    }

    #[test]
    fn run_describe_terminate_command_executer_docker_image() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");
        setup_env();
        let build_args = BuildEnclavesArgs {
            docker_uri: COMMAND_EXECUTER_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &build_args.docker_uri,
            &build_args.docker_dir,
            &build_args.output,
            &build_args.signing_certificate,
            &build_args.private_key,
            &build_args.img_name,
            &build_args.img_version,
            &build_args.metadata,
        )
        .expect("Docker build failed");

        let args = RunEnclavesArgs {
            enclave_cid: None,
            eif_path: build_args.output,
            cpu_ids: None,
            cpu_count: Some(2),
            memory_mib: 2046,
            debug_mode: true,
            attach_console: false,
            enclave_name: Some("testName".to_string()),
        };
        run_describe_terminate(args);
    }

    fn run_describe_terminate(args: RunEnclavesArgs) {
        setup_env();
        let req_enclave_cid = args.enclave_cid;
        let req_mem_size = args.memory_mib;
        let req_nr_cpus: u64 = args.cpu_count.unwrap().try_into().unwrap();
        let debug_mode = args.debug_mode;
        let mut enclave_manager = run_enclaves(&args, None)
            .expect("Run enclaves failed")
            .enclave_manager;
        let enclave_cid = enclave_manager.get_console_resources_enclave_cid().unwrap();
        let enclave_flags = enclave_manager
            .get_console_resources_enclave_flags()
            .unwrap();
        if let Some(req_enclave_cid) = req_enclave_cid {
            assert_eq!(req_enclave_cid, enclave_cid);
        }

        if debug_mode {
            assert_eq!(enclave_flags & NE_ENCLAVE_DEBUG_MODE, NE_ENCLAVE_DEBUG_MODE);
        } else {
            assert_eq!(enclave_flags & NE_ENCLAVE_DEBUG_MODE, 0);
        }

        let cid_copy = enclave_cid;

        let console = Console::new_nonblocking(
            VMADDR_CID_HYPERVISOR,
            u32::try_from(cid_copy).unwrap() + CID_TO_CONSOLE_PORT_OFFSET,
        )
        .expect("Failed to connect to the console");
        let mut buffer: Vec<u8> = Vec::new();
        let duration: Duration = Duration::from_secs(MAX_BOOT_TIMEOUT_SEC);
        console
            .read_to_buffer(&mut buffer, duration)
            .expect("Failed to check that the enclave booted");

        let contents = String::from_utf8(buffer).unwrap();
        let boot = contents.contains("nsm: loading out-of-tree module");

        assert_eq!(boot, true);

        let info = get_enclave_describe_info(&enclave_manager, false).unwrap();
        let replies: Vec<EnclaveDescribeInfo> = vec![info];
        let reply = &replies[0];
        let flags = &reply.flags;

        assert_eq!({ reply.enclave_cid }, enclave_cid);
        assert_eq!(reply.memory_mib, req_mem_size);
        assert_eq!({ reply.cpu_count }, req_nr_cpus);
        assert_eq!(reply.state, "RUNNING");
        if debug_mode {
            assert_eq!(flags, "DEBUG_MODE");
        } else {
            assert_eq!(flags, "NONE");
        }
        let _enclave_id = generate_enclave_id(0).expect("Describe enclaves failed");

        terminate_enclaves(&mut enclave_manager, None).expect("Terminate enclaves failed");

        let info = get_enclave_describe_info(&enclave_manager, false).unwrap();

        assert_eq!(info.enclave_cid, 0);
        assert_eq!(info.cpu_count, 0);
        assert_eq!(info.memory_mib, 0);
    }

    #[test]
    fn build_run_describe_terminate_simple_eif_image() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");
        setup_env();
        let build_args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &build_args.docker_uri,
            &build_args.docker_dir,
            &build_args.output,
            &build_args.signing_certificate,
            &build_args.private_key,
            &build_args.img_name,
            &build_args.img_version,
            &build_args.metadata,
        )
        .expect("Docker build failed");

        let run_args = RunEnclavesArgs {
            enclave_cid: None,
            eif_path: build_args.output,
            cpu_ids: None,
            cpu_count: Some(2),
            memory_mib: 128,
            debug_mode: true,
            attach_console: false,
            enclave_name: Some("testName".to_string()),
        };

        run_describe_terminate(run_args);
    }

    #[test]
    fn console_without_debug_mode() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");
        setup_env();
        let build_args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &build_args.docker_uri,
            &build_args.docker_dir,
            &build_args.output,
            &build_args.signing_certificate,
            &build_args.private_key,
            &build_args.img_name,
            &build_args.img_version,
            &build_args.metadata,
        )
        .expect("Docker build failed");

        let run_args = RunEnclavesArgs {
            enclave_cid: None,
            eif_path: build_args.output,
            cpu_ids: None,
            cpu_count: Some(2),
            memory_mib: 128,
            debug_mode: false,
            attach_console: false,
            enclave_name: Some("testName".to_string()),
        };

        let mut enclave_manager = run_enclaves(&run_args, None)
            .expect("Run enclaves failed")
            .enclave_manager;
        let enclave_cid = enclave_manager.get_console_resources_enclave_cid().unwrap();
        let enclave_flags = enclave_manager
            .get_console_resources_enclave_flags()
            .unwrap();

        if run_args.debug_mode {
            assert_eq!(enclave_flags & NE_ENCLAVE_DEBUG_MODE, NE_ENCLAVE_DEBUG_MODE);
        } else {
            assert_eq!(enclave_flags & NE_ENCLAVE_DEBUG_MODE, 0);
        };

        let info = get_enclave_describe_info(&enclave_manager, false).unwrap();
        let replies: Vec<EnclaveDescribeInfo> = vec![info];
        let _reply = &replies[0];

        assert_eq!(enclave_console(enclave_cid, None).is_err(), true);

        terminate_enclaves(&mut enclave_manager, None).expect("Terminate enclaves failed");
    }

    #[test]
    fn console_multiple_connect() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");
        setup_env();
        let build_args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &build_args.docker_uri,
            &build_args.docker_dir,
            &build_args.output,
            &build_args.signing_certificate,
            &build_args.private_key,
            &build_args.img_name,
            &build_args.img_version,
            &build_args.metadata,
        )
        .expect("Docker build failed");

        let run_args = RunEnclavesArgs {
            enclave_cid: None,
            eif_path: build_args.output,
            cpu_ids: None,
            cpu_count: Some(2),
            memory_mib: 128,
            debug_mode: true,
            attach_console: false,
            enclave_name: Some("testName".to_string()),
        };

        let mut enclave_manager = run_enclaves(&run_args, None)
            .expect("Run enclaves failed")
            .enclave_manager;
        let enclave_cid = enclave_manager.get_console_resources_enclave_cid().unwrap();
        let enclave_flags = enclave_manager
            .get_console_resources_enclave_flags()
            .unwrap();

        if run_args.debug_mode {
            assert_eq!(enclave_flags & NE_ENCLAVE_DEBUG_MODE, NE_ENCLAVE_DEBUG_MODE);
        } else {
            assert_eq!(enclave_flags & NE_ENCLAVE_DEBUG_MODE, 0);
        }

        let info = get_enclave_describe_info(&enclave_manager, false).unwrap();
        let replies: Vec<EnclaveDescribeInfo> = vec![info];
        let _reply = &replies[0];

        for _ in 0..3 {
            let console = Console::new(
                VMADDR_CID_HYPERVISOR,
                u32::try_from(enclave_cid).unwrap() + CID_TO_CONSOLE_PORT_OFFSET,
            )
            .expect("Failed to connect to the console");

            drop(console);

            std::thread::sleep(std::time::Duration::from_secs(2));
        }

        terminate_enclaves(&mut enclave_manager, None).expect("Terminate enclaves failed");
    }

    #[test]
    fn run_describe_terminate_simple_docker_image_loop() {
        for _ in 0..5 {
            run_describe_terminate_simple_docker_image();
        }
    }

    #[test]
    fn run_describe_terminate_loop() {
        for _ in 0..3 {
            run_describe_terminate_command_executer_docker_image();
            run_describe_terminate_simple_docker_image();
            run_describe_terminate_signed_enclave_image();
            run_describe_terminate_command_executer_docker_image();
            run_describe_terminate_signed_enclave_image();
        }
    }

    #[test]
    fn build_run_save_pcrs_describe() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");
        setup_env();
        let args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &args.docker_uri,
            &args.docker_dir,
            &args.output,
            &args.signing_certificate,
            &args.private_key,
            &args.img_name,
            &args.img_version,
            &args.metadata,
        )
        .expect("Docker build failed");

        setup_env();
        let run_args = RunEnclavesArgs {
            enclave_cid: None,
            eif_path: args.output,
            cpu_ids: None,
            cpu_count: Some(2),
            memory_mib: 128,
            debug_mode: true,
            attach_console: false,
            enclave_name: Some("testName".to_string()),
        };
        let run_result = run_enclaves(&run_args, None).expect("Run enclaves failed");
        let mut enclave_manager = run_result.enclave_manager;
        let mut describe_thread = run_result.describe_thread;

        assert!(describe_thread.is_some());

        let thread_result = describe_thread
            .take()
            .unwrap()
            .join()
            .expect("Failed to join thread.")
            .expect("Failed to save PCRs.");

        enclave_manager
            .set_measurements(thread_result.measurements)
            .expect("Failed to set measurements inside enclave handle.");

        get_enclave_describe_info(&enclave_manager, false).unwrap();
        let build_info = enclave_manager.get_measurements().unwrap();
        let measurements = build_info.measurements;

        assert_eq!(
            measurements.get("PCR0").unwrap(),
            sample_docker_pcrs::IMAGE_PCR
        );
        assert_eq!(
            measurements.get("PCR1").unwrap(),
            sample_docker_pcrs::KERNEL_PCR
        );
        assert_eq!(
            measurements.get("PCR2").unwrap(),
            sample_docker_pcrs::APP_PCR
        );

        let _enclave_id = generate_enclave_id(0).expect("Describe enclaves failed");
        terminate_enclaves(&mut enclave_manager, None).expect("Terminate enclaves failed");
    }

    fn create_metadata_json(dir: &TempDir) {
        let file_path = dir.path().join("meta.json");
        let mut meta_file = File::create(file_path).unwrap();
        let content = json!({
            "AppVersion": "3.2",
            "TestField": "Some info",
            "CustomField": "Added by user",
        });
        let json_bytes = serde_json::to_vec(&content).unwrap();
        meta_file.write_all(&json_bytes[..]).unwrap();
    }

    #[test]
    fn build_with_metadata_run_describe() {
        let dir = tempdir().unwrap();
        create_metadata_json(&dir);
        let eif_path = dir.path().join("test.eif");
        let meta_path = dir.path().join("meta.json");
        setup_env();
        let args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: Some("TestName".to_string()),
            img_version: Some("1.0".to_string()),
            metadata: Some(meta_path.to_str().unwrap().to_string()),
        };

        build_from_docker(
            &args.docker_uri,
            &args.docker_dir,
            &args.output,
            &args.signing_certificate,
            &args.private_key,
            &args.img_name,
            &args.img_version,
            &args.metadata,
        )
        .expect("Docker build failed");

        setup_env();
        let run_args = RunEnclavesArgs {
            enclave_cid: None,
            eif_path: args.output,
            cpu_ids: None,
            cpu_count: Some(2),
            memory_mib: 128,
            debug_mode: true,
            attach_console: false,
            enclave_name: Some("testName".to_string()),
        };
        let run_result = run_enclaves(&run_args, None).expect("Run enclaves failed");
        let mut enclave_manager = run_result.enclave_manager;
        let mut describe_thread = run_result.describe_thread;

        assert!(describe_thread.is_some());

        let thread_result = describe_thread
            .take()
            .unwrap()
            .join()
            .expect("Failed to join thread.")
            .expect("Failed to save PCRs.");

        enclave_manager
            .set_measurements(thread_result.measurements)
            .expect("Failed to set measurements inside enclave handle.");
        let metadata = thread_result.metadata.expect("Failed to fetch metadata");
        enclave_manager
            .set_metadata(metadata.clone())
            .expect("Failed to set metadata inside enclave handle.");

        assert_eq!(metadata.build_info.img_os, "Linux");
        assert_eq!(metadata.build_info.build_tool, env!("CARGO_PKG_NAME"));
        assert_eq!(
            metadata.build_info.build_tool_version,
            env!("CARGO_PKG_VERSION")
        );

        assert_eq!(
            *metadata
                .docker_info
                .get("RepoTags")
                .unwrap()
                .get(0)
                .unwrap(),
            json!("public.ecr.aws/aws-nitro-enclaves/hello:v1")
        );

        assert_eq!(
            *metadata.custom_info.get("AppVersion").unwrap(),
            json!("3.2")
        );
        assert_eq!(
            *metadata.custom_info.get("TestField").unwrap(),
            json!("Some info")
        );
        assert_eq!(
            *metadata.custom_info.get("CustomField").unwrap(),
            json!("Added by user")
        );

        let _enclave_id = generate_enclave_id(0).expect("Describe enclaves failed");
        terminate_enclaves(&mut enclave_manager, None).expect("Terminate enclaves failed");
    }

    #[test]
    fn build_run_default_enclave_name() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");
        setup_env();
        let args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &args.docker_uri,
            &args.docker_dir,
            &args.output,
            &args.signing_certificate,
            &args.private_key,
            &args.img_name,
            &args.img_version,
            &args.metadata,
        )
        .expect("Docker build failed");

        setup_env();
        let mut run_args = RunEnclavesArgs {
            enclave_cid: None,
            eif_path: args.output,
            cpu_ids: None,
            cpu_count: Some(2),
            memory_mib: 128,
            debug_mode: true,
            attach_console: false,
            enclave_name: None,
        };
        let names = Vec::new();
        run_args.enclave_name =
            Some(new_enclave_name(run_args.clone(), names).expect("Failed to set new name."));
        let run_result = run_enclaves(&run_args, None).expect("Run enclaves failed");
        let mut enclave_manager = run_result.enclave_manager;

        get_enclave_describe_info(&enclave_manager, false).unwrap();
        let enclave_name = enclave_manager.enclave_name.clone();

        // Assert that EIF name has been set
        assert_eq!(enclave_name, "test");

        terminate_enclaves(&mut enclave_manager, None).expect("Terminate enclaves failed");
    }

    #[test]
    fn new_enclave_names() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");

        let mut run_args = RunEnclavesArgs {
            enclave_cid: None,
            eif_path: eif_path.to_str().unwrap().to_string(),
            cpu_ids: None,
            cpu_count: Some(2),
            memory_mib: 128,
            debug_mode: true,
            attach_console: false,
            enclave_name: Some("enclaveName".to_string()),
        };
        let mut names = Vec::new();
        let name =
            new_enclave_name(run_args.clone(), names.clone()).expect("Failed to set new name.");
        names.push(name);

        run_args.enclave_name = Some("enclaveNameOther".to_string());
        let name =
            new_enclave_name(run_args.clone(), names.clone()).expect("Failed to set new name.");
        names.push(name);

        run_args.enclave_name = Some("enclaveName".to_string());
        let name =
            new_enclave_name(run_args.clone(), names.clone()).expect("Failed to set new name.");
        names.push(name);

        run_args.enclave_name = Some("enclaveName".to_string());
        let name = new_enclave_name(run_args, names.clone()).expect("Failed to set new name.");
        names.push(name);

        assert_eq!(
            names,
            vec![
                "enclaveName",
                "enclaveNameOther",
                "enclaveName_1",
                "enclaveName_2"
            ]
        );
    }

    #[test]
    fn build_describe_simple_eif() {
        let dir = tempdir().unwrap();
        let eif_path = dir.path().join("test.eif");
        setup_env();
        let args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path.to_str().unwrap().to_string(),
            signing_certificate: None,
            private_key: None,
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &args.docker_uri,
            &args.docker_dir,
            &args.output,
            &args.signing_certificate,
            &args.private_key,
            &args.img_name,
            &args.img_version,
            &args.metadata,
        )
        .expect("Docker build failed");

        let eif_info = describe_eif(args.output).unwrap();

        assert_eq!(eif_info.version, 4);
        assert_eq!(eif_info.is_signed, false);
        assert!(eif_info.cert_info.is_none());
        assert!(eif_info.crc_check);
        assert!(eif_info.sign_check.is_none());
    }

    #[test]
    fn build_describe_signed_simple_eif() {
        let dir = tempdir().unwrap();
        let dir_path = dir.path().to_str().unwrap();
        let eif_path = format!("{}/test.eif", dir_path);
        let cert_path = format!("{}/cert.pem", dir_path);
        let key_path = format!("{}/key.pem", dir_path);
        generate_signing_cert_and_key(&cert_path, &key_path);

        setup_env();
        let args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path,
            signing_certificate: Some(cert_path),
            private_key: Some(key_path),
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &args.docker_uri,
            &args.docker_dir,
            &args.output,
            &args.signing_certificate,
            &args.private_key,
            &args.img_name,
            &args.img_version,
            &args.metadata,
        )
        .expect("Docker build failed");

        let eif_info = describe_eif(args.output).unwrap();

        assert_eq!(eif_info.version, 4);
        assert_eq!(eif_info.is_signed, true);
        assert!(eif_info.cert_info.is_some());
        assert!(eif_info.crc_check);
        assert!(eif_info.sign_check.unwrap());
    }

    #[test]
    fn get_certificate_pcr() {
        let dir = tempdir().unwrap();
        let dir_path = dir.path().to_str().unwrap();
        let eif_path = format!("{}/test.eif", dir_path);
        let cert_path = format!("{}/cert.pem", dir_path);
        let key_path = format!("{}/key.pem", dir_path);
        generate_signing_cert_and_key(&cert_path, &key_path);

        setup_env();
        let args = BuildEnclavesArgs {
            docker_uri: SAMPLE_DOCKER.to_string(),
            docker_dir: None,
            output: eif_path,
            signing_certificate: Some(cert_path.clone()),
            private_key: Some(key_path),
            img_name: None,
            img_version: None,
            metadata: None,
        };

        build_from_docker(
            &args.docker_uri,
            &args.docker_dir,
            &args.output,
            &args.signing_certificate,
            &args.private_key,
            &args.img_name,
            &args.img_version,
            &args.metadata,
        )
        .expect("Docker build failed");

        // Describe EIF and get PCR8
        let eif_info = describe_eif(args.output).unwrap();
        // Hash signing certificate and verify that PCR8 is the same (identifying the certificate)
        let pcr = get_file_pcr(cert_path, PcrType::SigningCertificate).unwrap();

        assert_eq!(
            eif_info
                .build_info
                .measurements
                .get(&"PCR8".to_string())
                .unwrap(),
            pcr.get(&"PCR8".to_string()).unwrap(),
        );
    }
}


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

use libc::{VMADDR_CID_HOST, VMADDR_CID_LOCAL};
use std::fs::File;
use std::os::unix::io::{AsRawFd, RawFd};
use std::process::Command;

use nitro_cli::common::{NitroCliErrorEnum, NitroCliFailure, NitroCliResult};
use nitro_cli::enclave_proc::cpu_info::CpuInfo;
use nitro_cli::enclave_proc::resource_manager::{
    EnclaveStartInfo, MemoryRegion, NE_ADD_VCPU, NE_CREATE_VM, NE_SET_USER_MEMORY_REGION,
    NE_START_ENCLAVE,
};
use nitro_cli::enclave_proc::utils::MiB;

const ENCLAVE_MEM_2MB_CHUNKS: u64 = 48;
#[cfg(target_arch = "aarch64")]
const ENCLAVE_MEM_32MB_CHUNKS: u64 = 3;
pub const NE_DEVICE_PATH: &str = "/dev/nitro_enclaves";

/// This is similar to `MemoryRegion`, except it doesn't implement `Drop`.
#[allow(dead_code)]
pub struct EnclaveMemoryRegion {
    /// Flags to determine the usage for the memory region.
    flags: u64,
    /// The region's size in bytes.
    mem_size: u64,
    /// The region's virtual address.
    mem_addr: u64,
}

impl EnclaveMemoryRegion {
    fn new(flags: u64, mem_addr: u64, mem_size: u64) -> Self {
        EnclaveMemoryRegion {
            flags,
            mem_size,
            mem_addr,
        }
    }

    fn new_from(region: &MemoryRegion) -> Self {
        EnclaveMemoryRegion {
            flags: 0,
            mem_size: region.mem_size(),
            mem_addr: region.mem_addr(),
        }
    }
}

/// Class that covers communication with the NE driver.
pub struct NitroEnclavesDeviceDriver {
    // NE device file.
    file: File,
}

impl NitroEnclavesDeviceDriver {
    /// Open the file descriptor for communicating with the NE driver.
    pub fn new() -> NitroCliResult<Self> {
        Ok(NitroEnclavesDeviceDriver {
            file: File::open(NE_DEVICE_PATH).map_err(|e| {
                NitroCliFailure::new()
                    .add_subaction(format!("Could not open {}: {}", NE_DEVICE_PATH, e))
                    .set_error_code(NitroCliErrorEnum::FileOperationFailure)
                    .set_file_and_line(file!(), line!())
                    .add_info(vec![NE_DEVICE_PATH, "Open"])
            })?,
        })
    }

    /// Allocate an enclave slot and return an enclave fd.
    pub fn create_enclave(&mut self) -> NitroCliResult<NitroEnclave> {
        let mut slot_uid: u64 = 0;
        // This is safe because we are providing valid values.
        let enc_fd =
            unsafe { libc::ioctl(self.file.as_raw_fd(), NE_CREATE_VM as _, &mut slot_uid) };

        if enc_fd < 0 {
            return Err(NitroCliFailure::new()
                .add_subaction(format!(
                    "Could not create an enclave descriptor: {}",
                    enc_fd
                ))
                .set_error_code(NitroCliErrorEnum::IoctlFailure)
                .set_file_and_line(file!(), line!()));
        }

        if slot_uid == 0 {
            return Err(NitroCliFailure::new()
                .add_subaction("Obtained invalid slot ID".to_string())
                .set_error_code(NitroCliErrorEnum::IoctlFailure)
                .set_file_and_line(file!(), line!()));
        }

        Ok(NitroEnclave::new(enc_fd).unwrap())
    }
}

/// Class for managing a Nitro Enclave provided by NitroEnclavesDeviceDriver.
pub struct NitroEnclave {
    enc_fd: RawFd,
}

impl NitroEnclave {
    pub fn new(enc_fd: RawFd) -> NitroCliResult<Self> {
        Ok(NitroEnclave { enc_fd })
    }

    fn release(&mut self) {
        // Close enclave descriptor.
        let rc = unsafe { libc::close(self.enc_fd) };
        if rc < 0 {
            panic!("Could not close enclave descriptor: {}.", rc)
        }
    }

    pub fn add_mem_region(&mut self, mem_region: EnclaveMemoryRegion) -> NitroCliResult<()> {
        let rc = unsafe { libc::ioctl(self.enc_fd, NE_SET_USER_MEMORY_REGION as _, &mem_region) };
        if rc < 0 {
            return Err(NitroCliFailure::new()
                .add_subaction(format!("Could not add memory region: {}", rc))
                .set_error_code(NitroCliErrorEnum::IoctlSetMemoryRegionFailure)
                .set_file_and_line(file!(), line!()));
        }

        Ok(())
    }

    pub fn add_cpu(&mut self, cpu_id: u32) -> NitroCliResult<()> {
        let mut actual_cpu_id: u32 = cpu_id;
        let rc = unsafe { libc::ioctl(self.enc_fd, NE_ADD_VCPU as _, &mut actual_cpu_id) };
        if rc < 0 {
            return Err(NitroCliFailure::new()
                .add_subaction(format!("Could not add vCPU: {}", rc))
                .set_error_code(NitroCliErrorEnum::IoctlAddVcpuFailure)
                .set_file_and_line(file!(), line!()));
        }

        Ok(())
    }

    pub fn start(&mut self, start_info: EnclaveStartInfo) -> NitroCliResult<()> {
        let rc = unsafe { libc::ioctl(self.enc_fd, NE_START_ENCLAVE as _, &start_info) };
        if rc < 0 {
            return Err(NitroCliFailure::new()
                .add_subaction(format!("Could not start enclave: {}", rc))
                .set_error_code(NitroCliErrorEnum::IoctlEnclaveStartFailure)
                .set_file_and_line(file!(), line!()));
        }

        Ok(())
    }
}

impl Drop for NitroEnclave {
    fn drop(&mut self) {
        if self.enc_fd < 0 {
            return;
        }
        self.release();
    }
}

// Class for checking the dmesg logs.
pub struct CheckDmesg {
    recorded_line: usize,
}

impl CheckDmesg {
    pub fn new() -> NitroCliResult<Self> {
        Ok(CheckDmesg { recorded_line: 0 })
    }

    /// Obtain the log lines from dmesg.
    fn get_dmesg_lines(&mut self) -> NitroCliResult<Vec<String>> {
        let dmesg = Command::new("dmesg")
            .output()
            .expect("Failed to execute dmesg process");
        let message = String::from_utf8(dmesg.stdout).unwrap();
        let lines: Vec<String> = message.split('\n').map(|s| s.to_string()).collect();
        Ok(lines)
    }

    /// Record the current number of lines from dmesg.
    pub fn record_current_line(&mut self) -> NitroCliResult<()> {
        self.recorded_line = self.get_dmesg_lines().unwrap().len();
        Ok(())
    }

    /// Verify if dmesg number of lines changed from the last recorded line.
    pub fn expect_no_changes(&mut self) -> NitroCliResult<()> {
        let checks = vec![
            "WARNING",
            "BUG",
            "ERROR",
            "FAILURE",
            "nitro_enclaves",
            // NE PCI device identifier
            "pci 0000:00:02.0",
        ];
        let lines = self.get_dmesg_lines().unwrap();

        for i in self.recorded_line..lines.len() {
            let upper_line = lines[i].to_uppercase();
            for word in checks.iter() {
                if upper_line.contains(&word.to_uppercase()) {
                    return Err(NitroCliFailure::new()
                        .add_subaction(format!("Dmesg line: {} contains: {}", lines[i], word))
                        .set_error_code(NitroCliErrorEnum::IoctlFailure)
                        .set_file_and_line(file!(), line!()));
                }
            }
        }

        Ok(())
    }
}

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

    #[test]
    pub fn test_ne_dev_open() {
        let mut driver = NitroEnclavesDeviceDriver::new().expect("Failed to open NE device");
        let enc_fd = unsafe { libc::ioctl(driver.file.as_raw_fd(), NE_CREATE_VM as _, 0) };
        assert!(
            enc_fd < 0,
            "Should not have been able to create enclave descriptor"
        );

        // Test unexpected ioctl.
        let enc_fd =
            unsafe { libc::ioctl(driver.file.as_raw_fd(), NE_SET_USER_MEMORY_REGION as _, 0) };
        assert!(
            enc_fd < 0,
            "Should not have been able to create enclave with invalid ioctl"
        );

        let mut slot_alloc_num: u64 = 1;
        if let Ok(value) = std::env::var("NE_SLOT_ALLOC_NUM") {
            if let Ok(value) = value.parse::<u64>() {
                slot_alloc_num = value;
            }
        }

        let mut check_dmesg = CheckDmesg::new().expect("Failed to obtain dmesg object");
        check_dmesg
            .record_current_line()
            .expect("Failed to record current line");

        for _i in 0..slot_alloc_num {
            // Allocate Nitro Enclave slot and free it.
            let _enclave = driver.create_enclave().unwrap();
        }

        check_dmesg.expect_no_changes().unwrap();
    }

    #[test]
    pub fn test_enclave_memory() {
        let mut driver = NitroEnclavesDeviceDriver::new().expect("Failed to open NE device");
        let mut enclave = driver.create_enclave().unwrap();

        // Add invalid memory region.
        let result = enclave.add_mem_region(EnclaveMemoryRegion::new(0, 0, 2 * MiB as u64));
        assert_eq!(result.is_err(), true);

        // Create a memory region using hugetlbfs.
        let region = MemoryRegion::new(libc::MAP_HUGE_2MB).unwrap();

        // Add unaligned memory region.
        let result = enclave.add_mem_region(EnclaveMemoryRegion::new(
            0,
            region.mem_addr() + 1,
            region.mem_size(),
        ));
        assert_eq!(result.is_err(), true);

        // Add wrongly sized memory region of 1 MiB.
        let result = enclave.add_mem_region(EnclaveMemoryRegion::new(
            0,
            region.mem_addr(),
            region.mem_size() / 2,
        ));
        assert_eq!(result.is_err(), true);

        // Add wrongly sized memory region of double the memory size.
        let result = enclave.add_mem_region(EnclaveMemoryRegion::new(
            0,
            region.mem_addr(),
            region.mem_size() * 2,
        ));
        assert_eq!(result.is_err(), true);

        // Add wrongly sized memory region of max value multiple of 2 MiB.
        let result = enclave.add_mem_region(EnclaveMemoryRegion::new(
            0,
            region.mem_addr(),
            u64::max_value() - (2 * 1024 * 1024) + 1,
        ));
        assert_eq!(result.is_err(), true);

        // Add wrong memory region with address out of range.
        let result = enclave.add_mem_region(EnclaveMemoryRegion::new(
            0,
            region.mem_addr() + region.mem_size(),
            region.mem_size(),
        ));
        assert_eq!(result.is_err(), true);

        let mut check_dmesg = CheckDmesg::new().expect("Failed to obtain dmesg object");
        check_dmesg
            .record_current_line()
            .expect("Failed to record current line");

        // Correctly add the memory region.
        let region = MemoryRegion::new(libc::MAP_HUGE_2MB).unwrap();
        let result = enclave.add_mem_region(EnclaveMemoryRegion::new_from(&region));
        assert_eq!(result.is_err(), false);

        check_dmesg.expect_no_changes().unwrap();

        // Add the same memory region twice.
        let result = enclave.add_mem_region(EnclaveMemoryRegion::new_from(&region));
        assert_eq!(result.is_err(), true);

        // Add a memory region with invalid flags.
        let region = MemoryRegion::new(libc::MAP_HUGE_2MB).unwrap();
        let result = enclave.add_mem_region(EnclaveMemoryRegion::new(
            1024,
            region.mem_addr(),
            region.mem_size(),
        ));
        assert_eq!(result.is_err(), true);
    }

    #[test]
    pub fn test_enclave_vcpu() {
        let mut driver = NitroEnclavesDeviceDriver::new().expect("Failed to open NE device");
        let mut enclave = driver.create_enclave().unwrap();
        let cpu_info = CpuInfo::new().expect("Failed to obtain CpuInfo.");

        // Add an invalid cpu id.
        let result = enclave.add_cpu(u32::max_value());
        assert_eq!(result.is_err(), true);

        let mut candidates = cpu_info.get_cpu_candidates();
        // Instance does not have the appropriate number of cpus.
        if candidates.is_empty() {
            return;
        }

        let cpu_id = candidates.pop().unwrap();

        let mut check_dmesg = CheckDmesg::new().expect("Failed to obtain dmesg object");
        check_dmesg
            .record_current_line()
            .expect("Failed to record current line");

        // Insert the first valid cpu id.
        let result = enclave.add_cpu(cpu_id);
        assert_eq!(result.is_err(), false);

        check_dmesg.expect_no_changes().unwrap();

        // Try inserting the cpu twice.
        let result = enclave.add_cpu(cpu_id);
        assert_eq!(result.is_err(), true);

        check_dmesg
            .record_current_line()
            .expect("Failed to record current line");

        // Add all remaining cpus.
        for cpu in &candidates {
            let result = enclave.add_cpu(*cpu);
            assert_eq!(result.is_err(), false);
        }

        check_dmesg.expect_no_changes().unwrap();

        // Clear the enclave.
        drop(enclave);

        let mut enclave = driver.create_enclave().unwrap();

        check_dmesg
            .record_current_line()
            .expect("Failed to record current line");

        // Add an auto-chosen cpu from the pool.
        let result = enclave.add_cpu(0);
        assert_eq!(result.is_err(), false);

        check_dmesg.expect_no_changes().unwrap();

        check_dmesg
            .record_current_line()
            .expect("Failed to record current line");

        // Add all remaining auto-chosen cpus.
        for _i in 0..candidates.len() {
            let result = enclave.add_cpu(0);
            assert_eq!(result.is_err(), false);
        }

        check_dmesg.expect_no_changes().unwrap();

        // Add one more cpu than the maximum available in the pool.
        let result = enclave.add_cpu(0);
        assert_eq!(result.is_err(), true);
    }

    #[test]
    pub fn test_enclave_start() {
        let mut mem_regions = Vec::new();
        let mut driver = NitroEnclavesDeviceDriver::new().expect("Failed to open NE device");
        let mut enclave = driver.create_enclave().unwrap();

        // Start enclave without resources.
        let result = enclave.start(EnclaveStartInfo::default());
        assert_eq!(result.is_err(), true);

        // Allocate memory for the enclave.
        #[cfg(target_arch = "x86_64")]
        for _i in 0..ENCLAVE_MEM_2MB_CHUNKS {
            mem_regions.push(MemoryRegion::new(libc::MAP_HUGE_2MB).unwrap());
        }

        #[cfg(target_arch = "aarch64")]
        {
            let mut mem_2mb_chunks = ENCLAVE_MEM_2MB_CHUNKS;

            for _i in 0..ENCLAVE_MEM_32MB_CHUNKS {
                let region = MemoryRegion::new(libc::MAP_HUGE_32MB);

                if region.is_err() {
                    break;
                }

                mem_regions.push(region.unwrap());

                mem_2mb_chunks = mem_2mb_chunks - (32 / 2);
            }

            for _i in 0..mem_2mb_chunks {
                mem_regions.push(MemoryRegion::new(libc::MAP_HUGE_2MB).unwrap());
            }
        }

        // Add memory to the enclave.
        for region in &mut mem_regions {
            let result = enclave.add_mem_region(EnclaveMemoryRegion::new_from(region));
            assert_eq!(result.is_err(), false);
        }

        // Start the enclave without cpus.
        let result = enclave.start(EnclaveStartInfo::default());
        assert_eq!(result.is_err(), true);

        let cpu_info = CpuInfo::new().expect("Failed to obtain CpuInfo.");
        let candidates = cpu_info.get_cpu_candidates();
        // Instance does not have the appropriate number of cpus.
        if candidates.len() < 2 {
            return;
        }

        // Clear the enclave.
        drop(enclave);

        let mut enclave = driver.create_enclave().unwrap();

        for cpu in &candidates {
            let result = enclave.add_cpu(*cpu);
            assert_eq!(result.is_err(), false);
        }

        // Start enclave without memory.
        let result = enclave.start(EnclaveStartInfo::default());
        assert_eq!(result.is_err(), true);

        drop(enclave);

        let mut enclave = driver.create_enclave().unwrap();

        // Add memory to the enclave.
        for region in &mut mem_regions {
            let result = enclave.add_mem_region(EnclaveMemoryRegion::new_from(region));
            assert_eq!(result.is_err(), false);
        }

        // Add the first available cpu.
        let result = enclave.add_cpu(candidates[0]);
        assert_eq!(result.is_err(), false);

        // Start without cpu pair.
        #[cfg(target_arch = "aarch64")]
        let mut check_dmesg = CheckDmesg::new().expect("Failed to obtain dmesg object");
        #[cfg(target_arch = "aarch64")]
        check_dmesg
            .record_current_line()
            .expect("Failed to record current line");

        let result = enclave.start(EnclaveStartInfo::default());
        #[cfg(target_arch = "x86_64")]
        assert_eq!(result.is_err(), true);
        #[cfg(target_arch = "aarch64")]
        assert_eq!(result.is_err(), false);

        #[cfg(target_arch = "aarch64")]
        check_dmesg.expect_no_changes().unwrap();

        #[cfg(target_arch = "aarch64")]
        drop(enclave);

        #[cfg(target_arch = "aarch64")]
        let mut enclave = driver.create_enclave().unwrap();

        // Add memory to the enclave.
        #[cfg(target_arch = "aarch64")]
        for region in &mut mem_regions {
            let result = enclave.add_mem_region(EnclaveMemoryRegion::new_from(region));
            assert_eq!(result.is_err(), false);
        }

        // Add the first available cpu.
        #[cfg(target_arch = "aarch64")]
        let result = enclave.add_cpu(candidates[0]);
        #[cfg(target_arch = "aarch64")]
        assert_eq!(result.is_err(), false);

        // Add the first cpu pair.
        let result = enclave.add_cpu(candidates[1]);
        assert_eq!(result.is_err(), false);

        // Start with an invalid flag.
        let mut enclave_start_info = EnclaveStartInfo::default();
        enclave_start_info.flags = 1234;
        let result = enclave.start(enclave_start_info);
        assert_eq!(result.is_err(), true);

        // Start with an invalid CID.
        let mut enclave_start_info = EnclaveStartInfo::default();
        enclave_start_info.enclave_cid = VMADDR_CID_LOCAL as u64;
        let result = enclave.start(enclave_start_info);
        assert_eq!(result.is_err(), true);

        enclave_start_info.enclave_cid = VMADDR_CID_HOST as u64;
        let result = enclave.start(enclave_start_info);
        assert_eq!(result.is_err(), true);

        enclave_start_info.enclave_cid = u32::max_value() as u64;
        let result = enclave.start(enclave_start_info);
        assert_eq!(result.is_err(), true);

        enclave_start_info.enclave_cid = u32::max_value() as u64 + 1234_u64;
        let result = enclave.start(enclave_start_info);
        assert_eq!(result.is_err(), true);

        let mut check_dmesg = CheckDmesg::new().expect("Failed to obtain dmesg object");
        check_dmesg
            .record_current_line()
            .expect("Failed to record current line");

        // Start the enclave.
        let result = enclave.start(EnclaveStartInfo::default());
        assert_eq!(result.is_err(), false);

        check_dmesg.expect_no_changes().unwrap();

        // Try starting an already running enclave.
        let result = enclave.start(EnclaveStartInfo::default());
        assert_eq!(result.is_err(), true);

        // Try adding an already added memory region
        // after the enclave start.
        let result = enclave.add_mem_region(EnclaveMemoryRegion::new_from(&mem_regions[0]));
        assert_eq!(result.is_err(), true);

        // Try adding a new memory region after the enclave start.
        let result = enclave.add_mem_region(EnclaveMemoryRegion::new_from(
            &MemoryRegion::new(libc::MAP_HUGE_2MB).unwrap(),
        ));
        assert_eq!(result.is_err(), true);

        // Try adding an already added vcpu after enclave start.
        let result = enclave.add_cpu(candidates[0]);
        assert_eq!(result.is_err(), true);

        // Try adding a new vcpu after enclave start.
        if candidates.len() >= 3 {
            let result = enclave.add_cpu(candidates[2]);
            assert_eq!(result.is_err(), true);
        }
    }

    #[test]
    pub fn test_enclave_multiple_start() {
        let mut mem_regions = Vec::new();
        let mut driver = NitroEnclavesDeviceDriver::new().expect("Failed to open NE device");

        // Allocate memory for the enclave.
        #[cfg(target_arch = "x86_64")]
        for _i in 0..ENCLAVE_MEM_2MB_CHUNKS {
            mem_regions.push(MemoryRegion::new(libc::MAP_HUGE_2MB).unwrap());
        }

        #[cfg(target_arch = "aarch64")]
        {
            let mut mem_2mb_chunks = ENCLAVE_MEM_2MB_CHUNKS;

            for _i in 0..ENCLAVE_MEM_32MB_CHUNKS {
                let region = MemoryRegion::new(libc::MAP_HUGE_32MB);

                if region.is_err() {
                    break;
                }

                mem_regions.push(region.unwrap());

                mem_2mb_chunks = mem_2mb_chunks - (32 / 2);
            }

            for _i in 0..mem_2mb_chunks {
                mem_regions.push(MemoryRegion::new(libc::MAP_HUGE_2MB).unwrap());
            }
        }

        let cpu_info = CpuInfo::new().expect("Failed to obtain CpuInfo.");
        let candidates = cpu_info.get_cpu_candidates();
        // Instance does not have the appropriate number of cpus.
        if candidates.len() < 2 {
            return;
        }

        let mut start_num: u64 = 1;
        if let Ok(value) = std::env::var("NE_MULTIPLE_START_NUM") {
            if let Ok(value) = value.parse::<u64>() {
                start_num = value;
            }
        }

        for _i in 0..start_num {
            let mut enclave = driver.create_enclave().unwrap();

            // Add memory to the enclave.
            for region in &mut mem_regions {
                let result = enclave.add_mem_region(EnclaveMemoryRegion::new_from(region));
                assert_eq!(result.is_err(), false);
            }

            // Add cpus to the enclave.
            for cpu in &candidates {
                let result = enclave.add_cpu(*cpu);
                assert_eq!(result.is_err(), false);
            }

            // Start and stop the enclave
            let result = enclave.start(EnclaveStartInfo::default());
            assert_eq!(result.is_err(), false);
        }
    }
}


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

#[cfg(test)]
mod test_nitro_cli_args {
    use clap::{App, AppSettings, Arg, SubCommand};
    use nitro_cli::create_app;

    #[test]
    fn terminate_enclave_enclave_id_arg_is_required() {
        let app = create_app!();
        let args = vec!["nitro cli", "terminate-enclave"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn terminate_enclave_enclave_id_takes_value() {
        let app = create_app!();
        let args = vec!["nitro cli", "terminate-enclave", "--enclave-id"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn terminate_enclave_enclave_id_takes_one_value() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "terminate-enclave",
            "--enclave-id",
            "i-1234_enc123",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn terminate_enclave_enclave_id_takes_multiple_values() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "terminate-enclave",
            "--enclave-id",
            "1234",
            "135",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn terminate_enclave_name() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "terminate-enclave",
            "--enclave-name",
            "testName",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn terminate_enclave_name_is_required() {
        let app = create_app!();
        let args = vec!["nitro cli", "terminate-enclave", "--enclave-name"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn terminate_enclave_name_takes_multiple_values() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "terminate-enclave",
            "--enclave-name",
            "name1",
            "name2",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn describe_enclaves_correct_command() {
        let app = create_app!();
        let args = vec!["nitro cli", "describe-enclaves"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn describe_enclaves_request_metadata_correct() {
        let app = create_app!();
        let args = vec!["nitro cli", "describe-enclaves", "--metadata"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn describe_eif_correct_command() {
        let app = create_app!();
        let args = vec!["nitro cli", "describe-eif", "--eif-path", "dir/image.eif"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn describe_eif_without_path_arg() {
        let app = create_app!();
        let args = vec!["nitro cli", "describe-eif", "--eif-path"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn console_without_enclave_id_arg_is_required() {
        let app = create_app!();
        let args = vec!["nitro cli", "console"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn console_enclave_id_takes_value() {
        let app = create_app!();
        let args = vec!["nitro cli", "console", "--enclave-id"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn console_correct_command() {
        let app = create_app!();
        let args = vec!["nitro cli", "console", "--enclave-id", "i-1234_enc123"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn console_enclave_id_takes_one_value() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "console",
            "--enclave-id",
            "i-1234_enc123",
            "135",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn console_enclave_name() {
        let app = create_app!();
        let args = vec!["nitro cli", "console", "--enclave-name", "testName"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn console_enclave_name_is_required() {
        let app = create_app!();
        let args = vec!["nitro cli", "console", "--enclave-name"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn console_enclave_name_takes_multiple_values() {
        let app = create_app!();
        let args = vec!["nitro cli", "console", "--enclave-name", "name1", "name2"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn console_correct_disconnect_timeout_command() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "console",
            "--enclave-id",
            "i-1234_enc123",
            "--disconnect-timeout",
            "10",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn console_correct_disconnect_timeout_command_with_name() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "console",
            "--enclave-name",
            "testName",
            "--disconnect-timeout",
            "10",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn console_disconnect_timeout_takes_value() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "console",
            "--enclave-id",
            "i-1234_enc123",
            "--disconnect-timeout",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn build_enclave_docker_uri_arg_is_required() {
        let app = create_app!();
        let args = vec!["nitro cli", "build-enclave", "--output-file", "image.eif"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn build_enclave_docker_dir_arg_is_not_required() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "build-enclave",
            "--docker-uri",
            "dkr.ecr.us-east-1.amazonaws.com/stronghold-develss",
            "--output-file",
            "image.eif",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn build_enclave_output_arg_is_required() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "build-enclave",
            "--docker-uri",
            "dkr.ecr.us-east-1.amazonaws.com/stronghold-develss",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn build_enclave_correct_command() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "build-enclave",
            "--docker-uri",
            "dkr.ecr.us-east-1.amazonaws.com/stronghold-develss",
            "--docker-dir",
            "dir/",
            "--output-file",
            "image.eif",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn build_signed_enclave_correct_command() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "build-enclave",
            "--docker-uri",
            "dkr.ecr.us-east-1.amazonaws.com/stronghold-develss",
            "--docker-dir",
            "dir/",
            "--output-file",
            "image.eif",
            "--signing-certificate",
            "cert.pem",
            "--private-key",
            "key.pem",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn build_enclave_with_metadata_correct_command() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "build-enclave",
            "--docker-uri",
            "dkr.ecr.us-east-1.amazonaws.com/stronghold-develss",
            "--docker-dir",
            "dir/",
            "--output-file",
            "image.eif",
            "--name",
            "TestName",
            "--version",
            "4.0",
            "--metadata",
            "meta.json",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn build_enclave_with_metadata_file_is_required() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "build-enclave",
            "--docker-uri",
            "dkr.ecr.us-east-1.amazonaws.com/stronghold-develss",
            "--docker-dir",
            "dir/",
            "--output-file",
            "image.eif",
            "--name",
            "TestName",
            "--version",
            "4.0",
            "--metadata",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn run_enclave_correct_command_with_eif_path() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--cpu-ids",
            "10000",
            "10001",
            "--memory",
            "512",
            "--eif-path",
            "dir/image.eif",
            "--enclave-cid",
            "1234",
            "--debug-mode",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn run_enclave_cpu_ids_arg_is_required() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--memory",
            "512",
            "--eif-path",
            "dir/image.eif",
            "--enclave-cid",
            "12345",
            "--debug-mode",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn run_enclave_cpu_ids_takes_value() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--cpu-ids",
            "--memory",
            "512",
            "--eif-path",
            "dir/image.eif",
            "--enclave-cid",
            "12345",
            "--debug-mode",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn run_enclave_cpu_ids_takes_multiple_values() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--cpu-ids",
            "10000",
            "10001",
            "--memory",
            "512",
            "--eif-path",
            "dir/image.eif",
            "--enclave-cid",
            "12345",
            "--debug-mode",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn run_enclave_memory_arg_is_required() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--cpu-ids",
            "10000",
            "--eif-path",
            "dir/image.eif",
            "--enclave-cid",
            "12345",
            "--debug-mode",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn run_enclave_memory_takes_value() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--cpu-ids",
            "10000",
            "--memory",
            "--eif-path",
            "dir/image.eif",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn run_enclave_enclave_cid_takes_value() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--cpu-ids",
            "10000",
            "--memory",
            "512",
            "--eif-path",
            "dir/image.eif",
            "--enclave-cid",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn run_enclave_console_does_not_take_value() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--cpu-ids",
            "10000",
            "--memory",
            "512",
            "--eif-path",
            "dir/image.eif",
            "--enclave-cid",
            "12345",
            "--debug-mode",
            "123",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn run_enclave_console_eif_path_is_required() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--cpu-ids",
            "10000",
            "--memory",
            "512",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn run_enclave_eif_path_takes_value() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--cpu-ids",
            "10000",
            "10001",
            "--memory",
            "512",
            "--eif-path",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn run_enclave_config_does_not_take_value() {
        let app = create_app!();
        let args = vec!["nitro cli", "run-enclave", "--config"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn run_enclave_config_takes_multiple_values() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--config",
            "config1.json",
            "config2.json",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn run_enclave_try_to_overwrite_config() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--config",
            "config.json",
            "--cpu-count",
            "2",
            "--memory",
            "1024",
            "--eif-path",
            "dir/image.eif",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn run_enclave_correct_command_with_name() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--cpu-ids",
            "10000",
            "10001",
            "--memory",
            "512",
            "--eif-path",
            "dir/image.eif",
            "--enclave-cid",
            "1234",
            "--debug-mode",
            "--enclave-name",
            "testName",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }

    #[test]
    fn run_enclave_name_takes_value() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "run-enclave",
            "--cpu-ids",
            "10000",
            "10001",
            "--memory",
            "512",
            "--eif-path",
            "dir/image.eif",
            "--enclave-cid",
            "1234",
            "--debug-mode",
            "--enclave-name",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn pcr_input_takes_value() {
        let app = create_app!();
        let args = vec!["nitro cli", "pcr", "--input"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn pcr_certificate_takes_value() {
        let app = create_app!();
        let args = vec!["nitro cli", "pcr", "--signing-certificate"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn pcr_conflicting_arguments() {
        let app = create_app!();
        let args = vec![
            "nitro cli",
            "pcr",
            "--signing-certificate",
            "cert.pem",
            "--input",
            "test.bin",
        ];

        assert_eq!(app.get_matches_from_safe(args).is_err(), true)
    }

    #[test]
    fn pcr_certificate_correct() {
        let app = create_app!();
        let args = vec!["nitro cli", "pcr", "--signing-certificate", "cert.pem"];

        assert_eq!(app.get_matches_from_safe(args).is_err(), false)
    }
}


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

pub mod starter;


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

/// Simple proxy for translating vsock traffic to TCP traffic
/// Example of usage:
/// vsock-proxy 8000 127.0.0.1 9000
///
use clap::{App, AppSettings, Arg};
use env_logger::init;
use log::info;

use vsock_proxy::starter::{Proxy, VsockProxyResult};

fn main() -> VsockProxyResult<()> {
    init();

    let matches = App::new("Vsock-TCP proxy")
        .about("Vsock-TCP proxy")
        .setting(AppSettings::DisableVersion)
        .arg(
            Arg::with_name("ipv4")
                .short('4')
                .long("ipv4")
                .help("Force the proxy to use IPv4 addresses only.")
                .required(false),
        )
        .arg(
            Arg::with_name("ipv6")
                .short('6')
                .long("ipv6")
                .help("Force the proxy to use IPv6 addresses only.")
                .required(false)
                .conflicts_with("ipv4"),
        )
        .arg(
            Arg::with_name("workers")
                .short('w')
                .long("num_workers")
                .help("Set the maximum number of simultaneous\nconnections supported.")
                .required(false)
                .takes_value(true)
                .default_value("4"),
        )
        .arg(
            Arg::with_name("local_port")
                .help("Local Vsock port to listen for incoming connections.")
                .required(true),
        )
        .arg(
            Arg::with_name("remote_addr")
                .help("Address of the server to be proxyed.")
                .required(true),
        )
        .arg(
            Arg::with_name("remote_port")
                .help("Remote TCP port of the server to be proxyed.")
                .required(true),
        )
        .arg(
            Arg::with_name("config_file")
                .long("config")
                .help("YAML file containing the services that\ncan be forwarded.\n")
                .required(false)
                .takes_value(true)
                .default_value("/etc/nitro_enclaves/vsock-proxy.yaml"),
        )
        .get_matches();

    let local_port = matches
        .value_of("local_port")
        // This argument is required, so clap ensures it's available
        .unwrap();
    let local_port = local_port
        .parse::<u32>()
        .map_err(|_| "Local port is not valid")?;

    let only_4 = matches.is_present("ipv4");
    let only_6 = matches.is_present("ipv6");
    let remote_addr = matches
        .value_of("remote_addr")
        // This argument is required, so clap ensures it's available
        .unwrap();

    let remote_port = matches
        .value_of("remote_port")
        // This argument is required, so clap ensures it's available
        .unwrap();
    let remote_port = remote_port
        .parse::<u16>()
        .map_err(|_| "Remote port is not valid")?;

    let num_workers = matches
        .value_of("workers")
        // This argument has a default value, so it is available
        .unwrap();
    let num_workers = num_workers
        .parse::<usize>()
        .map_err(|_| "Number of workers is not valid")?;

    let config_file = matches.value_of("config_file");

    let proxy = Proxy::new(
        local_port,
        remote_addr,
        remote_port,
        num_workers,
        config_file,
        only_4,
        only_6,
    )
    .map_err(|err| format!("Could not create proxy: {}", err))?;

    let listener = proxy
        .sock_listen()
        .map_err(|err| format!("Could not listen for connections: {}", err))?;
    info!("Proxy is now in listening state");
    loop {
        proxy
            .sock_accept(&listener)
            .map_err(|err| format!("Could not accept connection: {}", err))?;
    }
}


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

/// Contains code for Proxy, a library used for translating vsock traffic to
/// TCP traffic
///
use dns_lookup::lookup_host;
use idna::domain_to_ascii;
use log::info;
use nix::sys::select::{select, FdSet};
use nix::sys::socket::SockType;
use std::fs::File;
use std::io::{Read, Write};
use std::net::{IpAddr, SocketAddr, TcpStream};
use std::os::unix::io::AsRawFd;
use threadpool::ThreadPool;
use vsock::{VsockAddr, VsockListener};
use yaml_rust::YamlLoader;

const BUFF_SIZE: usize = 8192;
pub const VSOCK_PROXY_CID: u32 = 3;
pub const VSOCK_PROXY_PORT: u32 = 8000;

/// The most common result type provided by VsockProxy operations.
pub type VsockProxyResult<T> = Result<T, String>;

/// Checks if the forwarded server is allowed, providing its IP on success.
pub fn check_allowlist(
    remote_host: &str,
    remote_port: u16,
    config_file: Option<&str>,
    only_4: bool,
    only_6: bool,
) -> VsockProxyResult<IpAddr> {
    if let Some(config_file) = config_file {
        let mut f = File::open(config_file).map_err(|_| "Could not open the file")?;

        let mut content = String::new();
        f.read_to_string(&mut content)
            .map_err(|_| "Could not read the file")?;

        let docs = YamlLoader::load_from_str(&content).map_err(|_| "Bad yaml format")?;
        let services = (&docs[0])["allowlist"]
            .as_vec()
            .ok_or("No allowlist field")?;

        // Obtain the remote server's IP address.
        let mut addrs = Proxy::parse_addr(remote_host, only_4, only_6)
            .map_err(|err| format!("Could not parse remote address: {}", err))?;
        let remote_addr = *addrs.get(0).ok_or("No IP address found")?;

        for raw_service in services {
            let addr = raw_service["address"].as_str().ok_or("No address field")?;
            let port = raw_service["port"]
                .as_i64()
                .ok_or("No port field or invalid type")?;
            let port = port as u16;

            // Start by matching against ports.
            if port != remote_port {
                continue;
            }

            // Attempt to match directly against the allowlisted hostname first.
            if addr == remote_host {
                info!("Matched with host name \"{}\" and port \"{}\"", addr, port);
                return Ok(remote_addr);
            }

            // If hostname matching failed, attempt to match against IPs.
            addrs = Proxy::parse_addr(addr, only_4, only_6)?;
            for addr in addrs.into_iter() {
                if addr == remote_addr {
                    info!("Matched with host IP \"{}\" and port \"{}\"", addr, port);
                    return Ok(remote_addr);
                }
            }
        }
    }
    Err("The given address and port are not allowed".to_string())
}

/// Configuration parameters for port listening and remote destination
pub struct Proxy {
    local_port: u32,
    remote_addr: IpAddr,
    remote_port: u16,
    pool: ThreadPool,
    sock_type: SockType,
}

impl Proxy {
    pub fn new(
        local_port: u32,
        remote_host: &str,
        remote_port: u16,
        num_workers: usize,
        config_file: Option<&str>,
        only_4: bool,
        only_6: bool,
    ) -> VsockProxyResult<Self> {
        if num_workers == 0 {
            return Err("Number of workers must not be 0".to_string());
        }
        info!("Checking allowlist configuration");
        let remote_addr = check_allowlist(remote_host, remote_port, config_file, only_4, only_6)
            .map_err(|err| format!("Error at checking the allowlist: {}", err))?;
        let pool = ThreadPool::new(num_workers);
        let sock_type = SockType::Stream;

        info!(
            "Using IP \"{:?}\" for the given server \"{}\"",
            remote_addr, remote_host
        );
        Ok(Proxy {
            local_port,
            remote_addr,
            remote_port,
            pool,
            sock_type,
        })
    }

    /// Resolve a DNS name (IDNA format) into an IP address (v4 or v6)
    pub fn parse_addr(addr: &str, only_4: bool, only_6: bool) -> VsockProxyResult<Vec<IpAddr>> {
        // IDNA parsing
        let addr = domain_to_ascii(addr).map_err(|_| "Could not parse domain name")?;

        // DNS lookup
        // It results in a vector of IPs (V4 and V6)
        let ips = match lookup_host(&addr) {
            Err(_) => {
                return Ok(vec![]);
            }
            Ok(v) => {
                if v.is_empty() {
                    return Ok(v);
                }
                v
            }
        };

        // If there is no restriction, choose randomly
        if !only_4 && !only_6 {
            return Ok(ips.into_iter().collect());
        }

        // Split the IPs in v4 and v6
        let (ips_v4, ips_v6): (Vec<_>, Vec<_>) = ips.into_iter().partition(IpAddr::is_ipv4);

        if only_4 && !ips_v4.is_empty() {
            Ok(ips_v4.into_iter().collect())
        } else if only_6 && !ips_v6.is_empty() {
            Ok(ips_v6.into_iter().collect())
        } else {
            Err("No accepted IP was found".to_string())
        }
    }

    /// Creates a listening socket
    /// Returns the file descriptor for it or the appropriate error
    pub fn sock_listen(&self) -> VsockProxyResult<VsockListener> {
        let sockaddr = VsockAddr::new(VSOCK_PROXY_CID, self.local_port);
        let listener = VsockListener::bind(&sockaddr)
            .map_err(|_| format!("Could not bind to {:?}", sockaddr))?;
        info!("Bound to {:?}", sockaddr);

        Ok(listener)
    }

    /// Accepts an incoming connection coming on listener and handles it on a
    /// different thread
    /// Returns the handle for the new thread or the appropriate error
    pub fn sock_accept(&self, listener: &VsockListener) -> VsockProxyResult<()> {
        let (mut client, client_addr) = listener
            .accept()
            .map_err(|_| "Could not accept connection")?;
        info!("Accepted connection on {:?}", client_addr);

        let sockaddr = SocketAddr::new(self.remote_addr, self.remote_port);
        let sock_type = self.sock_type;
        self.pool.execute(move || {
            let mut server = match sock_type {
                SockType::Stream => TcpStream::connect(sockaddr)
                    .map_err(|_| format!("Could not connect to {:?}", sockaddr)),
                _ => Err("Socket type not implemented".to_string()),
            }
            .expect("Could not create connection");
            info!("Connected client from {:?} to {:?}", client_addr, sockaddr);

            let client_socket = client.as_raw_fd();
            let server_socket = server.as_raw_fd();

            let mut disconnected = false;
            while !disconnected {
                let mut set = FdSet::new();
                set.insert(client_socket);
                set.insert(server_socket);

                select(None, Some(&mut set), None, None, None).expect("select");

                if set.contains(client_socket) {
                    disconnected = transfer(&mut client, &mut server);
                }
                if set.contains(server_socket) {
                    disconnected = transfer(&mut server, &mut client);
                }
            }
            info!("Client on {:?} disconnected", client_addr);
        });

        Ok(())
    }
}

/// Transfers a chunck of maximum 4KB from src to dst
/// If no error occurs, returns true if the source disconnects and false otherwise
fn transfer(src: &mut dyn Read, dst: &mut dyn Write) -> bool {
    let mut buffer = [0u8; BUFF_SIZE];

    let nbytes = src.read(&mut buffer);
    let nbytes = nbytes.unwrap_or(0);

    if nbytes == 0 {
        return true;
    }

    dst.write_all(&buffer[..nbytes]).is_err()
}

#[cfg(test)]
mod tests {
    use rand;
    use std::fs;
    use std::fs::File;
    use std::io::Write;
    use std::process::Command;

    use super::*;

    /// Test transfer function with more data than buffer
    #[test]
    fn test_transfer() {
        let data: Vec<u8> = (0..2 * BUFF_SIZE).map(|_| rand::random::<u8>()).collect();

        let _ret = fs::create_dir("tmp");
        let mut src = File::create("tmp/src").unwrap();
        let mut dst = File::create("tmp/dst").unwrap();

        let _ret = src.write_all(&data);

        let mut src = File::open("tmp/src").unwrap();
        while !transfer(&mut src, &mut dst) {}

        let status = Command::new("cmp")
            .arg("tmp/src")
            .arg("tmp/dst")
            .status()
            .expect("command");

        let _ret = fs::remove_dir_all("tmp");

        assert!(status.success());
    }
}


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

use std::io::{Read, Write};
use std::net::TcpListener;
use std::net::{IpAddr, Ipv4Addr};
use std::str;
use std::sync::mpsc;
use std::thread;
use tempfile::NamedTempFile;
use vsock::{VsockAddr, VsockStream};

use vsock_proxy::starter::Proxy;

fn vsock_connect(port: u32) -> VsockStream {
    let sockaddr = VsockAddr::new(vsock_proxy::starter::VSOCK_PROXY_CID, port);
    VsockStream::connect(&sockaddr).expect("Could not connect")
}

/// Test connection with both client and server sending each other messages
#[test]
fn test_tcp_connection() {
    // Proxy will translate from port 8000 vsock to localhost port 9000 TCP
    let addr = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)).to_string();
    let mut file = NamedTempFile::new().unwrap();
    file.write_all(
        b"allowlist:\n\
            - {address: 127.0.0.1, port: 9000}",
    )
    .unwrap();
    let proxy = Proxy::new(
        vsock_proxy::starter::VSOCK_PROXY_PORT,
        &addr,
        9000,
        2,
        file.path().to_str(),
        false,
        false,
    )
    .unwrap();

    let (tx, rx) = mpsc::channel();

    // Create a listening TCP server on port 9000
    let server_handle = thread::spawn(move || {
        let server = TcpListener::bind("127.0.0.1:9000").expect("server bind");
        tx.send(true).expect("server send event");
        let (mut stream, _) = server.accept().expect("server accept");

        // Read request
        let mut buf = [0; 13];
        stream.read_exact(&mut buf).expect("server read");
        let msg = str::from_utf8(&buf).expect("from_utf8");
        assert_eq!(msg, "client2server");

        // Write response
        stream.write_all(b"server2client").expect("server write");
    });

    let _ret = rx.recv().expect("main recv event");
    let (tx, rx) = mpsc::channel();

    // Start proxy in a different thread
    let ret = proxy.sock_listen();
    let listener = ret.expect("proxy listen");
    let proxy_handle = thread::spawn(move || {
        tx.send(true).expect("proxy send event");
        let _ret = proxy.sock_accept(&listener).expect("proxy accept");
    });

    let _ret = rx.recv().expect("main recv event");

    // Start client that connects to proxy on port 8000 vsock
    let client_handle = thread::spawn(move || {
        let mut stream = vsock_connect(vsock_proxy::starter::VSOCK_PROXY_PORT);

        // Write request
        stream.write_all(b"client2server").expect("client write");

        // Read response
        let mut buf = [0; 13];

        stream.read_exact(&mut buf).expect("client read");
        let msg = str::from_utf8(&buf).expect("from_utf8");
        assert_eq!(msg, "server2client");
    });

    server_handle.join().expect("Server panicked");
    proxy_handle.join().expect("Proxy panicked");
    client_handle.join().expect("Client panicked");
}


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

#![deny(warnings)]
/// Simple utility tool for building an Eif file
///  cargo run --example eif_build -- --help  should be self explanatory.
/// Example of usage:
/// cargo run --example eif_build --target-dir=~/vmm-build -- --kernel bzImage \
///    --cmdline "reboot=k initrd=0x2000000,3228672 root=/dev/ram0 panic=1 pci=off nomodules \
///               console=ttyS0 i8042.noaux i8042.nomux i8042.nopnp i8042.dumbkbd"
///   --ramdisk  initramfs_x86.txt_part1.cpio.gz
///   --ramdisk  initramfs_x86.txt_part2.cpio.gz
///   --output   eif.bin
///
use std::path::Path;

use aws_nitro_enclaves_image_format::defs::EifIdentityInfo;
use aws_nitro_enclaves_image_format::utils::identity::parse_custom_metadata;
use aws_nitro_enclaves_image_format::{
    generate_build_info,
    utils::{get_pcrs, EifBuilder, SignEnclaveInfo},
};
use clap::{App, Arg};
use serde_json::json;
use sha2::{Digest, Sha256, Sha384, Sha512};
use std::fmt::Debug;
use std::fs::OpenOptions;
use std::io::Write;

fn main() {
    let matches = App::new("Enclave image format builder")
        .about("Builds an eif file")
        .arg(
            Arg::with_name("kernel")
                .long("kernel")
                .value_name("FILE")
                .required(true)
                .help("Sets path to a bzImage/Image file for x86_64/aarch64 architecture")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("kernel_config")
                .long("kernel_config")
                .value_name("FILE")
                .required(true)
                .help("Sets path to a bzImage.config/Image.config file for x86_64/aarch64 architecture")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("cmdline")
                .long("cmdline")
                .help("Sets the cmdline")
                .value_name("String")
                .required(true)
                .takes_value(true),
        )
        .arg(
            Arg::with_name("output")
                .long("output")
                .help("Specify output file path")
                .value_name("FILE")
                .required(true)
                .takes_value(true),
        )
        .arg(
            Arg::with_name("ramdisk")
                .long("ramdisk")
                .value_name("FILE")
                .required(true)
                .help("Sets path to a ramdisk file representing a cpio.gz archive")
                .takes_value(true)
                .multiple(true)
                .number_of_values(1),
        )
        .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("sha256")
                .long("sha256")
                .help("Sets algorithm to be used for measuring the image")
                .group("measurement_alg"),
        )
        .arg(
            Arg::with_name("sha512")
                .long("sha512")
                .help("Sets algorithm to be used for measuring the image")
                .group("measurement_alg"),
        )
        .arg(
            Arg::with_name("sha384")
                .long("sha384")
                .help("Sets algorithm to be used for measuring the image")
                .group("measurement_alg"),
        )
        .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 kernel_path = matches
        .value_of("kernel")
        .expect("Kernel path is a mandatory option");

    let kernel_config_path = matches
        .value_of("kernel_config")
        .expect("Kernel config path is a mandatory option");

    let cmdline = matches
        .value_of("cmdline")
        .expect("Cmdline is a mandatory option");

    let sha512 = matches.is_present("sha512");
    let sha256 = matches.is_present("sha256");

    let ramdisks: Vec<&str> = matches
        .values_of("ramdisk")
        .expect("At least one ramdisk should be specified")
        .collect();

    let output_path = matches
        .value_of("output")
        .expect("Output file should be provided");

    let signing_certificate = matches.value_of("signing-certificate");

    let private_key = matches.value_of("private-key");

    let sign_info = match (signing_certificate, private_key) {
        (None, None) => None,
        (Some(cert_path), Some(key_path)) => {
            Some(SignEnclaveInfo::new(cert_path, key_path).expect("Could not read signing info"))
        }
        _ => panic!("Both signing-certificate and private-key parameters must be provided"),
    };

    let img_name = matches.value_of("image_name").map(|val| val.to_string());
    let img_version = matches.value_of("image_name").map(|val| val.to_string());
    let metadata_path = matches.value_of("metadata").map(|val| val.to_string());
    let metadata = match metadata_path {
        Some(ref path) => {
            parse_custom_metadata(path).expect("Can not parse specified metadata file")
        }
        None => json!(null),
    };

    let eif_info = EifIdentityInfo {
        img_name: img_name.unwrap_or_else(|| {
            // Set default value to kernel file name
            Path::new(kernel_path)
                .file_name()
                .expect("Valid kernel file path should be provided")
                .to_str()
                .unwrap()
                .to_string()
        }),
        img_version: img_version.unwrap_or_else(|| "1.0".to_string()),
        build_info: generate_build_info!(kernel_config_path).expect("Can not generate build info"),
        docker_info: json!(null),
        custom_info: metadata,
    };

    if sha512 {
        build_eif(
            kernel_path,
            cmdline,
            ramdisks,
            output_path,
            sign_info,
            Sha512::new(),
            eif_info,
        );
    } else if sha256 {
        build_eif(
            kernel_path,
            cmdline,
            ramdisks,
            output_path,
            sign_info,
            Sha256::new(),
            eif_info,
        );
    } else {
        build_eif(
            kernel_path,
            cmdline,
            ramdisks,
            output_path,
            sign_info,
            Sha384::new(),
            eif_info,
        );
    }
}

pub fn build_eif<T: Digest + Debug + Write + Clone>(
    kernel_path: &str,
    cmdline: &str,
    ramdisks: Vec<&str>,
    output_path: &str,
    sign_info: Option<SignEnclaveInfo>,
    hasher: T,
    eif_info: EifIdentityInfo,
) {
    let mut output_file = OpenOptions::new()
        .read(true)
        .create(true)
        .write(true)
        .truncate(true)
        .open(output_path)
        .expect("Could not create output file");

    let mut build = EifBuilder::new(
        Path::new(kernel_path),
        cmdline.to_string(),
        sign_info,
        hasher.clone(),
        0, // flags
        eif_info,
    );
    for ramdisk in ramdisks {
        build.add_ramdisk(Path::new(ramdisk));
    }

    build.write_to(&mut output_file);
    let signed = build.is_signed();
    println!("Output written into {}", output_path);
    build.measure();
    let measurements = get_pcrs(
        &mut build.image_hasher,
        &mut build.bootstrap_hasher,
        &mut build.customer_app_hasher,
        &mut build.certificate_hasher,
        hasher.clone(),
        signed,
    )
    .expect("Failed to get boot measurements.");
    println!("BootMeasurement: {:?}: {:?}", hasher, measurements);
}


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

//! AWS Nitro Enclaves Image Format API
//!
//! This is the library that provides the API for interacting with Enclave Image Format (EIF) files
//! used in AWS Nitro Enclaves.
//!
//! defs provides the APIs that define an EIF and related objects.
//! utils provides the APIs to consume and interact with an EIF and related objects.

pub mod defs;
pub mod utils;