3 files changed, 56 insertions, 19 deletions

diff --git a/modules/mbedtls/SCsub b/modules/mbedtls/SCsub
index d11d7a7ec7..0c6c703e16 100755
--- a/modules/mbedtls/SCsub
+++ b/modules/mbedtls/SCsub
@@ -91,8 +91,12 @@ if env['builtin_mbedtls']:
 
     thirdparty_dir = "#thirdparty/mbedtls/library/"
     thirdparty_sources = [thirdparty_dir + file for file in thirdparty_sources]
-    env_mbed_tls.add_source_files(env.modules_sources, thirdparty_sources)
+
     env_mbed_tls.Prepend(CPPPATH=["#thirdparty/mbedtls/include/"])
 
+    env_thirdparty = env_mbed_tls.Clone()
+    env_thirdparty.disable_warnings()
+    env_thirdparty.add_source_files(env.modules_sources, thirdparty_sources)
+
 # Module sources
 env_mbed_tls.add_source_files(env.modules_sources, "*.cpp")
diff --git a/modules/mbedtls/stream_peer_mbed_tls.cpp b/modules/mbedtls/stream_peer_mbed_tls.cpp
index 884c26ddfe..5c81f32e9e 100755
--- a/modules/mbedtls/stream_peer_mbed_tls.cpp
+++ b/modules/mbedtls/stream_peer_mbed_tls.cpp
@@ -29,8 +29,11 @@
 /*************************************************************************/
 
 #include "stream_peer_mbed_tls.h"
-#include "mbedtls/platform_util.h"
-#include "os/file_access.h"
+
+#include "core/io/stream_peer_tcp.h"
+#include "core/os/file_access.h"
+
+#include <mbedtls/platform_util.h>
 
 static void my_debug(void *ctx, int level,
 		const char *file, int line,
@@ -98,12 +101,16 @@ Error StreamPeerMbedTLS::_do_handshake() {
 	int ret = 0;
 	while ((ret = mbedtls_ssl_handshake(&ssl)) != 0) {
 		if (ret != MBEDTLS_ERR_SSL_WANT_READ && ret != MBEDTLS_ERR_SSL_WANT_WRITE) {
+			// An error occurred.
 			ERR_PRINTS("TLS handshake error: " + itos(ret));
 			_print_error(ret);
 			disconnect_from_stream();
 			status = STATUS_ERROR;
 			return FAILED;
-		} else if (!blocking_handshake) {
+		}
+
+		// Handshake is still in progress.
+		if (!blocking_handshake) {
 			// Will retry via poll later
 			return OK;
 		}
@@ -192,7 +199,12 @@ Error StreamPeerMbedTLS::put_partial_data(const uint8_t *p_data, int p_bytes, in
 
 	int ret = mbedtls_ssl_write(&ssl, p_data, p_bytes);
 	if (ret == MBEDTLS_ERR_SSL_WANT_READ || ret == MBEDTLS_ERR_SSL_WANT_WRITE) {
-		ret = 0; // non blocking io
+		// 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) {
 		_print_error(ret);
 		disconnect_from_stream();
@@ -234,6 +246,10 @@ Error StreamPeerMbedTLS::get_partial_data(uint8_t *p_buffer, int p_bytes, int &r
 	int ret = mbedtls_ssl_read(&ssl, 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) {
 		_print_error(ret);
 		disconnect_from_stream();
@@ -256,9 +272,22 @@ void StreamPeerMbedTLS::poll() {
 
 	int ret = mbedtls_ssl_read(&ssl, NULL, 0);
 
-	if (ret < 0 && ret != MBEDTLS_ERR_SSL_WANT_READ && ret != MBEDTLS_ERR_SSL_WANT_WRITE) {
+	if (ret == MBEDTLS_ERR_SSL_WANT_READ || ret == MBEDTLS_ERR_SSL_WANT_WRITE) {
+		// Nothing to read/write (non blocking IO)
+	} else if (ret == MBEDTLS_ERR_SSL_PEER_CLOSE_NOTIFY) {
+		// Clean close (disconnect)
+		disconnect_from_stream();
+		return;
+	} else if (ret < 0) {
 		_print_error(ret);
 		disconnect_from_stream();
+		return;
+	}
+
+	Ref<StreamPeerTCP> tcp = base;
+	if (tcp.is_valid() && tcp->get_status() != StreamPeerTCP::STATUS_CONNECTED) {
+		disconnect_from_stream();
+		return;
 	}
 }
 
@@ -282,6 +311,12 @@ void StreamPeerMbedTLS::disconnect_from_stream() {
 	if (status != STATUS_CONNECTED && status != STATUS_HANDSHAKING)
 		return;
 
+	Ref<StreamPeerTCP> tcp = base;
+	if (tcp.is_valid() && tcp->get_status() == StreamPeerTCP::STATUS_CONNECTED) {
+		// We are still connected on the socket, try to send close notity.
+		mbedtls_ssl_close_notify(&ssl);
+	}
+
 	_cleanup();
 }
 
@@ -317,15 +352,13 @@ void StreamPeerMbedTLS::initialize_ssl() {
 	mbedtls_debug_set_threshold(1);
 #endif
 
-	PoolByteArray cert_array = StreamPeerSSL::get_project_cert_array();
-
-	if (cert_array.size() > 0)
-		_load_certs(cert_array);
-
 	available = true;
 }
 
 void StreamPeerMbedTLS::finalize_ssl() {
 
+	available = false;
+	_create = NULL;
+	load_certs_func = NULL;
 	mbedtls_x509_crt_free(&cacert);
 }
diff --git a/modules/mbedtls/stream_peer_mbed_tls.h b/modules/mbedtls/stream_peer_mbed_tls.h
index 7f4e5a4513..abf87b79cc 100755
--- a/modules/mbedtls/stream_peer_mbed_tls.h
+++ b/modules/mbedtls/stream_peer_mbed_tls.h
@@ -31,14 +31,14 @@
 #ifndef STREAM_PEER_OPEN_SSL_H
 #define STREAM_PEER_OPEN_SSL_H
 
-#include "io/stream_peer_ssl.h"
-
-#include "mbedtls/config.h"
-#include "mbedtls/ctr_drbg.h"
-#include "mbedtls/debug.h"
-#include "mbedtls/entropy.h"
-#include "mbedtls/net.h"
-#include "mbedtls/ssl.h"
+#include "core/io/stream_peer_ssl.h"
+
+#include <mbedtls/config.h>
+#include <mbedtls/ctr_drbg.h>
+#include <mbedtls/debug.h>
+#include <mbedtls/entropy.h>
+#include <mbedtls/net.h>
+#include <mbedtls/ssl.h>
 
 #include <stdio.h>
 #include <stdlib.h>