diff --git a/library/pk.c b/library/pk.c
index 6cfc16b..c83d02b 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -119,11 +119,16 @@
     else
 #endif
 #if defined(POLARSSL_ECP_C)
-    if( type == POLARSSL_PK_ECKEY || type == POLARSSL_PK_ECKEY_DH )
+    if( type == POLARSSL_PK_ECKEY )
     {
         size = sizeof( ecp_keypair );
         info = &eckey_info;
     }
+    else if( type == POLARSSL_PK_ECKEY_DH )
+    {
+        size = sizeof( ecp_keypair );
+        info = &eckeydh_info;
+    }
     else
 #endif
 #if defined(POLARSSL_ECDSA_C)