diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index cb8a7d9..6c27dac 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -4052,26 +4052,23 @@
     if( psk_len > POLARSSL_PSK_MAX_LEN )
         return( POLARSSL_ERR_SSL_BAD_INPUT_DATA );
 
-    if( ssl->psk != NULL )
+    if( ssl->psk != NULL || ssl->psk_identity != NULL )
     {
         polarssl_free( ssl->psk );
         polarssl_free( ssl->psk_identity );
     }
 
-    ssl->psk_len = psk_len;
-    ssl->psk_identity_len = psk_identity_len;
-
-    ssl->psk = polarssl_malloc( ssl->psk_len );
-    ssl->psk_identity = polarssl_malloc( ssl->psk_identity_len );
-
-    if( ssl->psk == NULL )
-        return( POLARSSL_ERR_SSL_MALLOC_FAILED );
-    if( ssl->psk_identity == NULL )
+    if( ( ssl->psk = polarssl_malloc( psk_len ) ) == NULL ||
+        ( ssl->psk_identity = polarssl_malloc( psk_identity_len ) ) == NULL )
     {
         polarssl_free( ssl->psk );
+        ssl->psk = NULL;
         return( POLARSSL_ERR_SSL_MALLOC_FAILED );
     }
 
+    ssl->psk_len = psk_len;
+    ssl->psk_identity_len = psk_identity_len;
+
     memcpy( ssl->psk, psk, ssl->psk_len );
     memcpy( ssl->psk_identity, psk_identity, ssl->psk_identity_len );