diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index c7d1b76..3c3e80d 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1417,16 +1417,14 @@
 {
     psa_key_type_t type = attributes->core.type;
 
-    if (PSA_KEY_TYPE_IS_RSA(type) || PSA_KEY_TYPE_IS_ECC(type) ||
-        PSA_KEY_TYPE_IS_DH(type)) {
-        if (PSA_KEY_TYPE_IS_PUBLIC_KEY(type)) {
-            /* Exporting public -> public */
-            return psa_export_key_buffer_internal(
-                key_buffer, key_buffer_size,
-                data, data_size, data_length);
-        }
-
-        if (PSA_KEY_TYPE_IS_RSA(type)) {
+    if (PSA_KEY_TYPE_IS_PUBLIC_KEY(type) &&
+            (PSA_KEY_TYPE_IS_RSA(type) || PSA_KEY_TYPE_IS_ECC(type) ||
+             PSA_KEY_TYPE_IS_DH(type))) {
+        /* Exporting public -> public */
+        return psa_export_key_buffer_internal(
+            key_buffer, key_buffer_size,
+            data, data_size, data_length);
+    } else if (PSA_KEY_TYPE_IS_RSA(type)) {
 #if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
             defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
             return mbedtls_psa_rsa_export_public_key(attributes,
@@ -1440,7 +1438,7 @@
             return PSA_ERROR_NOT_SUPPORTED;
 #endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
         * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
-        } else if (PSA_KEY_TYPE_IS_ECC(type)) {
+    } else if (PSA_KEY_TYPE_IS_ECC(type)) {
 #if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
             defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
             return mbedtls_psa_ecp_export_public_key(attributes,
@@ -1454,26 +1452,19 @@
             return PSA_ERROR_NOT_SUPPORTED;
 #endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
         * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) */
-        }
+    } else if (PSA_KEY_TYPE_IS_DH(type)) {
 #if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_FFDH_KEY_PAIR) || \
         defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_FFDH_PUBLIC_KEY)
-        else if (PSA_KEY_TYPE_IS_DH(type)) {
             return mbedtls_psa_export_ffdh_public_key(attributes,
                                                       key_buffer,
                                                       key_buffer_size,
                                                       data, data_size,
                                                       data_length);
-        } else {
-            return PSA_ERROR_NOT_SUPPORTED;
-        }
 #else
         return PSA_ERROR_NOT_SUPPORTED;
 #endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_FFDH_KEY_PAIR) ||
         * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_FFDH_PUBLIC_KEY) */
     } else {
-        /* This shouldn't happen in the reference implementation, but
-           it is valid for a special-purpose implementation to omit
-           support for exporting certain key types. */
         return PSA_ERROR_NOT_SUPPORTED;
     }
 }
diff --git a/library/psa_crypto_ffdh.c b/library/psa_crypto_ffdh.c
index e15d15b..b6a9bf4 100644
--- a/library/psa_crypto_ffdh.c
+++ b/library/psa_crypto_ffdh.c
@@ -142,31 +142,36 @@
     status = mbedtls_psa_ffdh_set_prime_generator(
         PSA_BITS_TO_BYTES(attributes->core.bits), &P, &G);
 
-    if (status == PSA_SUCCESS) {
-        MBEDTLS_MPI_CHK(mbedtls_mpi_read_binary(&X, key_buffer,
-                                                key_buffer_size));
-
-        MBEDTLS_MPI_CHK(mbedtls_mpi_read_binary(&GY, peer_key,
-                                                peer_key_length));
-
-        /* Calculate shared secret public key: K = G^(XY) mod P = GY^X mod P */
-        MBEDTLS_MPI_CHK(mbedtls_mpi_exp_mod(&K, &GY, &X, &P, NULL));
-
-        MBEDTLS_MPI_CHK(mbedtls_mpi_write_binary(&K, shared_secret,
-                                                 calculated_shared_secret_size));
-
-        *shared_secret_length = calculated_shared_secret_size;
+    if(status != PSA_SUCCESS) {
+        goto cleanup;
     }
+
+    MBEDTLS_MPI_CHK(mbedtls_mpi_read_binary(&X, key_buffer,
+                                            key_buffer_size));
+
+    MBEDTLS_MPI_CHK(mbedtls_mpi_read_binary(&GY, peer_key,
+                                            peer_key_length));
+
+    /* Calculate shared secret public key: K = G^(XY) mod P = GY^X mod P */
+    MBEDTLS_MPI_CHK(mbedtls_mpi_exp_mod(&K, &GY, &X, &P, NULL));
+
+    MBEDTLS_MPI_CHK(mbedtls_mpi_write_binary(&K, shared_secret,
+                                                calculated_shared_secret_size));
+
+    *shared_secret_length = calculated_shared_secret_size;
+
+    ret = 0;
+
 cleanup:
     mbedtls_mpi_free(&P); mbedtls_mpi_free(&G);
     mbedtls_mpi_free(&X); mbedtls_mpi_free(&GY);
     mbedtls_mpi_free(&K);
 
-    if (status == PSA_SUCCESS && ret != 0) {
-        return mbedtls_to_psa_error(ret);
+    if(status == PSA_SUCCESS && ret != 0) {
+        status = mbedtls_to_psa_error(ret);
     }
 
-    return PSA_SUCCESS;
+    return status;
 }
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_FFDH */
 
@@ -188,21 +193,25 @@
     status = mbedtls_psa_ffdh_set_prime_generator(
         PSA_BITS_TO_BYTES(attributes->core.bits), &P, &G);
 
-    if (status == PSA_SUCCESS) {
-        MBEDTLS_MPI_CHK(mbedtls_mpi_read_binary(&X, key_buffer,
-                                                key_buffer_size));
-
-        MBEDTLS_MPI_CHK(mbedtls_mpi_exp_mod(&GX, &G, &X, &P, NULL));
-        MBEDTLS_MPI_CHK(mbedtls_mpi_write_binary(&GX, data, data_size));
-
-        *data_length = mbedtls_mpi_size(&GX);
+    if(status != PSA_SUCCESS) {
+        goto cleanup;
     }
+
+    MBEDTLS_MPI_CHK(mbedtls_mpi_read_binary(&X, key_buffer,
+                                            key_buffer_size));
+
+    MBEDTLS_MPI_CHK(mbedtls_mpi_exp_mod(&GX, &G, &X, &P, NULL));
+    MBEDTLS_MPI_CHK(mbedtls_mpi_write_binary(&GX, data, data_size));
+
+    *data_length = mbedtls_mpi_size(&GX);
+
+    ret = 0;
 cleanup:
     mbedtls_mpi_free(&P); mbedtls_mpi_free(&G);
     mbedtls_mpi_free(&X); mbedtls_mpi_free(&GX);
 
     if (status == PSA_SUCCESS && ret != 0) {
-        return mbedtls_to_psa_error(ret);
+        status = mbedtls_to_psa_error(ret);
     }
 
     return status;