3 files changed, 83 insertions, 0 deletions

diff --git a/core/io/stream_peer.cpp b/core/io/stream_peer.cpp
index 27f8d4e88f..8ab025dda1 100644
--- a/core/io/stream_peer.cpp
+++ b/core/io/stream_peer.cpp
@@ -410,6 +410,63 @@ void StreamPeer::_bind_methods() {
 
 ////////////////////////////////
 
+int StreamPeerExtension::get_available_bytes() const {
+	int count;
+	if (GDVIRTUAL_CALL(_get_available_bytes, count)) {
+		return count;
+	}
+	WARN_PRINT_ONCE("StreamPeerExtension::_get_available_bytes is unimplemented!");
+	return -1;
+}
+
+Error StreamPeerExtension::get_data(uint8_t *r_buffer, int p_bytes) {
+	int err;
+	int received = 0;
+	if (GDVIRTUAL_CALL(_get_data, r_buffer, p_bytes, &received, err)) {
+		return (Error)err;
+	}
+	WARN_PRINT_ONCE("StreamPeerExtension::_get_data is unimplemented!");
+	return FAILED;
+}
+
+Error StreamPeerExtension::get_partial_data(uint8_t *r_buffer, int p_bytes, int &r_received) {
+	int err;
+	if (GDVIRTUAL_CALL(_get_partial_data, r_buffer, p_bytes, &r_received, err)) {
+		return (Error)err;
+	}
+	WARN_PRINT_ONCE("StreamPeerExtension::_get_partial_data is unimplemented!");
+	return FAILED;
+}
+
+Error StreamPeerExtension::put_data(const uint8_t *p_data, int p_bytes) {
+	int err;
+	int sent = 0;
+	if (GDVIRTUAL_CALL(_put_data, p_data, p_bytes, &sent, err)) {
+		return (Error)err;
+	}
+	WARN_PRINT_ONCE("StreamPeerExtension::_put_data is unimplemented!");
+	return FAILED;
+}
+
+Error StreamPeerExtension::put_partial_data(const uint8_t *p_data, int p_bytes, int &r_sent) {
+	int err;
+	if (GDVIRTUAL_CALL(_put_data, p_data, p_bytes, &r_sent, err)) {
+		return (Error)err;
+	}
+	WARN_PRINT_ONCE("StreamPeerExtension::_put_partial_data is unimplemented!");
+	return FAILED;
+}
+
+void StreamPeerExtension::_bind_methods() {
+	GDVIRTUAL_BIND(_get_data, "r_buffer", "r_bytes", "r_received");
+	GDVIRTUAL_BIND(_get_partial_data, "r_buffer", "r_bytes", "r_received");
+	GDVIRTUAL_BIND(_put_data, "p_data", "p_bytes", "r_sent");
+	GDVIRTUAL_BIND(_put_partial_data, "p_data", "p_bytes", "r_sent");
+	GDVIRTUAL_BIND(_get_available_bytes);
+}
+
+////////////////////////////////
+
 void StreamPeerBuffer::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("seek", "position"), &StreamPeerBuffer::seek);
 	ClassDB::bind_method(D_METHOD("get_size"), &StreamPeerBuffer::get_size);
diff --git a/core/io/stream_peer.h b/core/io/stream_peer.h
index effc3850af..89432951c5 100644
--- a/core/io/stream_peer.h
+++ b/core/io/stream_peer.h
@@ -33,6 +33,10 @@
 
 #include "core/object/ref_counted.h"
 
+#include "core/object/gdvirtual.gen.inc"
+#include "core/object/script_language.h"
+#include "core/variant/native_ptr.h"
+
 class StreamPeer : public RefCounted {
 	GDCLASS(StreamPeer, RefCounted);
 	OBJ_CATEGORY("Networking");
@@ -58,6 +62,7 @@ public:
 
 	virtual int get_available_bytes() const = 0;
 
+	/* helpers */
 	void set_big_endian(bool p_big_endian);
 	bool is_big_endian_enabled() const;
 
@@ -92,6 +97,26 @@ public:
 	StreamPeer() {}
 };
 
+class StreamPeerExtension : public StreamPeer {
+	GDCLASS(StreamPeerExtension, StreamPeer);
+
+protected:
+	static void _bind_methods();
+
+public:
+	virtual Error put_data(const uint8_t *p_data, int p_bytes) override;
+	virtual Error put_partial_data(const uint8_t *p_data, int p_bytes, int &r_sent) override;
+	virtual Error get_data(uint8_t *p_buffer, int p_bytes) override;
+	virtual Error get_partial_data(uint8_t *p_buffer, int p_bytes, int &r_received) override;
+	virtual int get_available_bytes() const override;
+
+	GDVIRTUAL3R(int, _put_data, GDNativeConstPtr<const uint8_t>, int, GDNativePtr<int>);
+	GDVIRTUAL3R(int, _put_partial_data, GDNativeConstPtr<const uint8_t>, int, GDNativePtr<int>);
+	GDVIRTUAL3R(int, _get_data, GDNativePtr<uint8_t>, int, GDNativePtr<int>);
+	GDVIRTUAL3R(int, _get_partial_data, GDNativePtr<uint8_t>, int, GDNativePtr<int>);
+	GDVIRTUAL0RC(int, _get_available_bytes);
+};
+
 class StreamPeerBuffer : public StreamPeer {
 	GDCLASS(StreamPeerBuffer, StreamPeer);
 
diff --git a/core/register_core_types.cpp b/core/register_core_types.cpp
index 3db1434c73..d4b5f5beb9 100644
--- a/core/register_core_types.cpp
+++ b/core/register_core_types.cpp
@@ -169,6 +169,7 @@ void register_core_types() {
 	GDREGISTER_VIRTUAL_CLASS(IP);
 
 	GDREGISTER_VIRTUAL_CLASS(StreamPeer);
+	GDREGISTER_CLASS(StreamPeerExtension);
 	GDREGISTER_CLASS(StreamPeerBuffer);
 	GDREGISTER_CLASS(StreamPeerTCP);
 	GDREGISTER_CLASS(TCPServer);