6 files changed, 28 insertions, 13 deletions
diff --git a/modules/mbedtls/dtls_server_mbedtls.h b/modules/mbedtls/dtls_server_mbedtls.h
index 29370062c4..a6626c9f65 100644
--- a/modules/mbedtls/dtls_server_mbedtls.h
+++ b/modules/mbedtls/dtls_server_mbedtls.h
@@ -28,8 +28,8 @@
 /* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.                */
 /*************************************************************************/
 
-#ifndef MBED_DTLS_SERVER_H
-#define MBED_DTLS_SERVER_H
+#ifndef DTLS_SERVER_MBEDTLS_H
+#define DTLS_SERVER_MBEDTLS_H
 
 #include "core/io/dtls_server.h"
 #include "ssl_context_mbedtls.h"
@@ -54,4 +54,4 @@ public:
 	~DTLSServerMbedTLS();
 };
 
-#endif // MBED_DTLS_SERVER_H
+#endif // DTLS_SERVER_MBEDTLS_H
diff --git a/modules/mbedtls/register_types.cpp b/modules/mbedtls/register_types.cpp
index 1af978e70a..2d4a18b3fc 100644
--- a/modules/mbedtls/register_types.cpp
+++ b/modules/mbedtls/register_types.cpp
@@ -39,14 +39,22 @@
 #include "tests/test_crypto_mbedtls.h"
 #endif
 
-void register_mbedtls_types() {
+void initialize_mbedtls_module(ModuleInitializationLevel p_level) {
+	if (p_level != MODULE_INITIALIZATION_LEVEL_SCENE) {
+		return;
+	}
+
 	CryptoMbedTLS::initialize_crypto();
 	StreamPeerMbedTLS::initialize_ssl();
 	PacketPeerMbedDTLS::initialize_dtls();
 	DTLSServerMbedTLS::initialize();
 }
 
-void unregister_mbedtls_types() {
+void uninitialize_mbedtls_module(ModuleInitializationLevel p_level) {
+	if (p_level != MODULE_INITIALIZATION_LEVEL_SCENE) {
+		return;
+	}
+
 	DTLSServerMbedTLS::finalize();
 	PacketPeerMbedDTLS::finalize_dtls();
 	StreamPeerMbedTLS::finalize_ssl();
diff --git a/modules/mbedtls/register_types.h b/modules/mbedtls/register_types.h
index 4bc2cca118..ebe76f44f1 100644
--- a/modules/mbedtls/register_types.h
+++ b/modules/mbedtls/register_types.h
@@ -31,7 +31,9 @@
 #ifndef MBEDTLS_REGISTER_TYPES_H
 #define MBEDTLS_REGISTER_TYPES_H
 
-void register_mbedtls_types();
-void unregister_mbedtls_types();
+#include "modules/register_module_types.h"
+
+void initialize_mbedtls_module(ModuleInitializationLevel p_level);
+void uninitialize_mbedtls_module(ModuleInitializationLevel p_level);
 
 #endif // MBEDTLS_REGISTER_TYPES_H
diff --git a/modules/mbedtls/ssl_context_mbedtls.h b/modules/mbedtls/ssl_context_mbedtls.h
index dd49792abd..5883388311 100644
--- a/modules/mbedtls/ssl_context_mbedtls.h
+++ b/modules/mbedtls/ssl_context_mbedtls.h
@@ -28,8 +28,8 @@
 /* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.                */
 /*************************************************************************/
 
-#ifndef SSL_CONTEXT_MBED_TLS_H
-#define SSL_CONTEXT_MBED_TLS_H
+#ifndef SSL_CONTEXT_MBEDTLS_H
+#define SSL_CONTEXT_MBEDTLS_H
 
 #include "crypto_mbedtls.h"
 
@@ -90,4 +90,4 @@ public:
 	~SSLContextMbedTLS();
 };
 
-#endif // SSL_CONTEXT_MBED_TLS_H
+#endif // SSL_CONTEXT_MBEDTLS_H
diff --git a/modules/mbedtls/stream_peer_mbedtls.cpp b/modules/mbedtls/stream_peer_mbedtls.cpp
index 1818048877..92590fbcf6 100644
--- a/modules/mbedtls/stream_peer_mbedtls.cpp
+++ b/modules/mbedtls/stream_peer_mbedtls.cpp
@@ -298,6 +298,10 @@ StreamPeerMbedTLS::Status StreamPeerMbedTLS::get_status() const {
 	return status;
 }
 
+Ref<StreamPeer> StreamPeerMbedTLS::get_stream() const {
+	return base;
+}
+
 StreamPeerSSL *StreamPeerMbedTLS::_create_func() {
 	return memnew(StreamPeerMbedTLS);
 }
diff --git a/modules/mbedtls/stream_peer_mbedtls.h b/modules/mbedtls/stream_peer_mbedtls.h
index 98b91e65ab..68b07feea9 100644
--- a/modules/mbedtls/stream_peer_mbedtls.h
+++ b/modules/mbedtls/stream_peer_mbedtls.h
@@ -28,8 +28,8 @@
 /* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.                */
 /*************************************************************************/
 
-#ifndef STREAM_PEER_OPEN_SSL_H
-#define STREAM_PEER_OPEN_SSL_H
+#ifndef STREAM_PEER_MBEDTLS_H
+#define STREAM_PEER_MBEDTLS_H
 
 #include "core/io/stream_peer_ssl.h"
 #include "ssl_context_mbedtls.h"
@@ -57,6 +57,7 @@ public:
 	virtual Error accept_stream(Ref<StreamPeer> p_base, Ref<CryptoKey> p_key, Ref<X509Certificate> p_cert, Ref<X509Certificate> p_ca_chain = Ref<X509Certificate>());
 	virtual Error connect_to_stream(Ref<StreamPeer> p_base, bool p_validate_certs = false, const String &p_for_hostname = String(), Ref<X509Certificate> p_valid_cert = Ref<X509Certificate>());
 	virtual Status get_status() const;
+	virtual Ref<StreamPeer> get_stream() const;
 
 	virtual void disconnect_from_stream();
 
@@ -75,4 +76,4 @@ public:
 	~StreamPeerMbedTLS();
 };
 
-#endif // STREAM_PEER_SSL_H
+#endif // STREAM_PEER_MBEDTLS_H