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

#![deny(warnings)]

use sha2::Digest;
use std::fmt::Debug;
use std::io::Result as IoResult;
use std::io::Write;
use std::vec::Vec;

use serde::{Deserialize, Serialize};

/// EifHasher class
///
/// A simple serialization/deserialization friendly Hasher class.
/// The only reason this exists is that we can't serialize a Hasher
/// from sha2 crate, so we are going to use the following algorithm:
///
/// 1. Initialize digest with 0.
/// 2. Gather block_size bytes in block field.
/// 3. digest = Hash(Concatenate(digest, block))
/// 4. Goto step 2
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EifHasher<T: Digest + Debug + Write + Clone> {
    /// The bytes that have not been hashed yet, they get hashed
    /// once we gather block_size bytes.
    pub block: Vec<u8>,
    /// Intermediary digest for the blocks processed untill now.
    pub digest: Vec<u8>,
    /// The number of bytes we need to gather before hashing, it needs to
    /// be at least twice of the hasher output, since the hash of each output
    /// is fed back into the hasher the size of the block impacts the performance
    /// 0 means nothing is cached and all the bytes are feed to the hasher.
    pub block_size: usize,
    pub output_size: usize,
    #[serde(skip)]
    /// The hasher to be used, it is always in reset state, so we can skip
    /// serialization.
    pub hasher: T,
}

fn initial_digest(len: usize) -> Vec<u8> {
    vec![0; len]
}

impl<T: Digest + Debug + Write + Clone> EifHasher<T> {
    pub fn new(block_size: usize, mut hasher: T) -> Result<Self, String> {
        let output_size = hasher.finalize_reset().len();
        if block_size > 0 && output_size * 2 > block_size {
            return Err("Invalid block_size".to_string());
        }

        Ok(EifHasher {
            block: Vec::with_capacity(block_size),
            digest: initial_digest(output_size),
            block_size,
            output_size,
            hasher,
        })
    }

    /// EifHasher constructor with fixed block size.
    ///
    /// It is needed in order for all clients of this class to use the same
    /// block size if we want to get the same results.
    pub fn new_with_fixed_block_size(hasher: T) -> Result<Self, String> {
        /// This impacts the performance of the hasher, it is a sweet
        /// spot where we get decent performance, 200MB/s, and where are not
        /// forced to keep a large serialized state, 256 bytes for SHA256.
        pub const FIXED_BLOCK_SIZE_HASHER_OUPUT_RATIO: usize = 8;
        Self::new(
            hasher.clone().finalize_reset().len() * FIXED_BLOCK_SIZE_HASHER_OUPUT_RATIO,
            hasher,
        )
    }

    /// EifHasher constructor without cache.
    ///
    /// EIfHasher acts like passthrough passing all the bytes to the actual hasher.
    pub fn new_without_cache(hasher: T) -> Result<Self, String> {
        Self::new(0, hasher)
    }

    pub fn finalize_reset(&mut self) -> IoResult<Vec<u8>> {
        if self.block_size == 0 {
            return Ok(self.hasher.finalize_reset().to_vec());
        }
        if !self.block.is_empty() {
            self.consume_block()?;
        }
        let result = self.digest.clone();
        self.digest = initial_digest(self.output_size);
        Ok(result)
    }

    pub fn tpm_extend_finalize_reset(&mut self) -> IoResult<Vec<u8>> {
        let result = self.finalize_reset()?;
        let mut hasher = self.hasher.clone();

        hasher.write_all(&initial_digest(self.output_size))?;
        hasher.write_all(&result[..])?;
        Ok(hasher.finalize_reset().to_vec())
    }

    fn consume_block(&mut self) -> IoResult<()> {
        self.hasher.write_all(&self.digest[..])?;
        self.hasher.write_all(&self.block[..])?;
        self.block.clear();
        let result = self.hasher.finalize_reset();
        self.digest.copy_from_slice(&result[..]);
        Ok(())
    }
}

impl<T: Digest + Debug + Write + Clone> Write for EifHasher<T> {
    fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
        if self.block_size == 0 {
            self.hasher.write_all(buf)?;
            return Ok(buf.len());
        }
        let mut remaining = buf;
        while self.block.len() + remaining.len() >= self.block_size {
            let (for_hasher, for_next_iter) =
                remaining.split_at(self.block_size - self.block.len());
            self.block.extend_from_slice(for_hasher);
            self.consume_block()?;
            remaining = for_next_iter;
        }

        self.block.extend_from_slice(remaining);
        Ok(buf.len())
    }

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

#[cfg(test)]
mod tests {
    use super::{initial_digest, EifHasher};
    use sha2::{Digest, Sha256, Sha384, Sha512};
    use std::fmt::Debug;
    use std::io::Write;

    const INPUT_BLOCK_SIZE_SHA256: usize = 64;
    const INPUT_BLOCK_SIZE_SHA384: usize = 128;
    const INPUT_BLOCK_SIZE_SHA512: usize = 128;

    #[test]
    fn invalid_block_size() {
        let hasher = EifHasher::new(31, Sha256::new());
        assert_eq!(hasher.is_err(), true);

        let hasher = EifHasher::new(63, Sha512::new());
        assert_eq!(hasher.is_err(), true);

        let hasher = EifHasher::new(47, Sha384::new());
        assert_eq!(hasher.is_err(), true)
    }

    #[test]
    fn test_hash_less_values_than_block_size() {
        hash_less_values_than_block_size(Sha256::new(), INPUT_BLOCK_SIZE_SHA256);
        hash_less_values_than_block_size(Sha512::new(), INPUT_BLOCK_SIZE_SHA512);
        hash_less_values_than_block_size(Sha384::new(), INPUT_BLOCK_SIZE_SHA384);
    }

    fn hash_less_values_than_block_size<T: Digest + Debug + Write + Clone>(
        mut hasher_alg: T,
        block_size: usize,
    ) {
        let data = vec![78u8; block_size - 1];
        let output_size = hasher_alg.finalize_reset().len();
        let mut hasher = EifHasher::new(block_size, hasher_alg.clone()).unwrap();

        hasher.write_all(&data[..]).unwrap();
        hasher_alg
            .write_all(&initial_digest(output_size)[..])
            .unwrap();
        hasher_alg.write_all(&data[..]).unwrap();

        let mut hasher_clone = hasher.clone();
        let mut hasher_alg_clone = hasher_alg.clone();
        assert_eq!(
            hasher_alg.finalize_reset().to_vec(),
            hasher.finalize_reset().unwrap()
        );

        let result = hasher_alg_clone.finalize_reset();
        hasher_alg_clone
            .write_all(&initial_digest(output_size)[..])
            .unwrap();
        hasher_alg_clone.write_all(&result[..]).unwrap();

        assert_eq!(
            hasher_clone.tpm_extend_finalize_reset().unwrap(),
            hasher_alg_clone.finalize_reset().to_vec()
        );
    }

    #[test]
    fn test_hash_exact_block_size_values() {
        hash_exact_block_size_values(Sha256::new(), INPUT_BLOCK_SIZE_SHA256);
        hash_exact_block_size_values(Sha384::new(), INPUT_BLOCK_SIZE_SHA384);
        hash_exact_block_size_values(Sha512::new(), INPUT_BLOCK_SIZE_SHA512);
    }

    fn hash_exact_block_size_values<T: Digest + Debug + Write + Clone>(
        mut hasher_alg: T,
        block_size: usize,
    ) {
        let data = vec![78u8; block_size];
        let output_size = hasher_alg.finalize_reset().len();
        let mut hasher = EifHasher::new(block_size, hasher_alg.clone()).unwrap();

        hasher.write_all(&data).unwrap();
        hasher_alg
            .write_all(&initial_digest(output_size)[..])
            .unwrap();
        hasher_alg.write_all(&data[..block_size]).unwrap();

        let mut hasher_clone = hasher.clone();
        let mut hasher_alg_clone = hasher_alg.clone();

        assert_eq!(
            hasher_alg.finalize_reset().to_vec(),
            hasher.finalize_reset().unwrap()
        );

        let result = hasher_alg_clone.finalize_reset();
        hasher_alg_clone
            .write_all(&initial_digest(output_size)[..])
            .unwrap();
        hasher_alg_clone.write_all(&result[..]).unwrap();

        assert_eq!(
            hasher_clone.tpm_extend_finalize_reset().unwrap(),
            hasher_alg_clone.finalize_reset().to_vec()
        );
    }

    #[test]
    fn test_hash_more_values_than_block_size() {
        hash_more_values_than_block_size(Sha256::new(), INPUT_BLOCK_SIZE_SHA256);
        hash_more_values_than_block_size(Sha384::new(), INPUT_BLOCK_SIZE_SHA384);
        hash_more_values_than_block_size(Sha512::new(), INPUT_BLOCK_SIZE_SHA512);
    }

    fn hash_more_values_than_block_size<T: Digest + Debug + Write + Clone>(
        mut hasher_alg: T,
        block_size: usize,
    ) {
        let data = vec![78u8; block_size + block_size / 2 - 1];
        let output_size = hasher_alg.finalize_reset().len();
        let (data1, data2) = data.split_at(block_size);
        let mut hasher = EifHasher::new(block_size, hasher_alg.clone()).unwrap();

        hasher.write_all(&data).unwrap();

        hasher_alg.write_all(&initial_digest(output_size)).unwrap();
        hasher_alg.write_all(data1).unwrap();
        let result = hasher_alg.finalize_reset();
        hasher_alg.write_all(&result).unwrap();
        hasher_alg.write_all(data2).unwrap();

        let mut hasher_clone = hasher.clone();
        let mut hasher_alg_clone = hasher_alg.clone();

        assert_eq!(
            hasher_alg.finalize_reset().to_vec(),
            hasher.finalize_reset().unwrap()
        );

        let result = hasher_alg_clone.finalize_reset();
        hasher_alg_clone
            .write_all(&initial_digest(output_size)[..])
            .unwrap();
        hasher_alg_clone.write_all(&result[..]).unwrap();

        assert_eq!(
            hasher_clone.tpm_extend_finalize_reset().unwrap(),
            hasher_alg_clone.finalize_reset().to_vec()
        );
    }

    #[test]
    fn test_hash_with_writes_of_different_sizes() {
        hash_with_writes_of_different_sizes(Sha256::new(), INPUT_BLOCK_SIZE_SHA256);
        hash_with_writes_of_different_sizes(Sha384::new(), INPUT_BLOCK_SIZE_SHA512);
        hash_with_writes_of_different_sizes(Sha512::new(), INPUT_BLOCK_SIZE_SHA512);
    }

    fn hash_with_writes_of_different_sizes<T: Digest + Debug + Write + Clone>(
        hasher_alg: T,
        block_size: usize,
    ) {
        let data = vec![78u8; block_size * 256];
        let mut hasher_in_one_go = EifHasher::new(block_size, hasher_alg.clone()).unwrap();
        let mut hasher_in_random_chunks = EifHasher::new(block_size, hasher_alg.clone()).unwrap();
        let mut hasher_one_byte_at_atime = EifHasher::new(block_size, hasher_alg).unwrap();

        hasher_in_one_go.write_all(&data).unwrap();
        let mut remaining = &data[..];
        let mut iteration = 1;
        while !remaining.is_empty() {
            let chunk_size = std::cmp::max(1, (iteration % remaining.len()) % block_size);
            let (to_be_written, unhandled) = remaining.split_at(chunk_size);
            hasher_in_random_chunks.write_all(to_be_written).unwrap();

            remaining = unhandled;
            iteration += 1987;
        }

        for x in data {
            hasher_one_byte_at_atime.write_all(&[x]).unwrap();
        }

        let result1 = hasher_in_one_go.finalize_reset().unwrap();
        let result2 = hasher_in_random_chunks.finalize_reset().unwrap();
        let result3 = hasher_one_byte_at_atime.finalize_reset().unwrap();
        assert_eq!(result1, result2);
        assert_eq!(result1, result3);
    }

