From 398a123d469a4b567f13ef26c9d35855154f334e Mon Sep 17 00:00:00 2001
From: Samuel Chiang <sachiang@amazon.com>
Date: Fri, 31 Jan 2025 23:05:00 +0000
Subject: [PATCH] Move length check to more internal function

---
 crypto/evp_extra/p_pqdsa.c      | 14 ++++++--------
 crypto/evp_extra/p_pqdsa_asn1.c |  7 +++----
 crypto/pqdsa/internal.h         |  4 ++--
 crypto/pqdsa/pqdsa.c            | 18 ++++++++++++++----
 4 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/crypto/evp_extra/p_pqdsa.c b/crypto/evp_extra/p_pqdsa.c
index caed4935ad..57584765fe 100644
--- a/crypto/evp_extra/p_pqdsa.c
+++ b/crypto/evp_extra/p_pqdsa.c
@@ -280,13 +280,9 @@ EVP_PKEY *EVP_PKEY_pqdsa_new_raw_public_key(int nid, const uint8_t *in, size_t l
     goto err;
   }
 
-  const PQDSA *pqdsa =  PQDSA_KEY_get0_dsa(ret->pkey.pqdsa_key);
-  if (pqdsa->public_key_len != len) {
-    OPENSSL_PUT_ERROR(EVP, EVP_R_INVALID_BUFFER_SIZE);
-    goto err;
-  }
-
-  if (!PQDSA_KEY_set_raw_public_key(ret->pkey.pqdsa_key, in)) {
+  CBS cbs;
+  CBS_init(&cbs, in, len);
+  if (!PQDSA_KEY_set_raw_public_key(ret->pkey.pqdsa_key, &cbs)) {
     // PQDSA_KEY_set_raw_public_key sets the appropriate error.
     goto err;
   }
@@ -316,7 +312,9 @@ EVP_PKEY *EVP_PKEY_pqdsa_new_raw_private_key(int nid, const uint8_t *in, size_t
     goto err;
   }
 
-  if (!PQDSA_KEY_set_raw_private_key(ret->pkey.pqdsa_key, in)) {
+  CBS cbs;
+  CBS_init(&cbs, in, len);
+  if (!PQDSA_KEY_set_raw_private_key(ret->pkey.pqdsa_key, &cbs)) {
     // PQDSA_KEY_set_raw_private_key sets the appropriate error.
     goto err;
   }
diff --git a/crypto/evp_extra/p_pqdsa_asn1.c b/crypto/evp_extra/p_pqdsa_asn1.c
index e61ff6694f..5edcfeb2ba 100644
--- a/crypto/evp_extra/p_pqdsa_asn1.c
+++ b/crypto/evp_extra/p_pqdsa_asn1.c
@@ -98,12 +98,11 @@ static int pqdsa_pub_decode(EVP_PKEY *out, CBS *params, CBS *key) {
   }
   // Set the pqdsa params on |out| and check if the parsed length corresponds
   // with the nid's specified length.
-  if (!EVP_PKEY_pqdsa_set_params(out, OBJ_cbs2nid(params)) ||
-      CBS_len(key) != out->pkey.pqdsa_key->pqdsa->public_key_len) {
+  if (!EVP_PKEY_pqdsa_set_params(out, OBJ_cbs2nid(params))) {
     OPENSSL_PUT_ERROR(EVP, EVP_R_DECODE_ERROR);
     return 0;
   }
-  return PQDSA_KEY_set_raw_public_key(out->pkey.pqdsa_key, CBS_data(key));
+  return PQDSA_KEY_set_raw_public_key(out->pkey.pqdsa_key, key);
 }
 
 static int pqdsa_pub_encode(CBB *out, const EVP_PKEY *pkey) {
@@ -158,7 +157,7 @@ static int pqdsa_priv_decode(EVP_PKEY *out, CBS *params, CBS *key, CBS *pubkey)
   }
 
   // Set the private key
-  if (!PQDSA_KEY_set_raw_private_key(out->pkey.pqdsa_key, CBS_data(key))) {
+  if (!PQDSA_KEY_set_raw_private_key(out->pkey.pqdsa_key, key)) {
     OPENSSL_PUT_ERROR(EVP, EVP_R_DECODE_ERROR);
     return 0;
   }
diff --git a/crypto/pqdsa/internal.h b/crypto/pqdsa/internal.h
index 3ecd5334b5..7a3aabdd66 100644
--- a/crypto/pqdsa/internal.h
+++ b/crypto/pqdsa/internal.h
@@ -76,8 +76,8 @@ PQDSA_KEY *PQDSA_KEY_new(void);
 void PQDSA_KEY_free(PQDSA_KEY *key);
 int EVP_PKEY_pqdsa_set_params(EVP_PKEY *pkey, int nid);
 
-int PQDSA_KEY_set_raw_public_key(PQDSA_KEY *key, const uint8_t *in);
-int PQDSA_KEY_set_raw_private_key(PQDSA_KEY *key, const uint8_t *in);
+int PQDSA_KEY_set_raw_public_key(PQDSA_KEY *key, CBS *in);
+int PQDSA_KEY_set_raw_private_key(PQDSA_KEY *key, CBS *in);
 #if defined(__cplusplus)
 }  // extern C
 #endif
diff --git a/crypto/pqdsa/pqdsa.c b/crypto/pqdsa/pqdsa.c
index b32dae90f6..a32a5ab84d 100644
--- a/crypto/pqdsa/pqdsa.c
+++ b/crypto/pqdsa/pqdsa.c
@@ -66,8 +66,13 @@ const PQDSA *PQDSA_KEY_get0_dsa(PQDSA_KEY* key) {
   return key->pqdsa;
 }
 
-int PQDSA_KEY_set_raw_public_key(PQDSA_KEY *key, const uint8_t *in) {
-  key->public_key = OPENSSL_memdup(in, key->pqdsa->public_key_len);
+int PQDSA_KEY_set_raw_public_key(PQDSA_KEY *key, CBS *in) {
+  if (CBS_len(in) != key->pqdsa->public_key_len) {
+    OPENSSL_PUT_ERROR(EVP, EVP_R_INVALID_BUFFER_SIZE);
+    return 0;
+  }
+
+  key->public_key = OPENSSL_memdup(CBS_data(in), key->pqdsa->public_key_len);
   if (key->public_key == NULL) {
     return 0;
   }
@@ -75,8 +80,13 @@ int PQDSA_KEY_set_raw_public_key(PQDSA_KEY *key, const uint8_t *in) {
   return 1;
 }
 
-int PQDSA_KEY_set_raw_private_key(PQDSA_KEY *key, const uint8_t *in) {
-  key->private_key = OPENSSL_memdup(in, key->pqdsa->private_key_len);
+int PQDSA_KEY_set_raw_private_key(PQDSA_KEY *key, CBS *in) {
+  if (CBS_len(in) != key->pqdsa->private_key_len) {
+    OPENSSL_PUT_ERROR(EVP, EVP_R_INVALID_BUFFER_SIZE);
+    return 0;
+  }
+
+  key->private_key = OPENSSL_memdup(CBS_data(in), key->pqdsa->private_key_len);
   if (key->private_key == NULL) {
     return 0;
   }