From 3dafee6aa6704405cc43516ffc75431ae263ae48 Mon Sep 17 00:00:00 2001
From: "Dylan K. Taylor" <odigiman@gmail.com>
Date: Thu, 14 May 2020 10:58:37 +0100
Subject: [PATCH] NetworkSession: explicitly unregister effect manager hooks on
 dispose, close #3455

---
 src/entity/effect/EffectManager.php | 33 +++++++++++++++--------------
 src/network/mcpe/NetworkSession.php | 23 ++++++++++++++++++--
 2 files changed, 38 insertions(+), 18 deletions(-)

diff --git a/src/entity/effect/EffectManager.php b/src/entity/effect/EffectManager.php
index 9f05a4eb1..683c713bb 100644
--- a/src/entity/effect/EffectManager.php
+++ b/src/entity/effect/EffectManager.php
@@ -23,14 +23,13 @@ declare(strict_types=1);
 
 namespace pocketmine\entity\effect;
 
+use Ds\Set;
 use pocketmine\entity\Living;
 use pocketmine\event\entity\EntityEffectAddEvent;
 use pocketmine\event\entity\EntityEffectRemoveEvent;
 use pocketmine\utils\Color;
-use pocketmine\utils\Utils;
 use function abs;
 use function count;
-use function spl_object_id;
 
 class EffectManager{
 
@@ -46,19 +45,21 @@ class EffectManager{
 	protected $onlyAmbientEffects = false;
 
 	/**
-	 * @var \Closure[]
-	 * @phpstan-var (\Closure(EffectInstance, bool $replacesOldEffect) : void)[]
+	 * @var \Closure[]|Set
+	 * @phpstan-var Set<\Closure(EffectInstance, bool $replacesOldEffect) : void>
 	 */
-	protected $effectAddHooks = [];
+	protected $effectAddHooks;
 	/**
-	 * @var \Closure[]
-	 * @phpstan-var (\Closure(EffectInstance) : void)[]
+	 * @var \Closure[]|Set
+	 * @phpstan-var Set<\Closure(EffectInstance) : void>
 	 */
-	protected $effectRemoveHooks = [];
+	protected $effectRemoveHooks;
 
 	public function __construct(Living $entity){
 		$this->entity = $entity;
 		$this->bubbleColor = new Color(0, 0, 0, 0);
+		$this->effectAddHooks = new Set();
+		$this->effectRemoveHooks = new Set();
 	}
 
 	/**
@@ -222,18 +223,18 @@ class EffectManager{
 	}
 
 	/**
-	 * @phpstan-param \Closure(EffectInstance, bool $replacesOldEffect) : void $closure
+	 * @return \Closure[]|Set
+	 * @phpstan-return Set<\Closure(EffectInstance, bool $replacesOldEffect) : void>
 	 */
-	public function onEffectAdd(\Closure $closure) : void{
-		Utils::validateCallableSignature(function(EffectInstance $effect, bool $replacesOldEffect) : void{}, $closure);
-		$this->effectAddHooks[spl_object_id($closure)] = $closure;
+	public function getEffectAddHooks() : Set{
+		return $this->effectAddHooks;
 	}
 
 	/**
-	 * @phpstan-param \Closure(EffectInstance) : void $closure
+	 * @return \Closure[]|Set
+	 * @phpstan-return Set<\Closure(EffectInstance) : void>
 	 */
-	public function onEffectRemove(\Closure $closure) : void{
-		Utils::validateCallableSignature(function(EffectInstance $effect) : void{}, $closure);
-		$this->effectRemoveHooks[spl_object_id($closure)] = $closure;
+	public function getEffectRemoveHooks() : Set{
+		return $this->effectRemoveHooks;
 	}
 }
diff --git a/src/network/mcpe/NetworkSession.php b/src/network/mcpe/NetworkSession.php
index 1826c8618..ce295a065 100644
--- a/src/network/mcpe/NetworkSession.php
+++ b/src/network/mcpe/NetworkSession.php
@@ -23,6 +23,7 @@ declare(strict_types=1);
 
 namespace pocketmine\network\mcpe;
 
+use Ds\Set;
 use Mdanter\Ecc\Crypto\Key\PublicKeyInterface;
 use pocketmine\entity\Attribute;
 use pocketmine\entity\effect\EffectInstance;
@@ -166,6 +167,12 @@ class NetworkSession{
 	/** @var PacketSender */
 	private $sender;
 
+	/**
+	 * @var \Closure[]|Set
+	 * @phpstan-var Set<\Closure() : void>
+	 */
+	private $disposeHooks;
+
 	public function __construct(Server $server, NetworkSessionManager $manager, PacketPool $packetPool, PacketSender $sender, Compressor $compressor, string $ip, int $port){
 		$this->server = $server;
 		$this->manager = $manager;
@@ -179,6 +186,8 @@ class NetworkSession{
 		$this->compressor = $compressor;
 		$this->packetPool = $packetPool;
 
+		$this->disposeHooks = new Set();
+
 		$this->connectTime = time();
 
 		$this->setHandler(new LoginPacketHandler(
@@ -218,12 +227,18 @@ class NetworkSession{
 		$this->player = new $class($this->server, $this, $this->info, $this->authenticated);
 
 		$this->invManager = new InventoryManager($this->player, $this);
-		$this->player->getEffects()->onEffectAdd(function(EffectInstance $effect, bool $replacesOldEffect) : void{
+
+		$effectManager = $this->player->getEffects();
+		$effectManager->getEffectAddHooks()->add($effectAddHook = function(EffectInstance $effect, bool $replacesOldEffect) : void{
 			$this->onEntityEffectAdded($this->player, $effect, $replacesOldEffect);
 		});
-		$this->player->getEffects()->onEffectRemove(function(EffectInstance $effect) : void{
+		$effectManager->getEffectRemoveHooks()->add($effectRemoveHook = function(EffectInstance $effect) : void{
 			$this->onEntityEffectRemoved($this->player, $effect);
 		});
+		$this->disposeHooks->add(static function() use ($effectManager, $effectAddHook, $effectRemoveHook) : void{
+			$effectManager->getEffectAddHooks()->remove($effectAddHook);
+			$effectManager->getEffectRemoveHooks()->remove($effectRemoveHook);
+		});
 	}
 
 	public function getPlayer() : ?Player{
@@ -467,6 +482,10 @@ class NetworkSession{
 			$this->disconnectGuard = true;
 			$func();
 			$this->disconnectGuard = false;
+			foreach($this->disposeHooks as $callback){
+				$callback();
+			}
+			$this->disposeHooks->clear();
 			$this->setHandler(null);
 			$this->connected = false;
 			$this->manager->remove($this);