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
+10 -18
View File
@@ -27,7 +27,7 @@ import (
dissector "github.com/go-gost/tls-dissector"
xbypass "github.com/go-gost/x/bypass"
"github.com/go-gost/x/config"
ctxvalue "github.com/go-gost/x/ctx"
xctx "github.com/go-gost/x/ctx"
xio "github.com/go-gost/x/internal/io"
xnet "github.com/go-gost/x/internal/net"
xhttp "github.com/go-gost/x/internal/net/http"
@@ -161,15 +161,7 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
if clientIP := xhttp.GetClientIP(req); clientIP != nil {
ro.ClientIP = clientIP.String()
}
{
clientAddr := ro.RemoteAddr
if ro.ClientIP != "" {
if _, port, _ := net.SplitHostPort(ro.RemoteAddr); port != "" {
clientAddr = net.JoinHostPort(ro.ClientIP, port)
}
}
ctx = ctxvalue.ContextWithClientAddr(ctx, ctxvalue.ClientAddr(clientAddr))
ctx = xctx.ContextWithSrcAddr(ctx, &net.TCPAddr{IP: clientIP})
}
// http/2
@@ -187,8 +179,8 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
log := ho.Log
log.Debugf("connected to node %s(%s)", node.Name, node.Addr)
ro.Src = cc.LocalAddr().String()
ro.Dst = cc.RemoteAddr().String()
ro.SrcAddr = cc.LocalAddr().String()
ro.DstAddr = cc.RemoteAddr().String()
ro.Time = time.Time{}
shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriteCloser(br, conn, conn), cc, node, req, &pStats, &ho)
@@ -343,8 +335,8 @@ func (h *Sniffer) serveH2(ctx context.Context, conn net.Conn, ho *HandleOptions)
}
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()
log.Debugf("connected to node %s(%s)", node.Name, node.Addr)
return cc, nil
},
@@ -435,7 +427,7 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser,
err = errors.New("unauthorized")
return
}
ctx = ctxvalue.ContextWithClientID(ctx, ctxvalue.ClientID(id))
ctx = xctx.ContextWithClientID(ctx, xctx.ClientID(id))
}
if httpSettings.Host != "" {
@@ -797,7 +789,7 @@ func (h *Sniffer) HandleTLS(ctx context.Context, conn net.Conn, opts ...HandleOp
ro.TLS.Proto = clientHello.SupportedProtos[0]
}
ctx = ctxvalue.ContextWithClientAddr(ctx, ctxvalue.ClientAddr(ro.RemoteAddr))
// ctx = xctx.ContextWithClientAddr(ctx, xctx.ClientAddr(ro.RemoteAddr))
host := clientHello.ServerName
if host != "" {
@@ -821,8 +813,8 @@ func (h *Sniffer) HandleTLS(ctx context.Context, conn net.Conn, opts ...HandleOp
log := ho.Log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()})
log.Debugf("connected to node %s(%s)", node.Name, node.Addr)
ro.Src = cc.LocalAddr().String()
ro.Dst = cc.RemoteAddr().String()
ro.SrcAddr = cc.LocalAddr().String()
ro.DstAddr = cc.RemoteAddr().String()
if h.Certificate != nil && h.PrivateKey != nil &&
len(clientHello.SupportedProtos) > 0 && (clientHello.SupportedProtos[0] == "h2" || clientHello.SupportedProtos[0] == "http/1.1") {
+9 -23
View File
@@ -1,11 +1,11 @@
package mux
import (
"context"
"net"
"time"
"github.com/go-gost/core/metadata"
xnet "github.com/go-gost/x/internal/net"
"github.com/go-gost/x/ctx"
smux "github.com/xtaci/smux"
)
@@ -145,27 +145,13 @@ func (c *streamConn) Write(b []byte) (n int, err error) {
return c.stream.Write(b)
}
func (c *streamConn) Metadata() metadata.Metadata {
if md, ok := c.Conn.(metadata.Metadatable); ok {
return md.Metadata()
}
return nil
}
func (c *streamConn) SrcAddr() net.Addr {
if sc, ok := c.Conn.(xnet.SrcAddr); ok {
return sc.SrcAddr()
}
return nil
}
func (c *streamConn) DstAddr() net.Addr {
if sc, ok := c.Conn.(xnet.DstAddr); ok {
return sc.DstAddr()
}
return nil
}
func (c *streamConn) Close() error {
return c.stream.Close()
}
func (c *streamConn) Context() context.Context {
if sc, ok := c.Conn.(ctx.Context); ok {
return sc.Context()
}
return nil
}
+9 -16
View File
@@ -24,7 +24,7 @@ import (
"github.com/go-gost/core/recorder"
dissector "github.com/go-gost/tls-dissector"
xbypass "github.com/go-gost/x/bypass"
ctxvalue "github.com/go-gost/x/ctx"
xctx "github.com/go-gost/x/ctx"
xio "github.com/go-gost/x/internal/io"
xnet "github.com/go-gost/x/internal/net"
xhttp "github.com/go-gost/x/internal/net/http"
@@ -151,16 +151,9 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
if clientIP := xhttp.GetClientIP(req); clientIP != nil {
ro.ClientIP = clientIP.String()
ctx = xctx.ContextWithSrcAddr(ctx, &net.TCPAddr{IP: clientIP})
}
clientAddr := ro.RemoteAddr
if ro.ClientIP != "" {
if _, port, _ := net.SplitHostPort(ro.RemoteAddr); port != "" {
clientAddr = net.JoinHostPort(ro.ClientIP, port)
}
}
ctx = ctxvalue.ContextWithClientAddr(ctx, ctxvalue.ClientAddr(clientAddr))
// http/2
if req.Method == "PRI" && len(req.Header) == 0 && req.URL.Path == "*" && req.Proto == "HTTP/2.0" {
return h.serveH2(ctx, xnet.NewReadWriteConn(br, conn, conn), &ho)
@@ -194,8 +187,8 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
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()
ro.Time = time.Time{}
shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriteCloser(br, conn, conn), cc, req, ro, &pStats, log)
@@ -254,8 +247,8 @@ func (h *Sniffer) serveH2(ctx context.Context, conn net.Conn, ho *HandleOptions)
}
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 = tls.Client(cc, cfg)
return cc, nil
@@ -598,7 +591,7 @@ func (h *Sniffer) HandleTLS(ctx context.Context, conn net.Conn, opts ...HandleOp
ro.TLS.Proto = clientHello.SupportedProtos[0]
}
ctx = ctxvalue.ContextWithClientAddr(ctx, ctxvalue.ClientAddr(ro.RemoteAddr))
// ctx = xctx.ContextWithClientAddr(ctx, xctx.ClientAddr(ro.RemoteAddr))
host := clientHello.ServerName
if host != "" {
@@ -623,8 +616,8 @@ func (h *Sniffer) HandleTLS(ctx context.Context, conn net.Conn, opts ...HandleOp
defer cc.Close()
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()
if h.Certificate != nil && h.PrivateKey != nil &&
len(clientHello.SupportedProtos) > 0 && (clientHello.SupportedProtos[0] == "h2" || clientHello.SupportedProtos[0] == "http/1.1") {
+5 -19
View File
@@ -1,11 +1,11 @@
package tls
import (
"context"
"crypto/tls"
"net"
"github.com/go-gost/core/metadata"
xnet "github.com/go-gost/x/internal/net"
"github.com/go-gost/x/ctx"
)
type listener struct {
@@ -34,23 +34,9 @@ type tlsConn struct {
*tls.Conn
}
func (c *tlsConn) Metadata() metadata.Metadata {
if md, ok := c.NetConn().(metadata.Metadatable); ok {
return md.Metadata()
}
return nil
}
func (c *tlsConn) SrcAddr() net.Addr {
if sc, ok := c.NetConn().(xnet.SrcAddr); ok {
return sc.SrcAddr()
}
return nil
}
func (c *tlsConn) DstAddr() net.Addr {
if sc, ok := c.NetConn().(xnet.DstAddr); ok {
return sc.DstAddr()
func (c *tlsConn) Context() context.Context {
if sc, ok := c.NetConn().(ctx.Context); ok {
return sc.Context()
}
return nil
}
+18 -17
View File
@@ -1,12 +1,13 @@
package ws
import (
"context"
"net"
"sync"
"time"
ctx_pkg "github.com/go-gost/x/ctx"
xio "github.com/go-gost/x/internal/io"
xnet "github.com/go-gost/x/internal/net"
"github.com/gorilla/websocket"
)
@@ -14,28 +15,31 @@ type WebsocketConn interface {
net.Conn
WriteMessage(int, []byte) error
ReadMessage() (int, []byte, error)
xnet.SrcAddr
xio.CloseRead
xio.CloseWrite
}
type websocketConn struct {
*websocket.Conn
rb []byte
srcAddr net.Addr
mux sync.Mutex
rb []byte
ctx context.Context
mux sync.Mutex
}
func Conn(conn *websocket.Conn) WebsocketConn {
return &websocketConn{
Conn: conn,
ctx := context.Background()
if cc, ok := conn.NetConn().(ctx_pkg.Context); ok {
if cv := cc.Context(); cv != nil {
ctx = cv
}
}
return ContextConn(ctx, conn)
}
func ConnWithSrcAddr(conn *websocket.Conn, srcAddr net.Addr) WebsocketConn {
func ContextConn(ctx context.Context, conn *websocket.Conn) WebsocketConn {
return &websocketConn{
Conn: conn,
srcAddr: srcAddr,
Conn: conn,
ctx: ctx,
}
}
@@ -80,13 +84,6 @@ func (c *websocketConn) SetWriteDeadline(t time.Time) error {
return c.Conn.SetWriteDeadline(t)
}
func (c *websocketConn) SrcAddr() net.Addr {
if sa, ok := c.Conn.NetConn().(xnet.SrcAddr); ok {
return sa.SrcAddr()
}
return c.srcAddr
}
func (c *websocketConn) CloseRead() error {
if sc, ok := c.Conn.NetConn().(xio.CloseRead); ok {
return sc.CloseRead()
@@ -100,3 +97,7 @@ func (c *websocketConn) CloseWrite() error {
}
return xio.ErrUnsupported
}
func (c *websocketConn) Context() context.Context {
return c.ctx
}