diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 864e495..0aa2357 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -108,7 +108,7 @@
 
 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_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
     mbedtls_rsa_context * rsa = (mbedtls_rsa_context *) ctx;
@@ -119,6 +119,8 @@
 #endif /* SIZE_MAX > UINT_MAX */
 
     *sig_len = mbedtls_rsa_get_len( rsa );
+    if( sig_size < *sig_len )
+        return( MBEDTLS_ERR_PK_BUFFER_TOO_SMALL );
 
     return( mbedtls_rsa_pkcs1_sign( rsa, f_rng, p_rng,
                                     md_alg, (unsigned int) hash_len,
@@ -247,7 +249,7 @@
 
 static int ecdsa_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng );
 
 static int eckey_verify_wrap( void *ctx, mbedtls_md_type_t md_alg,
@@ -269,7 +271,7 @@
 
 static int eckey_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
@@ -278,7 +280,8 @@
     mbedtls_ecdsa_init( &ecdsa );
 
     if( ( ret = mbedtls_ecdsa_from_keypair( &ecdsa, ctx ) ) == 0 )
-        ret = ecdsa_sign_wrap( &ecdsa, md_alg, hash, hash_len, sig, sig_len,
+        ret = ecdsa_sign_wrap( &ecdsa, md_alg, hash, hash_len,
+                               sig, sig_size, sig_len,
                                f_rng, p_rng );
 
     mbedtls_ecdsa_free( &ecdsa );
@@ -295,7 +298,7 @@
 
 static int ecdsa_sign_rs_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
                    void *rs_ctx );
 
@@ -367,7 +370,7 @@
 
 static int eckey_sign_rs_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
                        void *rs_ctx )
 {
@@ -383,7 +386,7 @@
         MBEDTLS_MPI_CHK( mbedtls_ecdsa_from_keypair( &rs->ecdsa_ctx, ctx ) );
 
     MBEDTLS_MPI_CHK( ecdsa_sign_rs_wrap( &rs->ecdsa_ctx, md_alg,
-                                         hash, hash_len, sig, sig_len,
+                                         hash, hash_len, sig, sig_size, sig_len,
                                          f_rng, p_rng, &rs->ecdsa_rs ) );
 
 cleanup:
@@ -652,11 +655,13 @@
 
 static int ecdsa_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
     return( mbedtls_ecdsa_write_signature( (mbedtls_ecdsa_context *) ctx,
-                md_alg, hash, hash_len, sig, sig_len, f_rng, p_rng ) );
+                                           md_alg, hash, hash_len,
+                                           sig, sig_size, sig_len,
+                                           f_rng, p_rng ) );
 }
 
 #if defined(MBEDTLS_ECP_RESTARTABLE)
@@ -770,7 +775,7 @@
 
 static int rsa_alt_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
     mbedtls_rsa_alt_context *rsa_alt = (mbedtls_rsa_alt_context *) ctx;
@@ -783,6 +788,8 @@
     *sig_len = rsa_alt->key_len_func( rsa_alt->key );
     if( *sig_len > MBEDTLS_PK_SIGNATURE_MAX_SIZE )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+    if( *sig_len > sig_size )
+        return( MBEDTLS_ERR_PK_BUFFER_TOO_SMALL );
 
     return( rsa_alt->sign_func( rsa_alt->key, f_rng, p_rng,
                 md_alg, (unsigned int) hash_len, hash, sig ) );
@@ -822,7 +829,8 @@
 
     if( ( ret = rsa_alt_sign_wrap( (void *) prv, MBEDTLS_MD_NONE,
                                    hash, sizeof( hash ),
-                                   sig, &sig_len, f_rng, p_rng ) ) != 0 )
+                                   sig, sizeof( sig ), &sig_len,
+                                   f_rng, p_rng ) ) != 0 )
     {
         return( ret );
     }
@@ -1010,7 +1018,7 @@
 
 static int pk_opaque_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
 #if !defined(MBEDTLS_ECDSA_C)
@@ -1019,41 +1027,28 @@
     ((void) hash);
     ((void) hash_len);
     ((void) sig);
+    ((void) sig_size);
     ((void) sig_len);
     ((void) f_rng);
     ((void) p_rng);
     return( MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE );
 #else /* !MBEDTLS_ECDSA_C */
     const psa_key_id_t *key = (const psa_key_id_t *) ctx;
-    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_algorithm_t alg = PSA_ALG_ECDSA( mbedtls_psa_translate_md( md_alg ) );
-    size_t buf_len;
     psa_status_t status;
 
     /* PSA has its own RNG */
     (void) f_rng;
     (void) p_rng;
 
-    /* PSA needs an output buffer of known size, but our API doesn't provide
-     * that information. Assume that the buffer is large enough for a
-     * maximal-length signature with that key (otherwise the application is
-     * buggy anyway). */
-    status = psa_get_key_attributes( *key, &attributes );
-    if( status != PSA_SUCCESS )
-        return( mbedtls_psa_err_translate_pk( status ) );
-    buf_len = MBEDTLS_ECDSA_MAX_SIG_LEN( psa_get_key_bits( &attributes ) );
-    psa_reset_key_attributes( &attributes );
-    if( buf_len > MBEDTLS_PK_SIGNATURE_MAX_SIZE )
-        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
-
     /* make the signature */
     status = psa_sign_hash( *key, alg, hash, hash_len,
-                            sig, buf_len, sig_len );
+                            sig, sig_size, sig_len );
     if( status != PSA_SUCCESS )
         return( mbedtls_psa_err_translate_pk( status ) );
 
     /* transcode it to ASN.1 sequence */
-    return( pk_ecdsa_sig_asn1_from_psa( sig, sig_len, buf_len ) );
+    return( pk_ecdsa_sig_asn1_from_psa( sig, sig_len, sig_size ) );
 #endif /* !MBEDTLS_ECDSA_C */
 }