fix pending_alert issues

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
diff --git a/library/ssl_misc.h b/library/ssl_misc.h
index 5be5b03..a1128ed 100644
--- a/library/ssl_misc.h
+++ b/library/ssl_misc.h
@@ -1358,7 +1358,6 @@
 #define MBEDTLS_SSL_PEND_FATAL_ALERT( type, user_return_value )         \
             mbedtls_ssl_pend_fatal_alert( ssl, type, user_return_value )
 
-
 #if defined(MBEDTLS_SSL_DTLS_ANTI_REPLAY)
 void mbedtls_ssl_dtls_replay_reset( mbedtls_ssl_context *ssl );
 #endif
diff --git a/library/ssl_msg.c b/library/ssl_msg.c
index 9230bcd..1ea5e89 100644
--- a/library/ssl_msg.c
+++ b/library/ssl_msg.c
@@ -5641,24 +5641,37 @@
 
 /*
  * Send pending fatal alerts or warnings.
+ * 0, No alert message.
+ * !0, error from send_alert_message or handshake_step return
  */
 int mbedtls_ssl_handle_pending_alert( mbedtls_ssl_context *ssl )
 {
     int ret;
 
-    /* Send alert if requested */
-    if( ssl->send_alert != 0 )
-    {
-        /* Clear send_alert to avoid infinite loop */
-        ssl->send_alert = 0;
+    /* No pending alert, return success*/
+    if( ssl->send_alert == 0 )
+        return( 0 );
 
-        ret = mbedtls_ssl_send_alert_message( ssl,
-                                 MBEDTLS_SSL_ALERT_LEVEL_FATAL,
-                                 ssl->alert_type );
-        if( ret != 0 )
-            return( ret );
+    ret = mbedtls_ssl_send_alert_message( ssl,
+                                MBEDTLS_SSL_ALERT_LEVEL_FATAL,
+                                ssl->alert_type );
+
+    /* Success or send message fail, clear send_alert flag
+     * except WANT_WRITE. WANT_WRITE means need re-send message.
+     */
+    if( ret != MBEDTLS_ERR_SSL_WANT_WRITE )
+    {
+        ssl->send_alert = 0;
     }
-    return( 0 );
+
+    if( ret != 0 )
+    {
+        /* some errors on send alert message */
+        return( ret );
+    }
+
+    /* Assume alert_reason == handshake_step return */
+    return( ssl->alert_reason );
 }
 
 /*
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 7bb5f9f..c11810d 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -5205,10 +5205,15 @@
 
     if( ret != 0 )
     {
+        /* handshake_step return error. And it is same
+         * with alert_reason.
+         */
         int alert_ret;
         alert_ret = mbedtls_ssl_handle_pending_alert( ssl );
         if( alert_ret != 0 )
         {
+            /* If success send, ret == alert_ret.
+             */
             ret = alert_ret;
             goto cleanup;
         }