From 07989705d27191a8b30b8692171ff939293ab4e6 Mon Sep 17 00:00:00 2001
From: UUBulb <35923940+uubulb@users.noreply.github.com>
Date: Tue, 26 Nov 2024 09:18:13 +0800
Subject: [PATCH] fix: only close user/agentIo connect channel once (#21)

---
 cmd/dashboard/controller/fm.go       |  6 +++++-
 cmd/dashboard/controller/terminal.go |  6 +++++-
 service/rpc/io_stream.go             | 10 ++++++++--
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/cmd/dashboard/controller/fm.go b/cmd/dashboard/controller/fm.go
index bdcd7a2..9bd14b3 100644
--- a/cmd/dashboard/controller/fm.go
+++ b/cmd/dashboard/controller/fm.go
@@ -95,5 +95,9 @@ func fmStream(c *gin.Context) (any, error) {
 		return nil, newWsError("%v", err)
 	}
 
-	return nil, newWsError("%v", rpc.NezhaHandlerSingleton.StartStream(streamId, time.Second*10))
+	if err = rpc.NezhaHandlerSingleton.StartStream(streamId, time.Second*10); err != nil {
+		return nil, newWsError("%v", err)
+	}
+
+	return nil, newWsError("")
 }
diff --git a/cmd/dashboard/controller/terminal.go b/cmd/dashboard/controller/terminal.go
index 67466ba..2b54417 100644
--- a/cmd/dashboard/controller/terminal.go
+++ b/cmd/dashboard/controller/terminal.go
@@ -95,5 +95,9 @@ func terminalStream(c *gin.Context) (any, error) {
 		return nil, newWsError("%v", err)
 	}
 
-	return nil, newWsError("%v", rpc.NezhaHandlerSingleton.StartStream(streamId, time.Second*10))
+	if err = rpc.NezhaHandlerSingleton.StartStream(streamId, time.Second*10); err != nil {
+		return nil, newWsError("%v", err)
+	}
+
+	return nil, newWsError("")
 }
diff --git a/service/rpc/io_stream.go b/service/rpc/io_stream.go
index 2270672..313083b 100644
--- a/service/rpc/io_stream.go
+++ b/service/rpc/io_stream.go
@@ -15,6 +15,8 @@ type ioStreamContext struct {
 	agentIo          io.ReadWriteCloser
 	userIoConnectCh  chan struct{}
 	agentIoConnectCh chan struct{}
+	userIoChOnce     sync.Once
+	agentIoChOnce    sync.Once
 }
 
 type bp struct {
@@ -74,7 +76,9 @@ func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser)
 	}
 
 	stream.userIo = userIo
-	close(stream.userIoConnectCh)
+	stream.userIoChOnce.Do(func() {
+		close(stream.userIoConnectCh)
+	})
 
 	return nil
 }
@@ -86,7 +90,9 @@ func (s *NezhaHandler) AgentConnected(streamId string, agentIo io.ReadWriteClose
 	}
 
 	stream.agentIo = agentIo
-	close(stream.agentIoConnectCh)
+	stream.agentIoChOnce.Do(func() {
+		close(stream.agentIoConnectCh)
+	})
 
 	return nil
 }