    #[test]
    fn test_no_cache() {
        let data = vec![78u8; 127 * 256];
        let mut eif_hasher = EifHasher::new_without_cache(Sha384::new()).unwrap();
        let mut hasher = Sha384::new();

        hasher.write_all(&data[..]).unwrap();
        for value in data {
            eif_hasher.write(&[value]).unwrap();
        }

        let result1 = eif_hasher.finalize_reset().unwrap();
        let result2 = hasher.finalize_reset();
        assert_eq!(result1, result2.to_vec());
    }
}


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

#![deny(warnings)]
/// Definition Eif (Enclave Image format)
///
/// This crate is consumed by the following clients:
///    - eif_utils: needed by the eif_builder.
///    - eif_loader: needed to properly send an eif file over vsock.
///
/// With that in mind please be frugal with the dependencies for this
/// crate.
use byteorder::{BigEndian, ByteOrder};
use num_derive::FromPrimitive;
use num_traits::FromPrimitive;
use serde::{Deserialize, Serialize};
use std::mem::size_of;

pub const EIF_MAGIC: [u8; 4] = [46, 101, 105, 102]; // .eif in ascii
pub const MAX_NUM_SECTIONS: usize = 32;
/// EIF Header Flags
/// bits 1  : Architecture - toggled for aarch64, cleared for x86_64
/// bits 2-8: Unused
pub const EIF_HDR_ARCH_ARM64: u16 = 0x1;

/// Current EIF version to be incremented every time we change the format
/// of this structures, we assume changes are backwards compatible.
/// V1 -> V2: Add support to generate and check CRC.
/// V2 -> V3: Add the signature section.
/// V3 -> V4: Add the metadata section.
pub const CURRENT_VERSION: u16 = 4;

#[derive(Clone, Copy, Debug)]
pub struct EifHeader {
    /// Magic number used to identify this file format
    pub magic: [u8; 4],
    /// EIF version
    pub version: u16,
    /// EIF header flags
    pub flags: u16,
    /// Default enclave memory used to boot this eif file
    pub default_mem: u64,
    /// Default enclave cpus number used to boot this eif file
    pub default_cpus: u64,
    pub reserved: u16,
    pub num_sections: u16,
    pub section_offsets: [u64; MAX_NUM_SECTIONS],
    /// Sizes for each section, we need this information in the
    /// header because the vsock_loader needs to know how large are the
    /// ramdisks
    pub section_sizes: [u64; MAX_NUM_SECTIONS],
    pub unused: u32,
    /// crc32 IEEE used for validating the eif file is correct, it contains
    /// the crc for everything except the bytes representing this field.
    /// Needs to be the last field of the header.
    pub eif_crc32: u32,
}

impl EifHeader {
    pub fn from_be_bytes(bytes: &[u8]) -> Result<Self, String> {
        let mut pos = 0;

        let mut magic = [0u8; 4];
        magic.copy_from_slice(&bytes[pos..pos + size_of::<[u8; 4]>()]);
        pos += size_of::<[u8; 4]>();

        let version = BigEndian::read_u16(&bytes[pos..]);
        pos += size_of::<u16>();
        let flags = BigEndian::read_u16(&bytes[pos..]);
        pos += size_of::<u16>();
        let default_mem = BigEndian::read_u64(&bytes[pos..]);
        pos += size_of::<u64>();
        let default_cpus = BigEndian::read_u64(&bytes[pos..]);
        pos += size_of::<u64>();
        let reserved = BigEndian::read_u16(&bytes[pos..]);
        pos += size_of::<u16>();
        let num_sections = BigEndian::read_u16(&bytes[pos..]);
        pos += size_of::<u16>();

        let mut section_offsets = [0u64; MAX_NUM_SECTIONS];
        for item in section_offsets.iter_mut() {
            *item = BigEndian::read_u64(&bytes[pos..]);
            pos += size_of::<u64>();
        }

        let mut section_sizes = [0u64; MAX_NUM_SECTIONS];
        for item in section_sizes.iter_mut() {
            *item = BigEndian::read_u64(&bytes[pos..]);
            pos += size_of::<u64>();
        }

        let unused = BigEndian::read_u32(&bytes[pos..]);
        pos += size_of::<u32>();
        let eif_crc32 = BigEndian::read_u32(&bytes[pos..]);
        pos += size_of::<u32>();

        if bytes.len() != pos {
            return Err("Invalid EifHeader length".to_string());
        }

        Ok(EifHeader {
            magic,
            version,
            flags,
            default_mem,
            default_cpus,
            reserved,
            num_sections,
            section_offsets,
            section_sizes,
            unused,
            eif_crc32,
        })
    }

    pub fn to_be_bytes(&self) -> Vec<u8> {
        let mut buf = [0u8; Self::size()];
        let mut result = Vec::new();
        let mut pos = 0;

        buf[pos..pos + size_of::<[u8; 4]>()].copy_from_slice(&self.magic);
        pos += size_of::<[u8; 4]>();

        BigEndian::write_u16(&mut buf[pos..], self.version);
        pos += size_of::<u16>();
        BigEndian::write_u16(&mut buf[pos..], self.flags);
        pos += size_of::<u16>();
        BigEndian::write_u64(&mut buf[pos..], self.default_mem);
        pos += size_of::<u64>();
        BigEndian::write_u64(&mut buf[pos..], self.default_cpus);
        pos += size_of::<u64>();
        BigEndian::write_u16(&mut buf[pos..], self.reserved);
        pos += size_of::<u16>();
        BigEndian::write_u16(&mut buf[pos..], self.num_sections);
        pos += size_of::<u16>();

        for elem in self.section_offsets.iter() {
            BigEndian::write_u64(&mut buf[pos..], *elem);
            pos += size_of::<u64>();
        }

        for elem in self.section_sizes.iter() {
            BigEndian::write_u64(&mut buf[pos..], *elem);
            pos += size_of::<u64>();
        }

        BigEndian::write_u32(&mut buf[pos..], self.unused);
        pos += size_of::<u32>();
        BigEndian::write_u32(&mut buf[pos..], self.eif_crc32);

        result.extend_from_slice(&buf[..]);
        result
    }

    pub const fn size() -> usize {
        4 * size_of::<u16>()
            + 2 * size_of::<u32>()
            + 2 * size_of::<u64>()
            + size_of::<[u8; 4]>()
            + 2 * size_of::<[u64; MAX_NUM_SECTIONS]>()
    }
}

#[derive(Clone, Copy, Debug, FromPrimitive, PartialEq, Eq)]
#[repr(u16)]
pub enum EifSectionType {
    EifSectionInvalid,
    EifSectionKernel,
    EifSectionCmdline,
    EifSectionRamdisk,
    EifSectionSignature,
    EifSectionMetadata,
}

#[derive(Clone, Copy, Debug)]
pub struct EifSectionHeader {
    pub section_type: EifSectionType,
    pub flags: u16,
    pub section_size: u64,
}

impl EifSectionHeader {
    pub fn from_be_bytes(bytes: &[u8]) -> Result<Self, String> {
        let mut pos = 0;

        let section_type = BigEndian::read_u16(&bytes[pos..]);
        pos += size_of::<u16>();
        let flags = BigEndian::read_u16(&bytes[pos..]);
        pos += size_of::<u16>();
        let section_size = BigEndian::read_u64(&bytes[pos..]);
        pos += size_of::<u64>();

        if bytes.len() != pos {
            return Err("Invalid EifSectionHeader length".to_string());
        }

        Ok(EifSectionHeader {
            section_type: FromPrimitive::from_u16(section_type)
                .ok_or_else(|| "Invalid section type".to_string())?,
            flags,
            section_size,
        })
    }

    pub fn to_be_bytes(&self) -> Vec<u8> {
        let mut result = Vec::new();
        let mut buf = [0u8; Self::size()];
        let mut pos = 0;

        BigEndian::write_u16(&mut buf[pos..], self.section_type as u16);
        pos += size_of::<u16>();
        BigEndian::write_u16(&mut buf[pos..], self.flags);
        pos += size_of::<u16>();
        BigEndian::write_u64(&mut buf[pos..], self.section_size);

        result.extend_from_slice(&buf[..]);
        result
    }

    pub const fn size() -> usize {
        size_of::<EifSectionType>() + size_of::<u16>() + size_of::<u64>()
    }
}

