From 7a6ed68f2dbbbf4c2f4aee347a80df0182f5223a Mon Sep 17 00:00:00 2001
From: Anthony Hu <anthony@wolfssl.com>
Date: Fri, 3 Feb 2023 10:36:26 -0500
Subject: [PATCH] Ensure that i2d APIs for public keys gives appropriate data.

---
 src/pk.c           |  10 ++--
 src/ssl.c          | 116 ++++++++++++++++++++++++++++++++++++++++++++-
 tests/api.c        |  10 +++-
 wolfssl/internal.h |   2 +-
 4 files changed, 129 insertions(+), 9 deletions(-)

diff --git a/src/pk.c b/src/pk.c
index 1bd29d1ef..3e31f5202 100644
--- a/src/pk.c
+++ b/src/pk.c
@@ -1200,7 +1200,7 @@ WOLFSSL_RSA* wolfSSL_RSAPublicKey_dup(WOLFSSL_RSA *rsa)
 
 #endif /* WOLFSSL_KEY_GEN && !HAVE_USER_RSA */
 
-#if defined(WOLFSSL_KEY_GEN) && !defined(HAVE_USER_RSA)
+#ifndef HAVE_USER_RSA
 static int wolfSSL_RSA_To_Der_ex(WOLFSSL_RSA* rsa, byte** outBuf, int publicKey,
     void* heap);
 #endif
@@ -1307,7 +1307,7 @@ WOLFSSL_RSA *wolfSSL_d2i_RSAPrivateKey(WOLFSSL_RSA **out,
     return rsa;
 }
 
-#if defined(WOLFSSL_KEY_GEN) && !defined(HAVE_USER_RSA) && \
+#if defined(OPENSSL_EXTRA) && !defined(HAVE_USER_RSA) && \
     !defined(HAVE_FAST_RSA)
 /* Converts an internal RSA structure to DER format for the private key.
  *
@@ -1382,7 +1382,7 @@ int wolfSSL_i2d_RSAPublicKey(WOLFSSL_RSA *rsa, unsigned char **pp)
 
     return ret;
 }
-#endif /* defined(WOLFSSL_KEY_GEN) && !defined(HAVE_USER_RSA) &&
+#endif /* defined(OPENSSL_EXTRA) && !defined(HAVE_USER_RSA) &&
         * !defined(HAVE_FAST_RSA) */
 
 #endif /* OPENSSL_EXTRA */
@@ -1512,7 +1512,7 @@ WOLFSSL_RSA* wolfSSL_d2i_RSAPrivateKey_bio(WOLFSSL_BIO *bio, WOLFSSL_RSA **out)
 
 #ifdef OPENSSL_EXTRA
 
-#if defined(WOLFSSL_KEY_GEN) && !defined(HAVE_USER_RSA)
+#ifndef HAVE_USER_RSA
 /* Create a DER encoding of key.
  *
  * Not OpenSSL API.
@@ -1645,7 +1645,7 @@ static int wolfSSL_RSA_To_Der_ex(WOLFSSL_RSA* rsa, byte** outBuf, int publicKey,
     WOLFSSL_LEAVE("wolfSSL_RSA_To_Der", ret);
     return ret;
 }
-#endif /* WOLFSSL_KEY_GEN && !HAVE_USER_RSA */
+#endif /* !HAVE_USER_RSA */
 
 #endif /* OPENSSL_EXTRA */
 
diff --git a/src/ssl.c b/src/ssl.c
index 5e43f2fef..e28ca0bf5 100644
--- a/src/ssl.c
+++ b/src/ssl.c
@@ -10402,6 +10402,9 @@ WOLFSSL_EVP_PKEY* wolfSSL_d2i_PUBKEY(WOLFSSL_EVP_PKEY** out,
     return d2iGenericKey(out, in, inSz, 0);
 }
 
