From 60e891b08c0e3f3a0a4750573bfe9af4e111af61 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=90=91=E6=9E=97?= <as2252258@163.com>
Date: Mon, 14 Feb 2022 14:28:47 +0800
Subject: [PATCH] modify plugin name

---
 kiri-websocket-server/FdCollector.php | 46 +++++++++++++++++++++++++--
 kiri-websocket-server/Sender.php      | 33 ++++++++++++++++---
 2 files changed, 72 insertions(+), 7 deletions(-)

diff --git a/kiri-websocket-server/FdCollector.php b/kiri-websocket-server/FdCollector.php
index dce3fb9e..3fdffc6b 100644
--- a/kiri-websocket-server/FdCollector.php
+++ b/kiri-websocket-server/FdCollector.php
@@ -2,13 +2,55 @@
 
 namespace Kiri\Websocket;
 
+use JetBrains\PhpStorm\Pure;
+use Kiri\Annotation\Inject;
+use Kiri\Core\HashMap;
+use Swoole\Http\Response;
+
 class FdCollector
 {
 
 
-	public function set($fd)
-	{
+	#[Inject(HashMap::class)]
+	public HashMap $fds;
 
+
+	/**
+	 * @param int $fd
+	 * @param Response $response
+	 * @return void
+	 */
+	public function set(int $fd, Response $response)
+	{
+		$this->fds->put('fd_' . $fd, $response);
+	}
+
+
+	/**
+	 * @param int $fd
+	 * @return bool
+	 */
+	#[Pure] public function has(int $fd): bool
+	{
+		return $this->fds->has('fd_' . $fd);
+	}
+
+	/**
+	 * @param int $fd
+	 * @return ?Response
+	 */
+	#[Pure] public function get(int $fd): ?Response
+	{
+		return $this->fds->get('fd_' . $fd);
+	}
+
+	/**
+	 * @param int $fd
+	 * @return void
+	 */
+	public function remove(int $fd)
+	{
+		$this->fds->del('fd_' . $fd);
 	}
 
 }
diff --git a/kiri-websocket-server/Sender.php b/kiri-websocket-server/Sender.php
index f949ce5f..dc805e3f 100644
--- a/kiri-websocket-server/Sender.php
+++ b/kiri-websocket-server/Sender.php
@@ -37,9 +37,18 @@ class Sender implements WebSocketInterface
 	 */
 	public function push(int $fd, string $data, int $opcode = WEBSOCKET_OPCODE_TEXT, int $flags = SWOOLE_WEBSOCKET_FLAG_FIN): bool
 	{
-		if ($this->isEstablished($fd)) {
+		if (!$this->isEstablished($fd)) {
+			return false;
+		}
+		if ($this->server instanceof Server) {
 			return $this->server->push($fd, $data, $opcode, $flags);
 		}
+		$collector = Kiri::getContainer()->get(FdCollector::class);
+
+		$response = $collector->get($fd);
+		if (!empty($response)) {
+			return $response->push($data, $opcode, $flags);
+		}
 		return false;
 	}
 
@@ -69,10 +78,13 @@ class Sender implements WebSocketInterface
 	 */
 	public function disconnect(int $fd, int $code = SWOOLE_WEBSOCKET_CLOSE_NORMAL, string $reason = ''): bool
 	{
-		if ($this->isEstablished($fd)) {
+		if (!$this->isEstablished($fd)) {
+			return false;
+		}
+		if ($this->server instanceof Server) {
 			return $this->server->disconnect($fd, $code, $reason);
 		}
-		return false;
+		return $this->server->close($fd, $reason);
 	}
 
 
@@ -82,7 +94,14 @@ class Sender implements WebSocketInterface
 	 */
 	public function isEstablished(int $fd): bool
 	{
-		return $this->exist($fd) && $this->server->isEstablished($fd);
+		if (!$this->exist($fd)) {
+			return false;
+		}
+
+		if ($this->server instanceof Server) {
+			return $this->server->isEstablished($fd);
+		}
+		return true;
 	}
 
 
@@ -92,7 +111,11 @@ class Sender implements WebSocketInterface
 	 */
 	public function exist(int $fd): bool
 	{
-		return $this->server->exist($fd);
+		if ($this->server instanceof Server) {
+			return $this->server->exist($fd);
+		}
+		$collector = Kiri::getContainer()->get(FdCollector::class);
+		return $collector->has($fd);
 	}
 
 }