Merge pull request #22316 from Faless/poll_more_checks_easy_squash

Add checks for clean disconnect in HTTP/TCP/SSL.

core/io/http_client.cpp

@@ -375,6 +375,18 @@ Error HTTPClient::poll() {
 			}
 		} break;
 		case STATUS_CONNECTED: {
+			// Check if we are still connected
+			if (ssl) {
+				Ref<StreamPeerSSL> tmp = connection;
+				tmp->poll();
+				if (tmp->get_status() != StreamPeerSSL::STATUS_CONNECTED) {
+					status = STATUS_CONNECTION_ERROR;
+					return ERR_CONNECTION_ERROR;
+				}
+			} else if (tcp_connection->get_status() != StreamPeerTCP::STATUS_CONNECTED) {
+				status = STATUS_CONNECTION_ERROR;
+				return ERR_CONNECTION_ERROR;
+			}
 			// Connection established, requests can now be made
 			return OK;
 		} break;

core/io/stream_peer_ssl.cpp

@@ -128,6 +128,7 @@ void StreamPeerSSL::_bind_methods() {
 	ADD_PROPERTY(PropertyInfo(Variant::BOOL, "blocking_handshake"), "set_blocking_handshake_enabled", "is_blocking_handshake_enabled");
 
 	BIND_ENUM_CONSTANT(STATUS_DISCONNECTED);
+	BIND_ENUM_CONSTANT(STATUS_HANDSHAKING);
 	BIND_ENUM_CONSTANT(STATUS_CONNECTED);
 	BIND_ENUM_CONSTANT(STATUS_ERROR);
 	BIND_ENUM_CONSTANT(STATUS_ERROR_HOSTNAME_MISMATCH);

core/io/stream_peer_tcp.cpp

@@ -249,6 +249,23 @@ StreamPeerTCP::Status StreamPeerTCP::get_status() {
 
 	if (status == STATUS_CONNECTING) {
 		_poll_connection();
+	} else if (status == STATUS_CONNECTED) {
+		Error err;
+		err = _sock->poll(NetSocket::POLL_TYPE_IN, 0);
+		if (err == OK) {
+			// FIN received
+			if (_sock->get_available_bytes() == 0) {
+				disconnect_from_host();
+				return status;
+			}
+		}
+		// Also poll write
+		err = _sock->poll(NetSocket::POLL_TYPE_IN_OUT, 0);
+		if (err != OK && err != ERR_BUSY) {
+			// Got an error
+			disconnect_from_host();
+			status = STATUS_ERROR;
+		}
 	}
 
 	return status;

modules/mbedtls/stream_peer_mbed_tls.cpp

@@ -29,6 +29,7 @@
 /*************************************************************************/
 
 #include "stream_peer_mbed_tls.h"
+#include "core/io/stream_peer_tcp.h"
 #include "core/os/file_access.h"
 #include "mbedtls/platform_util.h"
 
@@ -98,12 +99,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 +197,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 +244,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 +270,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() != STATUS_CONNECTED) {
+		disconnect_from_stream();
+		return;
 	}
 }
 
@@ -282,6 +309,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() == STATUS_CONNECTED) {
+		// We are still connected on the socket, try to send close notity.
+		mbedtls_ssl_close_notify(&ssl);
+	}
+
 	_cleanup();
 }