/// Array containing the signatures of at least one PCR.
/// For now, it only contains the signature of PRC0.
pub type EifSignature = Vec<PcrSignature>;

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct PcrSignature {
    /// The PEM-formatted signing certificate
    pub signing_certificate: Vec<u8>,
    /// The serialized COSESign1 object generated using the byte array
    /// formed from RegisterIndex and RegisterValue as payload
    pub signature: Vec<u8>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct PcrInfo {
    /// The index of the PCR
    pub register_index: i32,
    /// The value of the PCR
    pub register_value: Vec<u8>,
}

impl PcrInfo {
    pub fn new(register_index: i32, register_value: Vec<u8>) -> Self {
        PcrInfo {
            register_index,
            register_value,
        }
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct EifBuildInfo {
    #[serde(rename = "BuildTime")]
    pub build_time: String,
    #[serde(rename = "BuildTool")]
    pub build_tool: String,
    #[serde(rename = "BuildToolVersion")]
    pub build_tool_version: String,
    #[serde(rename = "OperatingSystem")]
    pub img_os: String,
    #[serde(rename = "KernelVersion")]
    pub img_kernel: String,
}

/// Structure used for (de)serializing metadata when
/// writing or reading the metadata section of the EIF
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EifIdentityInfo {
    #[serde(rename = "ImageName")]
    pub img_name: String,
    #[serde(rename = "ImageVersion")]
    pub img_version: String,
    #[serde(rename = "BuildMetadata")]
    pub build_info: EifBuildInfo,
    #[serde(rename = "DockerInfo")]
    pub docker_info: serde_json::Value,
    #[serde(rename = "CustomMetadata")]
    pub custom_info: serde_json::Value,
}

#[cfg(test)]
mod tests {
    use super::{EifHeader, EifSectionHeader, EifSectionType};
    use super::{EIF_MAGIC, MAX_NUM_SECTIONS};

    #[test]
    fn test_eif_section_type() {
        assert_eq!(std::mem::size_of::<EifSectionType>(), 2);
    }

    #[test]
    fn test_eif_section_header_to_from_be_bytes() {
        let eif_section_header = EifSectionHeader {
            section_type: EifSectionType::EifSectionSignature,
            flags: 3,
            section_size: 123,
        };

        let bytes = eif_section_header.to_be_bytes();
        assert_eq!(bytes.len(), EifSectionHeader::size());

        let new_eif_section_header = EifSectionHeader::from_be_bytes(&bytes).unwrap();
        assert_eq!(
            eif_section_header.section_type,
            new_eif_section_header.section_type
        );
        assert_eq!(eif_section_header.flags, new_eif_section_header.flags);
        assert_eq!(
            eif_section_header.section_size,
            new_eif_section_header.section_size
        );
    }

    #[test]
    fn test_eif_header_to_from_be_bytes() {
        let eif_header = EifHeader {
            magic: EIF_MAGIC,
            version: 3,
            flags: 4,
            default_mem: 5,
            default_cpus: 6,
            reserved: 2,
            num_sections: 5,
            section_offsets: [12u64; MAX_NUM_SECTIONS],
            section_sizes: [13u64; MAX_NUM_SECTIONS],
            unused: 0,
            eif_crc32: 123,
        };

        let bytes = eif_header.to_be_bytes();
        assert_eq!(bytes.len(), EifHeader::size());

        let new_eif_header = EifHeader::from_be_bytes(&bytes).unwrap();
        assert_eq!(eif_header.magic, new_eif_header.magic);
        assert_eq!(eif_header.version, new_eif_header.version);
        assert_eq!(eif_header.flags, new_eif_header.flags);
        assert_eq!(eif_header.default_mem, new_eif_header.default_mem);
        assert_eq!(eif_header.default_cpus, new_eif_header.default_cpus);
        assert_eq!(eif_header.reserved, new_eif_header.reserved);
        assert_eq!(eif_header.num_sections, new_eif_header.num_sections);
        assert_eq!(eif_header.section_offsets, new_eif_header.section_offsets);
        assert_eq!(eif_header.section_sizes, new_eif_header.section_sizes);
        assert_eq!(eif_header.unused, new_eif_header.unused);
        assert_eq!(eif_header.eif_crc32, new_eif_header.eif_crc32);
    }

    #[test]
    fn test_eif_header_size() {
        assert_eq!(EifHeader::size(), 548);
    }

    #[test]
    fn test_eif_section_header_size() {
        assert_eq!(EifSectionHeader::size(), 12);
    }

    #[test]
    fn test_eif_header_from_be_bytes_invalid_length() {
        let bytes = [0u8; 550];
        assert!(EifHeader::from_be_bytes(&bytes).is_err());
    }

    #[test]
    fn test_eif_section_header_from_be_bytes_invalid_length() {
        let bytes = [0u8; 16];
        assert!(EifSectionHeader::from_be_bytes(&bytes).is_err());
    }

    #[test]
    fn test_eif_section_header_from_be_bytes_invalid_section_type() {
        let mut bytes = [0u8; 12];
        // As there are 6 EIF section types, set the enum index to 6
        // so we get an invalid section to cause the error.
        bytes[1] = 6;

        assert_eq!(EifSectionHeader::from_be_bytes(&bytes).is_err(), true);
    }
}

pub mod eif_hasher;


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

use crate::defs::eif_hasher::EifHasher;
use crate::defs::{
    EifHeader, EifIdentityInfo, EifSectionHeader, EifSectionType, PcrInfo, PcrSignature,
};
use aws_nitro_enclaves_cose::{crypto::Openssl, CoseSign1};
use crc::{crc32, Hasher32};
use openssl::pkey::PKey;
use serde::{Deserialize, Serialize};
use serde_cbor::{from_slice, to_vec};
use sha2::{Digest, Sha384};

use std::collections::BTreeMap;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom, Write};
use std::mem::size_of;

/// The information about the signing certificate to be provided for a `describe-eif` request.
#[derive(Clone, Serialize, Deserialize)]
pub struct SignCertificateInfo {
    #[serde(rename = "IssuerName")]
    /// Certificate's subject name.
    pub issuer_name: BTreeMap<String, String>,
    #[serde(rename = "Algorithm")]
    /// Certificate's signature algorithm
    pub algorithm: String,
    #[serde(rename = "NotBefore")]
    /// Not before validity period
    pub not_before: String,
    #[serde(rename = "NotAfter")]
    /// Not after validity period
    pub not_after: String,
    #[serde(rename = "Signature")]
    /// Certificate's signature in hex format: 'XX:XX:XX..'
    pub signature: String,
}

impl SignCertificateInfo {
    /// Create new signing certificate information structure
    pub fn new(
        issuer_name: BTreeMap<String, String>,
        algorithm: String,
        not_before: String,
        not_after: String,
        signature: String,
    ) -> Self {
        SignCertificateInfo {
            issuer_name,
            algorithm,
            not_before,
            not_after,
            signature,
        }
    }
}

/// Used for providing EIF info when requested by
/// 'describe-eif' or 'describe-enclaves' commands
pub struct EifReader {
    /// Deserialized EIF header
    pub header: EifHeader,
    /// Serialized signature section
    pub signature_section: Option<Vec<u8>>,
    /// Hash of the whole EifImage.
    pub image_hasher: EifHasher<Sha384>,
    /// Hash of the EifSections provided by Amazon
    /// Kernel + cmdline + First Ramdisk
    pub bootstrap_hasher: EifHasher<Sha384>,
    /// Hash of the remaining ramdisks.
    pub app_hasher: EifHasher<Sha384>,
    /// Hash the signing certificate
    pub cert_hasher: EifHasher<Sha384>,
    pub eif_crc: u32,
    pub sign_check: Option<bool>,
    /// Generated and custom EIF metadata
    pub metadata: Option<EifIdentityInfo>,
}

impl EifReader {
    /// Reads EIF and extracts sections to be written in the hashers based
    /// on section type. Also writes sections in the eif_crc, excluding the
    /// CRC from the header
    pub fn from_eif(eif_path: String) -> Result<Self, String> {
        let mut eif_crc = crc32::Digest::new_with_initial(crc32::IEEE, 0);
        let mut curr_seek = 0;
        let mut eif_file =
            File::open(eif_path).map_err(|e| format!("Failed to open the EIF file: {:?}", e))?;

        // Extract EIF header
        let mut header_buf = vec![0u8; EifHeader::size()];
        eif_file
            .read_exact(&mut header_buf)
            .map_err(|e| format!("Error while reading EIF header: {:?}", e))?;

        // Exclude last field of header which is CRC
        let len_without_crc = header_buf.len() - size_of::<u32>();
        eif_crc.write(&header_buf[..len_without_crc]);

        let header = EifHeader::from_be_bytes(&header_buf)
            .map_err(|e| format!("Error while parsing EIF header: {:?}", e))?;
        curr_seek += EifHeader::size();
        eif_file
            .seek(SeekFrom::Start(curr_seek as u64))
            .map_err(|e| format!("Failed to seek file from start: {:?}", e))?;

        let mut section_buf = vec![0u8; EifSectionHeader::size()];
        let mut image_hasher = EifHasher::new_without_cache(Sha384::new())
            .map_err(|e| format!("Could not create image_hasher: {:?}", e))?;
        let mut bootstrap_hasher = EifHasher::new_without_cache(Sha384::new())
            .map_err(|e| format!("Could not create bootstrap_hasher: {:?}", e))?;
        let mut app_hasher = EifHasher::new_without_cache(Sha384::new())
            .map_err(|e| format!("Could not create app_hasher: {:?}", e))?;
        let mut cert_hasher = EifHasher::new_without_cache(Sha384::new())
            .map_err(|e| format!("Could not create cert_hasher: {:?}", e))?;
        let mut ramdisk_idx = 0;
        let mut signature_section = None;
        let mut metadata = None;

        // Read all section headers and treat by type
        while eif_file
            .read_exact(&mut section_buf)
            .map_err(|e| format!("Error while reading EIF header: {:?}", e))
            .is_ok()
        {
            let section = EifSectionHeader::from_be_bytes(&section_buf)
                .map_err(|e| format!("Error extracting EIF section header: {:?}", e))?;
            eif_crc.write(&section_buf);

            let mut buf = vec![0u8; section.section_size as usize];
            curr_seek += EifSectionHeader::size();
            eif_file
                .seek(SeekFrom::Start(curr_seek as u64))
                .map_err(|e| format!("Failed to seek after EIF header: {:?}", e))?;
            eif_file
                .read_exact(&mut buf)
                .map_err(|e| format!("Error while reading kernel from EIF: {:?}", e))?;
            eif_crc.write(&buf);

            curr_seek += section.section_size as usize;
            eif_file
                .seek(SeekFrom::Start(curr_seek as u64))
                .map_err(|e| format!("Failed to seek after EIF section: {:?}", e))?;

            match section.section_type {
                EifSectionType::EifSectionKernel | EifSectionType::EifSectionCmdline => {
                    image_hasher.write_all(&buf).map_err(|e| {
                        format!("Failed to write EIF section to image_hasher: {:?}", e)
                    })?;
                    bootstrap_hasher.write_all(&buf).map_err(|e| {
                        format!("Failed to write EIF section to bootstrap_hasher: {:?}", e)
                    })?;
                }
                EifSectionType::EifSectionRamdisk => {
                    image_hasher.write_all(&buf).map_err(|e| {
                        format!("Failed to write ramdisk section to image_hasher: {:?}", e)
                    })?;
                    if ramdisk_idx == 0 {
                        bootstrap_hasher.write_all(&buf).map_err(|e| {
                            format!(
                                "Failed to write ramdisk section to bootstrap_hasher: {:?}",
                                e
                            )
                        })?;
                    } else {
                        app_hasher.write_all(&buf).map_err(|e| {
                            format!("Failed to write ramdisk section to app_hasher: {:?}", e)
                        })?;
                    }
                    ramdisk_idx += 1;
                }
                EifSectionType::EifSectionSignature => {
                    signature_section = Some(buf.clone());
                    // Deserialize PCR0 signature structure and write it to the hasher
                    let des_sign: Vec<PcrSignature> = from_slice(&buf[..])
                        .map_err(|e| format!("Error deserializing certificate: {:?}", e))?;

                    let cert = openssl::x509::X509::from_pem(&des_sign[0].signing_certificate)
                        .map_err(|e| format!("Error while digesting certificate: {:?}", e))?;
                    let cert_der = cert.to_der().map_err(|e| {
                        format!("Failed to deserialize signing certificate: {:?}", e)
                    })?;
                    cert_hasher.write_all(&cert_der).map_err(|e| {
                        format!("Failed to write signature section to cert_hasher: {:?}", e)
                    })?;
                }
                EifSectionType::EifSectionMetadata => {
                    metadata = serde_json::from_slice(&buf[..])
                        .map_err(|e| format!("Error deserializing metadata: {:?}", e))?;
                }
                EifSectionType::EifSectionInvalid => {
                    return Err("Found invalid EIF section".to_string());
                }
            }
        }

        Ok(EifReader {
            header,
            signature_section,
            image_hasher,
            bootstrap_hasher,
            app_hasher,
            cert_hasher,
            eif_crc: eif_crc.sum32(),
            sign_check: None,
            metadata,
        })
    }

    pub fn get_metadata(&self) -> Option<EifIdentityInfo> {
        self.metadata.clone()
    }

    /// Returns deserialized header section
    pub fn get_header(&self) -> EifHeader {
        self.header
    }

    /// Compare header CRC to the one we computed
    pub fn check_crc(&self) -> bool {
        self.header.eif_crc32 == self.eif_crc
    }

    /// Extract signature section from EIF and parse the signing certificate
    pub fn get_certificate_info(
        &mut self,
        measurements: BTreeMap<String, String>,
    ) -> Result<SignCertificateInfo, String> {
        let signature_buf = match &self.signature_section {
            Some(section) => section,
            None => {
                return Err("Signature section missing from EIF.".to_string());
            }
        };
        // Deserialize PCR0 signature structure and write it to the hasher
        let des_sign: Vec<PcrSignature> = from_slice(&signature_buf[..])
            .map_err(|e| format!("Error deserializing certificate: {:?}", e))?;

        let cert = openssl::x509::X509::from_pem(&des_sign[0].signing_certificate)
            .map_err(|e| format!("Error while digesting certificate: {:?}", e))?;

        // Parse issuer into a BTreeMap
        let mut issuer_name = BTreeMap::new();
        for (_, e) in cert.issuer_name().entries().enumerate() {
            issuer_name.insert(
                e.object().to_string(),
                format!("{:?}", e.data()).replace(&['\"'][..], ""),
            );
        }

        let algorithm = format!("{:#?}", cert.signature_algorithm().object());

        // Get measured PCR0 signature payload
        let pcr0 = match measurements.get("PCR0") {
            Some(pcr) => pcr,
            None => {
                return Err("Failed to get PCR0.".to_string());
            }
        };
        let pcr_info = PcrInfo::new(
            0,
            hex::decode(pcr0).map_err(|e| format!("Error while decoding PCR0: {:?}", e))?,
        );

        let measured_payload =
            to_vec(&pcr_info).map_err(|e| format!("Could not serialize PCR info: {:?}", e))?;

        // Extract public key from certificate and convert to PKey
        let public_key = &cert
            .public_key()
            .map_err(|e| format!("Failed to get public key: {:?}", e))?;
        let coses_key = PKey::public_key_from_pem(
            &public_key
                .public_key_to_pem()
                .map_err(|e| format!("Failed to serialize public key: {:?}", e))?[..],
        )
        .map_err(|e| format!("Failed to decode key nit elliptic key structure: {:?}", e))?;

        // Deserialize COSE signature and extract the payload using the public key
        let pcr_sign = CoseSign1::from_bytes(&des_sign[0].signature[..])
            .map_err(|e| format!("Failed to deserialize signature: {:?}", e))?;
        let coses_payload = pcr_sign
            .get_payload::<Openssl>(Some(coses_key.as_ref()))
            .map_err(|e| format!("Failed to get signature payload: {:?}", e))?;

        self.sign_check = Some(measured_payload == coses_payload);

        Ok(SignCertificateInfo {
            issuer_name,
            algorithm,
            not_before: format!("{:#?}", cert.not_before()),
            not_after: format!("{:#?}", cert.not_after()),
            // Change format from [\  XX, \  X, ..] to XX:XX:XX...
            signature: format!("{:02X?}", cert.signature().as_slice().to_vec())
                .replace(&vec!['[', ']'][..], "")
                .replace(", ", ":"),
        })
    }
}


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

use crate::defs::EifBuildInfo;
use chrono::offset::Utc;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::Path;

/// Utilities and helpers to fill EIF identity information

const MAX_META_FILE_SIZE: u64 = 4096;
const UNKNOWN_IMG_STR: &str = "Unknown";

/// Generate basic build info (build tool, version, time, image kernel info)
pub fn generate_build_info(
    build_tool: &str,
    build_tool_version: &str,
    img_config_path: &str,
) -> Result<EifBuildInfo, String> {
    let now = Utc::now();

    let config_file = File::open(img_config_path)
        .map_err(|e| format!("Failed to open kernel image config file: {}", e))?;
    let os_string = BufReader::new(config_file)
        .lines()
        .nth(2)
        .unwrap()
        .map_err(|e| format!("Failed to read kernel config file: {}", e))?;

    // Extract OS and version from line format:
    // ' # Linux/x86_64 4.14.177-104.253.amzn2.x86_64 Kernel Configuration '
    let sep: Vec<char> = vec![' ', '/', '-'];
    let os_words: Vec<&str> = os_string.split(&sep[..]).collect();

    Ok(EifBuildInfo {
        build_time: now.to_rfc3339(),
        build_tool: build_tool.to_string(),
        build_tool_version: build_tool_version.to_string(),
        img_os: os_words.get(1).unwrap_or(&UNKNOWN_IMG_STR).to_string(),
        img_kernel: os_words.get(3).unwrap_or(&UNKNOWN_IMG_STR).to_string(),
    })
}

/// Macro helper for generate_buid_info function to automatically pick up cargo info for build tool fields
#[macro_export]
macro_rules! generate_build_info {
    ($kernel_config_path:expr) => {
        $crate::utils::identity::generate_build_info(
            env!("CARGO_PKG_NAME"),
            env!("CARGO_PKG_VERSION"),
            $kernel_config_path,
        )
    };
}

/// Read user-provided metadata from a file in a JSON format
pub fn parse_custom_metadata(path: &str) -> Result<serde_json::Value, String> {
    if !Path::new(&path).is_file() {
        return Err("Specified path is not a file".to_string());
    }

    // Check file size
    let file_meta =
        std::fs::metadata(path).map_err(|e| format!("Failed to get file metadata: {}", e))?;
    if file_meta.len() > MAX_META_FILE_SIZE {
        return Err(format!(
            "Metadata file size exceeded limit of {}B",
            MAX_META_FILE_SIZE
        ));
    }

    // Get json Value
    let custom_file =
        File::open(path).map_err(|e| format!("Failed to open custom metadata file: {}", e))?;
    let json_value: serde_json::Value = serde_json::from_reader(custom_file)
        .map_err(|e| format!("Failed to deserialize json: {}", e))?;

    Ok(json_value)
}


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

use crate::defs::eif_hasher::EifHasher;
use crate::defs::{
    EifHeader, EifIdentityInfo, EifSectionHeader, EifSectionType, PcrInfo, PcrSignature, EIF_MAGIC,
    MAX_NUM_SECTIONS,
};
use aws_nitro_enclaves_cose::{crypto::Openssl, header_map::HeaderMap, CoseSign1};
use crc::{crc32, Hasher32};
use openssl::asn1::Asn1Time;
use openssl::pkey::PKey;
use serde::{Deserialize, Serialize};
use serde_cbor::{from_slice, to_vec};
use sha2::Digest;
use std::cmp::Ordering;
use std::collections::BTreeMap;

/// Contains code for EifBuilder a simple library used for building an EifFile
/// from a:
///    - kernel_file
///    - cmdline string
///    - ramdisks files.
///  TODO:
///     - Unittests.
///     - Add support to write default_mem & default_cpus, flags.
///     - Various validity checks: E.g: kernel is a bzImage.
use std::ffi::CString;
use std::fmt::Debug;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom, Write};
use std::mem::size_of;
use std::path::Path;

const DEFAULT_SECTIONS_COUNT: u16 = 3;

#[derive(Clone, Debug)]
pub struct SignEnclaveInfo {
    pub signing_certificate: Vec<u8>,
    pub private_key: Vec<u8>,
}

impl SignEnclaveInfo {
    pub fn new(cert_path: &str, key_path: &str) -> Result<Self, String> {
        let mut certificate_file = File::open(cert_path)
            .map_err(|err| format!("Could not open the certificate file: {:?}", err))?;
        let mut signing_certificate = Vec::new();
        certificate_file
            .read_to_end(&mut signing_certificate)
            .map_err(|err| format!("Could not read the certificate file: {:?}", err))?;

        let mut key_file = File::open(key_path)
            .map_err(|err| format!("Could not open the key file: {:?}", err))?;
        let mut private_key = Vec::new();
        key_file
            .read_to_end(&mut private_key)
            .map_err(|err| format!("Could not read the key file: {:?}", err))?;

        Ok(SignEnclaveInfo {
            signing_certificate,
            private_key,
        })
    }
}

/// Utility function to calculate PCRs, used at build and describe.
pub fn get_pcrs<T: Digest + Debug + Write + Clone>(
    image_hasher: &mut EifHasher<T>,
    bootstrap_hasher: &mut EifHasher<T>,
    app_hasher: &mut EifHasher<T>,
    cert_hasher: &mut EifHasher<T>,
    hasher: T,
    is_signed: bool,
) -> Result<BTreeMap<String, String>, String> {
    let mut measurements = BTreeMap::new();
    let image_hasher = hex::encode(
        image_hasher
            .tpm_extend_finalize_reset()
            .map_err(|e| format!("Could not get result for image_hasher: {:?}", e))?,
    );
    let bootstrap_hasher = hex::encode(
        bootstrap_hasher
            .tpm_extend_finalize_reset()
            .map_err(|e| format!("Could not get result for bootstrap_hasher: {:?}", e))?,
    );
    let app_hash = hex::encode(
        app_hasher
            .tpm_extend_finalize_reset()
            .map_err(|e| format!("Could not get result for app_hasher: {:?}", e))?,
    );

    // Hash certificate only if signing key is set, otherwise related PCR will be zero
    let cert_hash = if is_signed {
        Some(hex::encode(
            cert_hasher
                .tpm_extend_finalize_reset()
                .map_err(|e| format!("Could not get result for cert_hash: {:?}", e))?,
        ))
    } else {
        None
    };

    measurements.insert("HashAlgorithm".to_string(), format!("{:?}", hasher));
    measurements.insert("PCR0".to_string(), image_hasher);
    measurements.insert("PCR1".to_string(), bootstrap_hasher);
    measurements.insert("PCR2".to_string(), app_hash);
    if let Some(cert_hash) = cert_hash {
        measurements.insert("PCR8".to_string(), cert_hash);
    }

    Ok(measurements)
}

pub struct EifBuilder<T: Digest + Debug + Write + Clone> {
    kernel: File,
    cmdline: Vec<u8>,
    ramdisks: Vec<File>,
    sign_info: Option<SignEnclaveInfo>,
    signature: Option<Vec<u8>>,
    signature_size: u64,
    metadata: Vec<u8>,
    eif_hdr_flags: u16,
    default_mem: u64,
    default_cpus: u64,
    /// Hash of the whole EifImage.
    pub image_hasher: EifHasher<T>,
    /// Hash of the EifSections provided by Amazon
    /// Kernel + cmdline + First Ramdisk
    pub bootstrap_hasher: EifHasher<T>,
    /// Hash of the remaining ramdisks.
    pub customer_app_hasher: EifHasher<T>,
    /// Hash the signing certificate
    pub certificate_hasher: EifHasher<T>,
    hasher_template: T,
    eif_crc: crc32::Digest,
}

impl<T: Digest + Debug + Write + Clone> EifBuilder<T> {
    pub fn new(
        kernel_path: &Path,
        cmdline: String,
        sign_info: Option<SignEnclaveInfo>,
        hasher: T,
        flags: u16,
        eif_info: EifIdentityInfo,
    ) -> Self {
        let kernel_file = File::open(kernel_path).expect("Invalid kernel path");
        let cmdline = CString::new(cmdline).expect("Invalid cmdline");
        let metadata = serde_json::to_vec(&eif_info).expect("Could not serialize metadata: {}");
        EifBuilder {
            kernel: kernel_file,
            cmdline: cmdline.into_bytes(),
            ramdisks: Vec::new(),
            sign_info,
            signature: None,
            signature_size: 0,
            metadata,
            eif_hdr_flags: flags,
            default_mem: 1024 * 1024 * 1024,
            default_cpus: 2,
            image_hasher: EifHasher::new_without_cache(hasher.clone())
                .expect("Could not create image_hasher"),
            bootstrap_hasher: EifHasher::new_without_cache(hasher.clone())
                .expect("Could not create bootstrap_hasher"),
            customer_app_hasher: EifHasher::new_without_cache(hasher.clone())
                .expect("Could not create customer app hasher"),
            certificate_hasher: EifHasher::new_without_cache(hasher.clone())
                .expect("Could not create certificate hasher"),
            hasher_template: hasher,
            eif_crc: crc32::Digest::new_with_initial(crc32::IEEE, 0),
        }
    }

    pub fn is_signed(&mut self) -> bool {
        self.sign_info.is_some()
    }

    pub fn add_ramdisk(&mut self, ramdisk_path: &Path) {
        let ramdisk_file = File::open(ramdisk_path).expect("Invalid ramdisk path");
        self.ramdisks.push(ramdisk_file);
    }

    /// The first two sections are the kernel and the cmdline and the last is metadata.
    fn num_sections(&self) -> u16 {
        DEFAULT_SECTIONS_COUNT + self.ramdisks.len() as u16 + self.sign_info.iter().count() as u16
    }

    fn sections_offsets(&self) -> [u64; MAX_NUM_SECTIONS] {
        let mut result = [0; MAX_NUM_SECTIONS];
        result[0] = self.kernel_offset();
        result[1] = self.cmdline_offset();
        result[2] = self.metadata_offset();

        for i in 0..self.ramdisks.len() {
            result[i + DEFAULT_SECTIONS_COUNT as usize] = self.ramdisk_offset(i);
        }

        if self.sign_info.is_some() {
            result[DEFAULT_SECTIONS_COUNT as usize + self.ramdisks.len()] = self.signature_offset();
        }

        result
    }

    fn sections_sizes(&self) -> [u64; MAX_NUM_SECTIONS] {
        let mut result = [0; MAX_NUM_SECTIONS];

        result[0] = self.kernel_size();
        result[1] = self.cmdline_size();
        result[2] = self.metadata_size();

        for i in 0..self.ramdisks.len() {
            result[i + DEFAULT_SECTIONS_COUNT as usize] = self.ramdisk_size(&self.ramdisks[i]);
        }

        if self.sign_info.is_some() {
            result[DEFAULT_SECTIONS_COUNT as usize + self.ramdisks.len()] = self.signature_size();
        }

        result
    }

    fn eif_header_offset(&self) -> u64 {
        0
    }

    fn kernel_offset(&self) -> u64 {
        self.eif_header_offset() + EifHeader::size() as u64
    }

    fn kernel_size(&self) -> u64 {
        self.kernel.metadata().unwrap().len() as u64
    }

    fn cmdline_offset(&self) -> u64 {
        self.kernel_offset() + EifSectionHeader::size() as u64 + self.kernel_size()
    }

    fn cmdline_size(&self) -> u64 {
        self.cmdline.len() as u64
    }

    fn ramdisk_offset(&self, index: usize) -> u64 {
        self.metadata_offset()
            + self.metadata_size()
            + EifSectionHeader::size() as u64
            + self.ramdisks[0..index]
                .iter()
                .fold(0, |mut total_len, file| {
                    total_len += file.metadata().expect("Invalid ramdisk metadata").len()
                        + EifSectionHeader::size() as u64;
                    total_len
                })
    }

    fn ramdisk_size(&self, ramdisk: &File) -> u64 {
        ramdisk.metadata().unwrap().len() as u64
    }

    fn signature_offset(&self) -> u64 {
        let index = self.ramdisks.len() - 1;
        self.ramdisk_offset(index)
            + EifSectionHeader::size() as u64
            + self.ramdisk_size(&self.ramdisks[index])
    }

    fn signature_size(&self) -> u64 {
        self.signature_size
    }

    fn metadata_offset(&self) -> u64 {
        self.cmdline_offset() + EifSectionHeader::size() as u64 + self.cmdline_size()
    }

    fn metadata_size(&self) -> u64 {
        self.metadata.len() as u64
    }

    /// Generate the signature of a certain PCR.
    fn generate_pcr_signature(
        &mut self,
        register_index: i32,
        register_value: Vec<u8>,
    ) -> PcrSignature {
        let sign_info = self.sign_info.as_ref().unwrap();
        let signing_certificate = sign_info.signing_certificate.clone();
        let pcr_info = PcrInfo::new(register_index, register_value);

        let payload = to_vec(&pcr_info).expect("Could not serialize PCR info");
        let private_key = PKey::private_key_from_pem(&sign_info.private_key)
            .expect("Could not deserialize the PEM-formatted private key");

        let signature =
            CoseSign1::new::<Openssl>(&payload, &HeaderMap::new(), private_key.as_ref())
                .unwrap()
                .as_bytes(false)
                .unwrap();

        PcrSignature {
            signing_certificate,
            signature,
        }
    }

    /// Generate the signature of the EIF.
    /// eif_signature = [pcr0_signature]
    fn generate_eif_signature(&mut self, measurements: &BTreeMap<String, String>) {
        let pcr0_index = 0;
        let pcr0_value = hex::decode(measurements.get("PCR0").unwrap()).unwrap();
        let pcr0_signature = self.generate_pcr_signature(pcr0_index, pcr0_value);

        let eif_signature = vec![pcr0_signature];
        let serialized_signature =
            to_vec(&eif_signature).expect("Could not serialize the signature");
        self.signature_size = serialized_signature.len() as u64;
        self.signature = Some(serialized_signature)
    }

    pub fn header(&mut self) -> EifHeader {
        EifHeader {
            magic: EIF_MAGIC,
            version: crate::defs::CURRENT_VERSION,
            flags: self.eif_hdr_flags,
            default_mem: self.default_mem,
            default_cpus: self.default_cpus,
            reserved: 0,
            num_sections: self.num_sections(),
            section_offsets: self.sections_offsets(),
            section_sizes: self.sections_sizes(),
            unused: 0,
            eif_crc32: self.eif_crc.sum32(),
        }
    }

    /// Compute the crc for the whole enclave image, excluding the
    /// eif_crc32 field from the EIF header.
    pub fn compute_crc(&mut self) {
        let eif_header = self.header();
        let eif_buffer = eif_header.to_be_bytes();
        // The last field of the EifHeader is the CRC itself, so we need
        // to exclude it from contributing to the CRC.
        let len_without_crc = eif_buffer.len() - size_of::<u32>();
        self.eif_crc.write(&eif_buffer[..len_without_crc]);

        let eif_section = EifSectionHeader {
            section_type: EifSectionType::EifSectionKernel,
            flags: 0,
            section_size: self.kernel_size(),
        };

        let eif_buffer = eif_section.to_be_bytes();
        self.eif_crc.write(&eif_buffer[..]);
        let mut kernel_file = &self.kernel;

        kernel_file
            .seek(SeekFrom::Start(0))
            .expect("Could not seek kernel to beginning");
        let mut buffer = Vec::new();
        kernel_file
            .read_to_end(&mut buffer)
            .expect("Failed to read kernel content");

        self.eif_crc.write(&buffer[..]);

        let eif_section = EifSectionHeader {
            section_type: EifSectionType::EifSectionCmdline,
            flags: 0,
            section_size: self.cmdline_size(),
        };

        let eif_buffer = eif_section.to_be_bytes();
        self.eif_crc.write(&eif_buffer[..]);
        self.eif_crc.write(&self.cmdline[..]);

        let eif_section = EifSectionHeader {
            section_type: EifSectionType::EifSectionMetadata,
            flags: 0,
            section_size: self.metadata_size(),
        };

        let eif_buffer = eif_section.to_be_bytes();
        self.eif_crc.write(&eif_buffer[..]);
        self.eif_crc.write(&self.metadata[..]);

        for mut ramdisk in &self.ramdisks {
            let eif_section = EifSectionHeader {
                section_type: EifSectionType::EifSectionRamdisk,
                flags: 0,
                section_size: self.ramdisk_size(ramdisk),
            };

            let eif_buffer = eif_section.to_be_bytes();
            self.eif_crc.write(&eif_buffer[..]);

            ramdisk
                .seek(SeekFrom::Start(0))
                .expect("Could not seek kernel to begining");
            let mut buffer = Vec::new();
            ramdisk
                .read_to_end(&mut buffer)
                .expect("Failed to read kernel content");
            self.eif_crc.write(&buffer[..]);
        }

        if let Some(signature) = &self.signature {
            let eif_section = EifSectionHeader {
                section_type: EifSectionType::EifSectionSignature,
                flags: 0,
                section_size: self.signature_size(),
            };

            let eif_buffer = eif_section.to_be_bytes();
            self.eif_crc.write(&eif_buffer[..]);
            self.eif_crc.write(&signature[..]);
        }
    }

    pub fn write_header(&mut self, file: &mut File) {
        let eif_header = self.header();
        file.seek(SeekFrom::Start(self.eif_header_offset())).expect(
            "Could not seek while writing eif \
             header",
        );
        let eif_buffer = eif_header.to_be_bytes();
        file.write_all(&eif_buffer[..])
            .expect("Failed to write eif header");
    }

    pub fn write_kernel(&mut self, eif_file: &mut File) {
        let eif_section = EifSectionHeader {
            section_type: EifSectionType::EifSectionKernel,
            flags: 0,
            section_size: self.kernel_size(),
        };

        eif_file
            .seek(SeekFrom::Start(self.kernel_offset()))
            .expect("Could not seek while writing kernel section");
        let eif_buffer = eif_section.to_be_bytes();
        eif_file
            .write_all(&eif_buffer[..])
            .expect("Failed to write kernel header");
        let mut kernel_file = &self.kernel;

        kernel_file
            .seek(SeekFrom::Start(0))
            .expect("Could not seek kernel to begining");
        let mut buffer = Vec::new();
        kernel_file
            .read_to_end(&mut buffer)
            .expect("Failed to read kernel content");

        eif_file
            .write_all(&buffer[..])
            .expect("Failed to write kernel data");
    }

    pub fn write_cmdline(&mut self, eif_file: &mut File) {
        let eif_section = EifSectionHeader {
            section_type: EifSectionType::EifSectionCmdline,
            flags: 0,
            section_size: self.cmdline_size(),
        };

        eif_file
            .seek(SeekFrom::Start(self.cmdline_offset()))
            .expect(
                "Could not seek while writing
        cmdline section",
            );
        let eif_buffer = eif_section.to_be_bytes();
        eif_file
            .write_all(&eif_buffer[..])
            .expect("Failed to write cmdline header");

        eif_file
            .write_all(&self.cmdline[..])
            .expect("Failed write cmdline header");
    }

    pub fn write_metadata(&mut self, eif_file: &mut File) {
        let eif_section = EifSectionHeader {
            section_type: EifSectionType::EifSectionMetadata,
            flags: 0,
            section_size: self.metadata_size(),
        };

        eif_file
            .seek(SeekFrom::Start(self.metadata_offset()))
            .expect("Could not seek while writing metadata section");

        let eif_buffer = eif_section.to_be_bytes();
        eif_file
            .write_all(&eif_buffer[..])
            .expect("Failed to write metadata header");

        eif_file
            .write_all(&self.metadata)
            .expect("Failed to write metadata content");
    }

    pub fn write_ramdisks(&mut self, eif_file: &mut File) {
        for (index, mut ramdisk) in self.ramdisks.iter().enumerate() {
            let eif_section = EifSectionHeader {
                section_type: EifSectionType::EifSectionRamdisk,
                flags: 0,
                section_size: self.ramdisk_size(ramdisk),
            };

            eif_file
                .seek(SeekFrom::Start(self.ramdisk_offset(index)))
                .expect(
                    "Could not seek while writing
        kernel section",
                );
            let eif_buffer = eif_section.to_be_bytes();
            eif_file
                .write_all(&eif_buffer[..])
                .expect("Failed to write section header");

            ramdisk
                .seek(SeekFrom::Start(0))
                .expect("Could not seek ramdisk to beginning");
            let mut buffer = Vec::new();
            ramdisk
                .read_to_end(&mut buffer)
                .expect("Failed to read ramdisk content");
            eif_file
                .write_all(&buffer[..])
                .expect("Failed to write ramdisk data");
        }
    }

    pub fn write_signature(&mut self, eif_file: &mut File) {
        if let Some(signature) = &self.signature {
            let eif_section = EifSectionHeader {
                section_type: EifSectionType::EifSectionSignature,
                flags: 0,
                section_size: self.signature_size(),
            };

            eif_file
                .seek(SeekFrom::Start(self.signature_offset()))
                .expect("Could not seek while writing signature section");
            let eif_buffer = eif_section.to_be_bytes();
            eif_file
                .write_all(&eif_buffer[..])
                .expect("Failed to write signature header");

            eif_file
                .write_all(&signature[..])
                .expect("Failed write signature header");
        }
    }

    pub fn write_to(&mut self, output_file: &mut File) -> BTreeMap<String, String> {
        self.measure();
        let measurements = get_pcrs(
            &mut self.image_hasher,
            &mut self.bootstrap_hasher,
            &mut self.customer_app_hasher,
            &mut self.certificate_hasher,
            self.hasher_template.clone(),
            self.sign_info.is_some(),
        )
        .expect("Failed to get measurements");
        if self.sign_info.is_some() {
            self.generate_eif_signature(&measurements);
        }
        self.compute_crc();
        self.write_header(output_file);
        self.write_kernel(output_file);
        self.write_cmdline(output_file);
        self.write_metadata(output_file);
        self.write_ramdisks(output_file);
        self.write_signature(output_file);
        measurements
    }

    pub fn measure(&mut self) {
        let mut kernel_file = &self.kernel;
        kernel_file
            .seek(SeekFrom::Start(0))
            .expect("Could not seek kernel to beginning");
        let mut buffer = Vec::new();
        kernel_file
            .read_to_end(&mut buffer)
            .expect("Failed to read kernel content");
        self.image_hasher.write_all(&buffer[..]).unwrap();
        self.bootstrap_hasher.write_all(&buffer[..]).unwrap();

        self.image_hasher.write_all(&self.cmdline[..]).unwrap();
        self.bootstrap_hasher.write_all(&self.cmdline[..]).unwrap();

        for (index, mut ramdisk) in self.ramdisks.iter().enumerate() {
            ramdisk
                .seek(SeekFrom::Start(0))
                .expect("Could not seek kernel to beginning");
            let mut buffer = Vec::new();
            ramdisk
                .read_to_end(&mut buffer)
                .expect("Failed to read kernel content");
            self.image_hasher.write_all(&buffer[..]).unwrap();
            // The first ramdisk is provided by amazon and it contains the
            // code to bootstrap the docker container.
            if index == 0 {
                self.bootstrap_hasher.write_all(&buffer[..]).unwrap();
            } else {
                self.customer_app_hasher.write_all(&buffer[..]).unwrap();
            }
        }

        if let Some(sign_info) = self.sign_info.as_ref() {
            let cert = openssl::x509::X509::from_pem(&sign_info.signing_certificate[..]).unwrap();
            let cert_der = cert.to_der().unwrap();
            // This is equivalent to extend(cert.digest(sha384)), since hasher is going to
            // hash the DER certificate (cert.digest()) and then tpm_extend_finalize_reset
            // will do the extend.
            self.certificate_hasher.write_all(&cert_der).unwrap();
        }
    }
}

/// PCR Signature verifier that checks the validity of
/// the certificate used to sign the enclave
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct PcrSignatureChecker {
    signing_certificate: Vec<u8>,
    signature: Vec<u8>,
}

impl PcrSignatureChecker {
    pub fn new(pcr_signature: &PcrSignature) -> Self {
        PcrSignatureChecker {
            signing_certificate: pcr_signature.signing_certificate.clone(),
            signature: pcr_signature.signature.clone(),
        }
    }

    /// Reads EIF section headers and looks for a signature.
    /// Seek to the signature section, if present, and save the certificate and signature
    pub fn from_eif(eif_path: &str) -> Result<Self, String> {
        let mut signing_certificate = Vec::new();
        let mut signature = Vec::new();

        let mut curr_seek = 0;
        let mut eif_file =
            File::open(eif_path).map_err(|e| format!("Failed to open the EIF file: {:?}", e))?;

        // Skip header
        let mut header_buf = vec![0u8; EifHeader::size()];
        eif_file
            .read_exact(&mut header_buf)
            .map_err(|e| format!("Error while reading EIF header: {:?}", e))?;

        curr_seek += EifHeader::size();
        eif_file
            .seek(SeekFrom::Start(curr_seek as u64))
            .map_err(|e| format!("Failed to seek file from start: {:?}", e))?;

        let mut section_buf = vec![0u8; EifSectionHeader::size()];

        // Read all section headers and skip if different from signature section
        while eif_file.read_exact(&mut section_buf).is_ok() {
            let section = EifSectionHeader::from_be_bytes(&section_buf)
                .map_err(|e| format!("Error extracting EIF section header: {:?}", e))?;
            curr_seek += EifSectionHeader::size();

            if section.section_type == EifSectionType::EifSectionSignature {
                let mut buf = vec![0u8; section.section_size as usize];
                eif_file
                    .seek(SeekFrom::Start(curr_seek as u64))
                    .map_err(|e| format!("Failed to seek after EIF section header: {:?}", e))?;
                eif_file.read_exact(&mut buf).map_err(|e| {
                    format!("Error while reading signature section from EIF: {:?}", e)
                })?;

                // Deserialize PCR signature structure and save certificate and signature
                let des_sign: Vec<PcrSignature> = from_slice(&buf[..])
                    .map_err(|e| format!("Error deserializing certificate: {:?}", e))?;

                signing_certificate = des_sign[0].signing_certificate.clone();
                signature = des_sign[0].signature.clone();
            }

            curr_seek += section.section_size as usize;
            eif_file
                .seek(SeekFrom::Start(curr_seek as u64))
                .map_err(|e| format!("Failed to seek after EIF section: {:?}", e))?;
        }

        Ok(Self {
            signing_certificate,
            signature,
        })
    }

    pub fn is_empty(&self) -> bool {
        self.signing_certificate.len() == 0 && self.signature.len() == 0
    }

    /// Verifies the validity of the signing certificate
    pub fn verify(&mut self) -> Result<(), String> {
        let signature = CoseSign1::from_bytes(&self.signature[..])
            .map_err(|err| format!("Could not deserialize the signature: {:?}", err))?;
        let cert = openssl::x509::X509::from_pem(&self.signing_certificate[..])
            .map_err(|_| "Could not deserialize the signing certificate".to_string())?;
        let public_key = cert
            .public_key()
            .map_err(|_| "Could not get the public key from the signing certificate".to_string())?;

        // Verify the signature
        let result = signature
            .verify_signature::<Openssl>(public_key.as_ref())
            .map_err(|err| format!("Could not verify EIF signature: {:?}", err))?;
        if !result {
            return Err("The EIF signature is not valid".to_string());
        }

        // Verify that the signing certificate is not expired
        let current_time = Asn1Time::days_from_now(0).map_err(|err| err.to_string())?;
        if current_time
            .compare(cert.not_after())
            .map_err(|err| err.to_string())?
            == Ordering::Greater
            || current_time
                .compare(cert.not_before())
                .map_err(|err| err.to_string())?
                == Ordering::Less
        {
            return Err("The signing certificate is expired".to_string());
        }

        Ok(())
    }
}


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

extern crate cbindgen;

use cbindgen::Config;
use std::env;
use std::path::{Path, PathBuf};

fn main() {
    let crate_env = env::var("CARGO_MANIFEST_DIR").unwrap();
    let crate_path = Path::new(&crate_env);
    let config = Config::from_root_or_default(crate_path);
    let out_path = output_dir();
    cbindgen::Builder::new()
        .with_crate(crate_path.to_str().unwrap())
        .with_config(config)
        .generate()
        .expect("Unable to generate bindings")
        .write_to_file(out_path.join("nsm.h"));
}

/// Sets target to target/$PROFILE/
fn output_dir() -> PathBuf {
    env::var("OUT_DIR")
        .map(PathBuf::from)
        .map(|dir| dir.ancestors().nth(3).unwrap().to_owned())
        .unwrap()
}


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

#![allow(clippy::missing_safety_doc)]

//! ***NitroSecureModule wrappers for non-Rust callers***
//! # Overview
//! This module implements wrappers over the NSM Rust API which enable
//! access to the API for non-Rust callers (ex.: C/C++ etc.).

pub use aws_nitro_enclaves_nsm_api::api::{Digest, ErrorCode};
use aws_nitro_enclaves_nsm_api::api::{Request, Response};
use aws_nitro_enclaves_nsm_api::driver::{nsm_exit, nsm_init, nsm_process_request};
use serde_bytes::ByteBuf;
use std::ptr::copy_nonoverlapping;
use std::{cmp, slice};

#[repr(C)]
pub struct NsmDescription {
    pub version_major: u16,
    pub version_minor: u16,
    pub version_patch: u16,
    pub module_id: [u8; 100],
    pub module_id_len: u32,
    pub max_pcrs: u16,
    pub locked_pcrs: [u16; 64],
    pub locked_pcrs_len: u32,
    pub digest: Digest,
}

/// NSM library initialization function.
/// *Returns*: A descriptor for the opened device file.
#[no_mangle]
pub extern "C" fn nsm_lib_init() -> i32 {
    nsm_init()
}

/// NSM library exit function.
/// *Argument 1 (input)*: The descriptor for the opened device file, as
/// obtained from `nsm_init()`.
#[no_mangle]
pub extern "C" fn nsm_lib_exit(fd: i32) {
    nsm_exit(fd)
}

/// NSM `ExtendPCR` operation for non-Rust callers.
/// *Argument 1 (input)*: The descriptor to the NSM device file.
/// *Argument 2 (input)*: The index of the PCR to extend.
/// *Argument 3 (input)*: The raw data to extend the PCR with.
/// *Argument 4 (input)*: The length of the raw data, in bytes.
/// *Argument 5 (output)*: The data from the extended PCR.
/// *Argument 6 (input/output)*: The capacity of the extended PCR data
/// buffer as input, the actual size of the buffer as output.
/// *Returns*: The status of the operation.
#[no_mangle]
pub unsafe extern "C" fn nsm_extend_pcr(
    fd: i32,
    index: u16,
    data: *const u8,
    data_len: u32,
    pcr_data: *mut u8,
    pcr_data_len: &mut u32,
) -> ErrorCode {
    let data_vec = nsm_get_vec_from_raw(data, data_len);
    match data_vec {
        Some(_) => (),
        None => return ErrorCode::InvalidArgument,
    }

    let request = Request::ExtendPCR {
        index,
        data: data_vec.unwrap(),
    };

    match nsm_process_request(fd, request) {
        Response::ExtendPCR { data: pcr } => nsm_get_raw_from_vec(&pcr, pcr_data, pcr_data_len),
        Response::Error(err) => err,
        _ => ErrorCode::InvalidResponse,
    }
}

/// NSM `DescribePCR` operation for non-Rust callers.
/// *Argument 1 (input)*: The descriptor to the NSM device file.
/// *Argument 2 (input)*: The index of the PCR to be described.
/// *Argument 3 (output)*: The lock state of the PCR.
/// *Argument 4 (output)*: The buffer that will hold the PCR data.
/// *Argument 5 (input / output)*: The PCR data buffer capacity (as input)
/// and the actual size of the received data (as output).
/// *Returns*: The status of the operation.
#[no_mangle]
pub unsafe extern "C" fn nsm_describe_pcr(
    fd: i32,
    index: u16,
    lock: &mut bool,
    data: *mut u8,
    data_len: &mut u32,
) -> ErrorCode {
    let request = Request::DescribePCR { index };

    match nsm_process_request(fd, request) {
        Response::DescribePCR {
            lock: pcr_lock,
            data: pcr_data,
        } => {
            *lock = pcr_lock;
            nsm_get_raw_from_vec(&pcr_data, data, data_len)
        }
        Response::Error(err) => err,
        _ => ErrorCode::InvalidResponse,
    }
}

/// NSM `LockPCR` operation for non-Rust callers.
/// *Argument 1 (input)*: The descriptor to the NSM device file.
/// *Argument 2 (input)*: The PCR to be locked.
/// *Returns*: The status of the operation.
#[no_mangle]
pub extern "C" fn nsm_lock_pcr(fd: i32, index: u16) -> ErrorCode {
    let request = Request::LockPCR { index };

    match nsm_process_request(fd, request) {
        Response::LockPCR => ErrorCode::Success,
        Response::Error(err) => err,
        _ => ErrorCode::InvalidResponse,
    }
}

/// NSM `LockPCRs` operation for non-Rust callers.
/// *Argument 1 (input)*: The descriptor to the NSM device file.
/// *Argument 2 (input)*: The range value for `[0, range)` to be locked.
/// *Returns*: The status of the operation.
#[no_mangle]
pub extern "C" fn nsm_lock_pcrs(fd: i32, range: u16) -> ErrorCode {
    let request = Request::LockPCRs { range };

    match nsm_process_request(fd, request) {
        Response::LockPCRs => ErrorCode::Success,
        Response::Error(err) => err,
        _ => ErrorCode::InvalidResponse,
    }
}

/// NSM `Describe` operation for non-Rust callers.
/// *Argument 1 (input)*: The descriptor to the NSM device file.
/// *Argument 2 (output)*: The obtained raw NSM description.
/// *Returns*: The status of the operation.
#[no_mangle]
pub extern "C" fn nsm_get_description(fd: i32, nsm_description: &mut NsmDescription) -> ErrorCode {
    let request = Request::DescribeNSM;

    match nsm_process_request(fd, request) {
        Response::DescribeNSM {
            version_major,
            version_minor,
            version_patch,
            module_id,
            max_pcrs,
            locked_pcrs,
            digest,
        } => {
            nsm_description.version_major = version_major;
            nsm_description.version_minor = version_minor;
            nsm_description.version_patch = version_patch;
            nsm_description.max_pcrs = max_pcrs;

            match digest {
                Digest::SHA256 => {
                    nsm_description.digest = Digest::SHA256;
                }
                Digest::SHA384 => {
                    nsm_description.digest = Digest::SHA384;
                }
                Digest::SHA512 => {
                    nsm_description.digest = Digest::SHA512;
                }
            }
            nsm_description.locked_pcrs_len = locked_pcrs.len() as u32;

            for (i, val) in locked_pcrs.iter().enumerate() {
                nsm_description.locked_pcrs[i] = *val;
            }

            let module_id_len = cmp::min(nsm_description.module_id.len() - 1, module_id.len());
            nsm_description.module_id[0..module_id_len]
                .copy_from_slice(&module_id.as_bytes()[0..module_id_len]);
            nsm_description.module_id[module_id_len] = 0;
            nsm_description.module_id_len = module_id_len as u32;

            ErrorCode::Success
        }
        Response::Error(err) => err,
        _ => ErrorCode::InvalidResponse,
    }
}

/// Get an optional byte buffer from user data.
/// *Argument 1 (input)*: User data.
/// *Argument 2 (input)*: Size of the user data buffer.
/// *Returns*: The optional byte buffer.
unsafe fn get_byte_buf_from_user_data(data: *const u8, len: u32) -> Option<ByteBuf> {
    let data_vec = nsm_get_vec_from_raw(data, len);
    data_vec.map(ByteBuf::from)
}

/// NSM `GetAttestationDoc` operation for non-Rust callers.
/// *Argument 1 (input)*: The descriptor to the NSM device file.
/// *Argument 2 (input)*: User data.
/// *Argument 3 (input)*: The size of the user data buffer.
/// *Argument 4 (input)*: Nonce data.
/// *Argument 5 (input)*: The size of the nonce data buffer.
/// *Argument 6 (input)*: Public key data.
/// *Argument 7 (input)*: The size of the public key data buffer.
/// *Argument 8 (output)*: The obtained attestation document.
/// *Argument 9 (input / output)*: The document buffer capacity (as input)
/// and the size of the received document (as output).
/// *Returns*: The status of the operation.
#[no_mangle]
pub unsafe extern "C" fn nsm_get_attestation_doc(
    fd: i32,
    user_data: *const u8,
    user_data_len: u32,
    nonce_data: *const u8,
    nonce_len: u32,
    pub_key_data: *const u8,
    pub_key_len: u32,
    att_doc_data: *mut u8,
    att_doc_len: &mut u32,
) -> ErrorCode {
    let request = Request::Attestation {
        user_data: get_byte_buf_from_user_data(user_data, user_data_len),
        nonce: get_byte_buf_from_user_data(nonce_data, nonce_len),
        public_key: get_byte_buf_from_user_data(pub_key_data, pub_key_len),
    };

    match nsm_process_request(fd, request) {
        Response::Attestation {
            document: attestation_doc,
        } => nsm_get_raw_from_vec(&attestation_doc, att_doc_data, att_doc_len),
        Response::Error(err) => err,
        _ => ErrorCode::InvalidResponse,
    }
}

/// NSM `GetRandom` operation for non-Rust callers. Returns up to 256 bytes of random data.
/// *fd (input)*: A valid descriptor to the NSM device.
/// *buf (output)*: A valid buffer to place the random data in.
/// *buf_len (input / output)*: The length of the passed buffer and the length of the output data
///                             if the function finishes with ErrorCode::Success.
#[no_mangle]
pub unsafe extern "C" fn nsm_get_random(fd: i32, buf: *mut u8, buf_len: &mut usize) -> ErrorCode {
    if fd < 0 || buf.is_null() || buf_len == &0 {
        return ErrorCode::InvalidArgument;
    }
    match nsm_process_request(fd, Request::GetRandom) {
        Response::GetRandom { random } => {
            *buf_len = std::cmp::min(*buf_len, random.len());
            std::ptr::copy_nonoverlapping(random.as_ptr(), buf, *buf_len);
            ErrorCode::Success
        }
        Response::Error(err) => err,
        _ => ErrorCode::InvalidResponse,
    }
}

/// Obtain a vector from a raw C-style pointer and length.
/// *Argument 1 (input)*: The raw input pointer.
/// *Argument 2 (input)*: The length of the input buffer.
/// *Returns*: The corresponding Rust vector.
unsafe fn nsm_get_vec_from_raw<T: Clone>(data: *const T, data_len: u32) -> Option<Vec<T>> {
    if data.is_null() {
        return None;
    }

    let slice = slice::from_raw_parts(data, data_len as usize);
    Some(slice.to_vec())
}

/// Fill a raw buffer using the data from a vector.
/// *Argument 1 (input)*: The input vector's slice.
/// *Argument 2 (output)*: The raw buffer to be filled with the vector data.
/// *Argument 3 (input / output)*: The capacity of the output buffer as input and
/// the actual size of the written data as output.
/// *Returns*: The status of the operation.
unsafe fn nsm_get_raw_from_vec<T>(input: &[T], output: *mut T, output_size: &mut u32) -> ErrorCode {
    if output.is_null() {
        *output_size = 0;
        return ErrorCode::BufferTooSmall;
    }

    let result = if *output_size as usize >= input.len() {
        ErrorCode::Success
    } else {
        ErrorCode::BufferTooSmall
    };

    *output_size = cmp::min(*output_size, input.len() as u32);
    copy_nonoverlapping(input.as_ptr(), output, *output_size as usize);

    result
}


// Copyright 2019-2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
// Author: Andrei Trandafir <aatrand@amazon.com>

//! ***NitroSecureModule test for Rust API***
//! # Overview
//! This module implements a comprehensive run-time test for the
//! NSM Rust API.

use aws_nitro_enclaves_nsm_api::api::{Digest, Request, Response};
use aws_nitro_enclaves_nsm_api::driver::{nsm_exit, nsm_init, nsm_process_request};
use serde_bytes::ByteBuf;
use std::collections::BTreeSet;

const RESERVED_PCRS: u16 = 5;

/// Structure holding PCR status.
struct PcrData {
    lock: bool,
    data: Vec<u8>,
}

/// Structure holding the NSM description.
struct NsmDescription {
    version_major: u16,
    version_minor: u16,
    version_patch: u16,
    module_id: String,
    max_pcrs: u16,
    locked_pcrs: BTreeSet<u16>,
    digest: Digest,
}

/// Get the description of the NSM.
/// *Argument 1 (input)*: Context from `nsm_init()`.
/// *Returns*: A description structure.
fn get_nsm_description(ctx: i32) -> NsmDescription {
    let response = nsm_process_request(ctx, Request::DescribeNSM);
    match response {
        Response::DescribeNSM {
            version_major,
            version_minor,
            version_patch,
            module_id,
            max_pcrs,
            locked_pcrs,
            digest,
        } => NsmDescription {
            version_major,
            version_minor,
            version_patch,
            module_id,
            max_pcrs,
            locked_pcrs,
            digest,
        },
        _ => panic!(
            "[Error] Request::DescribeNSM got invalid response: {:?}",
            response
        ),
    }
}

/// Get the length of a PCR in bytes, based on digest.
/// *Argument 1 (input)*: The NSM description with digest information.
/// *Returns*: PCR length in bytes.
fn get_pcr_len(description: &NsmDescription) -> usize {
    match description.digest {
        Digest::SHA256 => 32,
        Digest::SHA384 => 48,
        Digest::SHA512 => 64,
    }
}

/// Test the initial state of the PCRs.
/// *Argument 1 (input)*: Context from `nsm_init()`.
/// *Argument 2 (input)*: The NSM description.
fn check_initial_pcrs(ctx: i32, description: &NsmDescription) {
    let expected_pcr_len = get_pcr_len(description);

    // First, get the description of all available PCRs.
    let pcr_data: Vec<PcrData> = (0..description.max_pcrs)
        .map(|pcr| {
            let response = nsm_process_request(ctx, Request::DescribePCR { index: pcr });
            match response {
                Response::DescribePCR { lock, data } => {
                    assert_eq!(
                        data.len(),
                        expected_pcr_len,
                        "[Error] Request::DescribePCR got invalid response length."
                    );
                    PcrData { lock, data }
                }
                _ => panic!(
                    "[Error] Request::DescribePCR got invalid response: {:?}",
                    response
                ),
            }
        })
        .collect();
    println!(
        "Checked Request::DescribePCR for PCRs [0..{}).",
        description.max_pcrs
    );

    let zeroed_pcr: Vec<u8> = vec![0; expected_pcr_len];
    let locked_pcrs: Vec<u16> = description.locked_pcrs.iter().cloned().collect();
    let locked_pcrs_ref: Vec<u16> = (0..16).collect();

    // PCRs [0..4) must not be empty (shound contain non-zero bytes).
    for (index, pcr) in pcr_data.iter().enumerate().take(RESERVED_PCRS as usize) {
        if index != 3 {
            assert_ne!(
                pcr.data, zeroed_pcr,
                "[Error] PCR {} must not be empty.",
                index
            );
        } else {
            assert_eq!(pcr.data, zeroed_pcr, "[Error] PCR {} must be empty.", index);
        }
    }
    println!("Checked that PCRs [0..{}) are not empty.", RESERVED_PCRS);

    for (index, pcr) in pcr_data.iter().enumerate() {
        println!("PCR {} has value {:?}", index, pcr.data);
    }

    // All other PCRs should be empty.
    for (index, pcr) in pcr_data.iter().enumerate().skip(RESERVED_PCRS as usize) {
        assert_eq!(pcr.data, zeroed_pcr, "[Error] PCR {} must be empty.", index);
    }
    println!(
        "Checked that PCRs [{}..{}) are empty.",
        RESERVED_PCRS, description.max_pcrs
    );

    // PCRs [0..16) should be locked.
    assert_eq!(
        locked_pcrs, locked_pcrs_ref,
        "[Error] Initial locked PCR list is invalid."
    );
    for pcr in 0..16 {
        assert!(
            pcr_data[pcr as usize].lock,
            "[Error] PCR {} must be locked.",
            pcr
        );
    }

    // All other PCRs should not be locked.
    for pcr in 16..description.max_pcrs {
        assert!(
            !pcr_data[pcr as usize].lock,
            "[Error] PCR {} must not be locked.",
            pcr
        );
    }

    println!(
        "Checked that PCRs [0..16) are locked and [16..{}) are not locked.",
        description.max_pcrs
    );
}

/// Check and modify the lock state of the PCRs.
/// *Argument 1 (input)*: Context from `nsm_init()`.
/// *Argument 2 (input)*: The NSM description.
fn check_pcr_locks(ctx: i32, description: &NsmDescription) {
    let dummy_data: Vec<u8> = vec![1, 2, 3];
    let expected_pcr_len = get_pcr_len(description);
    let zeroed_pcr: Vec<u8> = vec![0; expected_pcr_len];
    let mut range = description.max_pcrs;
    let mut response: Response;

    // Test that PCRs [0..16) cannot be locked.
    for index in 0..16 {
        response = nsm_process_request(ctx, Request::LockPCR { index });
        match response {
            Response::Error(_) => (),
            _ => panic!(
                "[Error] PCR {} expected to not be lockable, but got: {:?}",
                index, response
            ),
        }
    }

    println!("Checked Request::LockPCR for PCRs [0..16).");

    // Extend the remaining PCRs multiple times.
    for loop_idx in 0..10 {
        for index in 16..description.max_pcrs {
            let data_copy = dummy_data.clone();
            response = nsm_process_request(
                ctx,
                Request::ExtendPCR {
                    index,
                    data: data_copy,
                },
            );

            match response {
                Response::ExtendPCR { data } => {
                    assert_eq!(
                        data.len(),
                        expected_pcr_len,
                        "[Error] Request::ExtendPCR got invalid response."
                    );
                    assert_ne!(data, zeroed_pcr, "[Error] PCR {} must not be empty.", index);
                }
                _ => panic!(
                    "[Error] Request::ExtendPCR got invalid response: {:?}",
                    response
                ),
            }
        }

        println!(
            "[Loop: {}] Checked Request::ExtendedPCR for PCRs [16..{}).",
            loop_idx, description.max_pcrs
        );
    }

    // Lock all remaining PCRs.
    for index in 16..description.max_pcrs {
        response = nsm_process_request(ctx, Request::LockPCR { index });

        match response {
            Response::LockPCR => (),
            _ => panic!(
                "[Error] Request::LockPCR got invalid response: {:?}",
                response
            ),
        }
    }

    println!(
        "Checked Request::LockPCR for PCRs [16..{}).",
        description.max_pcrs
    );

    // Lock PCRs in a valid range.
    response = nsm_process_request(ctx, Request::LockPCRs { range });
    match response {
        Response::LockPCRs => (),
        _ => panic!(
            "[Error] Request::LockPCRs expected to succeed for [0..{}), but got: {:?}",
            range, response
        ),
    }

    // Lock PCRs in an invalid range.
    range += 1;
    response = nsm_process_request(ctx, Request::LockPCRs { range });
    match response {
        Response::Error(_) => (),
        _ => panic!(
            "[Error] Request::LockPCRs expected to fail for [0..{}), but got: {:?}",
            range, response
        ),
    }

    println!(
        "Checked Request::LockPCRs for ranges {} and {}.",
        range - 1,
        range
    );

    // Attempt to extend locked PCRs.
    for index in 0..description.max_pcrs {
        let data_copy = dummy_data.clone();
        response = nsm_process_request(
            ctx,
            Request::ExtendPCR {
                index,
                data: data_copy,
            },
        );

        match response {
            Response::Error(_) => (),
            _ => panic!(
                "[Error] Request::ExtendPCR expected to fail, but got: {:?}",
                response
            ),
        }
    }

    println!(
        "Checked Request::ExtendPCR for locked PCRs [0..{}).",
        description.max_pcrs
    );

    // Describe all PCRs multiple times.
    for loop_idx in 0..10 {
        for index in 0..description.max_pcrs {
            response = nsm_process_request(ctx, Request::DescribePCR { index });

            match response {
                Response::DescribePCR { lock, data } => {
                    assert_eq!(
                        data.len(),
                        expected_pcr_len,
                        "[Error] Request::DescribePCR got invalid response length."
                    );

                    match index {
                        3 => {
                            assert_eq!(data, zeroed_pcr, "[Error] PCR {} must be empty.", index)
                        }
                        RESERVED_PCRS..=15 => {
                            assert_eq!(data, zeroed_pcr, "[Error] PCR {} must be empty.", index)
                        }
                        _ => {
                            assert_ne!(data, zeroed_pcr, "[Error] PCR {} must not be empty.", index)
                        }
                    }
                    assert!(lock, "[Error] PCR {} must be locked.", index);
                }
                _ => panic!(
                    "[Error] Request::ExtendPCR got invalid response: {:?}",
                    response
                ),
            }
        }

        println!(
            "[Loop: {}] Checked Request::DescribePCR for PCRs [0..{}).",
            loop_idx, description.max_pcrs
        );
    }
}

/// Check a single attestation operation.
/// *Argument 1 (input)*: Context from `nsm_init()`.
/// *Argument 2 (input)*: Optional user data.
/// *Argument 3 (input)*: Optional nonce data.
/// *Argument 4 (input)*: Optional public key.
fn check_single_attestation(
    ctx: i32,
    user_data: Option<ByteBuf>,
    nonce: Option<ByteBuf>,
    public_key: Option<ByteBuf>,
) {
    let response = nsm_process_request(
        ctx,
        Request::Attestation {
            user_data,
            nonce,
            public_key,
        },
    );
    match response {
        Response::Attestation { document } => {
            assert_ne!(document.len(), 0, "[Error] Attestation document is empty.");
        }
        _ => panic!(
            "[Error] Request::Attestation got invalid response: {:?}",
            response
        ),
    }
}

/// Check multiple attestation operations.
/// *Argument 1 (input)*: Context from `nsm_init()`.
fn check_attestation(ctx: i32) {
    const DATA_LEN: usize = 1024;
    let dummy_data: Vec<u8> = vec![128; DATA_LEN];

    check_single_attestation(ctx, None, None, None);
    println!("Checked Request::Attestation without any data.");

    check_single_attestation(ctx, Some(ByteBuf::from(&dummy_data[..])), None, None);
    println!(
        "Checked Request::Attestation with user data ({} bytes).",
        DATA_LEN
    );

    check_single_attestation(
        ctx,
        Some(ByteBuf::from(&dummy_data[..])),
        Some(ByteBuf::from(&dummy_data[..])),
        None,
    );
    println!(
        "Checked Request::Attestation with user data and nonce ({} bytes each).",
        DATA_LEN
    );

    check_single_attestation(
        ctx,
        Some(ByteBuf::from(&dummy_data[..])),
        Some(ByteBuf::from(&dummy_data[..])),
        Some(ByteBuf::from(&dummy_data[..])),
    );
    println!(
        "Checked Request::Attestation with user data, nonce and public key ({} bytes each).",
        DATA_LEN
    );
}

fn check_random(ctx: i32) {
    let mut prev_random: Vec<u8> = vec![];

    for _ in 0..16 {
        match nsm_process_request(ctx, Request::GetRandom) {
            Response::GetRandom { random } => {
                assert!(!random.is_empty());
                assert!(prev_random != random);
                prev_random = random;
            }

            resp => panic!(
                "GetRandom: expecting Response::GetRandom, but got {:?} instead",
                resp
            ),
        }
    }
}

fn main() {
    println!("NSM test started.");

    let ctx = nsm_init();
    assert!(ctx >= 0, "[Error] NSM initialization returned {}.", ctx);

    let description = get_nsm_description(ctx);
    assert_eq!(
        description.max_pcrs, 32,
        "[Error] NSM PCR count is {}.",
        description.max_pcrs
    );
    assert!(
        !description.module_id.is_empty(),
        "[Error] NSM module ID is missing."
    );

    println!(
        "NSM description: [major: {}, minor: {}, patch: {}, module_id: {}, max_pcrs: {},
        locked_pcrs: {:?}, digest: {:?}].",
        description.version_major,
        description.version_minor,
        description.version_patch,
        description.module_id,
        description.max_pcrs,
        description.locked_pcrs,
        description.digest
    );

    check_single_attestation(ctx, None, None, None);
    println!("Checked Request::Attestation without any data.");

    check_initial_pcrs(ctx, &description);
    check_pcr_locks(ctx, &description);

    check_attestation(ctx);

    check_random(ctx);

    nsm_exit(ctx);
    println!("NSM test finished.");
}


