1 files changed, 17 insertions, 14 deletions

diff --git a/modules/websocket/lws_peer.h b/modules/websocket/lws_peer.h
index e96b38b168..3ded3810d1 100644
--- a/modules/websocket/lws_peer.h
+++ b/modules/websocket/lws_peer.h
@@ -37,6 +37,7 @@
 #include "core/ring_buffer.h"
 #include "libwebsockets.h"
 #include "lws_config.h"
+#include "packet_buffer.h"
 #include "websocket_peer.h"
 
 class LWSPeer : public WebSocketPeer {
@@ -44,33 +45,33 @@ class LWSPeer : public WebSocketPeer {
 	GDCIIMPL(LWSPeer, WebSocketPeer);
 
 private:
-	enum {
-		PACKET_BUFFER_SIZE = 65536 - 5 // 4 bytes for the size, 1 for the type
-	};
+	int _in_size;
+	uint8_t _is_string;
+	// Our packet info is just a boolean (is_string), using uint8_t for it.
+	PacketBuffer<uint8_t> _in_buffer;
+	PacketBuffer<uint8_t> _out_buffer;
+
+	PoolVector<uint8_t> _packet_buffer;
 
-	uint8_t packet_buffer[PACKET_BUFFER_SIZE];
 	struct lws *wsi;
 	WriteMode write_mode;
-	bool _was_string;
+
+	int close_code;
+	String close_reason;
 
 public:
 	struct PeerData {
 		uint32_t peer_id;
 		bool force_close;
-		RingBuffer<uint8_t> rbw;
-		RingBuffer<uint8_t> rbr;
-		mutable uint8_t input_buffer[PACKET_BUFFER_SIZE];
-		uint32_t in_size;
-		int in_count;
-		int out_count;
+		bool clean_close;
 	};
 
 	virtual int get_available_packet_count() const;
 	virtual Error get_packet(const uint8_t **r_buffer, int &r_buffer_size);
 	virtual Error put_packet(const uint8_t *p_buffer, int p_buffer_size);
-	virtual int get_max_packet_size() const { return PACKET_BUFFER_SIZE; };
+	virtual int get_max_packet_size() const { return _packet_buffer.size(); };
 
-	virtual void close();
+	virtual void close(int p_code = 1000, String p_reason = "");
 	virtual bool is_connected_to_host() const;
 	virtual IP_Address get_connected_host() const;
 	virtual uint16_t get_connected_port() const;
@@ -79,9 +80,11 @@ public:
 	virtual void set_write_mode(WriteMode p_mode);
 	virtual bool was_string_packet() const;
 
-	void set_wsi(struct lws *wsi);
+	void set_wsi(struct lws *wsi, unsigned int _in_buf_size, unsigned int _in_pkt_size, unsigned int _out_buf_size, unsigned int _out_pkt_size);
 	Error read_wsi(void *in, size_t len);
 	Error write_wsi();
+	void send_close_status(struct lws *wsi);
+	String get_close_reason(void *in, size_t len, int &r_code);
 
 	LWSPeer();
 	~LWSPeer();