add ClientAddr for websocket conn

chain/router.go

@@ -96,6 +96,11 @@ func (r *Router) dial(ctx context.Context, network, address string, log logger.L
 			defer cancel()
 		}
 
+		buf := ctxvalue.BufferFromContext(ctx)
+		if buf != nil {
+			buf.Reset()
+		}
+
 		var ipAddr string
 		ipAddr, err = xnet.Resolve(ctx, "ip", address, r.options.Resolver, r.options.HostMapper, log)
 		if err != nil {
@@ -103,12 +108,15 @@ func (r *Router) dial(ctx context.Context, network, address string, log logger.L
 			break
 		}
 
+		if buf != nil {
+			buf.Reset()
+		}
+
 		var route chain.Route
 		if r.options.Chain != nil {
 			route = r.options.Chain.Route(ctx, network, ipAddr, chain.WithHostRouteOption(address))
 		}
 
-		buf := ctxvalue.BufferFromContext(ctx)
 		if buf == nil {
 			buf = &bytes.Buffer{}
 		}

connector/tunnel/metadata.go

@@ -5,9 +5,9 @@ import (
 	"time"
 
 	mdata "github.com/go-gost/core/metadata"
-	mdutil "github.com/go-gost/x/metadata/util"
 	"github.com/go-gost/relay"
 	"github.com/go-gost/x/internal/util/mux"
+	mdutil "github.com/go-gost/x/metadata/util"
 	"github.com/google/uuid"
 )
 
@@ -56,6 +56,9 @@ func (c *tunnelConnector) parseMetadata(md mdata.Metadata) (err error) {
 	if c.md.muxCfg.Version == 0 {
 		c.md.muxCfg.Version = 2
 	}
+	if c.md.muxCfg.MaxStreamBuffer == 0 {
+		c.md.muxCfg.MaxStreamBuffer = 1048576
+	}
 
 	return
 }

handler/forward/local/handler.go

@@ -7,7 +7,6 @@ import (
 	"errors"
 	"fmt"
 	"net"
-	"strconv"
 	"time"
 
 	"github.com/go-gost/core/chain"
@@ -180,7 +179,7 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 		}
 	}
 