// Copyright 2019-2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
// Author: Andrei Trandafir <aatrand@amazon.com>
// Author: Andrei Cipu <acipu@amazon.com>

//! ***NitroSecureModule multithread test for Rust API***
//! # Overview
//! This module implements a aggresive run-time test for the
//! NSM Rust API.

use aws_nitro_enclaves_nsm_api::api::{Request, Response};
use aws_nitro_enclaves_nsm_api::driver::{nsm_exit, nsm_init, nsm_process_request};
use serde_bytes::ByteBuf;
use std::convert::TryInto;
use std::sync::atomic;
use std::sync::Arc;
use std::thread;
use std::time;
use threadpool::ThreadPool;

enum ErrorCode {
    AttestationDocumentEmpty = 1,
    AttestationInvalidResponse = 2,
}

/// *Argument 2 (input)*: The NSM description.
fn extend_pcr(ctx: i32, j: usize) {
    let pcr: u16 = ((16 + j) & 15).try_into().unwrap();
    let one: u8 = ((j >> 24) & 0xFF).try_into().unwrap();
    let two: u8 = ((j >> 16) & 0xFF).try_into().unwrap();
    let three: u8 = ((j >> 8) & 0xFF).try_into().unwrap();
    let four: u8 = (j & 0xFF).try_into().unwrap();
    let dummy_data: Vec<u8> = vec![one, two, three, four];
    let mut _response: Response;

    // Extend the remaining PCRs multiple times.
    for _loop_idx in 0..2 {
        let data_copy = dummy_data.clone();
        _response = nsm_process_request(
            ctx,
            Request::ExtendPCR {
                index: pcr,
                data: data_copy,
            },
        );
    }
}

