8 files changed, 151 insertions, 33 deletions

diff --git a/src/Network.cpp b/src/Network.cpp
index cf4fe15..9cb2097 100644
--- a/src/Network.cpp
+++ b/src/Network.cpp
@@ -1,5 +1,7 @@
 #include "Network.hpp"
 
+#include <zlib.h>
+
 Network::Network(std::string address, unsigned short port) {
 	try {
 		socket = new Socket(address, port);
@@ -13,6 +15,8 @@ Network::Network(std::string address, unsigned short port) {
 	} catch (std::exception &e) {
 		LOG(WARNING) << "Stream creation failed: " << e.what();
 	}
+
+
 }
 
 Network::~Network() {
@@ -20,40 +24,114 @@ Network::~Network() {
 	delete socket;
 }
 
-std::shared_ptr<Packet> Network::ReceivePacket(ConnectionState state) {
-	int packetSize = stream->ReadVarInt();
-	auto packetData = stream->ReadByteArray(packetSize);
-	StreamBuffer streamBuffer(packetData.data(), packetData.size());
-	int packetId = streamBuffer.ReadVarInt();
-	auto packet = ReceivePacketByPacketId(packetId, state, streamBuffer);
-	return packet;
+std::shared_ptr<Packet> Network::ReceivePacket(ConnectionState state, bool useCompression) {
+    if (useCompression) {
+        int packetLength = stream->ReadVarInt();
+        auto packetData = stream->ReadByteArray(packetLength);
+        StreamBuffer streamBuffer(packetData.data(), packetData.size());
+        
+        int dataLength = streamBuffer.ReadVarInt();
+        if (dataLength == 0) {
+            auto packetData = streamBuffer.ReadByteArray(packetLength - streamBuffer.GetReadedLength());
+            StreamBuffer streamBuffer(packetData.data(), packetData.size());
+            int packetId = streamBuffer.ReadVarInt();
+            auto packet = ReceivePacketByPacketId(packetId, state, streamBuffer);
+            return packet;
+        } else {
+            std::vector<unsigned char> compressedData = streamBuffer.ReadByteArray(packetLength - streamBuffer.GetReadedLength());
+            std::vector<unsigned char> uncompressedData;
+            uncompressedData.resize(dataLength);
+
+            z_stream stream;
+            stream.avail_in = compressedData.size();
+            stream.next_in = compressedData.data();
+            stream.avail_out = uncompressedData.size();
+            stream.next_out = uncompressedData.data();
+            stream.zalloc = Z_NULL;
+            stream.zfree = Z_NULL;
+            stream.opaque = Z_NULL;
+            if (inflateInit(&stream) != Z_OK)
+                throw std::runtime_error("Zlib decompression initalization error");
+
+            int status = inflate(&stream, Z_FINISH);
+            switch (status) {
+            case Z_STREAM_END:
+                break;
+            case Z_OK:                            
+            case Z_STREAM_ERROR:
+            case Z_BUF_ERROR:
+                throw std::runtime_error("Zlib decompression error: " + std::to_string(status));
+            }
+
+            if (inflateEnd(&stream) != Z_OK)
+                throw std::runtime_error("Zlib decompression end error");
+
+            StreamBuffer streamBuffer(uncompressedData.data(), uncompressedData.size());
+            int packetId = streamBuffer.ReadVarInt();
+            auto packet = ReceivePacketByPacketId(packetId, state, streamBuffer);
+            return packet;            
+        }
+    } else {
+        int packetSize = stream->ReadVarInt();
+        auto packetData = stream->ReadByteArray(packetSize);
+        StreamBuffer streamBuffer(packetData.data(), packetData.size());
+        int packetId = streamBuffer.ReadVarInt();
+        auto packet = ReceivePacketByPacketId(packetId, state, streamBuffer);
+        return packet;
+    }	
 }
 
-void Network::SendPacket(Packet &packet) {
-	StreamCounter packetSize;
-	packetSize.WriteVarInt(packet.GetPacketId());
-	packet.ToStream(&packetSize);
-	stream->WriteVarInt(packetSize.GetCountedSize());
-	stream->WriteVarInt(packet.GetPacketId());
-	packet.ToStream(stream);
+void Network::SendPacket(Packet &packet, int compressionThreshold) {
+    if (compressionThreshold >= 0) {
+        StreamCounter packetSize;
+        packetSize.WriteVarInt(packet.GetPacketId());
+        packetSize.WriteVarInt(0);
+        packet.ToStream(&packetSize);
+        if (packetSize.GetCountedSize() < compressionThreshold) {
+            stream->WriteVarInt(packetSize.GetCountedSize());
+            stream->WriteVarInt(0);
+            stream->WriteVarInt(packet.GetPacketId());
+            packet.ToStream(stream);
+        } else {
+            throw std::runtime_error("Compressing data");
+            /*StreamBuffer buffer(packetSize.GetCountedSize());
+            packet.ToStream(&buffer);
+
+            z_stream stream;*/
+        }
+    }
+    else {
+        StreamCounter packetSize;
+        packetSize.WriteVarInt(packet.GetPacketId());
+        packet.ToStream(&packetSize);
+        stream->WriteVarInt(packetSize.GetCountedSize());
+        stream->WriteVarInt(packet.GetPacketId());
+        packet.ToStream(stream);
+    }	
 }
 
 std::shared_ptr<Packet> Network::ReceivePacketByPacketId(int packetId, ConnectionState state, StreamInput &stream) {
 	std::shared_ptr < Packet > packet(nullptr);
 	switch (state) {
 		case Handshaking:
-			switch (packetId) {
-				case PacketNameHandshakingCB::Handshake:
-					packet = std::make_shared<PacketHandshake>();
-					break;
-			}
+            switch (packetId) {
+            case PacketNameHandshakingCB::Handshake:
+                packet = std::make_shared<PacketHandshake>();
+                break;
+            }
             break;
 		case Login:
-			switch (packetId) {
-				case PacketNameLoginCB::LoginSuccess:
-					packet = std::make_shared<PacketLoginSuccess>();
-					break;
-			}
+            switch (packetId) {
+            case PacketNameLoginCB::LoginSuccess:
+                packet = std::make_shared<PacketLoginSuccess>();
+                break;
+            case PacketNameLoginCB::SetCompression:
+                packet = std::make_shared<PacketSetCompression>();
+                break;
+            case PacketNameLoginCB::Disconnect:
+                packet = std::make_shared<PacketDisconnect>();
+                break;
+            }
 			break;
 		case Play:
 			packet = ParsePacketPlay((PacketNamePlayCB) packetId);
diff --git a/src/Network.hpp b/src/Network.hpp
index 6263e1b..f28f808 100644
--- a/src/Network.hpp
+++ b/src/Network.hpp
@@ -20,7 +20,7 @@ public:
 	Network(std::string address, unsigned short port);
 	~Network();
 
-	std::shared_ptr<Packet> ReceivePacket(ConnectionState state = Play);
-	void SendPacket(Packet &packet);
+	std::shared_ptr<Packet> ReceivePacket(ConnectionState state = Play, bool useCompression = false);
+	void SendPacket(Packet &packet, int compressionThreshold = -1);
 	std::shared_ptr<Packet> ParsePacketPlay(PacketNamePlayCB id);
 };
\ No newline at end of file
diff --git a/src/NetworkClient.cpp b/src/NetworkClient.cpp
index 027eeb5..74d7804 100644
--- a/src/NetworkClient.cpp
+++ b/src/NetworkClient.cpp
@@ -16,13 +16,25 @@ NetworkClient::NetworkClient(std::string address, unsigned short port, std::stri
 	loginStart.Username = "HelloOne";
 	network.SendPacket(loginStart);
 
-	auto response = std::static_pointer_cast<PacketLoginSuccess>(network.ReceivePacket(Login));
 
-    while (!response)
-        response = std::static_pointer_cast<PacketLoginSuccess>(network.ReceivePacket(Login));
+    auto packet = network.ReceivePacket(Login);
+
+    while (!packet)
+        packet = network.ReceivePacket(Login);
+
+    if (packet->GetPacketId() == PacketNameLoginCB::SetCompression) {
+        auto compPacket = std::static_pointer_cast<PacketSetCompression>(packet);
+        LOG(INFO) << "Compression threshold: " << compPacket->Threshold;
+        compressionThreshold = compPacket->Threshold;
+        packet.reset();
+        while (!packet)
+            packet = network.ReceivePacket(Login, compressionThreshold >= 0);
+    }
+
+	auto response = std::static_pointer_cast<PacketLoginSuccess>(packet);
 
 	if (response->Username != username) {
-		throw std::logic_error("Received username is not sended username");
+		throw std::logic_error("Received username is not sended username: "+response->Username+" != "+username);
 	}
 
 	state = Play;
@@ -62,11 +74,11 @@ void NetworkClient::NetworkLoop() {
 			toSendMutex.lock();
 			while (!toSend.empty()) {
 				if (toSend.front() != nullptr)
-					network.SendPacket(*toSend.front());
+					network.SendPacket(*toSend.front(), compressionThreshold);
 				toSend.pop();
 			}
 			toSendMutex.unlock();
-			auto packet = network.ReceivePacket(state);
+            auto packet = network.ReceivePacket(state, compressionThreshold >= 0);
 			if (packet.get() != nullptr) {
 				if (packet->GetPacketId() != PacketNamePlayCB::KeepAliveCB) {
 					toReceiveMutex.lock();
@@ -76,7 +88,7 @@ void NetworkClient::NetworkLoop() {
 					timeOfLastKeepAlivePacket = std::chrono::steady_clock::now();
 					auto packetKeepAlive = std::static_pointer_cast<PacketKeepAliveCB>(packet);
 					auto packetKeepAliveSB = std::make_shared<PacketKeepAliveSB>(packetKeepAlive->KeepAliveId);
-					network.SendPacket(*packetKeepAliveSB);
+					network.SendPacket(*packetKeepAliveSB, compressionThreshold);
 				}
 			}
 			using namespace std::chrono_literals;
diff --git a/src/NetworkClient.hpp b/src/NetworkClient.hpp
index 9acae84..acb3ecb 100644
--- a/src/NetworkClient.hpp
+++ b/src/NetworkClient.hpp
@@ -17,6 +17,7 @@ class NetworkClient {
 	bool isActive=true;
 	ConnectionState state;
 	void NetworkLoop();
+    int compressionThreshold = -1;
 public:
 	NetworkClient(std::string address, unsigned short port, std::string username);
 	~NetworkClient();
diff --git a/src/Packet.hpp b/src/Packet.hpp
index f71922f..0470015 100644
--- a/src/Packet.hpp
+++ b/src/Packet.hpp
@@ -1040,4 +1040,20 @@ struct PacketDisconnect : Packet {
     }
 
     std::string Reason;
+};
+
+struct PacketSetCompression : Packet {
+    void ToStream(StreamOutput *stream) override {
+
+    }
+
+    void FromStream(StreamInput *stream) override {
+        Threshold = stream->ReadVarInt();
+    }
+
+    int GetPacketId() override {
+        return PacketNameLoginCB::SetCompression;
+    }
+
+    int Threshold;
 };
\ No newline at end of file
diff --git a/src/Section.cpp b/src/Section.cpp
index e6b6d93..c105b06 100644
--- a/src/Section.cpp
+++ b/src/Section.cpp
@@ -112,6 +112,9 @@ BlockId Section::GetBlockId(Vector pos) const {
 
 unsigned char Section::GetBlockLight(Vector pos) const
 {
+    if (light.empty())
+        return 0;
+
     int blockNumber = pos.y * 256 + pos.z * 16 + pos.x;
     unsigned char lightValue = this->light[blockNumber / 2];
     return (blockNumber % 2 == 0) ? (lightValue & 0xF) : (lightValue >> 4);
@@ -119,6 +122,9 @@ unsigned char Section::GetBlockLight(Vector pos) const
 
 unsigned char Section::GetBlockSkyLight(Vector pos) const
 {
+    if (sky.empty())
+        return 0;
+
     int blockNumber = pos.y * 256 + pos.z * 16 + pos.x;
     unsigned char skyValue = this->sky[blockNumber / 2];
     return (blockNumber % 2 == 0) ? (skyValue & 0xF) : (skyValue >> 4);
diff --git a/src/Stream.cpp b/src/Stream.cpp
index 6f93c9c..c7935e6 100644
--- a/src/Stream.cpp
+++ b/src/Stream.cpp
@@ -343,6 +343,10 @@ std::vector<unsigned char> StreamBuffer::GetBuffer() {
 	return std::vector<unsigned char>(buffer, buffer + bufferLength);
 }
 
+size_t StreamBuffer::GetReadedLength() {
+    return bufferPtr - buffer;
+}
+
 void StreamCounter::WriteData(unsigned char *buffPtr, size_t buffLen) {
 	buffPtr++;
 	size += buffLen;
diff --git a/src/Stream.hpp b/src/Stream.hpp
index 5432383..e2ba5cf 100644
--- a/src/Stream.hpp
+++ b/src/Stream.hpp
@@ -91,6 +91,7 @@ public:
 	~StreamBuffer();
 
 	std::vector<unsigned char> GetBuffer();
+    size_t GetReadedLength();
 };
 
 class StreamCounter : public StreamOutput {