PacketBatch: added a getPackets() method which encapsulates some logic

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,25 +327,19 @@ 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);
+			foreach($stream->getPackets($this->packetPool, 500) as $packet){
-			}catch(BinaryDataException $e){
-				$this->logger->debug("Packet batch: " . base64_encode($stream->getBuffer()));
-				throw BadPacketException::wrap($e, "Packet batch decode error");
-			}
-
 				try{
-				$this->handleDataPacket($pk);
+					$this->handleDataPacket($packet);
 				}catch(BadPacketException $e){
-				$this->logger->debug($pk->getName() . ": " . base64_encode($pk->getSerializer()->getBuffer()));
+					$this->logger->debug($packet->getName() . ": " . base64_encode($packet->getSerializer()->getBuffer()));
-				throw BadPacketException::wrap($e, "Error processing " . $pk->getName());
+					throw BadPacketException::wrap($e, "Error processing " . $packet->getName());
 				}
 			}
+		}catch(PacketDecodeException $e){
+			$this->logger->logException($e);
+			throw BadPacketException::wrap($e, "Packet batch decode error");
+		}
 	}
 
 	/**

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.
 	 *

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);
+	}
+}