/// Check a single attestation operation.
/// *Argument 1 (input)*: Context from `nsm_init()`.
/// *Argument 2 (input)*: Optional user data.
/// *Argument 3 (input)*: Optional nonce data.
/// *Argument 4 (input)*: Optional public key.
/// Returns Ok(()) in case of success
fn check_single_attestation(
    ctx: i32,
    user_data: Option<ByteBuf>,
    nonce: Option<ByteBuf>,
    public_key: Option<ByteBuf>,
) -> Result<(), ErrorCode> {
    let response = nsm_process_request(
        ctx,
        Request::Attestation {
            user_data,
            nonce,
            public_key,
        },
    );
    match response {
        Response::Attestation { document } => {
            if document.is_empty() {
                println!("[Error] Attestation document is empty.");
                return Err(ErrorCode::AttestationDocumentEmpty);
            }
        }
        _ => {
            println!(
                "[Error] Request::Attestation got invalid response: {:?}",
                response
            );
            return Err(ErrorCode::AttestationInvalidResponse);
        }
    }
    Ok(())
}

/// Check multiple attestation operations.
/// *Argument 1 (input)*: Context from `nsm_init()`.
/// Returns Ok(()) in case of success
fn check_attestation(ctx: i32, lp: usize) -> Result<(), ErrorCode> {
    const DATA_LEN: usize = 1024;
    let dummy_data: Vec<u8> = vec![128; DATA_LEN];
    let mut now = time::Instant::now();

    check_single_attestation(ctx, None, None, None)?;
    println!(
        "attestation loop={} wo/data took {} ns",
        lp,
        now.elapsed().as_nanos()
    );
    now = time::Instant::now();

    check_single_attestation(ctx, Some(ByteBuf::from(&dummy_data[..])), None, None)?;
    println!(
        "attestation loop={} w/data took {} ns",
        lp,
        now.elapsed().as_nanos()
    );
    now = time::Instant::now();

    check_single_attestation(
        ctx,
        Some(ByteBuf::from(&dummy_data[..])),
        Some(ByteBuf::from(&dummy_data[..])),
        None,
    )?;
    println!(
        "attestation loop={} w/data, nonce took {} ns",
        lp,
        now.elapsed().as_nanos()
    );
    now = time::Instant::now();

    check_single_attestation(
        ctx,
        Some(ByteBuf::from(&dummy_data[..])),
        Some(ByteBuf::from(&dummy_data[..])),
        Some(ByteBuf::from(&dummy_data[..])),
    )?;
    println!(
        "attestation loop={} w/user_data, nonce, PK took {} ns",
        lp,
        now.elapsed().as_nanos()
    );

    Ok(())
}

