diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 56ce69c..0d1d91b 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -492,15 +492,14 @@
 
     if( ( end - *p ) < 1 )
     {
-        return( MBEDTLS_ERR_X509_INVALID_SIGNATURE +
-                MBEDTLS_ERR_ASN1_OUT_OF_DATA );
+        return( MBEDTLS_ERR_ASN1_OUT_OF_DATA );
     }
     tag_type = **p;
 
     if( ( ret = mbedtls_asn1_get_tag( p, end, &len_partial,
                 MBEDTLS_ASN1_CONSTRUCTED | MBEDTLS_ASN1_SEQUENCE ) ) != 0 )
     {
-        return( MBEDTLS_ERR_X509_INVALID_SIGNATURE + ret );
+        return( ret );
     }
 
     if( ( ret = mbedtls_asn1_get_tag( p, end, &len_partial, MBEDTLS_ASN1_INTEGER ) )
@@ -536,7 +535,7 @@
 
     // Check if both parts are of the same size
     if( len_partial != len_signature )
-        return( MBEDTLS_ERR_X509_INVALID_SIGNATURE );
+        return( MBEDTLS_ERR_PK_SIG_LEN_MISMATCH );
 
     memcpy( sig->p + len_partial, *p, len_partial );
     len_signature += len_partial;
@@ -638,8 +637,8 @@
 
     psa_type = PSA_KEY_TYPE_ECC_PUBLIC_KEY( curve );
 
-    if( extract_ecdsa_sig( &p, p + sig_len, &signature ) != 0 )
-        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+    if( ( ret = extract_ecdsa_sig( &p, p + sig_len, &signature ) ) != 0 )
+        return( ret );
 
     key_len = mbedtls_pk_write_pubkey_der( &key, buf, buf_len );
     if( key_len <= 0 )
@@ -648,17 +647,17 @@
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
     }
 
-    if( mbedtls_psa_get_free_key_slot( &key_slot ) != PSA_SUCCESS )
+    if( ( ret = mbedtls_psa_get_free_key_slot( &key_slot ) ) != PSA_SUCCESS )
     {
         mbedtls_free( signature.p );
-        return( MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE );
+        return( mbedtls_psa_err_translate_pk( ret ) );
     }
 
     psa_key_policy_init( &policy );
     psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_VERIFY, psa_sig_md );
-    if( psa_set_key_policy( key_slot, &policy ) != PSA_SUCCESS )
+    if( ( ret = psa_set_key_policy( key_slot, &policy ) ) != PSA_SUCCESS )
     {
-        ret = MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
+        ret = mbedtls_psa_err_translate_pk( ret );
         goto cleanup;
     }