diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 85728c3..a09877e 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5814,13 +5814,6 @@
         return PSA_ERROR_BAD_STATE;
     }
 
-    if (output_length > operation->capacity) {
-        operation->capacity = 0;
-        /* Go through the error path to wipe all confidential data now
-         * that the operation object is useless. */
-        status = PSA_ERROR_INSUFFICIENT_DATA;
-        goto exit;
-    }
     if (output_length == 0 && operation->capacity == 0) {
         /* Edge case: this is a finished operation, and 0 bytes
          * were requested. The right error in this case could
@@ -5832,6 +5825,14 @@
     }
 
     LOCAL_OUTPUT_ALLOC(output_external, output_length, output);
+    if (output_length > operation->capacity) {
+        operation->capacity = 0;
+        /* Go through the error path to wipe all confidential data now
+         * that the operation object is useless. */
+        status = PSA_ERROR_INSUFFICIENT_DATA;
+        goto exit;
+    }
+
     operation->capacity -= output_length;
 
 #if defined(BUILTIN_ALG_ANY_HKDF)
@@ -5872,8 +5873,6 @@
     }
 
 exit:
-    LOCAL_OUTPUT_FREE(output_external, output);
-
     if (status != PSA_SUCCESS) {
         /* Preserve the algorithm upon errors, but clear all sensitive state.
          * This allows us to differentiate between exhausted operations and
@@ -5884,6 +5883,8 @@
         operation->alg = alg;
         memset(output, '!', output_length);
     }
+
+    LOCAL_OUTPUT_FREE(output_external, output);
     return status;
 }