fn main() {
    println!("NSM test started.");

    let term = Arc::new(atomic::AtomicBool::new(false));
    signal_hook::flag::register(signal_hook::consts::SIGTERM, Arc::clone(&term))
        .expect("Failed to register signal hook");

    let ctx = nsm_init();
    assert!(ctx >= 0, "[Error] NSM initialization returned {}.", ctx);

    // 90 threads is the limit for ~200M of memory
    let index = 90;

    let exit_code = Arc::new(atomic::AtomicI32::new(0));
    let pool = ThreadPool::new(index);
    let mut j = 0;

    // Wait until the sigterm
    while !term.load(atomic::Ordering::Relaxed) {
        if exit_code.load(atomic::Ordering::Relaxed) != 0 {
            break;
        }

        let waiting = pool.queued_count();
        if waiting >= 1 {
            println!("{} waiting", waiting);
            thread::sleep(time::Duration::from_millis(100));
        }
        let exit_code_t = Arc::clone(&exit_code);
        pool.execute(move || {
            let exit_code = check_attestation(ctx, j);
            if let Err(e) = exit_code {
                if let Err(er) = exit_code_t.compare_exchange(
                    0,
                    e as i32,
                    atomic::Ordering::Relaxed,
                    atomic::Ordering::Relaxed,
                ) {
                    println!("{:?}", er);
                }
            }
        });
        j += 1;
        pool.execute(move || {
            extend_pcr(ctx, j);
        });
        j += 1;
    } //while

    pool.join();
    nsm_exit(ctx);

    println!(
        "NSM test finished. Exitcode: {}",
        exit_code.load(atomic::Ordering::Relaxed)
    );

    std::process::exit(exit_code.load(atomic::Ordering::Relaxed));
}


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

//! AWS Nitro Secure Module API
//!
//! This is the library that provides the API for the Nitro Secure Module used in AWS Nitro
//! Enclaves for management, attestation and entropy generation.
//!
//! nsm_io provides the API and CBOR encoding functionality.
//! nsm_driver provides the ioctl interface for the Nitro Secure Module driver.

pub mod api;
#[cfg(feature = "nix")]
pub mod driver;