fix websocket for tunnel

dialer/mtls/conn.go

@@ -3,20 +3,20 @@ package mtls
 import (
 	"net"
 
-	"github.com/xtaci/smux"
+	"github.com/go-gost/x/internal/util/mux"
 )
 
 type muxSession struct {
 	conn    net.Conn
-	session *smux.Session
+	session *mux.Session
 }
 
 func (session *muxSession) GetConn() (net.Conn, error) {
-	return session.session.OpenStream()
+	return session.session.GetConn()
 }
 
 func (session *muxSession) Accept() (net.Conn, error) {
-	return session.session.AcceptStream()
+	return session.session.Accept()
 }
 
 func (session *muxSession) Close() error {

dialer/mtls/dialer.go

@@ -11,8 +11,8 @@ import (
 	"github.com/go-gost/core/dialer"
 	"github.com/go-gost/core/logger"
 	md "github.com/go-gost/core/metadata"
+	"github.com/go-gost/x/internal/util/mux"
 	"github.com/go-gost/x/registry"
-	"github.com/xtaci/smux"
 )
 
 func init() {
@@ -130,25 +130,7 @@ func (d *mtlsDialer) initSession(ctx context.Context, conn net.Conn) (*muxSessio
 	conn = tlsConn
 
 	// stream multiplex
-	smuxConfig := smux.DefaultConfig()
+	session, err := mux.ClientSession(conn, d.md.muxCfg)
-	smuxConfig.KeepAliveDisabled = d.md.muxKeepAliveDisabled
-	if d.md.muxKeepAliveInterval > 0 {
-		smuxConfig.KeepAliveInterval = d.md.muxKeepAliveInterval
-	}
-	if d.md.muxKeepAliveTimeout > 0 {
-		smuxConfig.KeepAliveTimeout = d.md.muxKeepAliveTimeout
-	}
-	if d.md.muxMaxFrameSize > 0 {
-		smuxConfig.MaxFrameSize = d.md.muxMaxFrameSize
-	}
-	if d.md.muxMaxReceiveBuffer > 0 {
-		smuxConfig.MaxReceiveBuffer = d.md.muxMaxReceiveBuffer
-	}
-	if d.md.muxMaxStreamBuffer > 0 {
-		smuxConfig.MaxStreamBuffer = d.md.muxMaxStreamBuffer
-	}
-
-	session, err := smux.Client(conn, smuxConfig)
 	if err != nil {
 		return nil, err
 	}

dialer/mtls/metadata.go

@@ -5,39 +5,25 @@ import (
 
 	mdata "github.com/go-gost/core/metadata"
 	mdutil "github.com/go-gost/core/metadata/util"
+	"github.com/go-gost/x/internal/util/mux"
 )
 
 type metadata struct {
 	handshakeTimeout time.Duration
-
+	muxCfg           *mux.Config
-	muxKeepAliveDisabled bool
-	muxKeepAliveInterval time.Duration
-	muxKeepAliveTimeout  time.Duration
-	muxMaxFrameSize      int
-	muxMaxReceiveBuffer  int
-	muxMaxStreamBuffer   int
 }
 
 func (d *mtlsDialer) parseMetadata(md mdata.Metadata) (err error) {
-	const (
+	d.md.handshakeTimeout = mdutil.GetDuration(md, "handshakeTimeout")
-		handshakeTimeout = "handshakeTimeout"
-
-		muxKeepAliveDisabled = "muxKeepAliveDisabled"
-		muxKeepAliveInterval = "muxKeepAliveInterval"
-		muxKeepAliveTimeout  = "muxKeepAliveTimeout"
-		muxMaxFrameSize      = "muxMaxFrameSize"
-		muxMaxReceiveBuffer  = "muxMaxReceiveBuffer"
-		muxMaxStreamBuffer   = "muxMaxStreamBuffer"
-	)
-
-	d.md.handshakeTimeout = mdutil.GetDuration(md, handshakeTimeout)
-
-	d.md.muxKeepAliveDisabled = mdutil.GetBool(md, muxKeepAliveDisabled)
-	d.md.muxKeepAliveInterval = mdutil.GetDuration(md, muxKeepAliveInterval)
-	d.md.muxKeepAliveTimeout = mdutil.GetDuration(md, muxKeepAliveTimeout)
-	d.md.muxMaxFrameSize = mdutil.GetInt(md, muxMaxFrameSize)
-	d.md.muxMaxReceiveBuffer = mdutil.GetInt(md, muxMaxReceiveBuffer)
-	d.md.muxMaxStreamBuffer = mdutil.GetInt(md, muxMaxStreamBuffer)
 
+	d.md.muxCfg = &mux.Config{
+		Version:           mdutil.GetInt(md, "mux.version"),
+		KeepAliveInterval: mdutil.GetDuration(md, "mux.keepaliveInterval"),
+		KeepAliveDisabled: mdutil.GetBool(md, "mux.keepaliveDisabled"),
+		KeepAliveTimeout:  mdutil.GetDuration(md, "mux.keepaliveTimeout"),
+		MaxFrameSize:      mdutil.GetInt(md, "mux.maxFrameSize"),
+		MaxReceiveBuffer:  mdutil.GetInt(md, "mux.maxReceiveBuffer"),
+		MaxStreamBuffer:   mdutil.GetInt(md, "mux.maxStreamBuffer"),
+	}
 	return
 }

dialer/mws/conn.go

@@ -3,20 +3,20 @@ package mws
 import (
 	"net"
 
-	"github.com/xtaci/smux"
+	"github.com/go-gost/x/internal/util/mux"
 )
 
 type muxSession struct {
 	conn    net.Conn
-	session *smux.Session
+	session *mux.Session
 }
 
 func (session *muxSession) GetConn() (net.Conn, error) {
-	return session.session.OpenStream()
+	return session.session.GetConn()
 }
 
 func (session *muxSession) Accept() (net.Conn, error) {
-	return session.session.AcceptStream()
+	return session.session.Accept()
 }
 
 func (session *muxSession) Close() error {

dialer/mws/dialer.go

@@ -10,10 +10,10 @@ import (
 
 	"github.com/go-gost/core/dialer"
 	md "github.com/go-gost/core/metadata"
+	"github.com/go-gost/x/internal/util/mux"
 	ws_util "github.com/go-gost/x/internal/util/ws"
 	"github.com/go-gost/x/registry"
 	"github.com/gorilla/websocket"
-	"github.com/xtaci/smux"
 )
 
 func init() {
@@ -178,25 +178,7 @@ func (d *mwsDialer) initSession(ctx context.Context, host string, conn net.Conn)
 	}
 
 	// stream multiplex
-	smuxConfig := smux.DefaultConfig()
+	session, err := mux.ClientSession(conn, d.md.muxCfg)
-	smuxConfig.KeepAliveDisabled = d.md.muxKeepAliveDisabled
-	if d.md.muxKeepAliveInterval > 0 {
-		smuxConfig.KeepAliveInterval = d.md.muxKeepAliveInterval
-	}
-	if d.md.muxKeepAliveTimeout > 0 {
-		smuxConfig.KeepAliveTimeout = d.md.muxKeepAliveTimeout
-	}
-	if d.md.muxMaxFrameSize > 0 {
-		smuxConfig.MaxFrameSize = d.md.muxMaxFrameSize
-	}
-	if d.md.muxMaxReceiveBuffer > 0 {
-		smuxConfig.MaxReceiveBuffer = d.md.muxMaxReceiveBuffer
-	}
-	if d.md.muxMaxStreamBuffer > 0 {
-		smuxConfig.MaxStreamBuffer = d.md.muxMaxStreamBuffer
-	}
-
-	session, err := smux.Client(cc, smuxConfig)
 	if err != nil {
 		return nil, err
 	}

dialer/mws/metadata.go

@@ -6,6 +6,7 @@ import (
 
 	mdata "github.com/go-gost/core/metadata"
 	mdutil "github.com/go-gost/core/metadata/util"
+	"github.com/go-gost/x/internal/util/mux"
 )
 
 const (
@@ -23,15 +24,9 @@ type metadata struct {
 	writeBufferSize   int
 	enableCompression bool
 
-	muxKeepAliveDisabled bool
-	muxKeepAliveInterval time.Duration
-	muxKeepAliveTimeout  time.Duration
-	muxMaxFrameSize      int
-	muxMaxReceiveBuffer  int
-	muxMaxStreamBuffer   int
-
 	header            http.Header
 	keepaliveInterval time.Duration
+	muxCfg            *mux.Config
 }
 
 func (d *mwsDialer) parseMetadata(md mdata.Metadata) (err error) {
@@ -46,13 +41,6 @@ func (d *mwsDialer) parseMetadata(md mdata.Metadata) (err error) {
 		enableCompression = "enableCompression"
 
 		header = "header"
-
-		muxKeepAliveDisabled = "muxKeepAliveDisabled"
-		muxKeepAliveInterval = "muxKeepAliveInterval"
-		muxKeepAliveTimeout  = "muxKeepAliveTimeout"
-		muxMaxFrameSize      = "muxMaxFrameSize"
-		muxMaxReceiveBuffer  = "muxMaxReceiveBuffer"
-		muxMaxStreamBuffer   = "muxMaxStreamBuffer"
 	)
 
 	d.md.host = mdutil.GetString(md, host)
@@ -62,12 +50,15 @@ func (d *mwsDialer) parseMetadata(md mdata.Metadata) (err error) {
 		d.md.path = defaultPath
 	}
 
-	d.md.muxKeepAliveDisabled = mdutil.GetBool(md, muxKeepAliveDisabled)
+	d.md.muxCfg = &mux.Config{
-	d.md.muxKeepAliveInterval = mdutil.GetDuration(md, muxKeepAliveInterval)
+		Version:           mdutil.GetInt(md, "mux.version"),
-	d.md.muxKeepAliveTimeout = mdutil.GetDuration(md, muxKeepAliveTimeout)
+		KeepAliveInterval: mdutil.GetDuration(md, "mux.keepaliveInterval"),
-	d.md.muxMaxFrameSize = mdutil.GetInt(md, muxMaxFrameSize)
+		KeepAliveDisabled: mdutil.GetBool(md, "mux.keepaliveDisabled"),
-	d.md.muxMaxReceiveBuffer = mdutil.GetInt(md, muxMaxReceiveBuffer)
+		KeepAliveTimeout:  mdutil.GetDuration(md, "mux.keepaliveTimeout"),
-	d.md.muxMaxStreamBuffer = mdutil.GetInt(md, muxMaxStreamBuffer)
+		MaxFrameSize:      mdutil.GetInt(md, "mux.maxFrameSize"),
+		MaxReceiveBuffer:  mdutil.GetInt(md, "mux.maxReceiveBuffer"),
+		MaxStreamBuffer:   mdutil.GetInt(md, "mux.maxStreamBuffer"),
+	}
 
 	d.md.handshakeTimeout = mdutil.GetDuration(md, handshakeTimeout)
 	d.md.readHeaderTimeout = mdutil.GetDuration(md, readHeaderTimeout)

handler/forward/local/handler.go

@@ -271,6 +271,14 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriter, log l
 				return resp.Write(rw)
 			}
 
+			if req.Header.Get("Upgrade") == "websocket" {
+				err := xnet.CopyBuffer(cc, br, 8192)
+				if err == nil {
+					err = io.EOF
+				}
+				return err
+			}
+
 			res, err := http.ReadResponse(bufio.NewReader(cc), req)
 			if err != nil {
 				log.Warnf("read response from node %s(%s): %v", target.Name, target.Addr, err)

handler/forward/remote/handler.go

@@ -276,6 +276,14 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriter, remot
 				return resp.Write(rw)
 			}
 
+			if req.Header.Get("Upgrade") == "websocket" {
+				err := xnet.CopyBuffer(cc, br, 8192)
+				if err == nil {
+					err = io.EOF
+				}
+				return err
+			}
+
 			res, err := http.ReadResponse(bufio.NewReader(cc), req)
 			if err != nil {
 				log.Warnf("read response from node %s(%s): %v", target.Name, target.Addr, err)

internal/util/forward/forward.go

@@ -38,10 +38,6 @@ func Sniffing(ctx context.Context, rdw io.ReadWriter) (rw io.ReadWriter, host st
 		if err == nil {
 			host = r.Host
 			protocol = ProtoHTTP
-
-			if r.Header.Get("Upgrade") == "websocket" {
-				protocol = ProtoWebsocket
-			}
 			return
 		}
 	}
@@ -94,7 +90,6 @@ func isHTTP(s string) bool {
 
 const (
 	ProtoHTTP  = "http"
-	ProtoWebsocket = "ws"
 	ProtoTLS   = "tls"
 	ProtoSSHv2 = "SSH-2"
 )

listener/mtls/listener.go

@@ -12,11 +12,11 @@ import (
 	admission "github.com/go-gost/x/admission/wrapper"
 	xnet "github.com/go-gost/x/internal/net"
 	"github.com/go-gost/x/internal/net/proxyproto"
+	"github.com/go-gost/x/internal/util/mux"
 	climiter "github.com/go-gost/x/limiter/conn/wrapper"
 	limiter "github.com/go-gost/x/limiter/traffic/wrapper"
 	metrics "github.com/go-gost/x/metrics/wrapper"
 	"github.com/go-gost/x/registry"
-	"github.com/xtaci/smux"
 )
 
 func init() {
@@ -105,24 +105,7 @@ func (l *mtlsListener) listenLoop() {
 func (l *mtlsListener) mux(conn net.Conn) {
 	defer conn.Close()
 
-	smuxConfig := smux.DefaultConfig()
+	session, err := mux.ServerSession(conn, l.md.muxCfg)
-	smuxConfig.KeepAliveDisabled = l.md.muxKeepAliveDisabled
-	if l.md.muxKeepAliveInterval > 0 {
-		smuxConfig.KeepAliveInterval = l.md.muxKeepAliveInterval
-	}
-	if l.md.muxKeepAliveTimeout > 0 {
-		smuxConfig.KeepAliveTimeout = l.md.muxKeepAliveTimeout
-	}
-	if l.md.muxMaxFrameSize > 0 {
-		smuxConfig.MaxFrameSize = l.md.muxMaxFrameSize
-	}
-	if l.md.muxMaxReceiveBuffer > 0 {
-		smuxConfig.MaxReceiveBuffer = l.md.muxMaxReceiveBuffer
-	}
-	if l.md.muxMaxStreamBuffer > 0 {
-		smuxConfig.MaxStreamBuffer = l.md.muxMaxStreamBuffer
-	}
-	session, err := smux.Server(conn, smuxConfig)
 	if err != nil {
 		l.logger.Error(err)
 		return
@@ -130,7 +113,7 @@ func (l *mtlsListener) mux(conn net.Conn) {
 	defer session.Close()
 
 	for {
-		stream, err := session.AcceptStream()
+		stream, err := session.Accept()
 		if err != nil {
 			l.logger.Error("accept stream: ", err)
 			return
@@ -138,8 +121,6 @@ func (l *mtlsListener) mux(conn net.Conn) {
 
 		select {
 		case l.cqueue <- stream:
-		case <-stream.GetDieCh():
-			stream.Close()
 		default:
 			stream.Close()
 			l.logger.Warnf("connection queue is full, client %s discarded", stream.RemoteAddr())

listener/mtls/metadata.go

@@ -1,10 +1,9 @@
 package mtls
 
 import (
-	"time"
-
 	mdata "github.com/go-gost/core/metadata"
 	mdutil "github.com/go-gost/core/metadata/util"
+	"github.com/go-gost/x/internal/util/mux"
 )
 
 const (
@@ -12,41 +11,26 @@ const (
 )
 
 type metadata struct {
-	muxKeepAliveDisabled bool
+	muxCfg  *mux.Config
-	muxKeepAliveInterval time.Duration
-	muxKeepAliveTimeout  time.Duration
-	muxMaxFrameSize      int
-	muxMaxReceiveBuffer  int
-	muxMaxStreamBuffer   int
-
 	backlog int
 	mptcp   bool
 }
 
 func (l *mtlsListener) parseMetadata(md mdata.Metadata) (err error) {
-	const (
+	l.md.backlog = mdutil.GetInt(md, "backlog")
-		backlog = "backlog"
-
-		muxKeepAliveDisabled = "muxKeepAliveDisabled"
-		muxKeepAliveInterval = "muxKeepAliveInterval"
-		muxKeepAliveTimeout  = "muxKeepAliveTimeout"
-		muxMaxFrameSize      = "muxMaxFrameSize"
-		muxMaxReceiveBuffer  = "muxMaxReceiveBuffer"
-		muxMaxStreamBuffer   = "muxMaxStreamBuffer"
-	)
-
-	l.md.backlog = mdutil.GetInt(md, backlog)
 	if l.md.backlog <= 0 {
 		l.md.backlog = defaultBacklog
 	}
 
-	l.md.muxKeepAliveDisabled = mdutil.GetBool(md, muxKeepAliveDisabled)
+	l.md.muxCfg = &mux.Config{
-	l.md.muxKeepAliveInterval = mdutil.GetDuration(md, muxKeepAliveInterval)
+		Version:           mdutil.GetInt(md, "mux.version"),
-	l.md.muxKeepAliveTimeout = mdutil.GetDuration(md, muxKeepAliveTimeout)
+		KeepAliveInterval: mdutil.GetDuration(md, "mux.keepaliveInterval"),
-	l.md.muxMaxFrameSize = mdutil.GetInt(md, muxMaxFrameSize)
+		KeepAliveDisabled: mdutil.GetBool(md, "mux.keepaliveDisabled"),
-	l.md.muxMaxReceiveBuffer = mdutil.GetInt(md, muxMaxReceiveBuffer)
+		KeepAliveTimeout:  mdutil.GetDuration(md, "mux.keepaliveTimeout"),
-	l.md.muxMaxStreamBuffer = mdutil.GetInt(md, muxMaxStreamBuffer)
+		MaxFrameSize:      mdutil.GetInt(md, "mux.maxFrameSize"),
-
+		MaxReceiveBuffer:  mdutil.GetInt(md, "mux.maxReceiveBuffer"),
+		MaxStreamBuffer:   mdutil.GetInt(md, "mux.maxStreamBuffer"),
+	}
 	l.md.mptcp = mdutil.GetBool(md, "mptcp")
 
 	return

listener/mws/listener.go

@@ -14,13 +14,13 @@ import (
 	admission "github.com/go-gost/x/admission/wrapper"
 	xnet "github.com/go-gost/x/internal/net"
 	"github.com/go-gost/x/internal/net/proxyproto"
+	"github.com/go-gost/x/internal/util/mux"
 	ws_util "github.com/go-gost/x/internal/util/ws"
 	climiter "github.com/go-gost/x/limiter/conn/wrapper"
 	limiter "github.com/go-gost/x/limiter/traffic/wrapper"
 	metrics "github.com/go-gost/x/metrics/wrapper"
 	"github.com/go-gost/x/registry"
 	"github.com/gorilla/websocket"
-	"github.com/xtaci/smux"
 )
 
 func init() {
@@ -170,24 +170,7 @@ func (l *mwsListener) upgrade(w http.ResponseWriter, r *http.Request) {
 func (l *mwsListener) mux(conn net.Conn) {
 	defer conn.Close()
 
-	smuxConfig := smux.DefaultConfig()
+	session, err := mux.ServerSession(conn, l.md.muxCfg)
-	smuxConfig.KeepAliveDisabled = l.md.muxKeepAliveDisabled
-	if l.md.muxKeepAliveInterval > 0 {
-		smuxConfig.KeepAliveInterval = l.md.muxKeepAliveInterval
-	}
-	if l.md.muxKeepAliveTimeout > 0 {
-		smuxConfig.KeepAliveTimeout = l.md.muxKeepAliveTimeout
-	}
-	if l.md.muxMaxFrameSize > 0 {
-		smuxConfig.MaxFrameSize = l.md.muxMaxFrameSize
-	}
-	if l.md.muxMaxReceiveBuffer > 0 {
-		smuxConfig.MaxReceiveBuffer = l.md.muxMaxReceiveBuffer
-	}
-	if l.md.muxMaxStreamBuffer > 0 {
-		smuxConfig.MaxStreamBuffer = l.md.muxMaxStreamBuffer
-	}
-	session, err := smux.Server(conn, smuxConfig)
 	if err != nil {
 		l.logger.Error(err)
 		return
@@ -195,7 +178,7 @@ func (l *mwsListener) mux(conn net.Conn) {
 	defer session.Close()
 
 	for {
-		stream, err := session.AcceptStream()
+		stream, err := session.Accept()
 		if err != nil {
 			l.logger.Error("accept stream: ", err)
 			return
@@ -203,8 +186,6 @@ func (l *mwsListener) mux(conn net.Conn) {
 
 		select {
 		case l.cqueue <- stream:
-		case <-stream.GetDieCh():
-			stream.Close()
 		default:
 			stream.Close()
 			l.logger.Warnf("connection queue is full, client %s discarded", stream.RemoteAddr())

listener/mws/metadata.go

@@ -6,6 +6,7 @@ import (
 
 	mdata "github.com/go-gost/core/metadata"
 	mdutil "github.com/go-gost/core/metadata/util"
+	"github.com/go-gost/x/internal/util/mux"
 )
 
 const (
@@ -24,12 +25,7 @@ type metadata struct {
 	writeBufferSize   int
 	enableCompression bool
 
-	muxKeepAliveDisabled bool
+	muxCfg *mux.Config
-	muxKeepAliveInterval time.Duration
-	muxKeepAliveTimeout  time.Duration
-	muxMaxFrameSize      int
-	muxMaxReceiveBuffer  int
-	muxMaxStreamBuffer   int
 
 	mptcp bool
 }
@@ -70,12 +66,15 @@ func (l *mwsListener) parseMetadata(md mdata.Metadata) (err error) {
 	l.md.writeBufferSize = mdutil.GetInt(md, writeBufferSize)
 	l.md.enableCompression = mdutil.GetBool(md, enableCompression)
 
-	l.md.muxKeepAliveDisabled = mdutil.GetBool(md, muxKeepAliveDisabled)
+	l.md.muxCfg = &mux.Config{
-	l.md.muxKeepAliveInterval = mdutil.GetDuration(md, muxKeepAliveInterval)
+		Version:           mdutil.GetInt(md, "mux.version"),
-	l.md.muxKeepAliveTimeout = mdutil.GetDuration(md, muxKeepAliveTimeout)
+		KeepAliveInterval: mdutil.GetDuration(md, "mux.keepaliveInterval"),
-	l.md.muxMaxFrameSize = mdutil.GetInt(md, muxMaxFrameSize)
+		KeepAliveDisabled: mdutil.GetBool(md, "mux.keepaliveDisabled"),
-	l.md.muxMaxReceiveBuffer = mdutil.GetInt(md, muxMaxReceiveBuffer)
+		KeepAliveTimeout:  mdutil.GetDuration(md, "mux.keepaliveTimeout"),
-	l.md.muxMaxStreamBuffer = mdutil.GetInt(md, muxMaxStreamBuffer)
+		MaxFrameSize:      mdutil.GetInt(md, "mux.maxFrameSize"),
+		MaxReceiveBuffer:  mdutil.GetInt(md, "mux.maxReceiveBuffer"),
+		MaxStreamBuffer:   mdutil.GetInt(md, "mux.maxStreamBuffer"),
+	}
 
 	if mm := mdutil.GetStringMapString(md, header); len(mm) > 0 {
 		hd := http.Header{}