From 24547b43323d0d7676b11b578330ac2418e6cc4e Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Wed, 16 Oct 2024 20:46:45 +0800 Subject: [PATCH] add ClientAddr for websocket conn --- chain/router.go | 10 +++++++++- connector/tunnel/metadata.go | 5 ++++- handler/forward/local/handler.go | 27 +-------------------------- handler/http/handler.go | 9 ++++++++- handler/relay/handler.go | 9 ++++++++- handler/sni/handler.go | 9 ++++++++- handler/socks/v4/handler.go | 9 ++++++++- handler/socks/v5/handler.go | 9 ++++++++- handler/ss/handler.go | 9 ++++++++- handler/tunnel/metadata.go | 3 +++ internal/net/addr.go | 4 ++++ internal/util/ws/ws.go | 18 ++++++++++++++++-- listener/ws/listener.go | 8 +++++++- service/service.go | 6 ++++++ 14 files changed, 98 insertions(+), 37 deletions(-) diff --git a/chain/router.go b/chain/router.go index 69af36ec..cb5907b8 100644 --- a/chain/router.go +++ b/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{} } diff --git a/connector/tunnel/metadata.go b/connector/tunnel/metadata.go index e127f31d..20f47329 100644 --- a/connector/tunnel/metadata.go +++ b/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 } diff --git a/handler/forward/local/handler.go b/handler/forward/local/handler.go index e7fa2ec3..241d7113 100644 --- a/handler/forward/local/handler.go +++ b/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, - } - } -} diff --git a/handler/http/handler.go b/handler/http/handler.go index d09ee80e..b952cb20 100644 --- a/handler/http/handler.go +++ b/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(), diff --git a/handler/relay/handler.go b/handler/relay/handler.go index 9075624a..99ae2bf2 100644 --- a/handler/relay/handler.go +++ b/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(), diff --git a/handler/sni/handler.go b/handler/sni/handler.go index 4d6b500e..8df3f03c 100644 --- a/handler/sni/handler.go +++ b/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(), diff --git a/handler/socks/v4/handler.go b/handler/socks/v4/handler.go index c390131d..44cd52cb 100644 --- a/handler/socks/v4/handler.go +++ b/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(), diff --git a/handler/socks/v5/handler.go b/handler/socks/v5/handler.go index f57efd69..4f73464e 100644 --- a/handler/socks/v5/handler.go +++ b/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(), diff --git a/handler/ss/handler.go b/handler/ss/handler.go index 57b11264..5bfb8299 100644 --- a/handler/ss/handler.go +++ b/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(), diff --git a/handler/tunnel/metadata.go b/handler/tunnel/metadata.go index f260f9c5..052f5aa5 100644 --- a/handler/tunnel/metadata.go +++ b/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") diff --git a/internal/net/addr.go b/internal/net/addr.go index 9675946f..db84ca4e 100644 --- a/internal/net/addr.go +++ b/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 +} diff --git a/internal/util/ws/ws.go b/internal/util/ws/ws.go index a6c15126..d19e006a 100644 --- a/internal/util/ws/ws.go +++ b/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 - mux sync.Mutex + rb []byte + 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 +} diff --git a/listener/ws/listener.go b/listener/ws/listener.go index f5e299b5..bfdeb4c7 100644 --- a/listener/ws/listener.go +++ b/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()) diff --git a/service/service.go b/service/service.go index daf00b3d..4ca064e8 100644 --- a/service/service.go +++ b/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