From f990089545e830e8e9a75e43aa484d4953a42bde Mon Sep 17 00:00:00 2001
From: Justin W Smith <103147162+justsmth@users.noreply.github.com>
Date: Mon, 8 Jan 2024 15:22:42 -0500
Subject: [PATCH] Per PR feedback; only need to verify NID

---
 aws-lc-rs/src/agreement.rs   |  3 +++
 aws-lc-rs/src/ec.rs          | 34 ++++++++++++++++++++++++++--------
 aws-lc-rs/src/ec/key_pair.rs | 12 +++++-------
 3 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/aws-lc-rs/src/agreement.rs b/aws-lc-rs/src/agreement.rs
index 978e0cdd0ef..f7249bfa91c 100644
--- a/aws-lc-rs/src/agreement.rs
+++ b/aws-lc-rs/src/agreement.rs
@@ -294,6 +294,9 @@ impl PrivateKey {
         alg: &'static Algorithm,
         key_bytes: &[u8],
     ) -> Result<Self, KeyRejected> {
+        if key_bytes.len() != alg.id.private_key_len() {
+            return Err(KeyRejected::wrong_algorithm());
+        }
         let evp_pkey = if AlgorithmID::X25519 == alg.id {
             LcPtr::new(unsafe {
                 EVP_PKEY_new_raw_private_key(
diff --git a/aws-lc-rs/src/ec.rs b/aws-lc-rs/src/ec.rs
index 45f59153f99..fcfc710cafc 100644
--- a/aws-lc-rs/src/ec.rs
+++ b/aws-lc-rs/src/ec.rs
@@ -273,19 +273,37 @@ fn evp_pkey_from_public_key(
     Ok(pkey)
 }
 
-#[inline]
-unsafe fn validate_evp_key(
-    evp_pkey: &ConstPointer<EVP_PKEY>,
+fn verify_ec_key_nid(
+    ec_key: &ConstPointer<EC_KEY>,
     expected_curve_nid: i32,
 ) -> Result<(), KeyRejected> {
-    let ec_key = ConstPointer::new(EVP_PKEY_get0_EC_KEY(**evp_pkey))?;
-
-    let ec_group = ConstPointer::new(EC_KEY_get0_group(*ec_key))?;
-    let key_nid = EC_GROUP_get_curve_name(*ec_group);
+    let ec_group = ConstPointer::new(unsafe { EC_KEY_get0_group(**ec_key) })?;
+    let key_nid = unsafe { EC_GROUP_get_curve_name(*ec_group) };
 
     if key_nid != expected_curve_nid {
         return Err(KeyRejected::wrong_algorithm());
     }
+    Ok(())
+}
+
+#[inline]
+pub(crate) fn verify_evp_key_nid(
+    evp_pkey: &ConstPointer<EVP_PKEY>,
+    expected_curve_nid: i32,
+) -> Result<(), KeyRejected> {
+    let ec_key = ConstPointer::new(unsafe { EVP_PKEY_get0_EC_KEY(**evp_pkey) })?;
+    verify_ec_key_nid(&ec_key, expected_curve_nid)?;
+
+    Ok(())
+}
+
+#[inline]
+unsafe fn validate_evp_key(
+    evp_pkey: &ConstPointer<EVP_PKEY>,
+    expected_curve_nid: i32,
+) -> Result<(), KeyRejected> {
+    let ec_key = ConstPointer::new(EVP_PKEY_get0_EC_KEY(**evp_pkey))?;
+    verify_ec_key_nid(&ec_key, expected_curve_nid)?;
 
     #[cfg(not(feature = "fips"))]
     if 1 != EC_KEY_check_key(*ec_key) {
@@ -333,7 +351,7 @@ pub(crate) unsafe fn unmarshal_der_to_private_key(
             .try_into()
             .map_err(|_| KeyRejected::too_large())?,
     ))?;
-    validate_evp_key(&evp_pkey.as_const(), nid)?;
+    verify_evp_key_nid(&evp_pkey.as_const(), nid)?;
 
     Ok(evp_pkey)
 }
diff --git a/aws-lc-rs/src/ec/key_pair.rs b/aws-lc-rs/src/ec/key_pair.rs
index ddec2e68031..dae052d4cd6 100644
--- a/aws-lc-rs/src/ec/key_pair.rs
+++ b/aws-lc-rs/src/ec/key_pair.rs
@@ -13,7 +13,7 @@ use aws_lc::{EVP_DigestSign, EVP_DigestSignInit, EVP_PKEY_get0_EC_KEY, EVP_PKEY}
 use crate::buffer::Buffer;
 use crate::digest::digest_ctx::DigestContext;
 use crate::ec::{
-    evp_key_generate, validate_evp_key, EcdsaSignatureFormat, EcdsaSigningAlgorithm, PublicKey,
+    evp_key_generate, verify_evp_key_nid, EcdsaSignatureFormat, EcdsaSigningAlgorithm, PublicKey,
 };
 use crate::encoding::{AsBigEndian, AsDer, EcPrivateKeyBin, EcPrivateKeyRfc5915Der};
 use crate::error::{KeyRejected, Unspecified};
@@ -88,15 +88,13 @@ impl EcdsaKeyPair {
         alg: &'static EcdsaSigningAlgorithm,
         pkcs8: &[u8],
     ) -> Result<Self, KeyRejected> {
-        unsafe {
-            let evp_pkey = LcPtr::try_from(pkcs8)?;
+        let evp_pkey = LcPtr::try_from(pkcs8)?;
 
-            validate_evp_key(&evp_pkey.as_const(), alg.id.nid())?;
+        verify_evp_key_nid(&evp_pkey.as_const(), alg.id.nid())?;
 
-            let key_pair = Self::new(alg, evp_pkey)?;
+        let key_pair = Self::new(alg, evp_pkey)?;
 
-            Ok(key_pair)
-        }
+        Ok(key_pair)
     }
 
     /// Generates a new key pair and returns the key pair serialized as a