-	var target *chain.Node
+	target := &chain.Node{}
 	if h.hop != nil {
 		target = h.hop.Select(ctx,
 			hop.ProtocolSelectOption(proto),
@@ -252,27 +251,3 @@ func (h *forwardHandler) checkRateLimit(addr net.Addr) bool {
 
 	return true
 }
-
-func convertAddr(addr net.Addr) net.Addr {
-	host, sp, _ := net.SplitHostPort(addr.String())
-	ip := net.ParseIP(host)
-	port, _ := strconv.Atoi(sp)
-
-	if ip == nil || ip.Equal(net.IPv6zero) {
-		ip = net.IPv4zero
-	}
-
-	switch addr.Network() {
-	case "tcp", "tcp4", "tcp6":
-		return &net.TCPAddr{
-			IP:   ip,
-			Port: port,
-		}
-
-	default:
-		return &net.UDPAddr{
-			IP:   ip,
-			Port: port,
-		}
-	}
-}

handler/http/handler.go

@@ -111,7 +111,14 @@ func (h *httpHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler
 		Time:       start,
 		SID:        string(ctxvalue.SidFromContext(ctx)),
 	}
-	ro.ClientIP, _, _ = net.SplitHostPort(conn.RemoteAddr().String())
+
+	ro.ClientIP = conn.RemoteAddr().String()
+	if clientAddr := ctxvalue.ClientAddrFromContext(ctx); clientAddr != "" {
+		ro.ClientIP = string(clientAddr)
+	}
+	if h, _, _ := net.SplitHostPort(ro.ClientIP); h != "" {
+		ro.ClientIP = h
+	}
 
 	log := h.options.Logger.WithFields(map[string]any{
 		"remote": conn.RemoteAddr().String(),

handler/relay/handler.go

@@ -104,7 +104,14 @@ func (h *relayHandler) Handle(ctx context.Context, conn net.Conn, opts ...handle
 		Time:       start,
 		SID:        string(ctxvalue.SidFromContext(ctx)),
 	}
-	ro.ClientIP, _, _ = net.SplitHostPort(conn.RemoteAddr().String())
+
+	ro.ClientIP = conn.RemoteAddr().String()
+	if clientAddr := ctxvalue.ClientAddrFromContext(ctx); clientAddr != "" {
+		ro.ClientIP = string(clientAddr)
+	}
+	if h, _, _ := net.SplitHostPort(ro.ClientIP); h != "" {
+		ro.ClientIP = h
+	}
 
 	log := h.options.Logger.WithFields(map[string]any{
 		"remote": conn.RemoteAddr().String(),

handler/sni/handler.go

@@ -79,7 +79,14 @@ func (h *sniHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.
 		Time:       start,
 		SID:        string(ctxvalue.SidFromContext(ctx)),
 	}
-	ro.ClientIP, _, _ = net.SplitHostPort(conn.RemoteAddr().String())
+
+	ro.ClientIP = conn.RemoteAddr().String()
+	if clientAddr := ctxvalue.ClientAddrFromContext(ctx); clientAddr != "" {
+		ro.ClientIP = string(clientAddr)
+	}
+	if h, _, _ := net.SplitHostPort(ro.ClientIP); h != "" {
+		ro.ClientIP = h
+	}
 
 	log := h.options.Logger.WithFields(map[string]any{
 		"remote": conn.RemoteAddr().String(),

handler/socks/v4/handler.go

@@ -105,7 +105,14 @@ func (h *socks4Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl
 		Time:       start,
 		SID:        string(ctxvalue.SidFromContext(ctx)),
 	}
-	ro.ClientIP, _, _ = net.SplitHostPort(conn.RemoteAddr().String())
+
+	ro.ClientIP = conn.RemoteAddr().String()
+	if clientAddr := ctxvalue.ClientAddrFromContext(ctx); clientAddr != "" {
+		ro.ClientIP = string(clientAddr)
+	}
+	if h, _, _ := net.SplitHostPort(ro.ClientIP); h != "" {
+		ro.ClientIP = h
+	}
 
 	log := h.options.Logger.WithFields(map[string]any{
 		"remote": conn.RemoteAddr().String(),

handler/socks/v5/handler.go

@@ -105,7 +105,14 @@ func (h *socks5Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl
 		Time:       start,
 		SID:        string(ctxvalue.SidFromContext(ctx)),
 	}
-	ro.ClientIP, _, _ = net.SplitHostPort(conn.RemoteAddr().String())
+
+	ro.ClientIP = conn.RemoteAddr().String()
+	if clientAddr := ctxvalue.ClientAddrFromContext(ctx); clientAddr != "" {
+		ro.ClientIP = string(clientAddr)
+	}
+	if h, _, _ := net.SplitHostPort(ro.ClientIP); h != "" {
+		ro.ClientIP = h
+	}
 
 	log := h.options.Logger.WithFields(map[string]any{
 		"remote": conn.RemoteAddr().String(),

handler/ss/handler.go

@@ -89,7 +89,14 @@ func (h *ssHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.H
 		Time:       start,
 		SID:        string(ctxvalue.SidFromContext(ctx)),
 	}
-	ro.ClientIP, _, _ = net.SplitHostPort(conn.RemoteAddr().String())
+
+	ro.ClientIP = conn.RemoteAddr().String()
+	if clientAddr := ctxvalue.ClientAddrFromContext(ctx); clientAddr != "" {
+		ro.ClientIP = string(clientAddr)
+	}
+	if h, _, _ := net.SplitHostPort(ro.ClientIP); h != "" {
+		ro.ClientIP = h
+	}
 
 	log := h.options.Logger.WithFields(map[string]any{
 		"remote": conn.RemoteAddr().String(),

handler/tunnel/metadata.go

@@ -96,6 +96,9 @@ func (h *tunnelHandler) parseMetadata(md mdata.Metadata) (err error) {
 	if h.md.muxCfg.Version == 0 {
 		h.md.muxCfg.Version = 2
 	}
+	if h.md.muxCfg.MaxStreamBuffer == 0 {
+		h.md.muxCfg.MaxStreamBuffer = 1048576
+	}
 
 	h.md.observePeriod = mdutil.GetDuration(md, "observePeriod")
 

internal/net/addr.go

@@ -153,3 +153,7 @@ func (pr *PortRange) Parse(s string) error {
 func (pr *PortRange) Contains(port int) bool {
 	return port >= pr.Min && port <= pr.Max
 }
+
+type ClientAddr interface {
+	ClientAddr() net.Addr
+}

internal/util/ws/ws.go

@@ -5,6 +5,7 @@ import (
 	"sync"
 	"time"
 
+	xnet "github.com/go-gost/x/internal/net"
 	"github.com/gorilla/websocket"
 )
 
@@ -12,12 +13,14 @@ type WebsocketConn interface {
 	net.Conn
 	WriteMessage(int, []byte) error
 	ReadMessage() (int, []byte, error)
+	xnet.ClientAddr
 }
 
 type websocketConn struct {
 	*websocket.Conn
-	rb  []byte
+	rb         []byte
-	mux sync.Mutex
+	clientAddr net.Addr
+	mux        sync.Mutex
 }
 
 func Conn(conn *websocket.Conn) WebsocketConn {
@@ -26,6 +29,13 @@ func Conn(conn *websocket.Conn) WebsocketConn {
 	}
 }
 
+func ConnWithClientAddr(conn *websocket.Conn, clientAddr net.Addr) WebsocketConn {
+	return &websocketConn{
+		Conn:       conn,
+		clientAddr: clientAddr,
+	}
+}
+
 func (c *websocketConn) Read(b []byte) (n int, err error) {
 	if len(c.rb) == 0 {
 		_, c.rb, err = c.Conn.ReadMessage()
@@ -66,3 +76,7 @@ func (c *websocketConn) SetWriteDeadline(t time.Time) error {
 	defer c.mux.Unlock()
 	return c.Conn.SetWriteDeadline(t)
 }
+
+func (c *websocketConn) ClientAddr() net.Addr {
+	return c.clientAddr
+}

listener/ws/listener.go

@@ -14,6 +14,7 @@ import (
 	md "github.com/go-gost/core/metadata"
 	admission "github.com/go-gost/x/admission/wrapper"
 	xnet "github.com/go-gost/x/internal/net"
+	xhttp "github.com/go-gost/x/internal/net/http"
 	"github.com/go-gost/x/internal/net/proxyproto"
 	limiter_util "github.com/go-gost/x/internal/util/limiter"
 	ws_util "github.com/go-gost/x/internal/util/ws"
@@ -177,8 +178,13 @@ func (l *wsListener) upgrade(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	var clientAddr net.Addr
+	if clientIP := xhttp.GetClientIP(r); clientIP != nil {
+		clientAddr = &net.IPAddr{IP: clientIP}
+	}
+
 	select {
-	case l.cqueue <- ws_util.Conn(conn):
+	case l.cqueue <- ws_util.ConnWithClientAddr(conn, clientAddr):
 	default:
 		conn.Close()
 		l.logger.Warnf("connection queue is full, client %s discarded", conn.RemoteAddr())

service/service.go

@@ -19,6 +19,7 @@ import (
 	"github.com/go-gost/core/recorder"
 	"github.com/go-gost/core/service"
 	ctxvalue "github.com/go-gost/x/ctx"
+	xnet "github.com/go-gost/x/internal/net"
 	xmetrics "github.com/go-gost/x/metrics"
 	"github.com/rs/xid"
 )
@@ -188,6 +189,11 @@ func (s *defaultService) Serve() error {
 		}
 
 		clientAddr := conn.RemoteAddr().String()
+		if ca, ok := conn.(xnet.ClientAddr); ok {
+			if addr := ca.ClientAddr(); addr != nil {
+				clientAddr = addr.String()
+			}
+		}
 		clientIP := clientAddr
 		if h, _, _ := net.SplitHostPort(clientAddr); h != "" {
 			clientIP = h