From 0d1185de75f594e7f7587d4271e51cacdde68fde Mon Sep 17 00:00:00 2001
From: YAMAMOTO Takashi <yamamoto@midokura.com>
Date: Fri, 4 Feb 2022 12:25:14 +0900
Subject: [PATCH] mqtt_pal.c: Tweak mbedtls implementation

To follow the protocol suggested in
https://github.com/LiamBindle/MQTT-C/issues/110
---
 src/mqtt_pal.c | 35 +++++++++++++++++++++++++----------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/src/mqtt_pal.c b/src/mqtt_pal.c
index dbf4c9c..1aa0995 100644
--- a/src/mqtt_pal.c
+++ b/src/mqtt_pal.c
@@ -48,6 +48,7 @@ int _mqtt_pal_dummy;
 #include <mbedtls/ssl.h>
 
 ssize_t mqtt_pal_sendall(mqtt_pal_socket_handle fd, const void* buf, size_t len, int flags) {
+    enum MQTTErrors error = 0;
     size_t sent = 0;
     while(sent < len) {
         int rv = mbedtls_ssl_write(fd, buf + sent, len - sent);
@@ -64,15 +65,28 @@ ssize_t mqtt_pal_sendall(mqtt_pal_socket_handle fd, const void* buf, size_t len,
                 /* should call mbedtls_ssl_write later again */
                 break;
             }
-            return MQTT_ERROR_SOCKET_ERROR;
+            error = MQTT_ERROR_SOCKET_ERROR;
+            break;
+        }
+        if (rv == 0) {
+            /*
+             * is this possible?
+             * Note: we never send an empty TLS record.
+             */
+            error = MQTT_ERROR_SOCKET_ERROR;
+            break;
         }
         sent += (size_t) rv;
     }
+    if (sent == 0) {
+        return error;
+    }
     return sent;
 }
 
 ssize_t mqtt_pal_recvall(mqtt_pal_socket_handle fd, void* buf, size_t bufsz, int flags) {
     const void *const start = buf;
+    enum MQTTErrors error = 0;
     int rv;
     do {
         rv = mbedtls_ssl_read(fd, buf, bufsz);
@@ -80,13 +94,10 @@ ssize_t mqtt_pal_recvall(mqtt_pal_socket_handle fd, void* buf, size_t bufsz, int
             /*
              * Note: mbedtls_ssl_read returns 0 when the underlying
              * transport was closed without CloseNotify.
+             *
+             * Raise an error to trigger a reconnect.
              */
-            if (buf == start) {
-                /*
-                 * Raise an error to trigger a reconnect.
-                 */
-                return MQTT_ERROR_SOCKET_ERROR;
-            }
+            error = MQTT_ERROR_SOCKET_ERROR;
             break;
         }
         if (rv < 0) {
@@ -102,12 +113,16 @@ ssize_t mqtt_pal_recvall(mqtt_pal_socket_handle fd, void* buf, size_t bufsz, int
                 /* should call mbedtls_ssl_read later again */
                 break;
             }
-            return MQTT_ERROR_SOCKET_ERROR;
+            /* Note: MBEDTLS_ERR_SSL_PEER_CLOSE_NOTIFY is handled here. */
+            error = MQTT_ERROR_SOCKET_ERROR;
+            break;
         }
         buf = (char*)buf + rv;
         bufsz -= rv;
-    } while (rv > 0 && bufsz > 0);
-
+    } while (bufsz > 0);
+    if (buf == start) {
+        return error;
+    }
     return buf - start;
 }