+#if defined(OPENSSL_EXTRA) && !defined(NO_CERTS) && !defined(NO_ASN) && \
+    !defined(NO_PWDBASED)
+
 /* helper function to get raw pointer to DER buffer from WOLFSSL_EVP_PKEY */
 static int wolfSSL_EVP_PKEY_get_der(const WOLFSSL_EVP_PKEY* key, unsigned char** der)
 {
@@ -10439,9 +10442,11 @@ static int wolfSSL_EVP_PKEY_get_der(const WOLFSSL_EVP_PKEY* key, unsigned char**
 
 int wolfSSL_i2d_PUBKEY(const WOLFSSL_EVP_PKEY *key, unsigned char **der)
 {
-    return wolfSSL_EVP_PKEY_get_der(key, der);
+    return wolfSSL_i2d_PublicKey(key, der);
 }
 
+#endif /* OPENSSL_EXTRA && !NO_CERTS && !NO_ASN && !NO_PWDBASED */
+
 static WOLFSSL_EVP_PKEY* _d2i_PublicKey(int type, WOLFSSL_EVP_PKEY** out,
     const unsigned char **in, long inSz, int priv)
 {
@@ -22949,7 +22954,114 @@ int wolfSSL_i2d_PrivateKey(const WOLFSSL_EVP_PKEY* key, unsigned char** der)
 
 int wolfSSL_i2d_PublicKey(const WOLFSSL_EVP_PKEY *key, unsigned char **der)
 {
-    return wolfSSL_EVP_PKEY_get_der(key, der);
+#if !defined(NO_RSA) || defined(HAVE_ECC)
+#ifdef HAVE_ECC
+    unsigned char *local_der = NULL;
+    word32 local_derSz = 0;
+    unsigned char *pub_der = NULL;
+    ecc_key eccKey;
+    word32 inOutIdx = 0;
+#endif
+    word32 pub_derSz = 0;
+    int ret = 0;
+    int key_type = 0;
+
+    if (key == NULL) {
+        return WOLFSSL_FATAL_ERROR;
+    }
+
+    key_type = key->type;
+    if ((key_type != EVP_PKEY_EC) && (key_type != EVP_PKEY_RSA)) {
+        ret = WOLFSSL_FATAL_ERROR;
+    }
+
+#ifndef NO_RSA
+    if (key_type == EVP_PKEY_RSA) {
+        return wolfSSL_i2d_RSAPublicKey(key->rsa, der);
+    }
+#endif
+
+    /* Now that RSA is taken care of, we only need to consider the ECC case. */
+
+#ifdef HAVE_ECC
+
+    /* We need to get the DER, then convert it to a public key. But what we get
+     * might be a buffereed private key so we need to decode it and then encode
+     * the public part. */
+    if (ret == 0) {
+        local_derSz = wolfSSL_EVP_PKEY_get_der(key, &local_der);
+        if (local_derSz <= 0) {
+            ret = WOLFSSL_FATAL_ERROR;
+        }
+    }
+
+    if (ret == 0) {
+        ret = wc_ecc_init(&eccKey);
+    }
+
+    if (ret == 0) {
+        ret = wc_EccPublicKeyDecode(local_der, &inOutIdx, &eccKey, local_derSz);
+    }
+
+    if (ret == 0) {
+        pub_derSz = wc_EccPublicKeyDerSize(&eccKey, 0);
+        if (pub_derSz <= 0) {
+            ret = WOLFSSL_FAILURE;
+        }
+    }
+
+    if (ret == 0) {
+        pub_der = (unsigned char*)XMALLOC(pub_derSz, NULL,
+                                          DYNAMIC_TYPE_PUBLIC_KEY);
+        if (pub_der == NULL) {
+            WOLFSSL_MSG("Failed to allocate output buffer.");
+            ret = WOLFSSL_FATAL_ERROR;
+        }
+    }
+
+    if (ret == 0) {
+        pub_derSz = wc_EccPublicKeyToDer(&eccKey, pub_der, pub_derSz, 0);
+        if (pub_derSz <= 0) {
+            ret = WOLFSSL_FATAL_ERROR;
+        }
+    }
+
+    /* This block is for actually returning the DER of the public key */
+    if ((ret == 0) && (der != NULL)) {
+        if (*der == NULL) {
+            *der = (unsigned char*)XMALLOC(pub_derSz, NULL,
+                                           DYNAMIC_TYPE_PUBLIC_KEY);
+            if (*der == NULL) {
+                WOLFSSL_MSG("Failed to allocate output buffer.");
+                ret = WOLFSSL_FATAL_ERROR;
+            }
+
+            if (ret == 0) {
+                XMEMCPY(*der, pub_der, pub_derSz);
+            }
+        }
+        else {
+            XMEMCPY(*der, pub_der, pub_derSz);
+            *der += pub_derSz;
+        }
+    }
+
+    XFREE(pub_der, NULL, DYNAMIC_TYPE_PUBLIC_KEY);
+    XFREE(local_der, NULL, DYNAMIC_TYPE_PUBLIC_KEY);
+
+    wc_ecc_free(&eccKey);
+#else
+    ret = WOLFSSL_FATAL_ERROR;
+#endif /* HAVE_ECC */
+
+    if (ret == 0) {
+        return pub_derSz;
+    }
+
+    return ret;
+#else
+    return WOLFSSL_FATAL_ERROR;
+#endif /* !NO_RSA || HAVE_ECC */
 }
 #endif /* !NO_ASN && !NO_PWDBASED */
 
diff --git a/tests/api.c b/tests/api.c
index eaa4e0f63..d20e53c3e 100644
--- a/tests/api.c
+++ b/tests/api.c
@@ -47216,7 +47216,7 @@ static int test_wolfSSL_d2i_and_i2d_PublicKey(void)
 #if defined(OPENSSL_EXTRA) && !defined(NO_RSA)
     EVP_PKEY* pkey;
     const unsigned char* p;
-    unsigned char* der = NULL;
+    unsigned char *der = NULL, *tmp = NULL;
     int derLen;
 
     p = client_keypub_der_2048;
@@ -47229,6 +47229,14 @@ static int test_wolfSSL_d2i_and_i2d_PublicKey(void)
     AssertIntEQ(derLen, sizeof_client_keypub_der_2048);
     AssertIntEQ(XMEMCMP(der, client_keypub_der_2048, derLen), 0);
 
+    /* Do same test except with pre-allocated buffer to ensure the der pointer
+     * is advanced. */
+    tmp = der;
+    AssertIntGE((derLen = wolfSSL_i2d_PublicKey(pkey, &tmp)), 0);
+    AssertIntEQ(derLen, sizeof_client_keypub_der_2048);
+    AssertIntEQ(XMEMCMP(der, client_keypub_der_2048, derLen), 0);
+    AssertTrue(der + derLen == tmp);
+
     XFREE(der, HEAP_HINT, DYNAMIC_TYPE_OPENSSL);
     EVP_PKEY_free(pkey);
 
diff --git a/wolfssl/internal.h b/wolfssl/internal.h
index 435082a42..91739be75 100644
--- a/wolfssl/internal.h
+++ b/wolfssl/internal.h
@@ -6094,7 +6094,7 @@ WOLFSSL_LOCAL int EncryptDerKey(byte *der, int *derSz, const EVP_CIPHER* cipher,
 #endif
 #endif
 
-#if defined(WOLFSSL_KEY_GEN) && !defined(NO_RSA) && !defined(HAVE_USER_RSA)
+#if !defined(NO_RSA) && !defined(HAVE_USER_RSA)
 WOLFSSL_LOCAL int wolfSSL_RSA_To_Der(WOLFSSL_RSA* rsa, byte** outBuf,
     int publicKey, void* heap);
 #endif