diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 3385229..b48af39 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -2333,11 +2333,11 @@
 
     operation->is_sign = is_sign;
 
-    /* Get the output length for the algorithm and key combination. None of the
-     * currently supported algorithms have an output length dependent on actual
-     * key size, so setting it to a bogus value is currently OK. */
+    /* Get the output length for the algorithm and key combination */
     operation->mac_size = PSA_MAC_LENGTH(
-                            psa_get_key_type( &attributes ), 0, alg );
+                            psa_get_key_type( &attributes ),
+                            psa_get_key_bits( &attributes ),
+                            alg );
 
     if( operation->mac_size < 4 )
     {
@@ -2350,7 +2350,7 @@
     }
 
     if( operation->mac_size > PSA_MAC_LENGTH( psa_get_key_type( &attributes ),
-                                              0,
+                                              psa_get_key_bits( &attributes ),
                                               PSA_ALG_FULL_LENGTH_MAC( alg ) ) )
     {
         /* It's impossible to "truncate" to a larger length than the full length