sniffing websocket frame

This commit is contained in:
ginuerzh
2024-10-19 19:36:06 +08:00
parent a0cbee8817
commit 6b932e35bf
28 changed files with 1238 additions and 301 deletions
+222 -89
View File
@@ -2,6 +2,7 @@ package tunnel
import (
"bufio"
"bytes"
"context"
"errors"
"fmt"
@@ -24,9 +25,13 @@ import (
"github.com/go-gost/relay"
admission "github.com/go-gost/x/admission/wrapper"
ctxvalue "github.com/go-gost/x/ctx"
ctx_internal "github.com/go-gost/x/internal/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"
"github.com/go-gost/x/internal/net/proxyproto"
"github.com/go-gost/x/internal/util/sniffing"
ws_util "github.com/go-gost/x/internal/util/ws"
climiter "github.com/go-gost/x/limiter/conn/wrapper"
metrics "github.com/go-gost/x/metrics/wrapper"
stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
@@ -34,25 +39,6 @@ import (
"golang.org/x/net/http/httpguts"
)
const (
defaultBodySize = 1024 * 1024 // 1MB
)
type recorderObjectCtxKey struct{}
var (
ctxKeyRecorderObject = &recorderObjectCtxKey{}
)
func contextWithRecorderObject(ctx context.Context, ro *xrecorder.HandlerRecorderObject) context.Context {
return context.WithValue(ctx, ctxKeyRecorderObject, ro)
}
func recorderObjectFromContext(ctx context.Context) *xrecorder.HandlerRecorderObject {
v, _ := ctx.Value(ctxKeyRecorderObject).(*xrecorder.HandlerRecorderObject)
return v
}
type entrypoint struct {
node string
service string
@@ -62,6 +48,9 @@ type entrypoint struct {
log logger.Logger
recorder recorder.RecorderObject
transport http.RoundTripper
sniffingWebsocket bool
websocketSampleRate float64
}
func (ep *entrypoint) Handle(ctx context.Context, conn net.Conn) (err error) {
@@ -119,6 +108,74 @@ func (ep *entrypoint) Handle(ctx context.Context, conn net.Conn) (err error) {
return ep.handleHTTP(ctx, conn, ro, log)
}
func (h *entrypoint) dial(ctx context.Context, network, addr string) (conn net.Conn, err error) {
var tunnelID relay.TunnelID
if h.ingress != nil {
if rule := h.ingress.GetRule(ctx, addr); rule != nil {
tunnelID = parseTunnelID(rule.Endpoint)
}
}
log := ctxvalue.LoggerFromContext(ctx)
if log == nil {
log = h.log
}
log.Debugf("dial: new connection to host %s", addr)
if tunnelID.IsZero() {
return nil, fmt.Errorf("%w %s", ErrTunnelRoute, addr)
}
if ro := ctx_internal.RecorderObjectFromContext(ctx); ro != nil {
ro.ClientID = tunnelID.String()
}
if tunnelID.IsPrivate() {
return nil, fmt.Errorf("%w: tunnel %s is private for host %s", ErrPrivateTunnel, tunnelID, addr)
}
log = log.WithFields(map[string]any{
"tunnel": tunnelID.String(),
})
d := &Dialer{
node: h.node,
pool: h.pool,
sd: h.sd,
retry: 3,
timeout: 15 * time.Second,
log: log,
}
conn, node, cid, err := d.Dial(ctx, "tcp", tunnelID.String())
if err != nil {
return
}
log.Debugf("dial: connected to host %s, tunnel: %s, connector: %s", addr, tunnelID, cid)
if node == h.node {
clientAddr := ctxvalue.ClientAddrFromContext(ctx)
var features []relay.Feature
af := &relay.AddrFeature{}
af.ParseFrom(string(clientAddr))
features = append(features, af) // src address
af = &relay.AddrFeature{}
af.ParseFrom(addr)
features = append(features, af) // dst address
if _, err = (&relay.Response{
Version: relay.Version1,
Status: relay.StatusOK,
Features: features,
}).WriteTo(conn); err != nil {
conn.Close()
return nil, err
}
}
return conn, nil
}
func (ep *entrypoint) handleHTTP(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) (err error) {
pStats := stats.Stats{}
conn = stats_wrapper.WrapConn(conn, &pStats)
@@ -148,7 +205,7 @@ func (ep *entrypoint) handleHTTP(ctx context.Context, conn net.Conn, ro *xrecord
ro.Time = time.Time{}
if err := ep.httpRoundTrip(ctx, conn, req, ro, &pStats, log); err != nil {
if err := ep.httpRoundTrip(ctx, xio.NewReadWriter(br, conn), req, ro, &pStats, log); err != nil {
log.Error(err)
return err
}
@@ -169,7 +226,7 @@ func (ep *entrypoint) handleHTTP(ctx context.Context, conn net.Conn, ro *xrecord
log.Trace(string(dump))
}
if err := ep.httpRoundTrip(ctx, conn, req, ro, &pStats, log); err != nil {
if err := ep.httpRoundTrip(ctx, xio.NewReadWriter(br, conn), req, ro, &pStats, log); err != nil {
return err
}
}
@@ -250,17 +307,20 @@ func (ep *entrypoint) httpRoundTrip(ctx context.Context, rw io.ReadWriter, req *
var reqBody *xhttp.Body
if opts := ep.recorder.Options; opts != nil && opts.HTTPBody {
if req.Body != nil {
maxSize := opts.MaxBodySize
if maxSize <= 0 {
maxSize = defaultBodySize
bodySize := opts.MaxBodySize
if bodySize <= 0 {
bodySize = sniffing.DefaultBodySize
}
reqBody = xhttp.NewBody(req.Body, maxSize)
if bodySize > sniffing.MaxBodySize {
bodySize = sniffing.MaxBodySize
}
reqBody = xhttp.NewBody(req.Body, bodySize)
req.Body = reqBody
}
}
ctx = ctxvalue.ContextWithClientAddr(ctx, ctxvalue.ClientAddr(clientAddr))
ctx = contextWithRecorderObject(ctx, ro)
ctx = ctx_internal.ContextWithRecorderObject(ctx, ro)
ctx = ctxvalue.ContextWithLogger(ctx, log)
resp, err := ep.transport.RoundTrip(req.WithContext(ctx))
@@ -289,16 +349,19 @@ func (ep *entrypoint) httpRoundTrip(ctx context.Context, rw io.ReadWriter, req *
}
if resp.StatusCode == http.StatusSwitchingProtocols {
return ep.handleUpgradeResponse(rw, req, resp)
return ep.handleUpgradeResponse(ctx, rw, req, resp, ro, log)
}
var respBody *xhttp.Body
if opts := ep.recorder.Options; opts != nil && opts.HTTPBody {
maxSize := opts.MaxBodySize
if maxSize <= 0 {
maxSize = defaultBodySize
bodySize := opts.MaxBodySize
if bodySize <= 0 {
bodySize = sniffing.DefaultBodySize
}
respBody = xhttp.NewBody(resp.Body, maxSize)
if bodySize > sniffing.MaxBodySize {
bodySize = sniffing.MaxBodySize
}
respBody = xhttp.NewBody(resp.Body, bodySize)
resp.Body = respBody
}
@@ -321,7 +384,7 @@ func upgradeType(h http.Header) string {
return h.Get("Upgrade")
}
func (ep *entrypoint) handleUpgradeResponse(rw io.ReadWriter, req *http.Request, res *http.Response) error {
func (ep *entrypoint) handleUpgradeResponse(ctx context.Context, rw io.ReadWriter, req *http.Request, res *http.Response, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
reqUpType := upgradeType(req.Header)
resUpType := upgradeType(res.Header)
if !strings.EqualFold(reqUpType, resUpType) {
@@ -332,81 +395,151 @@ func (ep *entrypoint) handleUpgradeResponse(rw io.ReadWriter, req *http.Request,
if !ok {
return fmt.Errorf("internal error: 101 switching protocols response with non-writable body")
}
defer backConn.Close()
res.Body = nil
if err := res.Write(rw); err != nil {
return fmt.Errorf("response write: %v", err)
}
if reqUpType == "websocket" && ep.sniffingWebsocket {
return ep.sniffingWebsocketFrame(ctx, rw, backConn, ro, log)
}
return xnet.Transport(rw, backConn)
}
func (h *entrypoint) dial(ctx context.Context, network, addr string) (conn net.Conn, err error) {
var tunnelID relay.TunnelID
if h.ingress != nil {
if rule := h.ingress.GetRule(ctx, addr); rule != nil {
tunnelID = parseTunnelID(rule.Endpoint)
func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWriter, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
errc := make(chan error, 1)
sampleRate := ep.websocketSampleRate
if sampleRate <= 0 {
sampleRate = sniffing.DefaultSampleRate
}
if sampleRate > sniffing.MaxSampleRate {
sampleRate = sniffing.MaxSampleRate
}
d := time.Duration(1 / sampleRate * 1e9)
go func() {
ro2 := &xrecorder.HandlerRecorderObject{}
*ro2 = *ro
ro = ro2
ticker := time.NewTicker(d)
defer ticker.Stop()
buf := &bytes.Buffer{}
for {
start := time.Now()
err := ep.copyWebsocketFrame(cc, rw, buf, "client", ro)
select {
case <-ticker.C:
if err != nil {
ro.Err = err.Error()
}
ro.Duration = time.Since(start)
ro.Time = time.Now()
if err := ro.Record(ctx, ep.recorder.Recorder); err != nil {
log.Errorf("record: %v", err)
}
default:
}
if err != nil {
errc <- err
return
}
}
}
}()
log := ctxvalue.LoggerFromContext(ctx)
if log == nil {
log = h.log
}
log.Debugf("dial: new connection to host %s", addr)
go func() {
ro2 := &xrecorder.HandlerRecorderObject{}
*ro2 = *ro
ro = ro2
if tunnelID.IsZero() {
return nil, fmt.Errorf("%w %s", ErrTunnelRoute, addr)
}
ticker := time.NewTicker(d)
defer ticker.Stop()
if ro := recorderObjectFromContext(ctx); ro != nil {
ro.ClientID = tunnelID.String()
}
buf := &bytes.Buffer{}
for {
start := time.Now()
if tunnelID.IsPrivate() {
return nil, fmt.Errorf("%w: tunnel %s is private for host %s", ErrPrivateTunnel, tunnelID, addr)
}
err := ep.copyWebsocketFrame(rw, cc, buf, "server", ro)
select {
case <-ticker.C:
if err != nil {
ro.Err = err.Error()
}
ro.Duration = time.Since(start)
ro.Time = time.Now()
if err := ro.Record(ctx, ep.recorder.Recorder); err != nil {
log.Errorf("record: %v", err)
}
default:
}
log = log.WithFields(map[string]any{
"tunnel": tunnelID.String(),
})
d := &Dialer{
node: h.node,
pool: h.pool,
sd: h.sd,
retry: 3,
timeout: 15 * time.Second,
log: log,
}
conn, node, cid, err := d.Dial(ctx, "tcp", tunnelID.String())
if err != nil {
return
}
log.Debugf("dial: connected to host %s, tunnel: %s, connector: %s", addr, tunnelID, cid)
if node == h.node {
clientAddr := ctxvalue.ClientAddrFromContext(ctx)
var features []relay.Feature
af := &relay.AddrFeature{}
af.ParseFrom(string(clientAddr))
features = append(features, af) // src address
af = &relay.AddrFeature{}
af.ParseFrom(addr)
features = append(features, af) // dst address
if _, err = (&relay.Response{
Version: relay.Version1,
Status: relay.StatusOK,
Features: features,
}).WriteTo(conn); err != nil {
conn.Close()
return nil, err
if err != nil {
errc <- err
return
}
}
}()
<-errc
return nil
}
func (ep *entrypoint) copyWebsocketFrame(w io.Writer, r io.Reader, buf *bytes.Buffer, from string, ro *xrecorder.HandlerRecorderObject) (err error) {
fr := ws_util.Frame{}
if _, err = fr.ReadFrom(r); err != nil {
return err
}
return conn, nil
ws := &xrecorder.WebsocketRecorderObject{
From: from,
Fin: fr.Header.Fin,
Rsv1: fr.Header.Rsv1,
Rsv2: fr.Header.Rsv2,
Rsv3: fr.Header.Rsv3,
OpCode: int(fr.Header.OpCode),
Masked: fr.Header.Masked,
MaskKey: fr.Header.MaskKey,
Length: fr.Header.PayloadLength,
}
if opts := ep.recorder.Options; opts != nil && opts.HTTPBody {
bodySize := opts.MaxBodySize
if bodySize <= 0 {
bodySize = sniffing.DefaultBodySize
}
if bodySize > sniffing.MaxBodySize {
bodySize = sniffing.MaxBodySize
}
buf.Reset()
if _, err := io.Copy(buf, io.LimitReader(fr.Data, int64(bodySize))); err != nil {
return err
}
ws.Payload = buf.Bytes()
}
ro.Websocket = ws
length := uint64(fr.Header.Length()) + uint64(fr.Header.PayloadLength)
if from == "client" {
ro.InputBytes = length
ro.OutputBytes = 0
} else {
ro.InputBytes = 0
ro.OutputBytes = length
}
fr.Data = io.MultiReader(bytes.NewReader(buf.Bytes()), fr.Data)
if _, err := fr.WriteTo(w); err != nil {
return err
}
return nil
}
func (ep *entrypoint) handleConnect(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) (err error) {
+2
View File
@@ -90,6 +90,8 @@ func (h *tunnelHandler) Init(md md.Metadata) (err error) {
log: h.log.WithFields(map[string]any{
"kind": "entrypoint",
}),
sniffingWebsocket: h.md.sniffingWebsocket,
websocketSampleRate: h.md.sniffingWebsocketSampleRate,
}
h.ep.transport = &http.Transport{
DialContext: h.ep.dial,
+10 -5
View File
@@ -22,11 +22,13 @@ const (
type metadata struct {
readTimeout time.Duration
entryPoint string
entryPointID relay.TunnelID
entryPointProxyProtocol int
entryPointKeepalive bool
entryPointReadTimeout time.Duration
entryPoint string
entryPointID relay.TunnelID
entryPointProxyProtocol int
entryPointKeepalive bool
entryPointReadTimeout time.Duration
sniffingWebsocket bool
sniffingWebsocketSampleRate float64
directTunnel bool
tunnelTTL time.Duration
@@ -53,6 +55,9 @@ func (h *tunnelHandler) parseMetadata(md mdata.Metadata) (err error) {
h.md.entryPointReadTimeout = 15 * time.Second
}
h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket")
h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate")
h.md.tunnelTTL = mdutil.GetDuration(md, "tunnel.ttl")
if h.md.tunnelTTL <= 0 {
h.md.tunnelTTL = defaultTTL