From c5eccc8e1c8078d8ce6617d7e66914c6c8bc2cf1 Mon Sep 17 00:00:00 2001
From: "Dylan K. Taylor" <odigiman@gmail.com>
Date: Sat, 15 Apr 2017 19:40:06 +0100
Subject: [PATCH] Minor improvements to Batch encoding

---
 src/pocketmine/Server.php                     | 21 ++++++----------
 .../network/CompressBatchedTask.php           | 12 ++++++----
 .../network/mcpe/protocol/BatchPacket.php     | 24 +++++++++++++++++++
 3 files changed, 39 insertions(+), 18 deletions(-)

diff --git a/src/pocketmine/Server.php b/src/pocketmine/Server.php
index f9c6477d5..f0420bf9c 100644
--- a/src/pocketmine/Server.php
+++ b/src/pocketmine/Server.php
@@ -1771,17 +1771,11 @@ class Server{
 	 */
 	public function batchPackets(array $players, array $packets, $forceSync = false){
 		Timings::$playerNetworkTimer->startTiming();
-		$str = "";
+
+		$pk = new BatchPacket();
 
 		foreach($packets as $p){
-			if($p instanceof DataPacket){
-				if(!$p->isEncoded){
-					$p->encode();
-				}
-				$str .= Binary::writeUnsignedVarInt(strlen($p->buffer)) . $p->buffer;
-			}else{
-				$str .= Binary::writeUnsignedVarInt(strlen($p)) . $p;
-			}
+			$pk->addPacket($p);
 		}
 
 		$targets = [];
@@ -1792,18 +1786,17 @@ class Server{
 		}
 
 		if(!$forceSync and $this->networkCompressionAsync){
-			$task = new CompressBatchedTask($str, $targets, $this->networkCompressionLevel);
+			$task = new CompressBatchedTask($pk, $targets, $this->networkCompressionLevel);
 			$this->getScheduler()->scheduleAsyncTask($task);
 		}else{
-			$this->broadcastPacketsCallback(zlib_encode($str, ZLIB_ENCODING_DEFLATE, $this->networkCompressionLevel), $targets);
+			$pk->compress($this->networkCompressionLevel);
+			$this->broadcastPacketsCallback($pk, $targets);
 		}
 
 		Timings::$playerNetworkTimer->stopTiming();
 	}
 
-	public function broadcastPacketsCallback($data, array $identifiers){
-		$pk = new BatchPacket();
-		$pk->payload = $data;
+	public function broadcastPacketsCallback(BatchPacket $pk, array $identifiers){
 		$pk->encode();
 		$pk->isEncoded = true;
 
diff --git a/src/pocketmine/network/CompressBatchedTask.php b/src/pocketmine/network/CompressBatchedTask.php
index c8fdb426b..357790ae5 100644
--- a/src/pocketmine/network/CompressBatchedTask.php
+++ b/src/pocketmine/network/CompressBatchedTask.php
@@ -21,6 +21,7 @@
 
 namespace pocketmine\network;
 
+use pocketmine\network\mcpe\protocol\BatchPacket;
 use pocketmine\scheduler\AsyncTask;
 use pocketmine\Server;
 
@@ -31,15 +32,18 @@ class CompressBatchedTask extends AsyncTask{
 	public $final;
 	public $targets;
 
-	public function __construct($data, array $targets, $level = 7){
-		$this->data = $data;
+	public function __construct(BatchPacket $data, array $targets, $level = 7){
+		$this->data = serialize($data);
 		$this->targets = $targets;
 		$this->level = $level;
 	}
 
 	public function onRun(){
 		try{
-			$this->final = zlib_encode($this->data, ZLIB_ENCODING_DEFLATE, $this->level);
+			/** @var BatchPacket $pk */
+			$pk = unserialize($this->data);
+			$pk->compress($this->level);
+			$this->final = serialize($pk);
 			$this->data = null;
 		}catch(\Throwable $e){
 
@@ -47,6 +51,6 @@ class CompressBatchedTask extends AsyncTask{
 	}
 
 	public function onCompletion(Server $server){
-		$server->broadcastPacketsCallback($this->final, (array) $this->targets);
+		$server->broadcastPacketsCallback(unserialize($this->final), (array) $this->targets);
 	}
 }
diff --git a/src/pocketmine/network/mcpe/protocol/BatchPacket.php b/src/pocketmine/network/mcpe/protocol/BatchPacket.php
index 7034d438e..f68194152 100644
--- a/src/pocketmine/network/mcpe/protocol/BatchPacket.php
+++ b/src/pocketmine/network/mcpe/protocol/BatchPacket.php
@@ -25,11 +25,14 @@ namespace pocketmine\network\mcpe\protocol;
 
 
 use pocketmine\network\mcpe\NetworkSession;
+use pocketmine\utils\Binary;
+use pocketmine\utils\BinaryStream;
 
 class BatchPacket extends DataPacket{
 	const NETWORK_ID = 0xfe;
 
 	public $payload;
+	public $compressed = false;
 
 	public function canBeBatched() : bool{
 		return false;
@@ -45,9 +48,30 @@ class BatchPacket extends DataPacket{
 
 	public function encode(){
 		$this->reset();
+		assert($this->compressed);
 		$this->put($this->payload);
 	}
 
+	/**
+	 * @param DataPacket|string $packet
+	 */
+	public function addPacket($packet){
+		if($packet instanceof DataPacket){
+			if(!$packet->isEncoded){
+				$packet->encode();
+			}
+			$packet = $packet->buffer;
+		}
+
+		$this->payload .= Binary::writeUnsignedVarInt(strlen($packet)) . $packet;
+	}
+
+	public function compress(int $level = 7){
+		assert(!$this->compressed);
+		$this->payload = zlib_encode($this->payload, ZLIB_ENCODING_DEFLATE, $level);
+		$this->compressed = true;
+	}
+
 	public function handle(NetworkSession $session) : bool{
 		if(strlen($this->payload) < 2){
 			throw new \InvalidStateException("Not enough bytes in payload, expected zlib header");