Add support to use KMS for signing EIF files when building them

Cargo.toml

@@ -20,7 +20,15 @@ byteorder = "1.3"
 clap = "3.2"
 hex = "0.4"
 crc = "1.8"
-aws-nitro-enclaves-cose = "0.5"
+aws-nitro-enclaves-cose = { version = "0.5", features = ["key_kms"] }
 openssl = "0.10"
 serde_cbor = "0.11"
 chrono = { version = "0.4", default-features = false, features = ["clock"]}
+aws-sdk-kms = "<=1.20"
+aws-config = "<=1.1"
+aws-types = "<=1.1"
+aws-smithy-runtime = { version = "<=1.2" }
+tokio = { version = "1.20", features = ["rt-multi-thread"] }
+
+[dev-dependencies]
+tempfile = "3.5"

examples/eif_build.rs

@@ -1,4 +1,4 @@
-// Copyright 2019-2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// Copyright 2019-2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
 #![deny(warnings)]
@@ -18,9 +18,9 @@ use aws_nitro_enclaves_image_format::defs::EifIdentityInfo;
 use aws_nitro_enclaves_image_format::utils::identity::parse_custom_metadata;
 use aws_nitro_enclaves_image_format::{
     generate_build_info,
-    utils::{get_pcrs, EifBuilder, SignEnclaveInfo},
+    utils::{get_pcrs, EifBuilder, SignKeyData, SignKeyDataInfo, SignKeyInfo},
 };
-use clap::{App, Arg};
+use clap::{App, Arg, ArgGroup};
 use serde_json::json;
 use sha2::{Digest, Sha256, Sha384, Sha512};
 use std::fmt::Debug;
@@ -76,14 +76,34 @@ fn main() {
             Arg::with_name("signing-certificate")
                 .long("signing-certificate")
                 .help("Specify the path to the signing certificate")
-                .takes_value(true),
+                .takes_value(true)
+                .requires("signing-key"),
         )
         .arg(
             Arg::with_name("private-key")
                 .long("private-key")
                 .help("Specify the path to the private-key")
                 .takes_value(true),
         )
