add context for conn

This commit is contained in:
ginuerzh
2025-08-04 19:29:38 +08:00
parent ad5cf6fd61
commit b597467858
90 changed files with 889 additions and 917 deletions
+23 -28
View File
@@ -14,7 +14,8 @@ import (
mdata "github.com/go-gost/core/metadata"
"github.com/go-gost/core/observer/stats"
"github.com/go-gost/core/recorder"
ctxvalue "github.com/go-gost/x/ctx"
xctx "github.com/go-gost/x/ctx"
ictx "github.com/go-gost/x/internal/ctx"
xnet "github.com/go-gost/x/internal/net"
"github.com/go-gost/x/internal/net/proxyproto"
"github.com/go-gost/x/internal/util/forwarder"
@@ -87,27 +88,28 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
LocalAddr: conn.LocalAddr().String(),
Network: "tcp",
Time: start,
SID: string(ctxvalue.SidFromContext(ctx)),
SID: xctx.SidFromContext(ctx).String(),
}
if srcAddr := ctxvalue.SrcAddrFromContext(ctx); srcAddr != nil {
ro.ClientIP = srcAddr.String()
if srcAddr := xctx.SrcAddrFromContext(ctx); srcAddr != nil {
ro.ClientAddr = srcAddr.String()
}
log := h.options.Logger.WithFields(map[string]any{
"remote": conn.RemoteAddr().String(),
"local": conn.LocalAddr().String(),
"sid": ro.SID,
"client": ro.ClientIP,
})
log.Infof("%s <> %s", conn.RemoteAddr(), conn.LocalAddr())
network := "tcp"
if _, ok := conn.(net.PacketConn); ok {
network = "udp"
}
ro.Network = network
log := h.options.Logger.WithFields(map[string]any{
"network": ro.Network,
"remote": conn.RemoteAddr().String(),
"local": conn.LocalAddr().String(),
"client": ro.ClientAddr,
"sid": ro.SID,
})
log.Infof("%s <> %s", conn.RemoteAddr(), conn.LocalAddr())
pStats := xstats.Stats{}
conn = stats_wrapper.WrapConn(conn, &pStats)
@@ -133,13 +135,6 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
return rate_limiter.ErrRateLimit
}
var host string
if md, ok := conn.(mdata.Metadatable); ok {
if v := mdutil.GetString(md.Metadata(), "host"); v != "" {
host = v
}
}
var proto string
if network == "tcp" && h.md.sniffing {
if h.md.sniffingTimeout > 0 {
@@ -156,13 +151,13 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
dial := func(ctx context.Context, network, address string) (net.Conn, error) {
var buf bytes.Buffer
cc, err := h.options.Router.Dial(ctxvalue.ContextWithBuffer(ctx, &buf), "tcp", address)
cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "tcp", address)
ro.Route = buf.String()
cc = proxyproto.WrapClientConn(
h.md.proxyProtocol,
ctxvalue.SrcAddrFromContext(ctx),
ctxvalue.DstAddrFromContext(ctx),
xctx.SrcAddrFromContext(ctx),
xctx.DstAddrFromContext(ctx),
cc)
return cc, err
@@ -203,7 +198,7 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
}
var target *chain.Node
if host != "" {
if host := mdutil.GetString(ictx.MetadataFromContext(ctx), "host"); host != "" {
target = &chain.Node{
Addr: host,
}
@@ -239,7 +234,7 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
log.Debugf("%s >> %s", conn.RemoteAddr(), target.Addr)
var buf bytes.Buffer
cc, err := h.options.Router.Dial(ctxvalue.ContextWithBuffer(ctx, &buf), network, target.Addr)
cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), network, target.Addr)
ro.Route = buf.String()
if err != nil {
log.Error(err)
@@ -256,13 +251,13 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
}
log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()})
ro.Src = cc.LocalAddr().String()
ro.Dst = cc.RemoteAddr().String()
ro.SrcAddr = cc.LocalAddr().String()
ro.DstAddr = cc.RemoteAddr().String()
cc = proxyproto.WrapClientConn(
h.md.proxyProtocol,
ctxvalue.SrcAddrFromContext(ctx),
ctxvalue.DstAddrFromContext(ctx),
xctx.SrcAddrFromContext(ctx),
xctx.DstAddrFromContext(ctx),
cc)
t := time.Now()