go-gost/x
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add ClientAddr for websocket conn

Browse Source
This commit is contained in:
ginuerzh
2024-10-16 20:46:45 +08:00
parent ea179a0ee6
commit 24547b4332
14 changed files with 98 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
chain/router.go
+9 -1
View File
@@ -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
+4 -1
View File
@@ -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
+1 -26
View File
@@ -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
+8 -1
View File
@@ -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
+8 -1
View File
@@ -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
+8 -1
View File
@@ -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
+8 -1
View File
@@ -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
+8 -1
View File
@@ -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
+8 -1
View File
@@ -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
+3
View File
@@ -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
+4
View File
@@ -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
+14
View File
@@ -5,6 +5,7 @@ import (
"sync"
"time"
xnet "github.com/go-gost/x/internal/net"
"github.com/gorilla/websocket"
)
@@ -12,11 +13,13 @@ type WebsocketConn interface {
net.Conn
WriteMessage(int, []byte) error
ReadMessage() (int, []byte, error)
xnet.ClientAddr
}
type websocketConn struct {
*websocket.Conn
rb []byte
clientAddr net.Addr
mux sync.Mutex
}
@@ -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
+7 -1
View File
@@ -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
+6
View File
@@ -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