+        .arg(
+            Arg::with_name("kms-key-id")
+                .long("kms-key-id")
+                .help("Specify unique id of the KMS key")
+                .takes_value(true),
+        )
+        .arg(
+            Arg::with_name("kms-key-region")
+                .long("kms-key-region")
+                .help("Specify region in which the KMS key resides")
+                .takes_value(true)
+                .requires("kms-key-id")
+        )
+        .group(
+            ArgGroup::new("signing-key")
+                .args(&["kms-key-id", "private-key"])
+                .multiple(false)
+                .requires("signing-certificate")
+        )
         .arg(
             Arg::with_name("sha256")
                 .long("sha256")
@@ -150,14 +170,29 @@ fn main() {
 
     let private_key = matches.value_of("private-key");
 
-    let sign_info = match (signing_certificate, private_key) {
+    let kms_key_id = matches.value_of("kms-key-id");
+    let kms_key_region = matches.value_of("kms-key-region");
+
+    let sign_key_info = match (kms_key_id, private_key) {
         (None, None) => None,
-        (Some(cert_path), Some(key_path)) => {
-            Some(SignEnclaveInfo::new(cert_path, key_path).expect("Could not read signing info"))
-        }
-        _ => panic!("Both signing-certificate and private-key parameters must be provided"),
+        (Some(kms_id), None) => Some(SignKeyInfo::KmsKeyInfo {
+            id: kms_id.to_string(),
+            region: kms_key_region.map(str::to_string),
+        }),
+        (None, Some(key_path)) => Some(SignKeyInfo::LocalPrivateKeyInfo {
+            path: key_path.to_string(),
+        }),
+        _ => panic!("kms-key-id and private-key parameters are mutually exclusive"),
     };
 
+    let sign_key_data = sign_key_info.map(|key_info| {
+        SignKeyData::new(&SignKeyDataInfo {
+            cert_path: signing_certificate.unwrap().to_string(),
+            key_info: key_info,
+        })
+        .expect("Could not read signing info")
+    });
+
     let img_name = matches.value_of("image_name").map(|val| val.to_string());
     let img_version = matches.value_of("image_name").map(|val| val.to_string());
     let metadata_path = matches.value_of("metadata").map(|val| val.to_string());
@@ -190,7 +225,7 @@ fn main() {
             cmdline,
             ramdisks,
             output_path,
-            sign_info,
+            sign_key_data,
             Sha512::new(),
             eif_info,
         );
@@ -200,7 +235,7 @@ fn main() {
             cmdline,
             ramdisks,
             output_path,
-            sign_info,
+            sign_key_data,
             Sha256::new(),
             eif_info,
         );
@@ -210,7 +245,7 @@ fn main() {
             cmdline,
             ramdisks,
             output_path,
-            sign_info,
+            sign_key_data,
             Sha384::new(),
             eif_info,
         );
@@ -222,7 +257,7 @@ pub fn build_eif<T: Digest + Debug + Write + Clone>(
     cmdline: &str,
     ramdisks: Vec<&str>,
     output_path: &str,
-    sign_info: Option<SignEnclaveInfo>,
+    sign_info: Option<SignKeyData>,
     hasher: T,
     eif_info: EifIdentityInfo,
 ) {

src/utils/mod.rs

@@ -1,4 +1,4 @@
-// Copyright 2019-2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// Copyright 2019-2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 #![deny(warnings)]
 pub mod eif_reader;
@@ -9,7 +9,12 @@ 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 aws_config::BehaviorVersion;
+use aws_nitro_enclaves_cose::{
+    crypto::kms::KmsKey, crypto::Openssl, header_map::HeaderMap, CoseSign1,
+};
+use aws_sdk_kms::client::Client;
+use aws_types::region::Region;
 use crc::{crc32, Hasher32};
 use openssl::asn1::Asn1Time;
 use openssl::pkey::PKey;
@@ -34,35 +39,96 @@ use std::fs::File;
 use std::io::{Read, Seek, SeekFrom, Write};
 use std::mem::size_of;
 use std::path::Path;
+use tokio::runtime::Runtime;
 
 const DEFAULT_SECTIONS_COUNT: u16 = 3;
 
+// Signing key for eif images
+pub enum SignKey {
+    // Local private key
+    LocalPrivateKey(Vec<u8>),
+
+    // KMS signer implementation from Cose library
+    KmsKey(KmsKey),
+}
+
+// Full signining key data
+pub struct SignKeyData {
+    // x509 certificate
+    pub cert: Vec<u8>,
+
+    // Signing key itself
+    pub key: SignKey,
+}
+
+// Signing key details
+#[derive(Clone, Debug)]
+pub enum SignKeyInfo {
+    // Local private key file path
+    LocalPrivateKeyInfo { path: String },
+
+    // KMS key details
+    KmsKeyInfo { id: String, region: Option<String> },
+}
+
+// Full details of signing key
 #[derive(Clone, Debug)]
-pub struct SignEnclaveInfo {
-    pub signing_certificate: Vec<u8>,
-    pub private_key: Vec<u8>,
+pub struct SignKeyDataInfo {
+    // Path to the certificate file
+    pub cert_path: String,
+
+    // Details of signing key itself
+    pub key_info: SignKeyInfo,
 }
 
-impl SignEnclaveInfo {
-    pub fn new(cert_path: &str, key_path: &str) -> Result<Self, String> {
-        let mut certificate_file = File::open(cert_path)
+impl SignKeyData {
+    pub fn new(sign_info: &SignKeyDataInfo) -> Result<Self, String> {
+        let mut cert_file = File::open(&sign_info.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)
+        let mut cert = Vec::new();
+        cert_file
+            .read_to_end(&mut cert)
             .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))?;
+        let key = match &sign_info.key_info {
+            SignKeyInfo::LocalPrivateKeyInfo { path } => {
+                let mut key_file = File::open(path)
+                    .map_err(|err| format!("Could not open the key file: {:?}", err))?;
+                let mut key_data = Vec::new();
+                key_file
+                    .read_to_end(&mut key_data)
+                    .map_err(|err| format!("Could not read the key file: {:?}", err))?;
 
-        Ok(SignEnclaveInfo {
-            signing_certificate,
-            private_key,
-        })
+                SignKey::LocalPrivateKey(key_data)
+            }
+            SignKeyInfo::KmsKeyInfo { id, region } => {
+                let act = async {
+                    let mut config_loader = aws_config::defaults(BehaviorVersion::latest());
+                    if let Some(region_id) = region {
+                        config_loader = config_loader.region(Region::new(region_id.clone()));
+                    }
+
+                    let sdk_config = config_loader.load().await;
+                    if sdk_config.region().is_none() {
+                        return Err("AWS region for KMS is not specified".to_string());
+                    }
+
+                    let id_copy = id.clone();
+                    tokio::task::spawn_blocking(move || {
+                        let client = Client::new(&sdk_config);
+                        KmsKey::new_with_public_key(client, id_copy, None)
+                            .map_err(|e| e.to_string())
+                    })
+                    .await
+                    .unwrap()
+                };
+                let runtime = Runtime::new().unwrap();
+                let key = runtime.block_on(act)?;
+                SignKey::KmsKey(key)
+            }
+        };
+
+        Ok(SignKeyData { cert, key })
     }
 }
 
@@ -118,7 +184,7 @@ pub struct EifBuilder<T: Digest + Debug + Write + Clone> {
     kernel: File,
     cmdline: Vec<u8>,
     ramdisks: Vec<File>,
-    sign_info: Option<SignEnclaveInfo>,
+    sign_info: Option<SignKeyData>,
     signature: Option<Vec<u8>>,
     signature_size: u64,
     metadata: Vec<u8>,
@@ -142,7 +208,7 @@ impl<T: Digest + Debug + Write + Clone> EifBuilder<T> {
     pub fn new(
         kernel_path: &Path,
         cmdline: String,
-        sign_info: Option<SignEnclaveInfo>,
+        sign_info: Option<SignKeyData>,
         hasher: T,
         flags: u16,
         eif_info: EifIdentityInfo,
@@ -285,23 +351,27 @@ impl<T: Digest + Debug + Write + Clone> EifBuilder<T> {
         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 sign_info = self
+            .sign_info
+            .as_ref()
+            .expect("Signing key is expected to be provided");
         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();
+        let cose_sign = match &sign_info.key {
+            SignKey::LocalPrivateKey(key) => {
+                let pkey = PKey::private_key_from_pem(key)
+                    .expect("Could not deserialize the PEM-formatted private key");
+
+                CoseSign1::new::<Openssl>(&payload, &HeaderMap::new(), &pkey)
+            }
+            SignKey::KmsKey(key) => CoseSign1::new::<Openssl>(&payload, &HeaderMap::new(), key),
+        };
 
         PcrSignature {
-            signing_certificate,
-            signature,
+            signing_certificate: sign_info.cert.clone(),
+            signature: cose_sign.unwrap().as_bytes(false).unwrap(),
         }
     }
 
@@ -614,7 +684,7 @@ impl<T: Digest + Debug + Write + Clone> EifBuilder<T> {
         }
 
         if let Some(sign_info) = self.sign_info.as_ref() {
-            let cert = openssl::x509::X509::from_pem(&sign_info.signing_certificate[..]).unwrap();
+            let cert = openssl::x509::X509::from_pem(&sign_info.cert[..]).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