From 30591d047c82d9903efb3f5b049285e1131ff8d4 Mon Sep 17 00:00:00 2001
From: "Dylan K. Taylor" <odigiman@gmail.com>
Date: Wed, 1 Jul 2020 13:35:34 +0100
Subject: [PATCH] PacketBatch: added a getPackets() method which encapsulates
 some logic

---
 src/network/mcpe/NetworkSession.php           | 29 ++++-----
 .../mcpe/protocol/serializer/PacketBatch.php  | 14 +++++
 .../protocol/serializer/PacketBatchTest.php   | 62 +++++++++++++++++++
 3 files changed, 87 insertions(+), 18 deletions(-)
 create mode 100644 tests/phpunit/network/mcpe/protocol/serializer/PacketBatchTest.php

diff --git a/src/network/mcpe/NetworkSession.php b/src/network/mcpe/NetworkSession.php
index 0f9bfad87..9ce7ca057 100644
--- a/src/network/mcpe/NetworkSession.php
+++ b/src/network/mcpe/NetworkSession.php
@@ -98,7 +98,6 @@ use pocketmine\player\Player;
 use pocketmine\player\PlayerInfo;
 use pocketmine\Server;
 use pocketmine\timings\Timings;
-use pocketmine\utils\BinaryDataException;
 use pocketmine\utils\TextFormat;
 use pocketmine\utils\Utils;
 use pocketmine\world\Position;
@@ -328,24 +327,18 @@ class NetworkSession{
 			Timings::$playerNetworkReceiveDecompressTimer->stopTiming();
 		}
 
-		$count = 0;
-		while(!$stream->feof() and $this->connected){
-			if($count++ >= 500){
-				throw new BadPacketException("Too many packets in a single batch");
-			}
-			try{
-				$pk = $stream->getPacket($this->packetPool);
-			}catch(BinaryDataException $e){
-				$this->logger->debug("Packet batch: " . base64_encode($stream->getBuffer()));
-				throw BadPacketException::wrap($e, "Packet batch decode error");
-			}
-
-			try{
-				$this->handleDataPacket($pk);
-			}catch(BadPacketException $e){
-				$this->logger->debug($pk->getName() . ": " . base64_encode($pk->getSerializer()->getBuffer()));
-				throw BadPacketException::wrap($e, "Error processing " . $pk->getName());
+		try{
+			foreach($stream->getPackets($this->packetPool, 500) as $packet){
+				try{
+					$this->handleDataPacket($packet);
+				}catch(BadPacketException $e){
+					$this->logger->debug($packet->getName() . ": " . base64_encode($packet->getSerializer()->getBuffer()));
+					throw BadPacketException::wrap($e, "Error processing " . $packet->getName());
+				}
 			}
+		}catch(PacketDecodeException $e){
+			$this->logger->logException($e);
+			throw BadPacketException::wrap($e, "Packet batch decode error");
 		}
 	}
 
diff --git a/src/network/mcpe/protocol/serializer/PacketBatch.php b/src/network/mcpe/protocol/serializer/PacketBatch.php
index a242a2b6f..98a41fd04 100644
--- a/src/network/mcpe/protocol/serializer/PacketBatch.php
+++ b/src/network/mcpe/protocol/serializer/PacketBatch.php
@@ -24,6 +24,7 @@ declare(strict_types=1);
 namespace pocketmine\network\mcpe\protocol\serializer;
 
 use pocketmine\network\mcpe\protocol\Packet;
+use pocketmine\network\mcpe\protocol\PacketDecodeException;
 use pocketmine\network\mcpe\protocol\PacketPool;
 use pocketmine\utils\BinaryDataException;
 
@@ -48,6 +49,19 @@ class PacketBatch{
 		return $packetPool->getPacket($this->serializer->getString());
 	}
 
+	/**
+	 * @return \Generator|Packet[]
+	 * @phpstan-return \Generator<int, Packet, void, void>
+	 */
+	public function getPackets(PacketPool $packetPool, int $max) : \Generator{
+		for($c = 0; $c < $max and !$this->serializer->feof(); ++$c){
+			yield $c => $packetPool->getPacket($this->serializer->getString());
+		}
+		if(!$this->serializer->feof()){
+			throw new PacketDecodeException("Reached limit of $max packets in a single batch");
+		}
+	}
+
 	/**
 	 * Constructs a packet batch from the given list of packets.
 	 *
diff --git a/tests/phpunit/network/mcpe/protocol/serializer/PacketBatchTest.php b/tests/phpunit/network/mcpe/protocol/serializer/PacketBatchTest.php
new file mode 100644
index 000000000..217914d31
--- /dev/null
+++ b/tests/phpunit/network/mcpe/protocol/serializer/PacketBatchTest.php
@@ -0,0 +1,62 @@
+<?php
+
+/*
+ *
+ *  ____            _        _   __  __ _                  __  __ ____
+ * |  _ \ ___   ___| | _____| |_|  \/  (_)_ __   ___      |  \/  |  _ \
+ * | |_) / _ \ / __| |/ / _ \ __| |\/| | | '_ \ / _ \_____| |\/| | |_) |
+ * |  __/ (_) | (__|   <  __/ |_| |  | | | | | |  __/_____| |  | |  __/
+ * |_|   \___/ \___|_|\_\___|\__|_|  |_|_|_| |_|\___|     |_|  |_|_|
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * @author PocketMine Team
+ * @link http://www.pocketmine.net/
+ *
+ *
+*/
+
+declare(strict_types=1);
+
+namespace pocketmine\mcpe\protocol\serializer;
+
+use pocketmine\network\mcpe\protocol\Packet;
+use pocketmine\network\mcpe\protocol\PacketDecodeException;
+use pocketmine\network\mcpe\protocol\PacketPool;
+use pocketmine\network\mcpe\protocol\serializer\PacketBatch;
+use PHPUnit\Framework\TestCase;
+use pocketmine\network\mcpe\protocol\TestPacket;
+
+class PacketBatchTest extends TestCase{
+
+	public function testDecodeTooBig() : void{
+		$limit = 10;
+		$write = new PacketBatch();
+		for($i = 0; $i < $limit + 1; $i++){
+			$write->putPacket(new TestPacket());
+		}
+		$read = new PacketBatch($write->getBuffer());
+		$this->expectException(PacketDecodeException::class);
+		$readCount = 0;
+		foreach($read->getPackets(PacketPool::getInstance(), $limit) as $packet){
+			$readCount++;
+		}
+	}
+
+	public function testDecodeAtLimit() : void{
+		$limit = 10;
+		$write = new PacketBatch();
+		for($i = 0; $i < $limit; $i++){
+			$write->putPacket(new TestPacket());
+		}
+		$read = new PacketBatch($write->getBuffer());
+		$readCount = 0;
+		foreach($read->getPackets(PacketPool::getInstance(), $limit) as $packet){
+			$readCount++;
+		}
+		self::assertSame($limit, $readCount);
+	}
+}