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

add context for conn

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
internal/ctx/value.go
+64 -6
View File
@@ -1,22 +1,80 @@
package ctx
import (
"bytes"
"context"
"github.com/go-gost/core/logger"
"github.com/go-gost/core/metadata"
xctx "github.com/go-gost/x/ctx"
xrecorder "github.com/go-gost/x/recorder"
)
type bufferKey struct{}
func ContextWithBuffer(ctx context.Context, buffer *bytes.Buffer) context.Context {
return context.WithValue(ctx, bufferKey{}, buffer)
}
func BufferFromContext(ctx context.Context) *bytes.Buffer {
v, _ := ctx.Value(bufferKey{}).(*bytes.Buffer)
return v
}
type loggerKey struct{}
func ContextWithLogger(ctx context.Context, log logger.Logger) context.Context {
return context.WithValue(ctx, loggerKey{}, log)
}
func LoggerFromContext(ctx context.Context) logger.Logger {
v, _ := ctx.Value(loggerKey{}).(logger.Logger)
return v
}
type mdKey struct{}
func ContextWithMetadata(ctx context.Context, md metadata.Metadata) context.Context {
return context.WithValue(ctx, mdKey{}, md)
}
func MetadataFromContext(ctx context.Context) metadata.Metadata {
v, _ := ctx.Value(mdKey{}).(metadata.Metadata)
return v
}
type recorderObjectCtxKey struct{}
var (
ctxKeyRecorderObject = &recorderObjectCtxKey{}
)
func ContextWithRecorderObject(ctx context.Context, ro *xrecorder.HandlerRecorderObject) context.Context {
return context.WithValue(ctx, ctxKeyRecorderObject, ro)
return context.WithValue(ctx, recorderObjectCtxKey{}, ro)
}
func RecorderObjectFromContext(ctx context.Context) *xrecorder.HandlerRecorderObject {
v, _ := ctx.Value(ctxKeyRecorderObject).(*xrecorder.HandlerRecorderObject)
v, _ := ctx.Value(recorderObjectCtxKey{}).(*xrecorder.HandlerRecorderObject)
return v
}
func Copy(ctx context.Context) context.Context {
if ctx == nil {
return nil
}
ctx2 := context.Background()
if v := xctx.SrcAddrFromContext(ctx); v != nil {
ctx2 = xctx.ContextWithSrcAddr(ctx2, v)
}
if v := xctx.DstAddrFromContext(ctx); v != nil {
ctx2 = xctx.ContextWithDstAddr(ctx2, v)
}
if v := xctx.SidFromContext(ctx); v != "" {
ctx2 = xctx.ContextWithSid(ctx2, v)
}
if v := xctx.ClientIDFromContext(ctx); v != "" {
ctx2 = xctx.ContextWithClientID(ctx2, v)
}
if v := MetadataFromContext(ctx); v != nil {
ctx2 = ContextWithMetadata(ctx2, v)
}
return ctx2
}
internal/net/net.go
-8
View File
@@ -26,14 +26,6 @@ type RemoteAddr interface {
RemoteAddr() net.Addr
}
type SrcAddr interface {
SrcAddr() net.Addr
}
type DstAddr interface {
DstAddr() net.Addr
}
// tcpraw.TCPConn
type SetDSCP interface {
SetDSCP(int) error
internal/net/proxyproto/conn.go
+6 -8
View File
@@ -1,6 +1,7 @@
package proxyproto
import (
"context"
"net"
"strconv"
@@ -10,6 +11,11 @@ import (
type serverConn struct {
net.Conn
ctx context.Context
}
func (c *serverConn) Context() context.Context {
return c.ctx
}
func (c *serverConn) RemoteAddr() net.Addr {
@@ -26,14 +32,6 @@ func (c *serverConn) LocalAddr() net.Addr {
return c.Conn.LocalAddr()
}
func (c *serverConn) SrcAddr() net.Addr {
return c.Conn.RemoteAddr()
}
func (c *serverConn) DstAddr() net.Addr {
return c.Conn.LocalAddr()
}
func (c *serverConn) CloseRead() error {
if sc, ok := c.Conn.(xio.CloseRead); ok {
return sc.CloseRead()
internal/net/proxyproto/listener.go
+14 -1
View File
@@ -1,9 +1,11 @@
package proxyproto
import (
"context"
"net"
"time"
"github.com/go-gost/x/ctx"
proxyproto "github.com/pires/go-proxyproto"
)
@@ -16,7 +18,18 @@ func (ln *listener) Accept() (net.Conn, error) {
if err != nil {
return nil, err
}
return &serverConn{Conn: conn}, nil
innerCtx := context.Background()
if c, ok := conn.(ctx.Context); ok {
if v := c.Context(); v != nil {
innerCtx = v
}
}
innerCtx = ctx.ContextWithSrcAddr(innerCtx, conn.RemoteAddr())
innerCtx = ctx.ContextWithDstAddr(innerCtx, conn.LocalAddr())
return &serverConn{Conn: conn, ctx: innerCtx}, nil
}
func WrapListener(ppv int, ln net.Listener, readHeaderTimeout time.Duration) net.Listener {
internal/util/forwarder/sniffer.go
+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") {
internal/util/mux/mux.go
+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
}
internal/util/sniffing/sniffer.go
+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") {
internal/util/tls/listener.go
+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
}
internal/util/ws/ws.go
+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
}