diff --git a/library/pk.c b/library/pk.c
index d8b4c85..61544eb 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -38,15 +38,6 @@
 #include "polarssl/ecdsa.h"
 #endif
 
-#if defined(POLARSSL_MEMORY_C)
-#include "polarssl/memory.h"
-#else
-#define polarssl_malloc     malloc
-#define polarssl_free       free
-#endif
-
-#include <stdlib.h>
-
 /*
  * Initialise a pk_context
  */
@@ -114,7 +105,7 @@
     }
 
     if( ( info = pk_info_from_type( type ) ) == NULL )
-        return( POLARSSL_ERR_PK_TYPE_MISMATCH );
+        return( POLARSSL_ERR_PK_BAD_INPUT_DATA );
 
     if( ( ctx->pk_ctx = info->ctx_alloc_func() ) == NULL )
         return( POLARSSL_ERR_PK_MALLOC_FAILED );
@@ -144,7 +135,7 @@
                const unsigned char *sig, size_t sig_len )
 {
     if( ctx == NULL || ctx->pk_info == NULL )
-        return( POLARSSL_ERR_PK_TYPE_MISMATCH ); // TODO
+        return( POLARSSL_ERR_PK_BAD_INPUT_DATA );
 
     return( ctx->pk_info->verify_func( ctx->pk_ctx, hash, md_info, sig, sig_len ) );
 }
@@ -166,7 +157,7 @@
 int pk_debug( const pk_context *ctx, pk_debug_item *items )
 {
     if( ctx == NULL || ctx->pk_info == NULL )
-        return( POLARSSL_ERR_PK_TYPE_MISMATCH ); // TODO
+        return( POLARSSL_ERR_PK_BAD_INPUT_DATA );
 
     ctx->pk_info->debug_func( ctx->pk_ctx, items );
     return( 0 );