diff --git a/dialer/mws/dialer.go b/dialer/mws/dialer.go index b0a75f1b..a127ef6c 100644 --- a/dialer/mws/dialer.go +++ b/dialer/mws/dialer.go @@ -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{}) diff --git a/internal/net/pipe.go b/internal/net/pipe.go index b2c5fd00..94270648 100644 --- a/internal/net/pipe.go +++ b/internal/net/pipe.go @@ -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() { + errCh <- pipeHalf(ctx, rw1, rw2) + }() go func() { - defer wg.Done() - if err := pipeBuffer(rw1, rw2, bufferSize/2); err != nil { - ch <- err + errCh <- pipeHalf(ctx, rw2, rw1) + }() + + // 等待第一个错误或完成 + var firstErr error + for i := 0; i < 2; i++ { + select { + case err := <-errCh: + if firstErr == nil && err != nil { + firstErr = err + cancel() // 一个方向出错，取消另一个 + } + case <-ctx.Done(): + // 超时或主动取消 + forceClose(rw1, rw2) + return ctx.Err() } - }() - go func() { - defer wg.Done() - if err := pipeBuffer(rw2, rw1, bufferSize/2); err != nil { - ch <- err - } - }() - - done := make(chan struct{}) - go func() { - wg.Wait() - close(done) - }() - - select { - case <-done: - case <-ctx.Done(): - return nil } - select { - case err := <-ch: - return err - default: - } - - return nil + return firstErr } -func pipeBuffer(dst io.ReadWriteCloser, src io.ReadWriteCloser, bufferSize int) error { - buf := bufpool.Get(bufferSize) +// pipeHalf 单向管道传输 +func pipeHalf(ctx context.Context, src, dst io.ReadWriteCloser) error { + defer func() { + // 传输完成后执行TCP半关闭 + halfClose(src, dst) + }() + + 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，完全关闭 + } +} + +// forceClose 强制关闭两个连接 +func forceClose(conns ...io.ReadWriteCloser) { + for _, conn := range conns { + if conn != nil { + conn.Close() + } } - - return err }