1 files changed, 35 insertions, 27 deletions

diff --git a/modules/mbedtls/stream_peer_mbedtls.cpp b/modules/mbedtls/stream_peer_mbedtls.cpp
index b4200410fb..3af66f6d83 100644
--- a/modules/mbedtls/stream_peer_mbedtls.cpp
+++ b/modules/mbedtls/stream_peer_mbedtls.cpp
@@ -166,21 +166,24 @@ Error StreamPeerMbedTLS::put_partial_data(const uint8_t *p_data, int p_bytes, in
 		return OK;
 	}
 
-	int ret = mbedtls_ssl_write(tls_ctx->get_context(), p_data, p_bytes);
-	if (ret == MBEDTLS_ERR_SSL_WANT_READ || ret == MBEDTLS_ERR_SSL_WANT_WRITE) {
-		// Non blocking IO
-		ret = 0;
-	} else if (ret == MBEDTLS_ERR_SSL_PEER_CLOSE_NOTIFY) {
-		// Clean close
-		disconnect_from_stream();
-		return ERR_FILE_EOF;
-	} else if (ret <= 0) {
-		TLSContextMbedTLS::print_mbedtls_error(ret);
-		disconnect_from_stream();
-		return ERR_CONNECTION_ERROR;
-	}
+	do {
+		int ret = mbedtls_ssl_write(tls_ctx->get_context(), &p_data[r_sent], p_bytes - r_sent);
+		if (ret == MBEDTLS_ERR_SSL_WANT_READ || ret == MBEDTLS_ERR_SSL_WANT_WRITE) {
+			// Non blocking IO.
+			break;
+		} else if (ret == MBEDTLS_ERR_SSL_PEER_CLOSE_NOTIFY) {
+			// Clean close
+			disconnect_from_stream();
+			return ERR_FILE_EOF;
+		} else if (ret <= 0) {
+			TLSContextMbedTLS::print_mbedtls_error(ret);
+			disconnect_from_stream();
+			return ERR_CONNECTION_ERROR;
+		}
+		r_sent += ret;
+
+	} while (r_sent < p_bytes);
 
-	r_sent = ret;
 	return OK;
 }
 
@@ -209,20 +212,25 @@ Error StreamPeerMbedTLS::get_partial_data(uint8_t *p_buffer, int p_bytes, int &r
 
 	r_received = 0;
 
-	int ret = mbedtls_ssl_read(tls_ctx->get_context(), p_buffer, p_bytes);
-	if (ret == MBEDTLS_ERR_SSL_WANT_READ || ret == MBEDTLS_ERR_SSL_WANT_WRITE) {
-		ret = 0; // non blocking io
-	} else if (ret == MBEDTLS_ERR_SSL_PEER_CLOSE_NOTIFY) {
-		// Clean close
-		disconnect_from_stream();
-		return ERR_FILE_EOF;
-	} else if (ret <= 0) {
-		TLSContextMbedTLS::print_mbedtls_error(ret);
-		disconnect_from_stream();
-		return ERR_CONNECTION_ERROR;
-	}
+	do {
+		int ret = mbedtls_ssl_read(tls_ctx->get_context(), &p_buffer[r_received], p_bytes - r_received);
+		if (ret == MBEDTLS_ERR_SSL_WANT_READ || ret == MBEDTLS_ERR_SSL_WANT_WRITE) {
+			// Non blocking IO.
+			break;
+		} else if (ret == MBEDTLS_ERR_SSL_PEER_CLOSE_NOTIFY) {
+			// Clean close
+			disconnect_from_stream();
+			return ERR_FILE_EOF;
+		} else if (ret <= 0) {
+			TLSContextMbedTLS::print_mbedtls_error(ret);
+			disconnect_from_stream();
+			return ERR_CONNECTION_ERROR;
+		}
+
+		r_received += ret;
+
+	} while (r_received < p_bytes);
 
-	r_received = ret;
 	return OK;
 }