sniffing websocket frame

handler/forward/local/handler.go

@@ -157,6 +157,8 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 			return cc, err
 		}
 		sniffer := &forwarder.Sniffer{
+			Websocket:           h.md.sniffingWebsocket,
+			WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
 			Recorder:            h.recorder.Recorder,
 			RecorderOptions:     h.recorder.Options,
 			Certificate:         h.md.certificate,

handler/forward/local/metadata.go

@@ -14,10 +14,12 @@ import (
 
 type metadata struct {
 	readTimeout   time.Duration
+	httpKeepalive bool
 
 	sniffing                    bool
 	sniffingTimeout             time.Duration
-	httpKeepalive   bool
+	sniffingWebsocket           bool
+	sniffingWebsocketSampleRate float64
 
 	certificate *x509.Certificate
 	privateKey  crypto.PrivateKey
@@ -31,10 +33,12 @@ func (h *forwardHandler) parseMetadata(md mdata.Metadata) (err error) {
 		h.md.readTimeout = 15 * time.Second
 	}
 
+	h.md.httpKeepalive = mdutil.GetBool(md, "http.keepalive")
+
 	h.md.sniffing = mdutil.GetBool(md, "sniffing")
 	h.md.sniffingTimeout = mdutil.GetDuration(md, "sniffing.timeout")
-
+	h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket")
-	h.md.httpKeepalive = mdutil.GetBool(md, "http.keepalive")
+	h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate")
 
 	certFile := mdutil.GetString(md, "mitm.certFile", "mitm.caCertFile")
 	keyFile := mdutil.GetString(md, "mitm.keyFile", "mitm.caKeyFile")

handler/forward/remote/handler.go

@@ -158,6 +158,8 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 			return cc, err
 		}
 		sniffer := &forwarder.Sniffer{
+			Websocket:           h.md.sniffingWebsocket,
+			WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
 			Recorder:            h.recorder.Recorder,
 			RecorderOptions:     h.recorder.Options,
 			Certificate:         h.md.certificate,

handler/forward/remote/metadata.go

@@ -15,10 +15,12 @@ import (
 type metadata struct {
 	readTimeout   time.Duration
 	proxyProtocol int
+	httpKeepalive bool
 
 	sniffing                    bool
 	sniffingTimeout             time.Duration
-	httpKeepalive   bool
+	sniffingWebsocket           bool
+	sniffingWebsocketSampleRate float64
 
 	certificate *x509.Certificate
 	privateKey  crypto.PrivateKey
@@ -35,6 +37,8 @@ func (h *forwardHandler) parseMetadata(md mdata.Metadata) (err error) {
 
 	h.md.sniffing = mdutil.GetBool(md, "sniffing")
 	h.md.sniffingTimeout = mdutil.GetDuration(md, "sniffing.timeout")
+	h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket")
+	h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate")
 
 	h.md.httpKeepalive = mdutil.GetBool(md, "http.keepalive")
 

handler/http/handler.go

@@ -29,12 +29,15 @@ import (
 	"github.com/go-gost/core/recorder"
 	xbypass "github.com/go-gost/x/bypass"
 	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"
 	limiter_util "github.com/go-gost/x/internal/util/limiter"
 	"github.com/go-gost/x/internal/util/sniffing"
 	stats_util "github.com/go-gost/x/internal/util/stats"
 	tls_util "github.com/go-gost/x/internal/util/tls"
+	ws_util "github.com/go-gost/x/internal/util/ws"
 	rate_limiter "github.com/go-gost/x/limiter/rate"
 	traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper"
 	stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
@@ -43,21 +46,6 @@ import (
 	"golang.org/x/net/http/httpguts"
 )
 
-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
-}
-
 func init() {
 	registry.HandlerRegistry().Register("http", NewHandler)
 }
@@ -295,7 +283,8 @@ func (h *httpHandler) handleRequest(ctx context.Context, conn net.Conn, req *htt
 	}
 	ctx = ctxvalue.ContextWithClientID(ctx, ctxvalue.ClientID(clientID))
 
-	if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, network, addr) {
+	if h.options.Bypass != nil &&
+		h.options.Bypass.Contains(ctx, network, addr) {
 		resp.StatusCode = http.StatusForbidden
 
 		if log.IsLevelEnabled(logger.TraceLevel) {
@@ -338,7 +327,7 @@ func (h *httpHandler) handleRequest(ctx context.Context, conn net.Conn, req *htt
 		return h.handleProxy(ctx, conn, req, ro, log)
 	}
 
-	ctx = contextWithRecorderObject(ctx, ro)
+	ctx = ctx_internal.ContextWithRecorderObject(ctx, ro)
 	ctx = ctxvalue.ContextWithLogger(ctx, log)
 	cc, err := h.dial(ctx, "tcp", addr)
 
@@ -386,6 +375,8 @@ func (h *httpHandler) handleRequest(ctx context.Context, conn net.Conn, req *htt
 			return cc, nil
 		}
 		sniffer := &sniffing.Sniffer{
+			Websocket:           h.md.sniffingWebsocket,
+			WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
 			Recorder:            h.recorder.Recorder,
 			RecorderOptions:     h.recorder.Options,
 			Certificate:         h.md.certificate,
@@ -452,7 +443,7 @@ func (h *httpHandler) handleProxy(ctx context.Context, conn net.Conn, req *http.
 			log.Trace(string(dump))
 		}
 
-		if err := h.proxyRoundTrip(ctx, conn, req, ro, &pStats, log); err != nil {
+		if err := h.proxyRoundTrip(ctx, xio.NewReadWriter(br, conn), req, ro, &pStats, log); err != nil {
 			return err
 		}
 	}
@@ -537,20 +528,36 @@ func (h *httpHandler) proxyRoundTrip(ctx context.Context, rw io.ReadWriter, req
 	}
 	ro.HTTP.StatusCode = res.StatusCode
 
+	if h.options.Bypass != nil &&
+		h.options.Bypass.Contains(ctx, "tcp", host) {
+		res.StatusCode = http.StatusForbidden
+
+		if log.IsLevelEnabled(logger.TraceLevel) {
+			dump, _ := httputil.DumpResponse(res, false)
+			log.Trace(string(dump))
+		}
+		log.Debug("bypass: ", host)
+		res.Write(rw)
+		return xbypass.ErrBypass
+	}
+
 	var reqBody *xhttp.Body
 	if opts := h.recorder.Options; opts != nil && opts.HTTPBody {
 		if req.Body != nil {
-			maxSize := opts.MaxBodySize
+			bodySize := opts.MaxBodySize
-			if maxSize <= 0 {
+			if bodySize <= 0 {
-				maxSize = sniffing.DefaultBodySize
+				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 := h.transport.RoundTrip(req.WithContext(ctx))
@@ -584,16 +591,19 @@ func (h *httpHandler) proxyRoundTrip(ctx context.Context, rw io.ReadWriter, req
 	}
 
 	if resp.StatusCode == http.StatusSwitchingProtocols {
-		return h.handleUpgradeResponse(rw, req, resp)
+		return h.handleUpgradeResponse(ctx, rw, req, resp, ro, log)
 	}
 
 	var respBody *xhttp.Body
 	if opts := h.recorder.Options; opts != nil && opts.HTTPBody {
-		maxSize := opts.MaxBodySize
+		bodySize := opts.MaxBodySize
-		if maxSize <= 0 {
+		if bodySize <= 0 {
-			maxSize = sniffing.DefaultBodySize
+			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
 	}
 
@@ -621,7 +631,7 @@ func (h *httpHandler) dial(ctx context.Context, network, addr string) (conn net.
 
 	var buf bytes.Buffer
 	conn, err = h.options.Router.Dial(ctxvalue.ContextWithBuffer(ctx, &buf), network, addr)
-	if ro := recorderObjectFromContext(ctx); ro != nil {
+	if ro := ctx_internal.RecorderObjectFromContext(ctx); ro != nil {
 		ro.Route = buf.String()
 	}
 
@@ -635,7 +645,7 @@ func upgradeType(h http.Header) string {
 	return h.Get("Upgrade")
 }
 
-func (h *httpHandler) handleUpgradeResponse(rw io.ReadWriter, req *http.Request, res *http.Response) error {
+func (h *httpHandler) 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) {
@@ -652,9 +662,146 @@ func (h *httpHandler) handleUpgradeResponse(rw io.ReadWriter, req *http.Request,
 		return fmt.Errorf("response write: %v", err)
 	}
 
+	if reqUpType == "websocket" && h.md.sniffingWebsocket {
+		return h.sniffingWebsocketFrame(ctx, rw, backConn, ro, log)
+	}
+
 	return xnet.Transport(rw, backConn)
 }
 
+func (h *httpHandler) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWriter, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
+	errc := make(chan error, 1)
+
+	sampleRate := h.md.sniffingWebsocketSampleRate
+	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 := h.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, h.recorder.Recorder); err != nil {
+					log.Errorf("record: %v", err)
+				}
+			default:
+			}
+
+			if err != nil {
+				errc <- err
+				return
+			}
+		}
+	}()
+
+	go func() {
+		ro2 := &xrecorder.HandlerRecorderObject{}
+		*ro2 = *ro
+		ro = ro2
+
+		ticker := time.NewTicker(d)
+		defer ticker.Stop()
+
+		buf := &bytes.Buffer{}
+		for {
+			start := time.Now()
+
+			err := h.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, h.recorder.Recorder); err != nil {
+					log.Errorf("record: %v", err)
+				}
+			default:
+			}
+
+			if err != nil {
+				errc <- err
+				return
+			}
+		}
+	}()
+
+	<-errc
+	return nil
+}
+
+func (h *httpHandler) 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.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 (h *httpHandler) decodeServerName(s string) (string, error) {
 	b, err := base64.RawURLEncoding.DecodeString(s)
 	if err != nil {

handler/http/metadata.go

@@ -32,6 +32,8 @@ type metadata struct {
 
 	sniffing                    bool
 	sniffingTimeout             time.Duration
+	sniffingWebsocket           bool
+	sniffingWebsocketSampleRate float64
 
 	certificate *x509.Certificate
 	privateKey  crypto.PrivateKey
@@ -80,6 +82,8 @@ func (h *httpHandler) parseMetadata(md mdata.Metadata) error {
 
 	h.md.sniffing = mdutil.GetBool(md, "sniffing")
 	h.md.sniffingTimeout = mdutil.GetDuration(md, "sniffing.timeout")
+	h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket")
+	h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate")
 
 	certFile := mdutil.GetString(md, "mitm.certFile", "mitm.caCertFile")
 	keyFile := mdutil.GetString(md, "mitm.keyFile", "mitm.caKeyFile")

handler/redirect/tcp/handler.go

@@ -128,6 +128,7 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han
 	}
 
 	ro.Host = dstAddr.String()
+	ro.Dst = dstAddr.String()
 
 	log = log.WithFields(map[string]any{
 		"dst":  fmt.Sprintf("%s/%s", dstAddr, dstAddr.Network()),
@@ -184,6 +185,8 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han
 		}
 
 		sniffer := &sniffing.Sniffer{
+			Websocket:           h.md.sniffingWebsocket,
+			WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
 			Recorder:            h.recorder.Recorder,
 			RecorderOptions:     h.recorder.Options,
 			Certificate:         h.md.certificate,

handler/redirect/tcp/metadata.go

@@ -19,6 +19,8 @@ type metadata struct {
 	sniffing                    bool
 	sniffingTimeout             time.Duration
 	sniffingFallback            bool
+	sniffingWebsocket           bool
+	sniffingWebsocketSampleRate float64
 
 	certificate *x509.Certificate
 	privateKey  crypto.PrivateKey
@@ -36,6 +38,8 @@ func (h *redirectHandler) parseMetadata(md mdata.Metadata) (err error) {
 	h.md.sniffing = mdutil.GetBool(md, "sniffing")
 	h.md.sniffingTimeout = mdutil.GetDuration(md, "sniffing.timeout")
 	h.md.sniffingFallback = mdutil.GetBool(md, "sniffing.fallback")
+	h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket")
+	h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate")
 
 	certFile := mdutil.GetString(md, "mitm.certFile", "mitm.caCertFile")
 	keyFile := mdutil.GetString(md, "mitm.keyFile", "mitm.caKeyFile")

handler/relay/connect.go

@@ -166,6 +166,8 @@ func (h *relayHandler) handleConnect(ctx context.Context, conn net.Conn, network
 			return cc, nil
 		}
 		sniffer := &sniffing.Sniffer{
+			Websocket:           h.md.sniffingWebsocket,
+			WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
 			Recorder:            h.recorder.Recorder,
 			RecorderOptions:     h.recorder.Options,
 			Certificate:         h.md.certificate,

handler/relay/metadata.go

@@ -9,8 +9,8 @@ import (
 
 	"github.com/go-gost/core/bypass"
 	mdata "github.com/go-gost/core/metadata"
-	mdutil "github.com/go-gost/x/metadata/util"
 	"github.com/go-gost/x/internal/util/mux"
+	mdutil "github.com/go-gost/x/metadata/util"
 	"github.com/go-gost/x/registry"
 )
 
@@ -25,6 +25,8 @@ type metadata struct {
 
 	sniffing                    bool
 	sniffingTimeout             time.Duration
+	sniffingWebsocket           bool
+	sniffingWebsocketSampleRate float64
 
 	certificate *x509.Certificate
 	privateKey  crypto.PrivateKey
@@ -63,6 +65,8 @@ func (h *relayHandler) parseMetadata(md mdata.Metadata) (err error) {
 
 	h.md.sniffing = mdutil.GetBool(md, "sniffing")
 	h.md.sniffingTimeout = mdutil.GetDuration(md, "sniffing.timeout")
+	h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket")
+	h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate")
 
 	certFile := mdutil.GetString(md, "mitm.certFile", "mitm.caCertFile")
 	keyFile := mdutil.GetString(md, "mitm.keyFile", "mitm.caKeyFile")

handler/sni/handler.go

@@ -143,6 +143,8 @@ func (h *sniHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.
 	}
 
 	sniffer := &sniffing.Sniffer{
+		Websocket:           h.md.sniffingWebsocket,
+		WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
 		Recorder:            h.recorder.Recorder,
 		RecorderOptions:     h.recorder.Options,
 		Certificate:         h.md.certificate,

handler/sni/metadata.go

@@ -16,6 +16,9 @@ type metadata struct {
 	readTimeout time.Duration
 	hash        string
 
+	sniffingWebsocket           bool
+	sniffingWebsocketSampleRate float64
+
 	certificate *x509.Certificate
 	privateKey  crypto.PrivateKey
 	alpn        string
@@ -29,6 +32,9 @@ func (h *sniHandler) parseMetadata(md mdata.Metadata) (err error) {
 	}
 	h.md.hash = mdutil.GetString(md, "hash")
 
+	h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket")
+	h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate")
+
 	certFile := mdutil.GetString(md, "mitm.certFile", "mitm.caCertFile")
 	keyFile := mdutil.GetString(md, "mitm.keyFile", "mitm.caKeyFile")
 	if certFile != "" && keyFile != "" {

handler/socks/v4/handler.go

@@ -275,6 +275,8 @@ func (h *socks4Handler) handleConnect(ctx context.Context, conn net.Conn, req *g
 			return cc, nil
 		}
 		sniffer := &sniffing.Sniffer{
+			Websocket:           h.md.sniffingWebsocket,
+			WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
 			Recorder:            h.recorder.Recorder,
 			RecorderOptions:     h.recorder.Options,
 			Certificate:         h.md.certificate,

handler/socks/v4/metadata.go

@@ -19,6 +19,8 @@ type metadata struct {
 
 	sniffing                    bool
 	sniffingTimeout             time.Duration
+	sniffingWebsocket           bool
+	sniffingWebsocketSampleRate float64
 
 	certificate *x509.Certificate
 	privateKey  crypto.PrivateKey
@@ -37,6 +39,8 @@ func (h *socks4Handler) parseMetadata(md mdata.Metadata) (err error) {
 
 	h.md.sniffing = mdutil.GetBool(md, "sniffing")
 	h.md.sniffingTimeout = mdutil.GetDuration(md, "sniffing.timeout")
+	h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket")
+	h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate")
 
 	certFile := mdutil.GetString(md, "mitm.certFile", "mitm.caCertFile")
 	keyFile := mdutil.GetString(md, "mitm.keyFile", "mitm.caKeyFile")

handler/socks/v5/connect.go

@@ -103,6 +103,8 @@ func (h *socks5Handler) handleConnect(ctx context.Context, conn net.Conn, networ
 			return cc, nil
 		}
 		sniffer := &sniffing.Sniffer{
+			Websocket:           h.md.sniffingWebsocket,
+			WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
 			Recorder:            h.recorder.Recorder,
 			RecorderOptions:     h.recorder.Options,
 			Certificate:         h.md.certificate,

handler/socks/v5/metadata.go

@@ -9,8 +9,8 @@ import (
 
 	"github.com/go-gost/core/bypass"
 	mdata "github.com/go-gost/core/metadata"
-	mdutil "github.com/go-gost/x/metadata/util"
 	"github.com/go-gost/x/internal/util/mux"
+	mdutil "github.com/go-gost/x/metadata/util"
 	"github.com/go-gost/x/registry"
 )
 
@@ -27,6 +27,8 @@ type metadata struct {
 
 	sniffing                    bool
 	sniffingTimeout             time.Duration
+	sniffingWebsocket           bool
+	sniffingWebsocketSampleRate float64
 
 	certificate *x509.Certificate
 	privateKey  crypto.PrivateKey
@@ -67,6 +69,8 @@ func (h *socks5Handler) parseMetadata(md mdata.Metadata) (err error) {
 
 	h.md.sniffing = mdutil.GetBool(md, "sniffing")
 	h.md.sniffingTimeout = mdutil.GetDuration(md, "sniffing.timeout")
+	h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket")
+	h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate")
 
 	certFile := mdutil.GetString(md, "mitm.certFile", "mitm.caCertFile")
 	keyFile := mdutil.GetString(md, "mitm.keyFile", "mitm.caKeyFile")

handler/ss/handler.go

@@ -192,6 +192,8 @@ func (h *ssHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.H
 			return cc, nil
 		}
 		sniffer := &sniffing.Sniffer{
+			Websocket:           h.md.sniffingWebsocket,
+			WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
 			Recorder:            h.recorder.Recorder,
 			RecorderOptions:     h.recorder.Options,
 			Certificate:         h.md.certificate,

handler/ss/metadata.go

@@ -19,6 +19,8 @@ type metadata struct {
 
 	sniffing                    bool
 	sniffingTimeout             time.Duration
+	sniffingWebsocket           bool
+	sniffingWebsocketSampleRate float64
 
 	certificate *x509.Certificate
 	privateKey  crypto.PrivateKey
@@ -37,6 +39,8 @@ func (h *ssHandler) parseMetadata(md mdata.Metadata) (err error) {
 
 	h.md.sniffing = mdutil.GetBool(md, "sniffing")
 	h.md.sniffingTimeout = mdutil.GetDuration(md, "sniffing.timeout")
+	h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket")
+	h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate")
 
 	certFile := mdutil.GetString(md, "mitm.certFile", "mitm.caCertFile")
 	keyFile := mdutil.GetString(md, "mitm.keyFile", "mitm.caKeyFile")

handler/sshd/handler.go

@@ -180,6 +180,8 @@ func (h *forwardHandler) handleDirectForward(ctx context.Context, conn *sshd_uti
 			return cc, nil
 		}
 		sniffer := &sniffing.Sniffer{
+			Websocket:           h.md.sniffingWebsocket,
+			WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
 			Recorder:            h.recorder.Recorder,
 			RecorderOptions:     h.recorder.Options,
 			Certificate:         h.md.certificate,

handler/sshd/metadata.go

@@ -17,6 +17,8 @@ type metadata struct {
 
 	sniffing                    bool
 	sniffingTimeout             time.Duration
+	sniffingWebsocket           bool
+	sniffingWebsocketSampleRate float64
 
 	certificate *x509.Certificate
 	privateKey  crypto.PrivateKey
@@ -32,6 +34,8 @@ func (h *forwardHandler) parseMetadata(md mdata.Metadata) (err error) {
 
 	h.md.sniffing = mdutil.GetBool(md, "sniffing")
 	h.md.sniffingTimeout = mdutil.GetDuration(md, "sniffing.timeout")
+	h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket")
+	h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate")
 
 	certFile := mdutil.GetString(md, "mitm.certFile", "mitm.caCertFile")
 	keyFile := mdutil.GetString(md, "mitm.keyFile", "mitm.caKeyFile")

handler/tunnel/entrypoint.go

@@ -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
+			bodySize := opts.MaxBodySize
-			if maxSize <= 0 {
+			if bodySize <= 0 {
-				maxSize = defaultBodySize
+				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
+		bodySize := opts.MaxBodySize
-		if maxSize <= 0 {
+		if bodySize <= 0 {
-			maxSize = defaultBodySize
+			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) {
+func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWriter, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
-	var tunnelID relay.TunnelID
+	errc := make(chan error, 1)
-	if h.ingress != nil {
-		if rule := h.ingress.GetRule(ctx, addr); rule != nil {
-			tunnelID = parseTunnelID(rule.Endpoint)
-		}
-	}
 
-	log := ctxvalue.LoggerFromContext(ctx)
+	sampleRate := ep.websocketSampleRate
-	if log == nil {
+	if sampleRate <= 0 {
-		log = h.log
+		sampleRate = sniffing.DefaultSampleRate
 	}
-	log.Debugf("dial: new connection to host %s", addr)
+	if sampleRate > sniffing.MaxSampleRate {
-
+		sampleRate = sniffing.MaxSampleRate
-	if tunnelID.IsZero() {
-		return nil, fmt.Errorf("%w %s", ErrTunnelRoute, addr)
 	}
+	d := time.Duration(1 / sampleRate * 1e9)
 
-	if ro := recorderObjectFromContext(ctx); ro != nil {
+	go func() {
-		ro.ClientID = tunnelID.String()
+		ro2 := &xrecorder.HandlerRecorderObject{}
-	}
+		*ro2 = *ro
+		ro = ro2
 
-	if tunnelID.IsPrivate() {
+		ticker := time.NewTicker(d)
-		return nil, fmt.Errorf("%w: tunnel %s is private for host %s", ErrPrivateTunnel, tunnelID, addr)
+		defer ticker.Stop()
-	}
 
-	log = log.WithFields(map[string]any{
+		buf := &bytes.Buffer{}
-		"tunnel": tunnelID.String(),
+		for {
-	})
+			start := time.Now()
 
-	d := &Dialer{
+			err := ep.copyWebsocketFrame(cc, rw, buf, "client", ro)
-		node:    h.node,
+			select {
-		pool:    h.pool,
+			case <-ticker.C:
-		sd:      h.sd,
-		retry:   3,
-		timeout: 15 * time.Second,
-		log:     log,
-	}
-	conn, node, cid, err := d.Dial(ctx, "tcp", tunnelID.String())
 				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.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
 		}
+	}()
+
+	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(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:
 			}
 
-	return conn, nil
+			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
+	}
+
+	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) {

handler/tunnel/handler.go

@@ -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,

handler/tunnel/metadata.go

@@ -27,6 +27,8 @@ type metadata struct {
 	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

internal/ctx/value.go

@@ -0,0 +1,22 @@
+package ctx
+
+import (
+	"context"
+
+	xrecorder "github.com/go-gost/x/recorder"
+)
+
+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
+}

internal/util/forwarder/sniffer.go

@@ -32,16 +32,13 @@ import (
 	xhttp "github.com/go-gost/x/internal/net/http"
 	"github.com/go-gost/x/internal/util/sniffing"
 	tls_util "github.com/go-gost/x/internal/util/tls"
+	ws_util "github.com/go-gost/x/internal/util/ws"
 	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"
 )
 
-const (
-	// Default max body size to record.
-	DefaultBodySize = 1024 * 1024 // 1MB
-)
-
 var (
 	DefaultCertPool = tls_util.NewMemoryCertPool()
 )
@@ -102,6 +99,9 @@ func WithLog(log logger.Logger) HandleOption {
 }
 
 type Sniffer struct {
+	Websocket           bool
+	WebsocketSampleRate float64
+
 	Recorder        recorder.Recorder
 	RecorderOptions *recorder.Options
 
@@ -177,7 +177,7 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
 
 	ro.Time = time.Time{}
 
-	shouldClose, err := h.httpRoundTrip(ctx, conn, cc, node, req, &pStats, &ho)
+	shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriter(br, conn), cc, node, req, &pStats, &ho)
 	if err != nil || shouldClose {
 		return err
 	}
@@ -198,7 +198,7 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
 			log.Trace(string(dump))
 		}
 
-		if shouldClose, err := h.httpRoundTrip(ctx, conn, cc, node, req, &pStats, &ho); err != nil || shouldClose {
+		if shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriter(br, conn), cc, node, req, &pStats, &ho); err != nil || shouldClose {
 			return err
 		}
 	}
@@ -441,11 +441,14 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, node
 	var reqBody *xhttp.Body
 	if opts := h.RecorderOptions; opts != nil && opts.HTTPBody {
 		if req.Body != nil {
-			maxSize := opts.MaxBodySize
+			bodySize := opts.MaxBodySize
-			if maxSize <= 0 {
+			if bodySize <= 0 {
-				maxSize = DefaultBodySize
+				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
 		}
 	}
@@ -479,6 +482,11 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, node
 		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 {
@@ -495,11 +503,14 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, node
 
 	var respBody *xhttp.Body
 	if opts := h.RecorderOptions; opts != nil && opts.HTTPBody {
-		maxSize := opts.MaxBodySize
+		bodySize := opts.MaxBodySize
-		if maxSize <= 0 {
+		if bodySize <= 0 {
-			maxSize = DefaultBodySize
+			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
 	}
 
@@ -513,11 +524,166 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, node
 		ro.HTTP.Response.ContentLength = respBody.Length()
 	}
 
-	if resp.StatusCode == http.StatusSwitchingProtocols {
+	return resp.Close, nil
-		xnet.Transport(rw, cc)
+}
+
+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 io.ReadWriter, cc 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) {
+		return fmt.Errorf("backend tried to switch protocol %q when %q was requested", resUpType, reqUpType)
 	}
 
-	return resp.Close, nil
+	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)
+}
+
+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 = 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 := h.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, h.Recorder); err != nil {
+					log.Errorf("record: %v", err)
+				}
+			default:
+			}
+
+			if err != nil {
+				errc <- err
+				return
+			}
+		}
+	}()
+
+	go func() {
+		ro2 := &xrecorder.HandlerRecorderObject{}
+		*ro2 = *ro
+		ro = ro2
+
+		ticker := time.NewTicker(d)
+		defer ticker.Stop()
+
+		buf := &bytes.Buffer{}
+		for {
+			start := time.Now()
+
+			err := h.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, h.Recorder); err != nil {
+					log.Errorf("record: %v", err)
+				}
+			default:
+			}
+
+			if err != nil {
+				errc <- err
+				return
+			}
+		}
+	}()
+
+	<-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 = 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 (h *Sniffer) rewriteBody(resp *http.Response, rewrites ...chain.HTTPBodyRewriteSettings) error {
@@ -914,11 +1080,14 @@ func (h *h2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	var reqBody *xhttp.Body
 	if opts := h.recorderOptions; opts != nil && opts.HTTPBody {
 		if req.Body != nil {
-			maxSize := opts.MaxBodySize
+			bodySize := opts.MaxBodySize
-			if maxSize <= 0 {
+			if bodySize <= 0 {
-				maxSize = DefaultBodySize
+				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
 		}
 	}
@@ -949,11 +1118,14 @@ func (h *h2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	var respBody *xhttp.Body
 	if opts := h.recorderOptions; opts != nil && opts.HTTPBody {
-		maxSize := opts.MaxBodySize
+		bodySize := opts.MaxBodySize
-		if maxSize <= 0 {
+		if bodySize <= 0 {
-			maxSize = DefaultBodySize
+			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
 	}
 

internal/util/sniffing/sniffer.go

@@ -28,14 +28,22 @@ import (
 	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"
 	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"
 )
 
 const (
-	// Default max body size to record.
+	// DefaultBodySize is the default HTTP body or websocket frame size to record.
-	DefaultBodySize = 1024 * 1024 // 1MB
+	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 = 10.0
+	// MaxSampleRate is the maximum websocket sample rate (samples per second).
+	MaxSampleRate = 100.0
 )
 
 var (
@@ -84,6 +92,9 @@ func WithLog(log logger.Logger) HandleOption {
 }
 
 type Sniffer struct {
+	Websocket           bool
+	WebsocketSampleRate float64
+
 	Recorder        recorder.Recorder
 	RecorderOptions *recorder.Options
 
@@ -176,7 +187,7 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
 
 	ro.Time = time.Time{}
 
-	shouldClose, err := h.httpRoundTrip(ctx, conn, cc, req, ro, &pStats, log)
+	shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriter(br, conn), cc, req, ro, &pStats, log)
 	if err != nil || shouldClose {
 		return err
 	}
@@ -197,7 +208,7 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
 			log.Trace(string(dump))
 		}
 
-		if shouldClose, err := h.httpRoundTrip(ctx, conn, cc, req, ro, &pStats, log); err != nil || shouldClose {
+		if shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriter(br, conn), cc, req, ro, &pStats, log); err != nil || shouldClose {
 			return err
 		}
 	}
@@ -302,11 +313,14 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, req *
 	var reqBody *xhttp.Body
 	if opts := h.RecorderOptions; opts != nil && opts.HTTPBody {
 		if req.Body != nil {
-			maxSize := opts.MaxBodySize
+			bodySize := opts.MaxBodySize
-			if maxSize <= 0 {
+			if bodySize <= 0 {
-				maxSize = DefaultBodySize
+				bodySize = DefaultBodySize
 			}
-			reqBody = xhttp.NewBody(req.Body, maxSize)
+			if bodySize > MaxBodySize {
+				bodySize = MaxBodySize
+			}
+			reqBody = xhttp.NewBody(req.Body, bodySize)
 			req.Body = reqBody
 		}
 	}
@@ -323,7 +337,7 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, req *
 	xio.SetReadDeadline(cc, time.Now().Add(h.ReadTimeout))
 	resp, err := http.ReadResponse(bufio.NewReader(cc), req)
 	if err != nil {
-		log.Errorf("read response: %v", err)
+		err = fmt.Errorf("read response: %w", err)
 		return
 	}
 	defer resp.Body.Close()
@@ -338,6 +352,11 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, req *
 		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 {
@@ -349,16 +368,19 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, req *
 
 	var respBody *xhttp.Body
 	if opts := h.RecorderOptions; opts != nil && opts.HTTPBody {
-		maxSize := opts.MaxBodySize
+		bodySize := opts.MaxBodySize
-		if maxSize <= 0 {
+		if bodySize <= 0 {
-			maxSize = DefaultBodySize
+			bodySize = DefaultBodySize
 		}
-		respBody = xhttp.NewBody(resp.Body, maxSize)
+		if bodySize > MaxBodySize {
+			bodySize = MaxBodySize
+		}
+		respBody = xhttp.NewBody(resp.Body, bodySize)
 		resp.Body = respBody
 	}
 
 	if err = resp.Write(rw); err != nil {
-		log.Errorf("write response: %v", err)
+		err = fmt.Errorf("write response: %w", err)
 		return
 	}
 
@@ -367,11 +389,166 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, req *
 		ro.HTTP.Response.ContentLength = respBody.Length()
 	}
 
-	if resp.StatusCode == http.StatusSwitchingProtocols {
+	return resp.Close, nil
-		xnet.Transport(rw, cc)
+}
+
+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 io.ReadWriter, cc 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) {
+		return fmt.Errorf("backend tried to switch protocol %q when %q was requested", resUpType, reqUpType)
 	}
 
-	return resp.Close, nil
+	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)
+}
+
+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 > MaxSampleRate {
+		sampleRate = 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 := h.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, h.Recorder); err != nil {
+					log.Errorf("record: %v", err)
+				}
+			default:
+			}
+
+			if err != nil {
+				errc <- err
+				return
+			}
+		}
+	}()
+
+	go func() {
+		ro2 := &xrecorder.HandlerRecorderObject{}
+		*ro2 = *ro
+		ro = ro2
+
+		ticker := time.NewTicker(d)
+		defer ticker.Stop()
+
+		buf := &bytes.Buffer{}
+		for {
+			start := time.Now()
+
+			err := h.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, h.Recorder); err != nil {
+					log.Errorf("record: %v", err)
+				}
+			default:
+			}
+
+			if err != nil {
+				errc <- err
+				return
+			}
+		}
+	}()
+
+	<-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, conn net.Conn, opts ...HandleOption) error {
@@ -644,11 +821,15 @@ func (h *h2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	var reqBody *xhttp.Body
 	if opts := h.recorderOptions; opts != nil && opts.HTTPBody {
 		if req.Body != nil {
-			maxSize := opts.MaxBodySize
+			bodySize := opts.MaxBodySize
-			if maxSize <= 0 {
+			if bodySize <= 0 {
-				maxSize = DefaultBodySize
+				bodySize = DefaultBodySize
 			}
-			reqBody = xhttp.NewBody(req.Body, maxSize)
+			if bodySize > MaxBodySize {
+				bodySize = MaxBodySize
+			}
+
+			reqBody = xhttp.NewBody(req.Body, bodySize)
 			req.Body = reqBody
 		}
 	}
@@ -679,11 +860,14 @@ func (h *h2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	var respBody *xhttp.Body
 	if opts := h.recorderOptions; opts != nil && opts.HTTPBody {
-		maxSize := opts.MaxBodySize
+		bodySize := opts.MaxBodySize
-		if maxSize <= 0 {
+		if bodySize <= 0 {
-			maxSize = DefaultBodySize
+			bodySize = DefaultBodySize
 		}
-		respBody = xhttp.NewBody(resp.Body, maxSize)
+		if bodySize > MaxBodySize {
+			bodySize = MaxBodySize
+		}
+		respBody = xhttp.NewBody(resp.Body, bodySize)
 		resp.Body = respBody
 	}
 

internal/util/ws/frame.go

@@ -0,0 +1,196 @@
+package ws
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+	"math"
+)
+
+// OpCode represents a WebSocket opcode.
+type OpCode int
+
+// https://tools.ietf.org/html/rfc6455#section-11.8.
+const (
+	OpContinuation OpCode = iota
+	OpText
+	OpBinary
+	// 3 - 7 are reserved for further non-control frames.
+	_
+	_
+	_
+	_
+	_
+	OpClose
+	OpPing
+	OpPong
+	// 11-16 are reserved for further control frames.
+)
+
+// FrameHeader represents a WebSocket frame header.
+// See https://tools.ietf.org/html/rfc6455#section-5.2.
+type FrameHeader struct {
+	Fin    bool
+	Rsv1   bool
+	Rsv2   bool
+	Rsv3   bool
+	OpCode OpCode
+
+	PayloadLength int64
+
+	Masked  bool
+	MaskKey uint32
+}
+
+// ReadFrom reads a header from the reader.
+// See https://tools.ietf.org/html/rfc6455#section-5.2.
+func (h *FrameHeader) ReadFrom(r io.Reader) (n int64, err error) {
+	var buf [8]byte
+
+	// First byte. FIN/RSV1/RSV2/RSV3/OpCode(4bits)
+	nn, err := io.ReadFull(r, buf[:2])
+	n += int64(n)
+	if err != nil {
+		return
+	}
+
+	b := buf[0]
+	h.Fin = b&(1<<7) != 0
+	h.Rsv1 = b&(1<<6) != 0
+	h.Rsv2 = b&(1<<5) != 0
+	h.Rsv3 = b&(1<<4) != 0
+
+	h.OpCode = OpCode(b & 0xf)
+
+	b = buf[1]
+	h.Masked = b&(1<<7) != 0
+
+	payloadLength := b &^ (1 << 7)
+	switch {
+	case payloadLength < 126:
+		h.PayloadLength = int64(payloadLength)
+	case payloadLength == 126:
+		nn, err = io.ReadFull(r, buf[:2])
+		h.PayloadLength = int64(binary.BigEndian.Uint16(buf[:]))
+	case payloadLength == 127:
+		nn, err = io.ReadFull(r, buf[:])
+		h.PayloadLength = int64(binary.BigEndian.Uint64(buf[:]))
+	}
+	n += int64(nn)
+	if err != nil {
+		return
+	}
+
+	if h.PayloadLength < 0 {
+		err = fmt.Errorf("received negative payload length: %v", h.PayloadLength)
+		return
+	}
+
+	if h.Masked {
+		nn, err = io.ReadFull(r, buf[:4])
+		n += int64(nn)
+		if err != nil {
+			return
+		}
+		h.MaskKey = binary.LittleEndian.Uint32(buf[:])
+	}
+
+	return
+}
+
+func (h FrameHeader) Length() int {
+	n := 2
+	switch {
+	case h.PayloadLength > math.MaxUint16:
+		n += 8
+	case h.PayloadLength > 125:
+		n += 2
+	}
+
+	if h.Masked {
+		n += 4
+	}
+	return n
+}
+
+func (h *FrameHeader) WriteTo(w io.Writer) (n int64, err error) {
+	var buf [14]byte
+	pos := 0
+
+	var b byte
+	if h.Fin {
+		b |= 1 << 7
+	}
+	if h.Rsv1 {
+		b |= 1 << 6
+	}
+	if h.Rsv2 {
+		b |= 1 << 5
+	}
+	if h.Rsv3 {
+		b |= 1 << 4
+	}
+
+	b |= byte(h.OpCode)
+
+	buf[0] = b
+
+	lengthByte := byte(0)
+	if h.Masked {
+		lengthByte |= 1 << 7
+	}
+
+	switch {
+	case h.PayloadLength > math.MaxUint16:
+		lengthByte |= 127
+	case h.PayloadLength > 125:
+		lengthByte |= 126
+	case h.PayloadLength >= 0:
+		lengthByte |= byte(h.PayloadLength)
+	}
+	buf[1] = lengthByte
+	pos = 2
+
+	switch {
+	case h.PayloadLength > math.MaxUint16:
+		binary.BigEndian.PutUint64(buf[2:], uint64(h.PayloadLength))
+		pos += 8
+	case h.PayloadLength > 125:
+		binary.BigEndian.PutUint16(buf[2:], uint16(h.PayloadLength))
+		pos += 2
+	}
+
+	if h.Masked {
+		binary.LittleEndian.PutUint32(buf[pos:], h.MaskKey)
+		pos += 4
+	}
+
+	nn, err := w.Write(buf[:pos])
+	n = int64(nn)
+	return
+}
+
+type Frame struct {
+	Header FrameHeader
+	Data   io.Reader
+}
+
+func (fr *Frame) ReadFrom(r io.Reader) (n int64, err error) {
+	if n, err = fr.Header.ReadFrom(r); err != nil {
+		return
+	}
+
+	fr.Data = io.LimitReader(r, fr.Header.PayloadLength)
+	return
+}
+
+func (fr *Frame) WriteTo(w io.Writer) (n int64, err error) {
+	n, err = fr.Header.WriteTo(w)
+	if err != nil {
+		return
+	}
+
+	nn, err := io.Copy(w, fr.Data)
+	n += nn
+	return
+}

recorder/recorder.go

@@ -38,6 +38,19 @@ type HTTPRecorderObject struct {
 	Response   HTTPResponseRecorderObject `json:"response"`
 }
 
+type WebsocketRecorderObject struct {
+	From    string `json:"from"`
+	Fin     bool   `json:"fin"`
+	Rsv1    bool   `json:"rsv1"`
+	Rsv2    bool   `json:"rsv2"`
+	Rsv3    bool   `json:"rsv3"`
+	OpCode  int    `json:"opcode"`
+	Masked  bool   `json:"masked"`
+	MaskKey uint32 `json:"maskKey"`
+	Length  int64  `json:"length"`
+	Payload []byte `json:"payload"`
+}
+
 type TLSRecorderObject struct {
 	ServerName        string `json:"serverName"`
 	CipherSuite       string `json:"cipherSuite"`
@@ -65,10 +78,12 @@ type HandlerRecorderObject struct {
 	RemoteAddr  string                   `json:"remote"`
 	LocalAddr   string                   `json:"local"`
 	Host        string                   `json:"host"`
+	Dst         string                   `json:"dst"`
 	Proto       string                   `json:"proto,omitempty"`
 	ClientIP    string                   `json:"clientIP"`
 	ClientID    string                   `json:"clientID,omitempty"`
 	HTTP        *HTTPRecorderObject      `json:"http,omitempty"`
+	Websocket   *WebsocketRecorderObject `json:"websocket,omitempty"`
 	TLS         *TLSRecorderObject       `json:"tls,omitempty"`
 	DNS         *DNSRecorderObject       `json:"dns,omitempty"`
 	Route       string                   `json:"route,omitempty"`