diff --git a/internal/util/sniffing/sniffer.go b/internal/util/sniffing/sniffer.go index 7b42e0df..a0f79ede 100644 --- a/internal/util/sniffing/sniffer.go +++ b/internal/util/sniffing/sniffer.go @@ -1,58 +1,37 @@ package sniffing import ( - "bufio" - "bytes" "context" "crypto" "crypto/tls" "crypto/x509" - "encoding/hex" - "errors" - "fmt" - "io" - "math" "net" - "net/http" - "net/http/httputil" "strings" "time" "github.com/go-gost/core/bypass" "github.com/go-gost/core/logger" - "github.com/go-gost/core/observer/stats" "github.com/go-gost/core/recorder" - dissector "github.com/go-gost/tls-dissector" - xbypass "github.com/go-gost/x/bypass" - 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" tls_util "github.com/go-gost/x/internal/util/tls" - ws_util "github.com/go-gost/x/internal/util/ws" - xstats "github.com/go-gost/x/observer/stats" - stats_wrapper "github.com/go-gost/x/observer/stats/wrapper" xrecorder "github.com/go-gost/x/recorder" - "golang.org/x/net/http/httpguts" - "golang.org/x/net/http2" - "golang.org/x/time/rate" ) const ( + // DefaultReadTimeout is the default timeout for reading data from connections. DefaultReadTimeout = 30 * time.Second // DefaultBodySize is the default HTTP body or websocket frame size to record. DefaultBodySize = 64 * 1024 // 64KB // MaxBodySize is the maximum HTTP body or websocket frame size to record. MaxBodySize = 1024 * 1024 // 1MB - // DeafultSampleRate is the default websocket sample rate (samples per second). + // DefaultSampleRate is the default websocket sample rate (samples per second). DefaultSampleRate = 10.0 ) -var ( - DefaultCertPool = tls_util.NewMemoryCertPool() -) +// DefaultCertPool is the default in-memory certificate pool used for TLS MITM. +var DefaultCertPool = tls_util.NewMemoryCertPool() +// HandleOptions holds configuration options for sniffing handlers. type HandleOptions struct { service string dial func(ctx context.Context, network, address string) (net.Conn, error) @@ -63,44 +42,53 @@ type HandleOptions struct { log logger.Logger } +// HandleOption configures HandleOptions for sniffing handlers. type HandleOption func(opts *HandleOptions) +// WithService sets the service name for bypass lookups. func WithService(service string) HandleOption { return func(opts *HandleOptions) { opts.service = service } } +// WithDial sets the dial function used to establish upstream connections. func WithDial(dial func(ctx context.Context, network, address string) (net.Conn, error)) HandleOption { return func(opts *HandleOptions) { opts.dial = dial } } +// WithDialTLS sets the dial function used for TLS-wrapped upstream connections. func WithDialTLS(dialTLS func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error)) HandleOption { return func(opts *HandleOptions) { opts.dialTLS = dialTLS } } +// WithBypass sets the bypass rules for filtering requests by host. func WithBypass(bypass bypass.Bypass) HandleOption { return func(opts *HandleOptions) { opts.bypass = bypass } } +// WithRecorderObject sets the recorder object for capturing traffic metadata. func WithRecorderObject(ro *xrecorder.HandlerRecorderObject) HandleOption { return func(opts *HandleOptions) { opts.recorderObject = ro } } +// WithLog sets the logger for the handler. func WithLog(log logger.Logger) HandleOption { return func(opts *HandleOptions) { opts.log = log } } +// Sniffer handles HTTP and TLS traffic sniffing, recording, and MITM TLS +// termination for protocol-aware forwarding. type Sniffer struct { Websocket bool WebsocketSampleRate float64 @@ -118,800 +106,40 @@ type Sniffer struct { ReadTimeout time.Duration } -func (h *Sniffer) HandleHTTP(ctx context.Context, network string, conn net.Conn, opts ...HandleOption) error { - var ho HandleOptions - for _, opt := range opts { - opt(&ho) +// clampBodySize returns the effective body capture size from recorder options, +// bounded by [DefaultBodySize, MaxBodySize]. Returns 0 if body recording is +// disabled. +func clampBodySize(opts *recorder.Options) int { + if opts == nil || !opts.HTTPBody { + return 0 } - - if h.ReadTimeout <= 0 { - h.ReadTimeout = DefaultReadTimeout + size := opts.MaxBodySize + if size <= 0 { + size = DefaultBodySize } - - pStats := xstats.Stats{} - conn = stats_wrapper.WrapConn(conn, &pStats) - - br := bufio.NewReader(conn) - req, err := http.ReadRequest(br) - if err != nil { - return err - } - - log := ho.log - if log.IsLevelEnabled(logger.TraceLevel) { - dump, _ := httputil.DumpRequest(req, false) - log.Trace(string(dump)) - } - - ro := ho.recorderObject - ro.HTTP = &xrecorder.HTTPRecorderObject{ - Host: req.Host, - Proto: req.Proto, - Scheme: req.URL.Scheme, - Method: req.Method, - URI: req.RequestURI, - Request: xrecorder.HTTPRequestRecorderObject{ - ContentLength: req.ContentLength, - Header: req.Header.Clone(), - }, - } - - if clientIP := xhttp.GetClientIP(req); clientIP != nil { - ro.ClientIP = clientIP.String() - ctx = xctx.ContextWithSrcAddr(ctx, &net.TCPAddr{IP: clientIP}) - } - - // http/2 - if req.Method == "PRI" && len(req.Header) == 0 && req.URL.Path == "*" && req.Proto == "HTTP/2.0" { - return h.serveH2(ctx, network, xnet.NewReadWriteConn(br, conn, conn), &ho) - } - - host := req.Host - if host != "" { - if _, _, err := net.SplitHostPort(host); err != nil { - host = net.JoinHostPort(strings.Trim(host, "[]"), "80") - } - ro.Host = host - - log = log.WithFields(map[string]any{ - "host": host, - }) - - if ho.bypass != nil && ho.bypass.Contains(ctx, network, host, bypass.WithService(ho.service)) { - return xbypass.ErrBypass - } - } - - dial := ho.dial - if dial == nil { - dial = (&net.Dialer{}).DialContext - } - cc, err := dial(ctx, network, host) - if err != nil { - return err - } - defer cc.Close() - - log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "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) - if err != nil || shouldClose { - return err - } - - for { - pStats.Reset() - - req, err := http.ReadRequest(br) - if err != nil { - if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) { - return nil - } - return err - } - - if log.IsLevelEnabled(logger.TraceLevel) { - dump, _ := httputil.DumpRequest(req, false) - log.Trace(string(dump)) - } - - if shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriteCloser(br, conn, conn), cc, req, ro, &pStats, log); err != nil || shouldClose { - return err - } + if size > MaxBodySize { + size = MaxBodySize } + return size } -func (h *Sniffer) serveH2(ctx context.Context, network string, conn net.Conn, ho *HandleOptions) error { - const expectedBody = "SM\r\n\r\n" - - buf := make([]byte, len(expectedBody)) - n, err := io.ReadFull(conn, buf) - if err != nil { - return fmt.Errorf("h2: error reading client preface: %s", err) - } - if string(buf[:n]) != expectedBody { - return errors.New("h2: invalid client preface") - } - - ro := ho.recorderObject - log := ho.log - - ro.Time = time.Time{} - - tr := &http2.Transport{ - DialTLSContext: func(ctx context.Context, nw, addr string, cfg *tls.Config) (net.Conn, error) { - if dial := ho.dialTLS; dial != nil { - return dial(ctx, network, addr, cfg) - } - - cc, err := (&net.Dialer{}).DialContext(ctx, network, addr) - if err != nil { - return nil, err - } - - log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()}) - ro.SrcAddr = cc.LocalAddr().String() - ro.DstAddr = cc.RemoteAddr().String() - - cc = tls.Client(cc, cfg) - return cc, nil - }, - } - defer tr.CloseIdleConnections() - - (&http2.Server{}).ServeConn(conn, &http2.ServeConnOpts{ - Context: ctx, - SawClientPreface: true, - Handler: &h2Handler{ - transport: tr, - recorder: h.Recorder, - recorderOptions: h.RecorderOptions, - recorderObject: ro, - log: log, - }, - }) - return nil -} - -func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, req *http.Request, ro *xrecorder.HandlerRecorderObject, pStats stats.Stats, log logger.Logger) (close bool, err error) { - close = true - - ro2 := &xrecorder.HandlerRecorderObject{} - *ro2 = *ro - ro = ro2 - - ro.Time = time.Now() - log.Infof("%s <-> %s", ro.RemoteAddr, req.Host) - defer func() { - if err != nil { - ro.Err = err.Error() - } - ro.InputBytes = pStats.Get(stats.KindInputBytes) - ro.OutputBytes = pStats.Get(stats.KindOutputBytes) - ro.Duration = time.Since(ro.Time) - if err := ro.Record(ctx, h.Recorder); err != nil { - log.Errorf("record: %v", err) - } - - log.WithFields(map[string]any{ - "duration": time.Since(ro.Time), - "inputBytes": ro.InputBytes, - "outputBytes": ro.OutputBytes, - }).Infof("%s >-< %s", ro.RemoteAddr, req.Host) - }() - - ro.HTTP = &xrecorder.HTTPRecorderObject{ - Host: req.Host, - Proto: req.Proto, - Scheme: req.URL.Scheme, - Method: req.Method, - URI: req.RequestURI, - Request: xrecorder.HTTPRequestRecorderObject{ - ContentLength: req.ContentLength, - Header: req.Header.Clone(), - }, - } - - // HTTP/1.0 - if req.ProtoMajor == 1 && req.ProtoMinor == 0 { - if strings.ToLower(req.Header.Get("Connection")) == "keep-alive" { - req.Header.Del("Connection") - } else { - req.Header.Set("Connection", "close") - } - } - - var reqBody *xhttp.Body - if opts := h.RecorderOptions; opts != nil && opts.HTTPBody { - if req.Body != nil { - bodySize := opts.MaxBodySize - if bodySize <= 0 { - bodySize = DefaultBodySize - } - if bodySize > MaxBodySize { - bodySize = MaxBodySize - } - reqBody = xhttp.NewBody(req.Body, bodySize) - req.Body = reqBody - } - } - - err = req.Write(cc) - - if reqBody != nil { - ro.HTTP.Request.Body = reqBody.Content() - ro.HTTP.Request.ContentLength = reqBody.Length() - } - - if err != nil { - return - } - - br := bufio.NewReader(cc) - var resp *http.Response - for { - xio.SetReadDeadline(cc, time.Now().Add(h.ReadTimeout)) - resp, err = http.ReadResponse(br, req) - if err != nil { - err = fmt.Errorf("read response: %v", err) - return - } - if resp.StatusCode == http.StatusContinue { - resp.Write(rw) - resp.Body.Close() - continue - } - - break - } - defer resp.Body.Close() - xio.SetReadDeadline(cc, time.Time{}) - - ro.HTTP.StatusCode = resp.StatusCode - ro.HTTP.Response.Header = resp.Header - ro.HTTP.Response.ContentLength = resp.ContentLength - - if log.IsLevelEnabled(logger.TraceLevel) { - dump, _ := httputil.DumpResponse(resp, false) - log.Trace(string(dump)) - } - - if resp.StatusCode == http.StatusSwitchingProtocols { - h.handleUpgradeResponse(ctx, rw, cc, req, resp, ro, log) - return - } - - // HTTP/1.0 - if req.ProtoMajor == 1 && req.ProtoMinor == 0 { - if !resp.Close { - resp.Header.Set("Connection", "keep-alive") - } - resp.ProtoMajor = req.ProtoMajor - resp.ProtoMinor = req.ProtoMinor - } - - var respBody *xhttp.Body - if opts := h.RecorderOptions; opts != nil && opts.HTTPBody { - bodySize := opts.MaxBodySize - if bodySize <= 0 { - bodySize = DefaultBodySize - } - if bodySize > MaxBodySize { - bodySize = MaxBodySize - } - respBody = xhttp.NewBody(resp.Body, bodySize) - resp.Body = respBody - } - - err = resp.Write(rw) - - if respBody != nil { - ro.HTTP.Response.Body = respBody.Content() - ro.HTTP.Response.ContentLength = respBody.Length() - } - - if err != nil { - err = fmt.Errorf("write response: %w", err) - return - } - - if resp.ContentLength >= 0 { - close = resp.Close - } - - return -} - -func upgradeType(h http.Header) string { - if !httpguts.HeaderValuesContainsToken(h["Connection"], "Upgrade") { - return "" - } - return h.Get("Upgrade") -} - -func (h *Sniffer) handleUpgradeResponse(ctx context.Context, rw, cc io.ReadWriteCloser, 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) { - return fmt.Errorf("backend tried to switch protocol %q when %q was requested", resUpType, reqUpType) - } - - res.Body = nil - if err := res.Write(rw); err != nil { - return fmt.Errorf("response write: %v", err) - } - - if reqUpType == "websocket" && h.Websocket { - return h.sniffingWebsocketFrame(ctx, rw, cc, ro, log) - } - - // return xnet.Transport(rw, cc) - return xnet.Pipe(ctx, rw, cc) -} - -func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWriter, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error { - errc := make(chan error, 1) - - sampleRate := h.WebsocketSampleRate - if sampleRate == 0 { - sampleRate = DefaultSampleRate - } - if sampleRate < 0 { - sampleRate = math.MaxFloat64 - } - - go func() { - ro2 := &xrecorder.HandlerRecorderObject{} - *ro2 = *ro - ro := ro2 - - limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate)) - - buf := &bytes.Buffer{} - for { - start := time.Now() - - if err := h.copyWebsocketFrame(cc, rw, buf, "client", ro); err != nil { - errc <- err - return - } - - if limiter.Allow() { - ro.Duration = time.Since(start) - ro.Time = time.Now() - if err := ro.Record(ctx, h.Recorder); err != nil { - log.Errorf("record: %v", err) - } - } - } - }() - - go func() { - ro2 := &xrecorder.HandlerRecorderObject{} - *ro2 = *ro - ro := ro2 - - limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate)) - - buf := &bytes.Buffer{} - for { - start := time.Now() - - if err := h.copyWebsocketFrame(rw, cc, buf, "server", ro); err != nil { - errc <- err - return - } - - if limiter.Allow() { - ro.Duration = time.Since(start) - ro.Time = time.Now() - if err := ro.Record(ctx, h.Recorder); err != nil { - log.Errorf("record: %v", err) - } - } - } - }() - - <-errc - return nil -} - -func (h *Sniffer) 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 - } - - 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 := h.RecorderOptions; opts != nil && opts.HTTPBody { - bodySize := opts.MaxBodySize - if bodySize <= 0 { - bodySize = DefaultBodySize - } - if bodySize > MaxBodySize { - bodySize = 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 (h *Sniffer) HandleTLS(ctx context.Context, network string, conn net.Conn, opts ...HandleOption) error { - var ho HandleOptions - for _, opt := range opts { - opt(&ho) - } - - if h.ReadTimeout <= 0 { - h.ReadTimeout = DefaultReadTimeout - } - - buf := new(bytes.Buffer) - clientHello, err := dissector.ParseClientHello(io.TeeReader(conn, buf)) - if err != nil { - return err - } - - log := ho.log - - ro := ho.recorderObject - ro.TLS = &xrecorder.TLSRecorderObject{ - ServerName: clientHello.ServerName, - ClientHello: hex.EncodeToString(buf.Bytes()), - } - if len(clientHello.SupportedProtos) > 0 { - ro.TLS.Proto = clientHello.SupportedProtos[0] - } - - // ctx = xctx.ContextWithClientAddr(ctx, xctx.ClientAddr(ro.RemoteAddr)) - - host := clientHello.ServerName - if host != "" { - if _, _, err := net.SplitHostPort(host); err != nil { - host = net.JoinHostPort(strings.Trim(host, "[]"), "443") - } - ro.Host = host - - if ho.bypass != nil && ho.bypass.Contains(ctx, network, host, bypass.WithService(ho.service)) { - return xbypass.ErrBypass - } - } - - dial := ho.dial - if dial == nil { - dial = (&net.Dialer{}).DialContext - } - cc, err := dial(ctx, network, host) - if err != nil { - return err - } - defer cc.Close() - - log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "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") { - if host == "" { - host = ro.Host - } - if h.MitmBypass == nil || !h.MitmBypass.Contains(ctx, network, host, bypass.WithService(ho.service)) { - return h.terminateTLS(ctx, network, xnet.NewReadWriteConn(io.MultiReader(buf, conn), conn, conn), cc, clientHello, &ho) - } - } - - if _, err := buf.WriteTo(cc); err != nil { - return err - } - - xio.SetReadDeadline(cc, time.Now().Add(h.ReadTimeout)) - serverHello, err := dissector.ParseServerHello(io.TeeReader(cc, buf)) - xio.SetReadDeadline(cc, time.Time{}) - - if serverHello != nil { - ro.TLS.CipherSuite = tls_util.CipherSuite(serverHello.CipherSuite).String() - ro.TLS.CompressionMethod = serverHello.CompressionMethod - if serverHello.Proto != "" { - ro.TLS.Proto = serverHello.Proto - } - if serverHello.Version > 0 { - ro.TLS.Version = tls_util.Version(serverHello.Version).String() - } - } - - if buf.Len() > 0 { - ro.TLS.ServerHello = hex.EncodeToString(buf.Bytes()) - } - - if _, err := buf.WriteTo(conn); err != nil { - return err - } - - log.Infof("%s <-> %s", ro.RemoteAddr, ro.Host) - // xnet.Transport(conn, cc) - xnet.Pipe(ctx, conn, cc) - log.WithFields(map[string]any{ - "duration": time.Since(ro.Time), - }).Infof("%s >-< %s", ro.RemoteAddr, ro.Host) - - return err -} - -func (h *Sniffer) terminateTLS(ctx context.Context, network string, conn, cc net.Conn, clientHello *dissector.ClientHelloInfo, ho *HandleOptions) error { - ro := ho.recorderObject - log := ho.log - - nextProtos := clientHello.SupportedProtos - if h.NegotiatedProtocol != "" { - nextProtos = []string{h.NegotiatedProtocol} - } - - cfg := &tls.Config{ - ServerName: clientHello.ServerName, - NextProtos: nextProtos, - CipherSuites: clientHello.CipherSuites, - } - if cfg.ServerName == "" { - cfg.InsecureSkipVerify = true - } - clientConn := tls.Client(cc, cfg) - if err := clientConn.HandshakeContext(ctx); err != nil { - return err - } - - cs := clientConn.ConnectionState() - ro.TLS.CipherSuite = tls_util.CipherSuite(cs.CipherSuite).String() - ro.TLS.Proto = cs.NegotiatedProtocol - ro.TLS.Version = tls_util.Version(cs.Version).String() - - host := cfg.ServerName +// normalizeHost ensures host contains a port component. If host is already +// in host:port form it is returned unchanged; otherwise defaultPort is appended. +func normalizeHost(host, defaultPort string) string { if host == "" { - if host = cs.PeerCertificates[0].Subject.CommonName; host == "" { - host = ro.Host - } + return host } - if h, _, _ := net.SplitHostPort(host); h != "" { - host = h + if _, _, err := net.SplitHostPort(host); err != nil { + return net.JoinHostPort(strings.Trim(host, "[]"), defaultPort) } - - negotiatedProtocol := cs.NegotiatedProtocol - if h.NegotiatedProtocol != "" { - negotiatedProtocol = h.NegotiatedProtocol - } - nextProtos = nil - if negotiatedProtocol != "" { - nextProtos = []string{negotiatedProtocol} - } - - // cache the tls server handshake record. - wb := &bytes.Buffer{} - conn = xnet.NewReadWriteConn(conn, io.MultiWriter(wb, conn), conn) - - serverConn := tls.Server(conn, &tls.Config{ - NextProtos: nextProtos, - GetCertificate: func(chi *tls.ClientHelloInfo) (*tls.Certificate, error) { - certPool := h.CertPool - if certPool == nil { - certPool = DefaultCertPool - } - serverName := chi.ServerName - if serverName == "" { - serverName = host - } - cert, err := certPool.Get(serverName) - if cert != nil { - pool := x509.NewCertPool() - pool.AddCert(h.Certificate) - if _, err = cert.Verify(x509.VerifyOptions{ - DNSName: serverName, - Roots: pool, - }); err != nil { - log.Warnf("verify cached certificate for %s: %v", serverName, err) - cert = nil - } - } - if cert == nil { - cert, err = tls_util.GenerateCertificate(serverName, 7*24*time.Hour, h.Certificate, h.PrivateKey) - certPool.Put(serverName, cert) - } - if err != nil { - return nil, err - } - - return &tls.Certificate{ - Certificate: [][]byte{cert.Raw}, - PrivateKey: h.PrivateKey, - }, nil - }, - }) - err := serverConn.HandshakeContext(ctx) - if record, _ := dissector.ReadRecord(wb); record != nil { - wb.Reset() - record.WriteTo(wb) - ro.TLS.ServerHello = hex.EncodeToString(wb.Bytes()) - } - if err != nil { - return err - } - - opts := []HandleOption{ - WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { - return clientConn, nil - }), - WithDialTLS(func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error) { - return clientConn, nil - }), - WithRecorderObject(ro), - WithLog(log), - } - return h.HandleHTTP(ctx, network, serverConn, opts...) + return host } -type h2Handler struct { - transport http.RoundTripper - recorder recorder.Recorder - recorderOptions *recorder.Options - recorderObject *xrecorder.HandlerRecorderObject - log logger.Logger -} - -func (h *h2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { - log := h.log - - ro := &xrecorder.HandlerRecorderObject{} - *ro = *h.recorderObject - ro.Time = time.Now() - - var err error - log.Infof("%s <-> %s", ro.RemoteAddr, r.Host) - defer func() { - ro.Duration = time.Since(ro.Time) - if err != nil { - ro.Err = err.Error() - } - if err := ro.Record(r.Context(), h.recorder); err != nil { - log.Errorf("record: %v", err) - } - - log.WithFields(map[string]any{ - "duration": time.Since(ro.Time), - }).Infof("%s >-< %s", ro.RemoteAddr, r.Host) - }() - - if clientIP := xhttp.GetClientIP(r); clientIP != nil { - ro.ClientIP = clientIP.String() - } - ro.HTTP = &xrecorder.HTTPRecorderObject{ - Host: r.Host, - Proto: r.Proto, - Scheme: "https", - Method: r.Method, - URI: r.RequestURI, - Request: xrecorder.HTTPRequestRecorderObject{ - ContentLength: r.ContentLength, - Header: r.Header.Clone(), - }, - } - - if log.IsLevelEnabled(logger.TraceLevel) { - dump, _ := httputil.DumpRequest(r, false) - log.Trace(string(dump)) - } - - url := r.URL - url.Scheme = "https" - url.Host = r.Host - req := &http.Request{ - Method: r.Method, - URL: url, - Host: r.Host, - Header: r.Header, - Body: r.Body, - ContentLength: r.ContentLength, - Trailer: r.Trailer, - } - - var reqBody *xhttp.Body - if opts := h.recorderOptions; opts != nil && opts.HTTPBody { - if req.Body != nil { - bodySize := opts.MaxBodySize - if bodySize <= 0 { - bodySize = DefaultBodySize - } - if bodySize > MaxBodySize { - bodySize = MaxBodySize - } - - reqBody = xhttp.NewBody(req.Body, bodySize) - req.Body = reqBody - } - } - - resp, err := h.transport.RoundTrip(req.WithContext(r.Context())) - if reqBody != nil { - ro.HTTP.Request.Body = reqBody.Content() - ro.HTTP.Request.ContentLength = reqBody.Length() - } - if err != nil { - log.Error(err) - w.WriteHeader(http.StatusServiceUnavailable) - return - } - defer resp.Body.Close() - - ro.HTTP.StatusCode = resp.StatusCode - ro.HTTP.Response.Header = resp.Header - ro.HTTP.Response.ContentLength = resp.ContentLength - - if log.IsLevelEnabled(logger.TraceLevel) { - dump, _ := httputil.DumpResponse(resp, false) - log.Trace(string(dump)) - } - - h.setHeader(w, resp.Header) - w.WriteHeader(resp.StatusCode) - - var respBody *xhttp.Body - if opts := h.recorderOptions; opts != nil && opts.HTTPBody { - bodySize := opts.MaxBodySize - if bodySize <= 0 { - bodySize = DefaultBodySize - } - if bodySize > MaxBodySize { - bodySize = MaxBodySize - } - respBody = xhttp.NewBody(resp.Body, bodySize) - resp.Body = respBody - } - - io.Copy(w, resp.Body) - - if respBody != nil { - ro.HTTP.Response.Body = respBody.Content() - ro.HTTP.Response.ContentLength = respBody.Length() - } -} - -func (h *h2Handler) setHeader(w http.ResponseWriter, header http.Header) { - for k, v := range header { - for i := range v { - w.Header().Add(k, v[i]) - } - } +// effectiveReadTimeout returns the read timeout from the Sniffer's ReadTimeout, +// falling back to DefaultReadTimeout. +func (h *Sniffer) effectiveReadTimeout() time.Duration { + if h.ReadTimeout > 0 { + return h.ReadTimeout + } + return DefaultReadTimeout } diff --git a/internal/util/sniffing/sniffer_h2.go b/internal/util/sniffing/sniffer_h2.go new file mode 100644 index 00000000..15efb423 --- /dev/null +++ b/internal/util/sniffing/sniffer_h2.go @@ -0,0 +1,189 @@ +package sniffing + +import ( + "context" + "crypto/tls" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httputil" + "time" + + "github.com/go-gost/core/logger" + "github.com/go-gost/core/recorder" + xhttp "github.com/go-gost/x/internal/net/http" + xrecorder "github.com/go-gost/x/recorder" + "golang.org/x/net/http2" +) + +// serveH2 handles HTTP/2 connections by reading the client preface and +// proxying requests through an http2.Transport. +func (h *Sniffer) serveH2(ctx context.Context, network string, conn net.Conn, ho *HandleOptions) error { + const expectedBody = "SM\r\n\r\n" + + buf := make([]byte, len(expectedBody)) + n, err := io.ReadFull(conn, buf) + if err != nil { + return fmt.Errorf("h2: error reading client preface: %s", err) + } + if string(buf[:n]) != expectedBody { + return errors.New("h2: invalid client preface") + } + + ro := ho.recorderObject + log := ho.log + + ro.Time = time.Time{} + + tr := &http2.Transport{ + DialTLSContext: func(ctx context.Context, nw, addr string, cfg *tls.Config) (net.Conn, error) { + if dial := ho.dialTLS; dial != nil { + return dial(ctx, network, addr, cfg) + } + + cc, err := (&net.Dialer{}).DialContext(ctx, network, addr) + if err != nil { + return nil, err + } + + log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()}) + ro.SrcAddr = cc.LocalAddr().String() + ro.DstAddr = cc.RemoteAddr().String() + + cc = tls.Client(cc, cfg) + return cc, nil + }, + } + defer tr.CloseIdleConnections() + + (&http2.Server{}).ServeConn(conn, &http2.ServeConnOpts{ + Context: ctx, + SawClientPreface: true, + Handler: &h2Handler{ + transport: tr, + recorder: h.Recorder, + recorderOptions: h.RecorderOptions, + recorderObject: ro, + log: log, + }, + }) + return nil +} + +// h2Handler is an http.Handler that proxies HTTP/2 requests through an +// http2.Transport while recording request and response metadata. +type h2Handler struct { + transport http.RoundTripper + recorder recorder.Recorder + recorderOptions *recorder.Options + recorderObject *xrecorder.HandlerRecorderObject + log logger.Logger +} + +func (h *h2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + log := h.log + + ro := &xrecorder.HandlerRecorderObject{} + *ro = *h.recorderObject + ro.Time = time.Now() + + var err error + log.Infof("%s <-> %s", ro.RemoteAddr, r.Host) + defer func() { + ro.Duration = time.Since(ro.Time) + if err != nil { + ro.Err = err.Error() + } + if rerr := ro.Record(r.Context(), h.recorder); rerr != nil { + log.Errorf("record: %v", rerr) + } + + log.WithFields(map[string]any{ + "duration": time.Since(ro.Time), + }).Infof("%s >-< %s", ro.RemoteAddr, r.Host) + }() + + if clientIP := xhttp.GetClientIP(r); clientIP != nil { + ro.ClientIP = clientIP.String() + } + ro.HTTP = &xrecorder.HTTPRecorderObject{ + Host: r.Host, + Proto: r.Proto, + Scheme: "https", + Method: r.Method, + URI: r.RequestURI, + Request: xrecorder.HTTPRequestRecorderObject{ + ContentLength: r.ContentLength, + Header: r.Header.Clone(), + }, + } + + if log.IsLevelEnabled(logger.TraceLevel) { + dump, _ := httputil.DumpRequest(r, false) + log.Trace(string(dump)) + } + + url := r.URL + url.Scheme = "https" + url.Host = r.Host + req := &http.Request{ + Method: r.Method, + URL: url, + Host: r.Host, + Header: r.Header, + Body: r.Body, + ContentLength: r.ContentLength, + Trailer: r.Trailer, + } + + var reqBody *xhttp.Body + if bodySize := clampBodySize(h.recorderOptions); bodySize > 0 && req.Body != nil { + reqBody = xhttp.NewBody(req.Body, bodySize) + req.Body = reqBody + } + + resp, respErr := h.transport.RoundTrip(req.WithContext(r.Context())) + if reqBody != nil { + ro.HTTP.Request.Body = reqBody.Content() + ro.HTTP.Request.ContentLength = reqBody.Length() + } + err = respErr + if err != nil { + log.Error(err) + w.WriteHeader(http.StatusServiceUnavailable) + return + } + defer resp.Body.Close() + + ro.HTTP.StatusCode = resp.StatusCode + ro.HTTP.Response.Header = resp.Header + ro.HTTP.Response.ContentLength = resp.ContentLength + + if log.IsLevelEnabled(logger.TraceLevel) { + dump, _ := httputil.DumpResponse(resp, false) + log.Trace(string(dump)) + } + + h.setHeader(w, resp.Header) + w.WriteHeader(resp.StatusCode) + + if bodySize := clampBodySize(h.recorderOptions); bodySize > 0 { + respBody := xhttp.NewBody(resp.Body, bodySize) + resp.Body = respBody + io.Copy(w, resp.Body) + ro.HTTP.Response.Body = respBody.Content() + ro.HTTP.Response.ContentLength = respBody.Length() + } else { + io.Copy(w, resp.Body) + } +} + +func (h *h2Handler) setHeader(w http.ResponseWriter, header http.Header) { + for k, v := range header { + for i := range v { + w.Header().Add(k, v[i]) + } + } +} diff --git a/internal/util/sniffing/sniffer_http.go b/internal/util/sniffing/sniffer_http.go new file mode 100644 index 00000000..afcde444 --- /dev/null +++ b/internal/util/sniffing/sniffer_http.go @@ -0,0 +1,282 @@ +package sniffing + +import ( + "bufio" + "context" + "errors" + "io" + "net" + "net/http" + "net/http/httputil" + "strings" + "time" + + "github.com/go-gost/core/bypass" + "github.com/go-gost/core/logger" + "github.com/go-gost/core/observer/stats" + xbypass "github.com/go-gost/x/bypass" + 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" + xstats "github.com/go-gost/x/observer/stats" + stats_wrapper "github.com/go-gost/x/observer/stats/wrapper" + xrecorder "github.com/go-gost/x/recorder" +) + +// HandleHTTP sniffs and proxies an HTTP connection. It reads the initial +// request, applies bypass rules, and forwards traffic to the upstream. +func (h *Sniffer) HandleHTTP(ctx context.Context, network string, conn net.Conn, opts ...HandleOption) error { + var ho HandleOptions + for _, opt := range opts { + opt(&ho) + } + + readTimeout := h.effectiveReadTimeout() + + pStats := xstats.Stats{} + conn = stats_wrapper.WrapConn(conn, &pStats) + + br := bufio.NewReader(conn) + req, err := http.ReadRequest(br) + if err != nil { + return err + } + + log := ho.log + if log.IsLevelEnabled(logger.TraceLevel) { + dump, _ := httputil.DumpRequest(req, false) + log.Trace(string(dump)) + } + + ro := ho.recorderObject + ro.HTTP = &xrecorder.HTTPRecorderObject{ + Host: req.Host, + Proto: req.Proto, + Scheme: req.URL.Scheme, + Method: req.Method, + URI: req.RequestURI, + Request: xrecorder.HTTPRequestRecorderObject{ + ContentLength: req.ContentLength, + Header: req.Header.Clone(), + }, + } + + if clientIP := xhttp.GetClientIP(req); clientIP != nil { + ro.ClientIP = clientIP.String() + ctx = xctx.ContextWithSrcAddr(ctx, &net.TCPAddr{IP: clientIP}) + } + + // http/2 + if req.Method == "PRI" && len(req.Header) == 0 && req.URL.Path == "*" && req.Proto == "HTTP/2.0" { + return h.serveH2(ctx, network, xnet.NewReadWriteConn(br, conn, conn), &ho) + } + + host := normalizeHost(req.Host, "80") + if host != "" { + ro.Host = host + + log = log.WithFields(map[string]any{ + "host": host, + }) + + if ho.bypass != nil && ho.bypass.Contains(ctx, network, host, bypass.WithService(ho.service)) { + return xbypass.ErrBypass + } + } + + dial := ho.dial + if dial == nil { + dial = (&net.Dialer{}).DialContext + } + cc, err := dial(ctx, network, host) + if err != nil { + return err + } + defer cc.Close() + + log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "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, readTimeout, ro, &pStats, log) + if err != nil || shouldClose { + return err + } + + for { + pStats.Reset() + + req, err := http.ReadRequest(br) + if err != nil { + if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) { + return nil + } + return err + } + + if log.IsLevelEnabled(logger.TraceLevel) { + dump, _ := httputil.DumpRequest(req, false) + log.Trace(string(dump)) + } + + if shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriteCloser(br, conn, conn), cc, req, readTimeout, ro, &pStats, log); err != nil || shouldClose { + return err + } + } +} + +// httpRoundTrip forwards a single HTTP request/response pair and records +// traffic metadata. Returns whether the connection should be closed. +func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, req *http.Request, readTimeout time.Duration, ro *xrecorder.HandlerRecorderObject, pStats stats.Stats, log logger.Logger) (close bool, err error) { + close = true + + ro2 := &xrecorder.HandlerRecorderObject{} + *ro2 = *ro + ro = ro2 + + ro.Time = time.Now() + log.Infof("%s <-> %s", ro.RemoteAddr, req.Host) + defer func() { + if err != nil { + ro.Err = err.Error() + } + ro.InputBytes = pStats.Get(stats.KindInputBytes) + ro.OutputBytes = pStats.Get(stats.KindOutputBytes) + ro.Duration = time.Since(ro.Time) + if rerr := ro.Record(ctx, h.Recorder); rerr != nil { + log.Errorf("record: %v", rerr) + } + + log.WithFields(map[string]any{ + "duration": time.Since(ro.Time), + "inputBytes": ro.InputBytes, + "outputBytes": ro.OutputBytes, + }).Infof("%s >-< %s", ro.RemoteAddr, req.Host) + }() + + ro.HTTP = &xrecorder.HTTPRecorderObject{ + Host: req.Host, + Proto: req.Proto, + Scheme: req.URL.Scheme, + Method: req.Method, + URI: req.RequestURI, + Request: xrecorder.HTTPRequestRecorderObject{ + ContentLength: req.ContentLength, + Header: req.Header.Clone(), + }, + } + + // HTTP/1.0 + if req.ProtoMajor == 1 && req.ProtoMinor == 0 { + if strings.ToLower(req.Header.Get("Connection")) == "keep-alive" { + req.Header.Del("Connection") + } else { + req.Header.Set("Connection", "close") + } + } + + var reqBody *xhttp.Body + if bodySize := clampBodySize(h.RecorderOptions); bodySize > 0 && req.Body != nil { + reqBody = xhttp.NewBody(req.Body, bodySize) + req.Body = reqBody + } + + err = req.Write(cc) + + if reqBody != nil { + ro.HTTP.Request.Body = reqBody.Content() + ro.HTTP.Request.ContentLength = reqBody.Length() + } + + if err != nil { + return + } + + br := bufio.NewReader(cc) + var resp *http.Response + for { + xio.SetReadDeadline(cc, time.Now().Add(readTimeout)) + resp, err = http.ReadResponse(br, req) + if err != nil { + err = wrapErr("read response", err) + return + } + if resp.StatusCode == http.StatusContinue { + resp.Write(rw) + resp.Body.Close() + continue + } + + break + } + defer resp.Body.Close() + xio.SetReadDeadline(cc, time.Time{}) + + ro.HTTP.StatusCode = resp.StatusCode + ro.HTTP.Response.Header = resp.Header + ro.HTTP.Response.ContentLength = resp.ContentLength + + if log.IsLevelEnabled(logger.TraceLevel) { + dump, _ := httputil.DumpResponse(resp, false) + log.Trace(string(dump)) + } + + if resp.StatusCode == http.StatusSwitchingProtocols { + h.handleUpgradeResponse(ctx, rw, cc, req, resp, ro, log) + return + } + + // HTTP/1.0 + if req.ProtoMajor == 1 && req.ProtoMinor == 0 { + if !resp.Close { + resp.Header.Set("Connection", "keep-alive") + } + resp.ProtoMajor = req.ProtoMajor + resp.ProtoMinor = req.ProtoMinor + } + + var respBody *xhttp.Body + if bodySize := clampBodySize(h.RecorderOptions); bodySize > 0 { + respBody = xhttp.NewBody(resp.Body, bodySize) + resp.Body = respBody + } + + err = resp.Write(rw) + + if respBody != nil { + ro.HTTP.Response.Body = respBody.Content() + ro.HTTP.Response.ContentLength = respBody.Length() + } + + if err != nil { + err = wrapErr("write response", err) + return + } + + if resp.ContentLength >= 0 { + close = resp.Close + } + + return +} + +// wrapErr formats an error with a context message. +func wrapErr(msg string, err error) error { + return &wrappedErr{msg: msg, err: err} +} + +type wrappedErr struct { + msg string + err error +} + +func (e *wrappedErr) Error() string { + return e.msg + ": " + e.err.Error() +} + +func (e *wrappedErr) Unwrap() error { + return e.err +} diff --git a/internal/util/sniffing/sniffer_rewrite.go b/internal/util/sniffing/sniffer_rewrite.go new file mode 100644 index 00000000..da3075e4 --- /dev/null +++ b/internal/util/sniffing/sniffer_rewrite.go @@ -0,0 +1,45 @@ +package sniffing + +import ( + "context" + "fmt" + "io" + "net/http" + "strings" + + "github.com/go-gost/core/logger" + xnet "github.com/go-gost/x/internal/net" + xrecorder "github.com/go-gost/x/recorder" + "golang.org/x/net/http/httpguts" +) + +// upgradeType extracts the Upgrade protocol from an HTTP header if the +// Connection header contains the "Upgrade" token. +func upgradeType(h http.Header) string { + if !httpguts.HeaderValuesContainsToken(h["Connection"], "Upgrade") { + return "" + } + return h.Get("Upgrade") +} + +// handleUpgradeResponse handles HTTP 101 Switching Protocols responses. +// It validates the upgrade type, writes the response, and either begins +// WebSocket frame sniffing or falls through to bidirectional copy. +func (h *Sniffer) handleUpgradeResponse(ctx context.Context, rw, cc io.ReadWriteCloser, 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) { + return fmt.Errorf("backend tried to switch protocol %q when %q was requested", resUpType, reqUpType) + } + + res.Body = nil + if err := res.Write(rw); err != nil { + return fmt.Errorf("response write: %v", err) + } + + if reqUpType == "websocket" && h.Websocket { + return h.sniffingWebsocketFrame(ctx, rw, cc, ro, log) + } + + return xnet.Pipe(ctx, rw, cc) +} diff --git a/internal/util/sniffing/sniffer_test.go b/internal/util/sniffing/sniffer_test.go new file mode 100644 index 00000000..4c8d788f --- /dev/null +++ b/internal/util/sniffing/sniffer_test.go @@ -0,0 +1,682 @@ +package sniffing + +import ( + "bufio" + "bytes" + "context" + "crypto/tls" + "errors" + "io" + "net" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/go-gost/core/bypass" + "github.com/go-gost/core/recorder" + xlogger "github.com/go-gost/x/logger" + xrecorder "github.com/go-gost/x/recorder" +) + +// ============================================================================= +// Test Helpers (local mocks) +// ============================================================================= + +type noopRecorder struct{} + +func (n *noopRecorder) Record(_ context.Context, _ []byte, _ ...recorder.RecordOption) error { + return nil +} + +type mockBypass struct { + contains bool + whitelist bool +} + +func (m *mockBypass) Contains(_ context.Context, _, _ string, _ ...bypass.Option) bool { + return m.contains +} + +func (m *mockBypass) IsWhitelist() bool { return m.whitelist } + +// ============================================================================= +// Pure Function Tests +// ============================================================================= + +func TestClampBodySize(t *testing.T) { + tests := []struct { + name string + opts *recorder.Options + want int + }{ + {"nil opts", nil, 0}, + {"HTTPBody disabled", &recorder.Options{HTTPBody: false, MaxBodySize: 1000}, 0}, + {"zero MaxBodySize defaults to DefaultBodySize", &recorder.Options{HTTPBody: true, MaxBodySize: 0}, DefaultBodySize}, + {"negative MaxBodySize defaults", &recorder.Options{HTTPBody: true, MaxBodySize: -1}, DefaultBodySize}, + {"valid MaxBodySize within bounds", &recorder.Options{HTTPBody: true, MaxBodySize: 1024}, 1024}, + {"MaxBodySize capped at MaxBodySize", &recorder.Options{HTTPBody: true, MaxBodySize: 10 * 1024 * 1024}, MaxBodySize}, + {"exact MaxBodySize", &recorder.Options{HTTPBody: true, MaxBodySize: MaxBodySize}, MaxBodySize}, + {"exact DefaultBodySize", &recorder.Options{HTTPBody: true, MaxBodySize: DefaultBodySize}, DefaultBodySize}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := clampBodySize(tt.opts); got != tt.want { + t.Errorf("clampBodySize() = %d, want %d", got, tt.want) + } + }) + } +} + +func TestNormalizeHost(t *testing.T) { + tests := []struct { + name string + host string + defaultPort string + want string + }{ + {"empty host", "", "443", ""}, + {"host with port", "example.com:8080", "443", "example.com:8080"}, + {"host without port, default 80", "example.com", "80", "example.com:80"}, + {"host without port, default 443", "example.com", "443", "example.com:443"}, + {"IPv4 with port", "127.0.0.1:9090", "80", "127.0.0.1:9090"}, + {"IPv4 without port", "127.0.0.1", "80", "127.0.0.1:80"}, + {"IPv6 bracketed with port", "[::1]:8080", "443", "[::1]:8080"}, + {"IPv6 bracketed without port", "[::1]", "443", "[::1]:443"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := normalizeHost(tt.host, tt.defaultPort) + if got != tt.want { + t.Errorf("normalizeHost(%q, %q) = %q, want %q", tt.host, tt.defaultPort, got, tt.want) + } + }) + } +} + +func TestEffectiveReadTimeout(t *testing.T) { + tests := []struct { + name string + snifferTO time.Duration + want time.Duration + }{ + {"zero defaults to DefaultReadTimeout", 0, DefaultReadTimeout}, + {"sniffer timeout set", 10 * time.Second, 10 * time.Second}, + {"another sniffer timeout", 5 * time.Second, 5 * time.Second}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + h := &Sniffer{ReadTimeout: tt.snifferTO} + got := h.effectiveReadTimeout() + if got != tt.want { + t.Errorf("effectiveReadTimeout() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestUpgradeType(t *testing.T) { + tests := []struct { + name string + header http.Header + want string + }{ + { + "websocket upgrade", + http.Header{"Connection": {"Upgrade"}, "Upgrade": {"websocket"}}, + "websocket", + }, + { + "h2c upgrade", + http.Header{"Connection": {"Upgrade"}, "Upgrade": {"h2c"}}, + "h2c", + }, + { + "no upgrade header", + http.Header{"Connection": {"keep-alive"}}, + "", + }, + { + "empty header", + http.Header{}, + "", + }, + { + "connection has upgrade but no upgrade header", + http.Header{"Connection": {"Upgrade"}}, + "", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := upgradeType(tt.header); got != tt.want { + t.Errorf("upgradeType() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestWrapErr(t *testing.T) { + err := wrapErr("read response", io.EOF) + if err == nil { + t.Fatal("expected error, got nil") + } + if err.Error() != "read response: EOF" { + t.Errorf("error = %q, want %q", err.Error(), "read response: EOF") + } + if !errors.Is(err, io.EOF) { + t.Error("errors.Is should unwrap to EOF") + } +} + +// ============================================================================= +// HandleOptions Tests +// ============================================================================= + +func TestWithService(t *testing.T) { + opts := &HandleOptions{} + WithService("mysvc")(opts) + if opts.service != "mysvc" { + t.Errorf("service = %q, want %q", opts.service, "mysvc") + } +} + +func TestWithDial(t *testing.T) { + opts := &HandleOptions{} + dial := func(ctx context.Context, network, address string) (net.Conn, error) { + return nil, nil + } + WithDial(dial)(opts) + if opts.dial == nil { + t.Error("dial should be set") + } +} + +func TestWithDialTLS(t *testing.T) { + opts := &HandleOptions{} + dialTLS := func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error) { + return nil, nil + } + WithDialTLS(dialTLS)(opts) + if opts.dialTLS == nil { + t.Error("dialTLS should be set") + } +} + +func TestWithBypass(t *testing.T) { + opts := &HandleOptions{} + bp := &mockBypass{contains: true} + WithBypass(bp)(opts) + if opts.bypass != bp { + t.Errorf("bypass not set") + } +} + +func TestWithRecorderObject(t *testing.T) { + opts := &HandleOptions{} + ro := &xrecorder.HandlerRecorderObject{} + WithRecorderObject(ro)(opts) + if opts.recorderObject != ro { + t.Errorf("recorderObject not set") + } +} + +func TestWithLog(t *testing.T) { + opts := &HandleOptions{} + log := xlogger.Nop() + WithLog(log)(opts) + if opts.log != log { + t.Errorf("log not set") + } +} + +// ============================================================================= +// HandleHTTP Integration Tests +// ============================================================================= + +func TestHandleHTTP_BasicProxy(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte("OK from upstream")) + })) + defer upstream.Close() + + h := &Sniffer{ + ReadTimeout: 5 * time.Second, + Recorder: &noopRecorder{}, + } + ro := &xrecorder.HandlerRecorderObject{} + + clientConn, serverConn := net.Pipe() + defer clientConn.Close() + defer serverConn.Close() + + errCh := make(chan error, 1) + go func() { + errCh <- h.HandleHTTP(context.Background(), "tcp", serverConn, + WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { + return net.Dial("tcp", upstream.Listener.Addr().String()) + }), + WithRecorderObject(ro), + WithLog(xlogger.Nop()), + ) + }() + + req, _ := http.NewRequest("GET", "http://example.com/", nil) + if err := req.Write(clientConn); err != nil { + t.Fatal(err) + } + + resp, err := http.ReadResponse(bufio.NewReader(clientConn), req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusOK) + } + + clientConn.Close() + if err := <-errCh; err != nil { + t.Logf("HandleHTTP returned: %v", err) + } + + if ro.SrcAddr == "" { + t.Error("SrcAddr should be populated") + } + if ro.DstAddr == "" { + t.Error("DstAddr should be populated") + } +} + +func TestHandleHTTP_HTTP2Detection(t *testing.T) { + h := &Sniffer{ + ReadTimeout: 5 * time.Second, + Recorder: &noopRecorder{}, + } + ro := &xrecorder.HandlerRecorderObject{ + RemoteAddr: "127.0.0.1:12345", + } + + clientConn, serverConn := net.Pipe() + defer clientConn.Close() + defer serverConn.Close() + + errCh := make(chan error, 1) + go func() { + errCh <- h.HandleHTTP(context.Background(), "tcp", serverConn, + WithRecorderObject(ro), + WithLog(xlogger.Nop()), + ) + }() + + clientConn.Write([]byte("PRI * HTTP/2.0\r\n\r\n")) + + clientConn.Close() + + err := <-errCh + if err == nil { + t.Log("HandleHTTP returned nil (expected error from incomplete h2 preface)") + } +} + +func TestHandleHTTP_DialError(t *testing.T) { + h := &Sniffer{ + ReadTimeout: 5 * time.Second, + Recorder: &noopRecorder{}, + } + ro := &xrecorder.HandlerRecorderObject{} + + clientConn, serverConn := net.Pipe() + defer clientConn.Close() + defer serverConn.Close() + + errCh := make(chan error, 1) + go func() { + errCh <- h.HandleHTTP(context.Background(), "tcp", serverConn, + WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { + return nil, io.ErrUnexpectedEOF + }), + WithRecorderObject(ro), + WithLog(xlogger.Nop()), + ) + }() + + req, _ := http.NewRequest("GET", "http://example.com/", nil) + req.Write(clientConn) + clientConn.Close() // unblock server-side ReadRequest loop + + err := <-errCh + if err == nil { + t.Error("expected error from failed dial, got nil") + } +} + +func TestHandleHTTP_HTTP10Request(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte("HTTP/1.0 OK")) + })) + defer upstream.Close() + + h := &Sniffer{ + ReadTimeout: 5 * time.Second, + Recorder: &noopRecorder{}, + } + ro := &xrecorder.HandlerRecorderObject{} + + clientConn, serverConn := net.Pipe() + defer clientConn.Close() + defer serverConn.Close() + + errCh := make(chan error, 1) + go func() { + errCh <- h.HandleHTTP(context.Background(), "tcp", serverConn, + WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { + return net.Dial("tcp", upstream.Listener.Addr().String()) + }), + WithRecorderObject(ro), + WithLog(xlogger.Nop()), + ) + }() + + clientConn.Write([]byte("GET / HTTP/1.0\r\nHost: example.com\r\nConnection: keep-alive\r\n\r\n")) + + resp, err := http.ReadResponse(bufio.NewReader(clientConn), nil) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusOK) + } + + clientConn.Close() + <-errCh +} + +// ============================================================================= +// copyWebsocketFrame Tests +// ============================================================================= + +func TestCopyWebsocketFrame_Basic(t *testing.T) { + h := &Sniffer{ + Recorder: &noopRecorder{}, + RecorderOptions: &recorder.Options{HTTPBody: false}, + } + + payload := []byte("hi") + mask := []byte{0xAA, 0xBB, 0xCC, 0xDD} + maskedPayload := make([]byte, len(payload)) + for i := range payload { + maskedPayload[i] = payload[i] ^ mask[i%4] + } + + var frame bytes.Buffer + frame.WriteByte(0x81) // FIN + text opcode + frame.WriteByte(0x82) // MASK + len=2 + frame.Write(mask) + frame.Write(maskedPayload) + + w := &bytes.Buffer{} + ro := &xrecorder.HandlerRecorderObject{} + + err := h.copyWebsocketFrame(w, &frame, &bytes.Buffer{}, "client", ro) + if err != nil { + t.Fatal(err) + } + + if ro.Websocket == nil { + t.Fatal("Websocket recorder object should be populated") + } + if ro.Websocket.OpCode != 1 { + t.Errorf("opcode = %d, want 1 (text)", ro.Websocket.OpCode) + } + if !ro.Websocket.Masked { + t.Error("client frame should be marked as masked") + } + if ro.InputBytes == 0 { + t.Error("InputBytes should be non-zero for client direction") + } +} + +func TestCopyWebsocketFrame_WithBodyRecording(t *testing.T) { + h := &Sniffer{ + Recorder: &noopRecorder{}, + RecorderOptions: &recorder.Options{HTTPBody: true, MaxBodySize: 1024}, + } + + payload := []byte("hello ws") + var frame bytes.Buffer + frame.WriteByte(0x81) // FIN + text + frame.WriteByte(byte(len(payload))) // MASK=0, len + frame.Write(payload) + + w := &bytes.Buffer{} + ro := &xrecorder.HandlerRecorderObject{} + + err := h.copyWebsocketFrame(w, &frame, &bytes.Buffer{}, "server", ro) + if err != nil { + t.Fatal(err) + } + + if ro.Websocket == nil { + t.Fatal("Websocket recorder object should be populated") + } + if string(ro.Websocket.Payload) != "hello ws" { + t.Errorf("payload = %q, want %q", ro.Websocket.Payload, "hello ws") + } + if ro.Websocket.Masked { + t.Error("server frame should be marked as unmasked") + } + if ro.OutputBytes == 0 { + t.Error("OutputBytes should be non-zero for server direction") + } +} + +func TestCopyWebsocketFrame_EmptyPayload(t *testing.T) { + h := &Sniffer{ + Recorder: &noopRecorder{}, + RecorderOptions: &recorder.Options{HTTPBody: true, MaxBodySize: 1024}, + } + + var frame bytes.Buffer + frame.WriteByte(0x81) // FIN + text + frame.WriteByte(0x00) // MASK=0, len=0 + + w := &bytes.Buffer{} + ro := &xrecorder.HandlerRecorderObject{} + + err := h.copyWebsocketFrame(w, &frame, &bytes.Buffer{}, "server", ro) + if err != nil { + t.Fatal(err) + } + + if ro.Websocket == nil { + t.Fatal("Websocket recorder object should be populated") + } + if ro.Websocket.Length != 0 { + t.Errorf("payload length = %d, want 0", ro.Websocket.Length) + } +} + +func TestCopyWebsocketFrame_ServerWithoutBodyRecording(t *testing.T) { + h := &Sniffer{ + Recorder: &noopRecorder{}, + RecorderOptions: &recorder.Options{HTTPBody: false}, + } + + payload := []byte("data") + var frame bytes.Buffer + frame.WriteByte(0x82) // FIN + binary + frame.WriteByte(byte(len(payload))) // MASK=0, len + frame.Write(payload) + + w := &bytes.Buffer{} + ro := &xrecorder.HandlerRecorderObject{} + + err := h.copyWebsocketFrame(w, &frame, &bytes.Buffer{}, "server", ro) + if err != nil { + t.Fatal(err) + } + + if ro.Websocket == nil { + t.Fatal("Websocket recorder object should be populated") + } + if ro.Websocket.Payload != nil { + t.Error("payload should be nil when body recording is disabled") + } + if ro.Websocket.OpCode != 2 { + t.Errorf("opcode = %d, want 2 (binary)", ro.Websocket.OpCode) + } + if ro.OutputBytes == 0 { + t.Error("OutputBytes should be non-zero for server direction") + } + if ro.InputBytes != 0 { + t.Error("InputBytes should be zero for server direction") + } +} + +func TestCopyWebsocketFrame_MaskedFrame(t *testing.T) { + h := &Sniffer{ + Recorder: &noopRecorder{}, + RecorderOptions: &recorder.Options{HTTPBody: false}, + } + + payload := []byte("secret") + mask := []byte{0x12, 0x34, 0x56, 0x78} + maskedPayload := make([]byte, len(payload)) + for i := range payload { + maskedPayload[i] = payload[i] ^ mask[i%4] + } + + var frame bytes.Buffer + frame.WriteByte(0x81) // FIN + text + frame.WriteByte(0x80 | byte(len(payload))) // MASK=1, len + frame.Write(mask) + frame.Write(maskedPayload) + + w := &bytes.Buffer{} + ro := &xrecorder.HandlerRecorderObject{} + + err := h.copyWebsocketFrame(w, &frame, &bytes.Buffer{}, "client", ro) + if err != nil { + t.Fatal(err) + } + + if ro.Websocket == nil { + t.Fatal("Websocket recorder object should be populated") + } + if !ro.Websocket.Masked { + t.Error("client frame should be masked") + } + written := w.Bytes() + if len(written) < 2+len(payload) { + t.Fatalf("written frame too short: %d bytes", len(written)) + } + if written[0] != 0x81 { + t.Errorf("first byte = 0x%02x, want 0x81", written[0]) + } + if written[1]&0x80 != 0x80 { + t.Error("mask bit should be preserved") + } +} + +// ============================================================================= +// handleUpgradeResponse Tests +// ============================================================================= + +func TestHandleUpgradeResponse_TypeMismatch(t *testing.T) { + h := &Sniffer{} + req, _ := http.NewRequest("GET", "http://example.com/", nil) + req.Header.Set("Connection", "Upgrade") + req.Header.Set("Upgrade", "websocket") + + res := &http.Response{ + Header: http.Header{ + "Connection": {"Upgrade"}, + "Upgrade": {"h2c"}, + }, + } + + clientConn, serverConn := net.Pipe() + defer clientConn.Close() + defer serverConn.Close() + + err := h.handleUpgradeResponse(context.Background(), serverConn, clientConn, req, res, + &xrecorder.HandlerRecorderObject{}, xlogger.Nop()) + if err == nil { + t.Fatal("expected type mismatch error, got nil") + } +} + +func TestHandleUpgradeResponse_NonWebsocket(t *testing.T) { + h := &Sniffer{Websocket: false} + + req, _ := http.NewRequest("GET", "http://example.com/", nil) + req.Header.Set("Connection", "Upgrade") + req.Header.Set("Upgrade", "h2c") + + res := &http.Response{ + StatusCode: http.StatusSwitchingProtocols, + ProtoMajor: 1, + ProtoMinor: 1, + Header: http.Header{ + "Connection": {"Upgrade"}, + "Upgrade": {"h2c"}, + }, + } + + errCh := make(chan error, 1) + clientConn, serverConn := net.Pipe() + defer clientConn.Close() + defer serverConn.Close() + + go func() { + errCh <- h.handleUpgradeResponse(context.Background(), serverConn, clientConn, req, res, + &xrecorder.HandlerRecorderObject{}, xlogger.Nop()) + }() + + br := bufio.NewReader(clientConn) + readResp, readErr := http.ReadResponse(br, req) + if readErr != nil { + t.Fatalf("reading upgrade response: %v", readErr) + } + if readResp.StatusCode != http.StatusSwitchingProtocols { + t.Errorf("status = %d, want %d", readResp.StatusCode, http.StatusSwitchingProtocols) + } + + clientConn.Close() + serverConn.Close() + <-errCh +} + +// ============================================================================= +// serveH2 unit test (client preface) +// ============================================================================= + +func TestServeH2_InvalidPreface(t *testing.T) { + h := &Sniffer{ + ReadTimeout: 5 * time.Second, + Recorder: &noopRecorder{}, + } + ho := &HandleOptions{ + log: xlogger.Nop(), + recorderObject: &xrecorder.HandlerRecorderObject{}, + } + + clientConn, serverConn := net.Pipe() + defer clientConn.Close() + defer serverConn.Close() + + go func() { + clientConn.Write([]byte("XXXXXX")) + }() + + err := h.serveH2(context.Background(), "tcp", serverConn, ho) + if err == nil { + t.Error("expected error from invalid h2 preface, got nil") + } +} diff --git a/internal/util/sniffing/sniffer_tls.go b/internal/util/sniffing/sniffer_tls.go new file mode 100644 index 00000000..7924e9ca --- /dev/null +++ b/internal/util/sniffing/sniffer_tls.go @@ -0,0 +1,237 @@ +package sniffing + +import ( + "bytes" + "context" + "crypto/tls" + "crypto/x509" + "encoding/hex" + "io" + "net" + "time" + + "github.com/go-gost/core/bypass" + dissector "github.com/go-gost/tls-dissector" + xbypass "github.com/go-gost/x/bypass" + xio "github.com/go-gost/x/internal/io" + xnet "github.com/go-gost/x/internal/net" + tls_util "github.com/go-gost/x/internal/util/tls" + xrecorder "github.com/go-gost/x/recorder" +) + +// HandleTLS sniffs and proxies a TLS connection. It parses the ClientHello +// for SNI-based routing, optionally performs MITM TLS termination for HTTP +// content inspection, and records TLS handshake metadata. +func (h *Sniffer) HandleTLS(ctx context.Context, network string, conn net.Conn, opts ...HandleOption) error { + var ho HandleOptions + for _, opt := range opts { + opt(&ho) + } + + readTimeout := h.effectiveReadTimeout() + + buf := new(bytes.Buffer) + clientHello, err := dissector.ParseClientHello(io.TeeReader(conn, buf)) + if err != nil { + return err + } + + log := ho.log + + ro := ho.recorderObject + ro.TLS = &xrecorder.TLSRecorderObject{ + ServerName: clientHello.ServerName, + ClientHello: hex.EncodeToString(buf.Bytes()), + } + if len(clientHello.SupportedProtos) > 0 { + ro.TLS.Proto = clientHello.SupportedProtos[0] + } + + host := normalizeHost(clientHello.ServerName, "443") + if host != "" { + ro.Host = host + + if ho.bypass != nil && ho.bypass.Contains(ctx, network, host, bypass.WithService(ho.service)) { + return xbypass.ErrBypass + } + } + + dial := ho.dial + if dial == nil { + dial = (&net.Dialer{}).DialContext + } + cc, err := dial(ctx, network, host) + if err != nil { + return err + } + defer cc.Close() + + log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "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") { + if host == "" { + host = ro.Host + } + if h.MitmBypass == nil || !h.MitmBypass.Contains(ctx, network, host, bypass.WithService(ho.service)) { + return h.terminateTLS(ctx, network, xnet.NewReadWriteConn(io.MultiReader(buf, conn), conn, conn), cc, clientHello, &ho) + } + } + + if _, err := buf.WriteTo(cc); err != nil { + return err + } + + xio.SetReadDeadline(cc, time.Now().Add(readTimeout)) + serverHello, serverHelloErr := dissector.ParseServerHello(io.TeeReader(cc, buf)) + xio.SetReadDeadline(cc, time.Time{}) + + if serverHello != nil { + ro.TLS.CipherSuite = tls_util.CipherSuite(serverHello.CipherSuite).String() + ro.TLS.CompressionMethod = serverHello.CompressionMethod + if serverHello.Proto != "" { + ro.TLS.Proto = serverHello.Proto + } + if serverHello.Version > 0 { + ro.TLS.Version = tls_util.Version(serverHello.Version).String() + } + } + + if buf.Len() > 0 { + ro.TLS.ServerHello = hex.EncodeToString(buf.Bytes()) + } + + if _, err := buf.WriteTo(conn); err != nil { + return err + } + + log.Infof("%s <-> %s", ro.RemoteAddr, ro.Host) + xnet.Pipe(ctx, conn, cc) + log.WithFields(map[string]any{ + "duration": time.Since(ro.Time), + }).Infof("%s >-< %s", ro.RemoteAddr, ro.Host) + + if serverHelloErr != nil { + return serverHelloErr + } + return nil +} + +// terminateTLS performs MITM TLS termination: handshakes with the upstream +// server as a client, then with the downstream client as a server using a +// dynamically generated certificate. The decrypted traffic is then handled +// as HTTP. +func (h *Sniffer) terminateTLS(ctx context.Context, network string, conn, cc net.Conn, clientHello *dissector.ClientHelloInfo, ho *HandleOptions) error { + ro := ho.recorderObject + log := ho.log + + nextProtos := clientHello.SupportedProtos + if h.NegotiatedProtocol != "" { + nextProtos = []string{h.NegotiatedProtocol} + } + + cfg := &tls.Config{ + ServerName: clientHello.ServerName, + NextProtos: nextProtos, + CipherSuites: clientHello.CipherSuites, + } + if cfg.ServerName == "" { + cfg.InsecureSkipVerify = true + } + clientConn := tls.Client(cc, cfg) + if err := clientConn.HandshakeContext(ctx); err != nil { + return err + } + + cs := clientConn.ConnectionState() + ro.TLS.CipherSuite = tls_util.CipherSuite(cs.CipherSuite).String() + ro.TLS.Proto = cs.NegotiatedProtocol + ro.TLS.Version = tls_util.Version(cs.Version).String() + + host := cfg.ServerName + if host == "" { + if len(cs.PeerCertificates) > 0 { + host = cs.PeerCertificates[0].Subject.CommonName + } + if host == "" { + host = ro.Host + } + } + if h, _, _ := net.SplitHostPort(host); h != "" { + host = h + } + + negotiatedProtocol := cs.NegotiatedProtocol + if h.NegotiatedProtocol != "" { + negotiatedProtocol = h.NegotiatedProtocol + } + nextProtos = nil + if negotiatedProtocol != "" { + nextProtos = []string{negotiatedProtocol} + } + + // cache the tls server handshake record. + wb := &bytes.Buffer{} + conn = xnet.NewReadWriteConn(conn, io.MultiWriter(wb, conn), conn) + + serverConn := tls.Server(conn, &tls.Config{ + NextProtos: nextProtos, + GetCertificate: func(chi *tls.ClientHelloInfo) (*tls.Certificate, error) { + certPool := h.CertPool + if certPool == nil { + certPool = DefaultCertPool + } + serverName := chi.ServerName + if serverName == "" { + serverName = host + } + cert, cerr := certPool.Get(serverName) + if cert != nil { + pool := x509.NewCertPool() + pool.AddCert(h.Certificate) + if _, cerr = cert.Verify(x509.VerifyOptions{ + DNSName: serverName, + Roots: pool, + }); cerr != nil { + log.Warnf("verify cached certificate for %s: %v", serverName, cerr) + cert = nil + } + } + if cert == nil { + cert, cerr = tls_util.GenerateCertificate(serverName, 7*24*time.Hour, h.Certificate, h.PrivateKey) + certPool.Put(serverName, cert) + } + if cerr != nil { + return nil, cerr + } + + return &tls.Certificate{ + Certificate: [][]byte{cert.Raw}, + PrivateKey: h.PrivateKey, + }, nil + }, + }) + handshakeErr := serverConn.HandshakeContext(ctx) + if record, _ := dissector.ReadRecord(wb); record != nil { + wb.Reset() + record.WriteTo(wb) + ro.TLS.ServerHello = hex.EncodeToString(wb.Bytes()) + } + if handshakeErr != nil { + return handshakeErr + } + + opts := []HandleOption{ + WithDial(func(ctx context.Context, network, address string) (net.Conn, error) { + return clientConn, nil + }), + WithDialTLS(func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error) { + return clientConn, nil + }), + WithRecorderObject(ro), + WithLog(log), + } + return h.HandleHTTP(ctx, network, serverConn, opts...) +} diff --git a/internal/util/sniffing/sniffer_ws.go b/internal/util/sniffing/sniffer_ws.go new file mode 100644 index 00000000..e043631e --- /dev/null +++ b/internal/util/sniffing/sniffer_ws.go @@ -0,0 +1,130 @@ +package sniffing + +import ( + "bytes" + "context" + "io" + "math" + "time" + + "github.com/go-gost/core/logger" + ws_util "github.com/go-gost/x/internal/util/ws" + xrecorder "github.com/go-gost/x/recorder" + "golang.org/x/time/rate" +) + +// sniffingWebsocketFrame copies WebSocket frames between rw and cc while +// recording frame metadata. It runs two goroutines for bidirectional copy +// and waits for the first error. +func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWriter, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error { + errc := make(chan error, 1) + + sampleRate := h.WebsocketSampleRate + if sampleRate == 0 { + sampleRate = DefaultSampleRate + } + if sampleRate < 0 { + sampleRate = math.MaxFloat64 + } + + go func() { + ro2 := &xrecorder.HandlerRecorderObject{} + *ro2 = *ro + ro := ro2 + + limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate)) + + buf := &bytes.Buffer{} + for { + start := time.Now() + + if err := h.copyWebsocketFrame(cc, rw, buf, "client", ro); err != nil { + errc <- err + return + } + + if limiter.Allow() { + ro.Duration = time.Since(start) + ro.Time = time.Now() + if err := ro.Record(ctx, h.Recorder); err != nil { + log.Errorf("record: %v", err) + } + } + } + }() + + go func() { + ro2 := &xrecorder.HandlerRecorderObject{} + *ro2 = *ro + ro := ro2 + + limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate)) + + buf := &bytes.Buffer{} + for { + start := time.Now() + + if err := h.copyWebsocketFrame(rw, cc, buf, "server", ro); err != nil { + errc <- err + return + } + + if limiter.Allow() { + ro.Duration = time.Since(start) + ro.Time = time.Now() + if err := ro.Record(ctx, h.Recorder); err != nil { + log.Errorf("record: %v", err) + } + } + } + }() + + <-errc + return nil +} + +// copyWebsocketFrame reads one WebSocket frame from r, records its metadata, +// and writes it to w. +func (h *Sniffer) 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 + } + + 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 bodySize := clampBodySize(h.RecorderOptions); bodySize > 0 { + 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 +}