Merge pull request #20195 from Faless/ssl_handshake

Non-blocking SSL handshake

6 files changed, 110 insertions, 38 deletions

diff --git a/core/io/http_client.cpp b/core/io/http_client.cpp
index fcbb22b5de..f1620f1493 100644
--- a/core/io/http_client.cpp
+++ b/core/io/http_client.cpp
@@ -279,6 +279,7 @@ void HTTPClient::close() {
 	chunk_left = 0;
 	read_until_eof = false;
 	response_num = 0;
+	handshaking = false;
 }
 
 Error HTTPClient::poll() {
@@ -327,16 +328,40 @@ Error HTTPClient::poll() {
 				} break;
 				case StreamPeerTCP::STATUS_CONNECTED: {
 					if (ssl) {
-						Ref<StreamPeerSSL> ssl = StreamPeerSSL::create();
-						Error err = ssl->connect_to_stream(tcp_connection, ssl_verify_host, conn_host);
-						if (err != OK) {
+						Ref<StreamPeerSSL> ssl;
+						if (!handshaking) {
+							// Connect the StreamPeerSSL and start handshaking
+							ssl = Ref<StreamPeerSSL>(StreamPeerSSL::create());
+							ssl->set_blocking_handshake_enabled(false);
+							Error err = ssl->connect_to_stream(tcp_connection, ssl_verify_host, conn_host);
+							if (err != OK) {
+								close();
+								status = STATUS_SSL_HANDSHAKE_ERROR;
+								return ERR_CANT_CONNECT;
+							}
+							connection = ssl;
+							handshaking = true;
+						} else {
+							// We are already handshaking, which means we can use your already active SSL connection
+							ssl = static_cast<Ref<StreamPeerSSL> >(connection);
+							ssl->poll(); // Try to finish the handshake
+						}
+
+						if (ssl->get_status() == StreamPeerSSL::STATUS_CONNECTED) {
+							// Handshake has been successfull
+							handshaking = false;
+							status = STATUS_CONNECTED;
+							return OK;
+						} else if (ssl->get_status() != StreamPeerSSL::STATUS_HANDSHAKING) {
+							// Handshake has failed
 							close();
 							status = STATUS_SSL_HANDSHAKE_ERROR;
 							return ERR_CANT_CONNECT;
 						}
-						connection = ssl;
+						// ... we will need to poll more for handshake to finish
+					} else {
+						status = STATUS_CONNECTED;
 					}
-					status = STATUS_CONNECTED;
 					return OK;
 				} break;
 				case StreamPeerTCP::STATUS_ERROR:
@@ -669,6 +694,7 @@ HTTPClient::HTTPClient() {
 	response_num = 0;
 	ssl = false;
 	blocking = false;
+	handshaking = false;
 	read_chunk_size = 4096;
 }
 
diff --git a/core/io/http_client.h b/core/io/http_client.h
index 38ec82ce8c..82b56b01db 100644
--- a/core/io/http_client.h
+++ b/core/io/http_client.h
@@ -165,6 +165,7 @@ private:
 	bool ssl;
 	bool ssl_verify_host;
 	bool blocking;
+	bool handshaking;
 
 	Vector<uint8_t> response_str;
 
diff --git a/core/io/stream_peer_ssl.cpp b/core/io/stream_peer_ssl.cpp
index 012ba78c6d..c71af6b641 100644
--- a/core/io/stream_peer_ssl.cpp
+++ b/core/io/stream_peer_ssl.cpp
@@ -52,6 +52,14 @@ bool StreamPeerSSL::is_available() {
 	return available;
 }
 
+void StreamPeerSSL::set_blocking_handshake_enabled(bool p_enabled) {
+	blocking_handshake = p_enabled;
+}
+
+bool StreamPeerSSL::is_blocking_handshake_enabled() const {
+	return blocking_handshake;
+}
+
 PoolByteArray StreamPeerSSL::get_project_cert_array() {
 
 	PoolByteArray out;
@@ -84,16 +92,21 @@ PoolByteArray StreamPeerSSL::get_project_cert_array() {
 void StreamPeerSSL::_bind_methods() {
 
 	ClassDB::bind_method(D_METHOD("poll"), &StreamPeerSSL::poll);
-	ClassDB::bind_method(D_METHOD("accept_stream", "stream"), &StreamPeerSSL::accept_stream);
+	ClassDB::bind_method(D_METHOD("accept_stream"), &StreamPeerSSL::accept_stream);
 	ClassDB::bind_method(D_METHOD("connect_to_stream", "stream", "validate_certs", "for_hostname"), &StreamPeerSSL::connect_to_stream, DEFVAL(false), DEFVAL(String()));
 	ClassDB::bind_method(D_METHOD("get_status"), &StreamPeerSSL::get_status);
 	ClassDB::bind_method(D_METHOD("disconnect_from_stream"), &StreamPeerSSL::disconnect_from_stream);
+	ClassDB::bind_method(D_METHOD("set_blocking_handshake_enabled", "enabled"), &StreamPeerSSL::set_blocking_handshake_enabled);
+	ClassDB::bind_method(D_METHOD("is_blocking_handshake_enabled"), &StreamPeerSSL::is_blocking_handshake_enabled);
+
+	ADD_PROPERTY(PropertyInfo(Variant::BOOL, "blocking_handshake"), "set_blocking_handshake_enabled", "is_blocking_handshake_enabled");
 
 	BIND_ENUM_CONSTANT(STATUS_DISCONNECTED);
 	BIND_ENUM_CONSTANT(STATUS_CONNECTED);
-	BIND_ENUM_CONSTANT(STATUS_ERROR_NO_CERTIFICATE);
+	BIND_ENUM_CONSTANT(STATUS_ERROR);
 	BIND_ENUM_CONSTANT(STATUS_ERROR_HOSTNAME_MISMATCH);
 }
 
 StreamPeerSSL::StreamPeerSSL() {
+	blocking_handshake = true;
 }
diff --git a/core/io/stream_peer_ssl.h b/core/io/stream_peer_ssl.h
index 77301a7c87..870704e875 100644
--- a/core/io/stream_peer_ssl.h
+++ b/core/io/stream_peer_ssl.h
@@ -49,14 +49,20 @@ protected:
 	friend class Main;
 	static bool initialize_certs;
 
+	bool blocking_handshake;
+
 public:
 	enum Status {
 		STATUS_DISCONNECTED,
+		STATUS_HANDSHAKING,
 		STATUS_CONNECTED,
-		STATUS_ERROR_NO_CERTIFICATE,
+		STATUS_ERROR,
 		STATUS_ERROR_HOSTNAME_MISMATCH
 	};
 
+	void set_blocking_handshake_enabled(bool p_enabled);
+	bool is_blocking_handshake_enabled() const;
+
 	virtual void poll() = 0;
 	virtual Error accept_stream(Ref<StreamPeer> p_base) = 0;
 	virtual Error connect_to_stream(Ref<StreamPeer> p_base, bool p_validate_certs = false, const String &p_for_hostname = String()) = 0;
diff --git a/modules/mbedtls/stream_peer_mbed_tls.cpp b/modules/mbedtls/stream_peer_mbed_tls.cpp
index a63e53ec1f..884c26ddfe 100755
--- a/modules/mbedtls/stream_peer_mbed_tls.cpp
+++ b/modules/mbedtls/stream_peer_mbed_tls.cpp
@@ -29,6 +29,8 @@
 /*************************************************************************/
 
 #include "stream_peer_mbed_tls.h"
+#include "mbedtls/platform_util.h"
+#include "os/file_access.h"
 
 static void my_debug(void *ctx, int level,
 		const char *file, int line,
@@ -81,6 +83,36 @@ int StreamPeerMbedTLS::bio_recv(void *ctx, unsigned char *buf, size_t len) {
 	return got;
 }
 
+void StreamPeerMbedTLS::_cleanup() {
+
+	mbedtls_ssl_free(&ssl);
+	mbedtls_ssl_config_free(&conf);
+	mbedtls_ctr_drbg_free(&ctr_drbg);
+	mbedtls_entropy_free(&entropy);
+
+	base = Ref<StreamPeer>();
+	status = STATUS_DISCONNECTED;
+}
+
+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) {
+			ERR_PRINTS("TLS handshake error: " + itos(ret));
+			_print_error(ret);
+			disconnect_from_stream();
+			status = STATUS_ERROR;
+			return FAILED;
+		} else if (!blocking_handshake) {
+			// Will retry via poll later
+			return OK;
+		}
+	}
+
+	status = STATUS_CONNECTED;
+	return OK;
+}
+
 Error StreamPeerMbedTLS::connect_to_stream(Ref<StreamPeer> p_base, bool p_validate_certs, const String &p_for_hostname) {
 
 	base = p_base;
@@ -95,6 +127,7 @@ Error StreamPeerMbedTLS::connect_to_stream(Ref<StreamPeer> p_base, bool p_valida
 	ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy, NULL, 0);
 	if (ret != 0) {
 		ERR_PRINTS(" failed\n  ! mbedtls_ctr_drbg_seed returned an error" + itos(ret));
+		_cleanup();
 		return FAILED;
 	}
 
@@ -112,29 +145,24 @@ Error StreamPeerMbedTLS::connect_to_stream(Ref<StreamPeer> p_base, bool p_valida
 
 	mbedtls_ssl_set_bio(&ssl, this, bio_send, bio_recv, NULL);
 
-	while ((ret = mbedtls_ssl_handshake(&ssl)) != 0) {
-		if (ret != MBEDTLS_ERR_SSL_WANT_READ && ret != MBEDTLS_ERR_SSL_WANT_WRITE) {
-			ERR_PRINTS("TLS handshake error: " + itos(ret));
-			_print_error(ret);
-			status = STATUS_ERROR_HOSTNAME_MISMATCH;
-			return FAILED;
-		}
-	}
+	status = STATUS_HANDSHAKING;
 
-	connected = true;
-	status = STATUS_CONNECTED;
+	if ((ret = _do_handshake()) != OK) {
+		status = STATUS_ERROR_HOSTNAME_MISMATCH;
+		return FAILED;
+	}
 
 	return OK;
 }
 
 Error StreamPeerMbedTLS::accept_stream(Ref<StreamPeer> p_base) {
 
-	return ERR_UNAVAILABLE;
+	return OK;
 }
 
 Error StreamPeerMbedTLS::put_data(const uint8_t *p_data, int p_bytes) {
 
-	ERR_FAIL_COND_V(!connected, ERR_UNCONFIGURED);
+	ERR_FAIL_COND_V(status != STATUS_CONNECTED, ERR_UNCONFIGURED);
 
 	Error err;
 	int sent = 0;
@@ -155,7 +183,7 @@ Error StreamPeerMbedTLS::put_data(const uint8_t *p_data, int p_bytes) {
 
 Error StreamPeerMbedTLS::put_partial_data(const uint8_t *p_data, int p_bytes, int &r_sent) {
 
-	ERR_FAIL_COND_V(!connected, ERR_UNCONFIGURED);
+	ERR_FAIL_COND_V(status != STATUS_CONNECTED, ERR_UNCONFIGURED);
 
 	r_sent = 0;
 
@@ -177,7 +205,7 @@ Error StreamPeerMbedTLS::put_partial_data(const uint8_t *p_data, int p_bytes, in
 
 Error StreamPeerMbedTLS::get_data(uint8_t *p_buffer, int p_bytes) {
 
-	ERR_FAIL_COND_V(!connected, ERR_UNCONFIGURED);
+	ERR_FAIL_COND_V(status != STATUS_CONNECTED, ERR_UNCONFIGURED);
 
 	Error err;
 
@@ -199,7 +227,7 @@ Error StreamPeerMbedTLS::get_data(uint8_t *p_buffer, int p_bytes) {
 
 Error StreamPeerMbedTLS::get_partial_data(uint8_t *p_buffer, int p_bytes, int &r_received) {
 
-	ERR_FAIL_COND_V(!connected, ERR_UNCONFIGURED);
+	ERR_FAIL_COND_V(status != STATUS_CONNECTED, ERR_UNCONFIGURED);
 
 	r_received = 0;
 
@@ -218,27 +246,30 @@ Error StreamPeerMbedTLS::get_partial_data(uint8_t *p_buffer, int p_bytes, int &r
 
 void StreamPeerMbedTLS::poll() {
 
-	ERR_FAIL_COND(!connected);
+	ERR_FAIL_COND(status != STATUS_CONNECTED && status != STATUS_HANDSHAKING);
 	ERR_FAIL_COND(!base.is_valid());
 
+	if (status == STATUS_HANDSHAKING) {
+		_do_handshake();
+		return;
+	}
+
 	int ret = mbedtls_ssl_read(&ssl, NULL, 0);
 
 	if (ret < 0 && ret != MBEDTLS_ERR_SSL_WANT_READ && ret != MBEDTLS_ERR_SSL_WANT_WRITE) {
 		_print_error(ret);
 		disconnect_from_stream();
-		return;
 	}
 }
 
 int StreamPeerMbedTLS::get_available_bytes() const {
 
-	ERR_FAIL_COND_V(!connected, 0);
+	ERR_FAIL_COND_V(status != STATUS_CONNECTED, 0);
 
 	return mbedtls_ssl_get_bytes_avail(&ssl);
 }
 StreamPeerMbedTLS::StreamPeerMbedTLS() {
 
-	connected = false;
 	status = STATUS_DISCONNECTED;
 }
 
@@ -248,17 +279,10 @@ StreamPeerMbedTLS::~StreamPeerMbedTLS() {
 
 void StreamPeerMbedTLS::disconnect_from_stream() {
 
-	if (!connected)
+	if (status != STATUS_CONNECTED && status != STATUS_HANDSHAKING)
 		return;
 
-	mbedtls_ssl_free(&ssl);
-	mbedtls_ssl_config_free(&conf);
-	mbedtls_ctr_drbg_free(&ctr_drbg);
-	mbedtls_entropy_free(&entropy);
-
-	base = Ref<StreamPeer>();
-	connected = false;
-	status = STATUS_DISCONNECTED;
+	_cleanup();
 }
 
 StreamPeerMbedTLS::Status StreamPeerMbedTLS::get_status() const {
diff --git a/modules/mbedtls/stream_peer_mbed_tls.h b/modules/mbedtls/stream_peer_mbed_tls.h
index 2b96a194a1..7f4e5a4513 100755
--- a/modules/mbedtls/stream_peer_mbed_tls.h
+++ b/modules/mbedtls/stream_peer_mbed_tls.h
@@ -48,8 +48,6 @@ private:
 	Status status;
 	String hostname;
 
-	bool connected;
-
 	Ref<StreamPeer> base;
 
 	static StreamPeerSSL *_create_func();
@@ -57,9 +55,11 @@ private:
 
 	static int bio_recv(void *ctx, unsigned char *buf, size_t len);
 	static int bio_send(void *ctx, const unsigned char *buf, size_t len);
+	void _cleanup();
 
 protected:
 	static mbedtls_x509_crt cacert;
+
 	mbedtls_entropy_context entropy;
 	mbedtls_ctr_drbg_context ctr_drbg;
 	mbedtls_ssl_context ssl;
@@ -67,6 +67,8 @@ protected:
 
 	static void _bind_methods();
 
+	Error _do_handshake();
+
 public:
 	virtual void poll();
 	virtual Error accept_stream(Ref<StreamPeer> p_base);