add SetReadDeadline in Pipe, to prevent memory not auto release in mws/mtcp

2025-12-29 18:16:46 +08:00
parent 7625973ca1
commit 180145189f
2 changed files with 115 additions and 47 deletions
@@ -130,7 +130,7 @@ func (d *mwsDialer) Handshake(ctx context.Context, conn net.Conn, options ...dia
 		if host == "" {
 			host = opts.Addr
 		}
-		s, err := d.initSession(ctx, host, conn, log)
+		s, err := d.initSession(ctx, host, conn, log, opts.Addr)
 		if err != nil {
 			log.Error(err)
 			conn.Close()
@@ -151,7 +151,7 @@ func (d *mwsDialer) Handshake(ctx context.Context, conn net.Conn, options ...dia
 	return cc, nil
 }

-func (d *mwsDialer) initSession(ctx context.Context, host string, conn net.Conn, log logger.Logger) (*muxSession, error) {
+func (d *mwsDialer) initSession(ctx context.Context, host string, conn net.Conn, log logger.Logger, optAddr string) (*muxSession, error) {
 	dialer := websocket.Dialer{
 		HandshakeTimeout:  d.md.handshakeTimeout,
 		ReadBufferSize:    d.md.readBufferSize,
@@ -191,7 +191,7 @@ func (d *mwsDialer) initSession(ctx context.Context, host string, conn net.Conn,
 			c.SetReadDeadline(time.Now().Add(d.md.keepaliveInterval * 2))
 			return nil
 		})
-		go d.keepAlive(cc)
+		go d.keepAlive(cc, optAddr)
 	}

 	// stream multiplex
@@ -203,7 +203,7 @@ func (d *mwsDialer) initSession(ctx context.Context, host string, conn net.Conn,
 	return &muxSession{conn: cc, session: session}, nil
 }

-func (d *mwsDialer) keepAlive(conn ws_util.WebsocketConn) {
+func (d *mwsDialer) keepAlive(conn ws_util.WebsocketConn, optAddr string) {
 	ticker := time.NewTicker(d.md.keepaliveInterval)
 	defer ticker.Stop()

@@ -211,6 +211,11 @@ func (d *mwsDialer) keepAlive(conn ws_util.WebsocketConn) {
 		d.options.Logger.Debug("send ping")
 		conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
 		if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+
+			d.sessionMutex.Lock()
+			delete(d.sessions, optAddr)
+			d.sessionMutex.Unlock()
+			conn.Close()
 			return
 		}
 		conn.SetWriteDeadline(time.Time{})
@@ -3,7 +3,6 @@ package net
 import (
 	"context"
 	"io"
-	"sync"
 	"time"

 	"github.com/go-gost/core/common/bufpool"
@@ -13,69 +12,133 @@ import (
 const (
 	// tcpWaitTimeout implements a TCP half-close timeout.
 	tcpWaitTimeout = 10 * time.Second
+	readTimeout    = 30 * time.Second
 )

+// Pipe 在两个连接之间建立双向数据通道
 func Pipe(ctx context.Context, rw1, rw2 io.ReadWriteCloser) error {
-	wg := sync.WaitGroup{}
-	wg.Add(2)
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()

-	ch := make(chan error, 2)
+	errCh := make(chan error, 2)

+	// 启动两个方向的传输
 	go func() {
-		defer wg.Done()
-		if err := pipeBuffer(rw1, rw2, bufferSize/2); err != nil {
-			ch <- err
-		}
-	}()
-	go func() {
-		defer wg.Done()
-		if err := pipeBuffer(rw2, rw1, bufferSize/2); err != nil {
-			ch <- err
-		}
+		errCh <- pipeHalf(ctx, rw1, rw2)
 	}()

-	done := make(chan struct{})
 	go func() {
-		wg.Wait()
-		close(done)
+		errCh <- pipeHalf(ctx, rw2, rw1)
 	}()

+	// 等待第一个错误或完成
+	var firstErr error
+	for i := 0; i < 2; i++ {
 		select {
-	case <-done:
+		case err := <-errCh:
+			if firstErr == nil && err != nil {
+				firstErr = err
+				cancel() // 一个方向出错，取消另一个
+			}
 		case <-ctx.Done():
-		return nil
+			// 超时或主动取消
+			forceClose(rw1, rw2)
+			return ctx.Err()
+		}
 	}

-	select {
-	case err := <-ch:
-		return err
-	default:
+	return firstErr
 }

-	return nil
-}
+// pipeHalf 单向管道传输
+func pipeHalf(ctx context.Context, src, dst io.ReadWriteCloser) error {
+	defer func() {
+		// 传输完成后执行TCP半关闭
+		halfClose(src, dst)
+	}()

-func pipeBuffer(dst io.ReadWriteCloser, src io.ReadWriteCloser, bufferSize int) error {
-	buf := bufpool.Get(bufferSize)
+	buf := bufpool.Get(bufferSize / 2)
 	defer bufpool.Put(buf)

-	_, err := io.CopyBuffer(dst, src, buf)

-	// Do the upload/download side TCP half-close.
+	// 创建带超时的读取器
+	reader := &readDeadliner{
+		Reader: src,
+		ctx:    ctx,
+	}
+
+	// 循环读取并写入，每次读取都有超时
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			// 设置读取超时
+			if rd, ok := src.(interface{ SetReadDeadline(time.Time) error }); ok {
+				rd.SetReadDeadline(time.Now().Add(readTimeout))
+			}
+
+			// 读取数据
+			nr, err := reader.Read(buf)
+			if err != nil {
+				if err != io.EOF {
+					return err
+				}
+				return nil // 正常结束
+			}
+
+			// 写入数据
+			_, err = dst.Write(buf[:nr])
+			if err != nil {
+				return err
+			}
+		}
+	}
+}
+
+// readDeadliner 包装读取器，支持上下文取消
+type readDeadliner struct {
+	io.Reader
+	ctx context.Context
+}
+
+func (r *readDeadliner) Read(p []byte) (int, error) {
+	// 检查上下文是否已取消
+	select {
+	case <-r.ctx.Done():
+		return 0, r.ctx.Err()
+	default:
+	}
+	return r.Reader.Read(p)
+}
+
+// halfClose 执行TCP半关闭
+func halfClose(src, dst io.ReadWriteCloser) {
+	// 关闭读取端
 	if cr, ok := src.(xio.CloseRead); ok {
 		cr.CloseRead()
 	}

+	// 关闭写入端，尝试半关闭
 	if cw, ok := dst.(xio.CloseWrite); ok {
-		if e := cw.CloseWrite(); e == xio.ErrUnsupported {
-			dst.Close()
+		if err := cw.CloseWrite(); err == xio.ErrUnsupported {
+			dst.Close() // 不支持半关闭，完全关闭
 		} else {
-			// Set TCP half-close timeout.
-			xio.SetReadDeadline(dst, time.Now().Add(tcpWaitTimeout))
+			// 设置半关闭超时
+			if rd, ok := dst.(interface{ SetReadDeadline(time.Time) error }); ok {
+				rd.SetReadDeadline(time.Now().Add(tcpWaitTimeout))
+			}
 		}
 	} else {
-		dst.Close()
+		dst.Close() // 不支持CloseWrite，完全关闭
+	}
 }

-	return err
+// forceClose 强制关闭两个连接
+func forceClose(conns ...io.ReadWriteCloser) {
+	for _, conn := range conns {
+		if conn != nil {
+			conn.Close()
+		}
+	}
 }