diff --git a/library/pk.c b/library/pk.c
index b6dd99d..369472a 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -518,6 +518,7 @@
                                          f_rng, p_rng, NULL ) );
 }
 
+#if defined(MBEDTLS_PSA_CRYPTO_C)
 /*
  * Make a signature with options
  */
@@ -529,12 +530,9 @@
                          int (*f_rng)(void *, unsigned char *, size_t),
                          void *p_rng )
 {
-    PK_VALIDATE_RET( ctx != NULL );
-    PK_VALIDATE_RET( ( md_alg == MBEDTLS_MD_NONE && hash_len == 0 ) ||
-                     hash != NULL );
-    PK_VALIDATE_RET( sig != NULL );
 
     *sig_len = 0;
+
     if( ctx->pk_info == NULL )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 
@@ -543,43 +541,22 @@
 
     if( type != MBEDTLS_PK_RSASSA_PSS )
     {
-        return( mbedtls_pk_sign_restartable( ctx, md_alg, hash, hash_len,
-                                             sig, sig_size, sig_len,
-                                             f_rng, p_rng, NULL ) );
+        return( mbedtls_pk_sign( ctx, md_alg, hash, hash_len,
+                                 sig, sig_size, sig_len, f_rng, p_rng ) );
     }
 
-#if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_PKCS1_V21)
-    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+#if defined(MBEDTLS_X509_RSASSA_PSS_SUPPORT) && defined(MBEDTLS_USE_PSA_CRYPTO)
 
-#if SIZE_MAX > UINT_MAX
-    if( md_alg == MBEDTLS_MD_NONE && UINT_MAX < hash_len )
-        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
-#endif /* SIZE_MAX > UINT_MAX */
-
-    if( sig_size < mbedtls_pk_get_len( ctx ) )
-        return( MBEDTLS_ERR_RSA_VERIFY_FAILED );
-
-    if( mbedtls_rsa_set_padding( mbedtls_pk_rsa( *ctx ),
-                                 MBEDTLS_RSA_PKCS_V21,
-                                 md_alg ) != 0 )
-    {
-        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
-    }
-
-    *sig_len = mbedtls_pk_get_len( ctx );
-    ret = mbedtls_rsa_rsassa_pss_sign_ext( mbedtls_pk_rsa( *ctx ),
-                                           f_rng, p_rng,
-                                           md_alg,
-                                           (unsigned int) hash_len,
-                                           hash,MBEDTLS_RSA_SALT_LEN_ANY,
-                                           sig
-                                         );
-    return( ret );
-#else
+    return( mbedtls_pk_psa_sign_ext( PSA_ALG_RSA_PSS_ANY_SALT(
+                                        mbedtls_psa_translate_md( md_alg ) ),
+                                     ctx->pk_ctx, hash, hash_len,
+                                     sig, sig_size, sig_len ) );
+#else /* MBEDTLS_X509_RSASSA_PSS_SUPPORT */
     return( MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE );
-#endif /* MBEDTLS_RSA_C && MBEDTLS_PKCS1_V21 */
+#endif /* !MBEDTLS_X509_RSASSA_PSS_SUPPORT */
 
 }
+#endif /* MBEDTLS_PSA_CRYPTO_C */
 
 /*
  * Decrypt message
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 92e9bf4..b910242 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -192,12 +192,12 @@
 }
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-static int rsa_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
-                   const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t sig_size, size_t *sig_len,
-                   int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
+int mbedtls_pk_psa_sign_ext( psa_algorithm_t psa_alg_md, void *pk_ctx,
+                             const unsigned char *hash, size_t hash_len,
+                             unsigned char *sig, size_t sig_size,
+                             size_t *sig_len )
 {
-    mbedtls_rsa_context * rsa = (mbedtls_rsa_context *) ctx;
+    mbedtls_rsa_context * rsa = (mbedtls_rsa_context *) pk_ctx;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
@@ -206,16 +206,6 @@
     int key_len;
     unsigned char buf[MBEDTLS_PK_RSA_PRV_DER_MAX_BYTES];
     mbedtls_pk_info_t pk_info = mbedtls_rsa_info;
-    psa_algorithm_t psa_alg_md =
-        PSA_ALG_RSA_PKCS1V15_SIGN( mbedtls_psa_translate_md( md_alg ) );
-
-    ((void) f_rng);
-    ((void) p_rng);
-
-#if SIZE_MAX > UINT_MAX
-    if( md_alg == MBEDTLS_MD_NONE && UINT_MAX < hash_len )
-        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
-#endif /* SIZE_MAX > UINT_MAX */
 
     *sig_len = mbedtls_rsa_get_len( rsa );
     if( sig_size < *sig_len )
@@ -224,11 +214,10 @@
     /* mbedtls_pk_write_key_der() expects a full PK context;
      * re-construct one to make it happy */
     key.pk_info = &pk_info;
-    key.pk_ctx = ctx;
+    key.pk_ctx = pk_ctx;
     key_len = mbedtls_pk_write_key_der( &key, buf, sizeof( buf ) );
     if( key_len <= 0 )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
-
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN_HASH );
     psa_set_key_algorithm( &attributes, psa_alg_md );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_RSA_KEY_PAIR );
@@ -241,7 +230,6 @@
         ret = mbedtls_pk_error_from_psa( status );
         goto cleanup;
     }
-
     status = psa_sign_hash( key_id, psa_alg_md, hash, hash_len,
                             sig, sig_size, sig_len );
     if( status != PSA_SUCCESS )
@@ -256,9 +244,31 @@
     status = psa_destroy_key( key_id );
     if( ret == 0 && status != PSA_SUCCESS )
         ret = mbedtls_pk_error_from_psa( status );
-
     return( ret );
 }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+static int rsa_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
+                   const unsigned char *hash, size_t hash_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
+                   int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
+{
+
+    ((void) f_rng);
+    ((void) p_rng);
+    ((void) md_alg);
+
+#if SIZE_MAX > UINT_MAX
+    if( md_alg == MBEDTLS_MD_NONE && UINT_MAX < hash_len )
+        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+#endif /* SIZE_MAX > UINT_MAX */
+
+    return( mbedtls_pk_psa_sign_ext( PSA_ALG_RSA_PKCS1V15_SIGN(
+                                        mbedtls_psa_translate_md( md_alg ) ),
+                                     ctx, hash, hash_len,
+                                     sig, sig_size, sig_len ) );
+}
 #else
 static int rsa_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
diff --git a/library/pk_wrap.h b/library/pk_wrap.h
index ca0d8d8..eead322 100644
--- a/library/pk_wrap.h
+++ b/library/pk_wrap.h
@@ -145,6 +145,15 @@
 #if defined(PSA_WANT_KEY_TYPE_RSA_PUBLIC_KEY)
 int mbedtls_pk_error_from_psa_rsa( psa_status_t status );
 #endif
-#endif
+
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+int mbedtls_pk_psa_sign_ext( psa_algorithm_t psa_alg_md, void *ctx,
+                             const unsigned char *hash, size_t hash_len,
+                             unsigned char *sig, size_t sig_size,
+                             size_t *sig_len );
+
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #endif /* MBEDTLS_PK_WRAP_H */