优化代码

server/global/gateway/tunnel.go

@@ -10,23 +10,26 @@ import (
 )
 
 type Tunnel struct {
-	ID         string // 唯一标识
-	LocalHost  string // 本地监听地址
-	LocalPort  int    // 本地端口
-	RemoteHost string // 远程连接地址
-	RemotePort int    // 远程端口
-	Gateway    *Gateway
-	ctx        context.Context
-	cancel     context.CancelFunc
-	listener   net.Listener
-	err        error
+	ID                string // 唯一标识
+	LocalHost         string // 本地监听地址
+	LocalPort         int    // 本地端口
+	RemoteHost        string // 远程连接地址
+	RemotePort        int    // 远程端口
+	Gateway           *Gateway
+	ctx               context.Context
+	cancel            context.CancelFunc
+	listener          net.Listener
+	localConnections  []net.Conn
+	remoteConnections []net.Conn
 }
 
 func (r *Tunnel) Open() {
 	localAddr := fmt.Sprintf("%s:%d", r.LocalHost, r.LocalPort)
+
 	go func() {
 		<-r.ctx.Done()
 		_ = r.listener.Close()
+		log.Debugf("SSH 隧道 %v 关闭", localAddr)
 	}()
 	for {
 		log.Debugf("等待客户端访问 %v", localAddr)
@@ -35,6 +38,7 @@ func (r *Tunnel) Open() {
 			log.Debugf("接受连接失败 %v, 退出循环", err.Error())
 			return
 		}
+		r.localConnections = append(r.localConnections, localConn)
 
 		log.Debugf("客户端 %v 连接至 %v", localConn.RemoteAddr().String(), localAddr)
 		remoteAddr := fmt.Sprintf("%s:%d", r.RemoteHost, r.RemotePort)
@@ -44,27 +48,27 @@ func (r *Tunnel) Open() {
 			log.Debugf("连接远程主机 %v 失败", remoteAddr)
 			return
 		}
+		r.remoteConnections = append(r.remoteConnections, remoteConn)
 
 		log.Debugf("连接远程主机 %v 成功", remoteAddr)
-		go copyConn(r.ctx, localConn, remoteConn)
-		go copyConn(r.ctx, remoteConn, localConn)
+		go copyConn(localConn, remoteConn)
+		go copyConn(remoteConn, localConn)
 		log.Debugf("转发数据 [%v]->[%v]", localAddr, remoteAddr)
 	}
 }
 
-func (r Tunnel) Close() {
+func (r *Tunnel) Close() {
+	for i := range r.localConnections {
+		_ = r.localConnections[i].Close()
+	}
+	r.localConnections = nil
+	for i := range r.remoteConnections {
+		_ = r.remoteConnections[i].Close()
+	}
+	r.remoteConnections = nil
 	r.cancel()
 }
 
-func copyConn(ctx context.Context, writer, reader net.Conn) {
+func copyConn(writer, reader net.Conn) {
 	_, _ = io.Copy(writer, reader)
-	for {
-		select {
-		case <-ctx.Done():
-			_ = writer.Close()
-			_ = reader.Close()
-			return
-		}
-	}
-
 }