diff --git a/handler/forward/local/handler.go b/handler/forward/local/handler.go index 775b7e48..50b53939 100644 --- a/handler/forward/local/handler.go +++ b/handler/forward/local/handler.go @@ -157,14 +157,16 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand return cc, err } sniffer := &forwarder.Sniffer{ - Recorder: h.recorder.Recorder, - RecorderOptions: h.recorder.Options, - Certificate: h.md.certificate, - PrivateKey: h.md.privateKey, - NegotiatedProtocol: h.md.alpn, - CertPool: h.certPool, - MitmBypass: h.md.mitmBypass, - ReadTimeout: h.md.readTimeout, + Websocket: h.md.sniffingWebsocket, + WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, + Recorder: h.recorder.Recorder, + RecorderOptions: h.recorder.Options, + Certificate: h.md.certificate, + PrivateKey: h.md.privateKey, + NegotiatedProtocol: h.md.alpn, + CertPool: h.certPool, + MitmBypass: h.md.mitmBypass, + ReadTimeout: h.md.readTimeout, } conn = xnet.NewReadWriteConn(br, conn, conn) diff --git a/handler/forward/local/metadata.go b/handler/forward/local/metadata.go index 2b76d418..13ecb87a 100644 --- a/handler/forward/local/metadata.go +++ b/handler/forward/local/metadata.go @@ -13,11 +13,13 @@ import ( ) type metadata struct { - readTimeout time.Duration + readTimeout time.Duration + httpKeepalive bool - sniffing bool - sniffingTimeout time.Duration - httpKeepalive bool + sniffing bool + sniffingTimeout time.Duration + 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.httpKeepalive = mdutil.GetBool(md, "http.keepalive") + 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") diff --git a/handler/forward/remote/handler.go b/handler/forward/remote/handler.go index 261e156c..4ac4461a 100644 --- a/handler/forward/remote/handler.go +++ b/handler/forward/remote/handler.go @@ -158,14 +158,16 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand return cc, err } sniffer := &forwarder.Sniffer{ - Recorder: h.recorder.Recorder, - RecorderOptions: h.recorder.Options, - Certificate: h.md.certificate, - PrivateKey: h.md.privateKey, - NegotiatedProtocol: h.md.alpn, - CertPool: h.certPool, - MitmBypass: h.md.mitmBypass, - ReadTimeout: h.md.readTimeout, + Websocket: h.md.sniffingWebsocket, + WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, + Recorder: h.recorder.Recorder, + RecorderOptions: h.recorder.Options, + Certificate: h.md.certificate, + PrivateKey: h.md.privateKey, + NegotiatedProtocol: h.md.alpn, + CertPool: h.certPool, + MitmBypass: h.md.mitmBypass, + ReadTimeout: h.md.readTimeout, } conn = xnet.NewReadWriteConn(br, conn, conn) diff --git a/handler/forward/remote/metadata.go b/handler/forward/remote/metadata.go index cdddc42c..6b1e5417 100644 --- a/handler/forward/remote/metadata.go +++ b/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 + sniffing bool + sniffingTimeout time.Duration + 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") diff --git a/handler/http/handler.go b/handler/http/handler.go index 0a6d77f0..7f2e8318 100644 --- a/handler/http/handler.go +++ b/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,14 +375,16 @@ func (h *httpHandler) handleRequest(ctx context.Context, conn net.Conn, req *htt return cc, nil } sniffer := &sniffing.Sniffer{ - Recorder: h.recorder.Recorder, - RecorderOptions: h.recorder.Options, - Certificate: h.md.certificate, - PrivateKey: h.md.privateKey, - NegotiatedProtocol: h.md.alpn, - CertPool: h.certPool, - MitmBypass: h.md.mitmBypass, - ReadTimeout: h.md.readTimeout, + Websocket: h.md.sniffingWebsocket, + WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, + Recorder: h.recorder.Recorder, + RecorderOptions: h.recorder.Options, + Certificate: h.md.certificate, + PrivateKey: h.md.privateKey, + NegotiatedProtocol: h.md.alpn, + CertPool: h.certPool, + MitmBypass: h.md.mitmBypass, + ReadTimeout: h.md.readTimeout, } conn = xnet.NewReadWriteConn(br, conn, conn) @@ -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 - if maxSize <= 0 { - maxSize = sniffing.DefaultBodySize + bodySize := opts.MaxBodySize + if bodySize <= 0 { + bodySize = sniffing.DefaultBodySize } - reqBody = xhttp.NewBody(req.Body, maxSize) + if bodySize > sniffing.MaxBodySize { + bodySize = sniffing.MaxBodySize + } + reqBody = xhttp.NewBody(req.Body, bodySize) req.Body = reqBody } } ctx = ctxvalue.ContextWithClientAddr(ctx, ctxvalue.ClientAddr(clientAddr)) - ctx = contextWithRecorderObject(ctx, ro) + ctx = ctx_internal.ContextWithRecorderObject(ctx, ro) ctx = ctxvalue.ContextWithLogger(ctx, log) resp, err := 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 - if maxSize <= 0 { - maxSize = sniffing.DefaultBodySize + bodySize := opts.MaxBodySize + if bodySize <= 0 { + bodySize = sniffing.DefaultBodySize } - respBody = xhttp.NewBody(resp.Body, maxSize) + if bodySize > sniffing.MaxBodySize { + bodySize = sniffing.MaxBodySize + } + respBody = xhttp.NewBody(resp.Body, bodySize) resp.Body = respBody } @@ -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 { diff --git a/handler/http/metadata.go b/handler/http/metadata.go index 385b4024..f0060f3e 100644 --- a/handler/http/metadata.go +++ b/handler/http/metadata.go @@ -30,8 +30,10 @@ type metadata struct { observePeriod time.Duration proxyAgent string - sniffing bool - sniffingTimeout time.Duration + 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") diff --git a/handler/redirect/tcp/handler.go b/handler/redirect/tcp/handler.go index 42c3d5b8..b3465602 100644 --- a/handler/redirect/tcp/handler.go +++ b/handler/redirect/tcp/handler.go @@ -128,9 +128,10 @@ 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()), + "dst": fmt.Sprintf("%s/%s", dstAddr, dstAddr.Network()), "host": dstAddr.String(), }) @@ -184,14 +185,16 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han } sniffer := &sniffing.Sniffer{ - Recorder: h.recorder.Recorder, - RecorderOptions: h.recorder.Options, - Certificate: h.md.certificate, - PrivateKey: h.md.privateKey, - NegotiatedProtocol: h.md.alpn, - CertPool: h.certPool, - MitmBypass: h.md.mitmBypass, - ReadTimeout: h.md.readTimeout, + Websocket: h.md.sniffingWebsocket, + WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, + Recorder: h.recorder.Recorder, + RecorderOptions: h.recorder.Options, + Certificate: h.md.certificate, + PrivateKey: h.md.privateKey, + NegotiatedProtocol: h.md.alpn, + CertPool: h.certPool, + MitmBypass: h.md.mitmBypass, + ReadTimeout: h.md.readTimeout, } conn = xnet.NewReadWriteConn(br, conn, conn) diff --git a/handler/redirect/tcp/metadata.go b/handler/redirect/tcp/metadata.go index 1bb7e10d..77865c0e 100644 --- a/handler/redirect/tcp/metadata.go +++ b/handler/redirect/tcp/metadata.go @@ -16,9 +16,11 @@ type metadata struct { readTimeout time.Duration tproxy bool - sniffing bool - sniffingTimeout time.Duration - sniffingFallback bool + 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") diff --git a/handler/relay/connect.go b/handler/relay/connect.go index bf81059a..fa83dbb7 100644 --- a/handler/relay/connect.go +++ b/handler/relay/connect.go @@ -166,14 +166,16 @@ func (h *relayHandler) handleConnect(ctx context.Context, conn net.Conn, network return cc, nil } sniffer := &sniffing.Sniffer{ - Recorder: h.recorder.Recorder, - RecorderOptions: h.recorder.Options, - Certificate: h.md.certificate, - PrivateKey: h.md.privateKey, - NegotiatedProtocol: h.md.alpn, - CertPool: h.certPool, - MitmBypass: h.md.mitmBypass, - ReadTimeout: h.md.readTimeout, + Websocket: h.md.sniffingWebsocket, + WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, + Recorder: h.recorder.Recorder, + RecorderOptions: h.recorder.Options, + Certificate: h.md.certificate, + PrivateKey: h.md.privateKey, + NegotiatedProtocol: h.md.alpn, + CertPool: h.certPool, + MitmBypass: h.md.mitmBypass, + ReadTimeout: h.md.readTimeout, } conn = xnet.NewReadWriteConn(br, conn, conn) diff --git a/handler/relay/metadata.go b/handler/relay/metadata.go index 85bf74e0..90cc4824 100644 --- a/handler/relay/metadata.go +++ b/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" ) @@ -23,8 +23,10 @@ type metadata struct { muxCfg *mux.Config observePeriod time.Duration - sniffing bool - sniffingTimeout time.Duration + 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") diff --git a/handler/sni/handler.go b/handler/sni/handler.go index 0913796d..582b46a4 100644 --- a/handler/sni/handler.go +++ b/handler/sni/handler.go @@ -143,14 +143,16 @@ func (h *sniHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler. } sniffer := &sniffing.Sniffer{ - Recorder: h.recorder.Recorder, - RecorderOptions: h.recorder.Options, - Certificate: h.md.certificate, - PrivateKey: h.md.privateKey, - NegotiatedProtocol: h.md.alpn, - CertPool: h.certPool, - MitmBypass: h.md.mitmBypass, - ReadTimeout: h.md.readTimeout, + Websocket: h.md.sniffingWebsocket, + WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, + Recorder: h.recorder.Recorder, + RecorderOptions: h.recorder.Options, + Certificate: h.md.certificate, + PrivateKey: h.md.privateKey, + NegotiatedProtocol: h.md.alpn, + CertPool: h.certPool, + MitmBypass: h.md.mitmBypass, + ReadTimeout: h.md.readTimeout, } conn = xnet.NewReadWriteConn(br, conn, conn) switch proto { diff --git a/handler/sni/metadata.go b/handler/sni/metadata.go index 5438d836..06852bc4 100644 --- a/handler/sni/metadata.go +++ b/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 != "" { diff --git a/handler/socks/v4/handler.go b/handler/socks/v4/handler.go index f9686c72..10b477b5 100644 --- a/handler/socks/v4/handler.go +++ b/handler/socks/v4/handler.go @@ -275,14 +275,16 @@ func (h *socks4Handler) handleConnect(ctx context.Context, conn net.Conn, req *g return cc, nil } sniffer := &sniffing.Sniffer{ - Recorder: h.recorder.Recorder, - RecorderOptions: h.recorder.Options, - Certificate: h.md.certificate, - PrivateKey: h.md.privateKey, - NegotiatedProtocol: h.md.alpn, - CertPool: h.certPool, - MitmBypass: h.md.mitmBypass, - ReadTimeout: h.md.readTimeout, + Websocket: h.md.sniffingWebsocket, + WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, + Recorder: h.recorder.Recorder, + RecorderOptions: h.recorder.Options, + Certificate: h.md.certificate, + PrivateKey: h.md.privateKey, + NegotiatedProtocol: h.md.alpn, + CertPool: h.certPool, + MitmBypass: h.md.mitmBypass, + ReadTimeout: h.md.readTimeout, } conn = xnet.NewReadWriteConn(br, conn, conn) diff --git a/handler/socks/v4/metadata.go b/handler/socks/v4/metadata.go index 26cf1c68..31308aac 100644 --- a/handler/socks/v4/metadata.go +++ b/handler/socks/v4/metadata.go @@ -17,8 +17,10 @@ type metadata struct { hash string observePeriod time.Duration - sniffing bool - sniffingTimeout time.Duration + 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") diff --git a/handler/socks/v5/connect.go b/handler/socks/v5/connect.go index a1d938f5..1e83c344 100644 --- a/handler/socks/v5/connect.go +++ b/handler/socks/v5/connect.go @@ -103,14 +103,16 @@ func (h *socks5Handler) handleConnect(ctx context.Context, conn net.Conn, networ return cc, nil } sniffer := &sniffing.Sniffer{ - Recorder: h.recorder.Recorder, - RecorderOptions: h.recorder.Options, - Certificate: h.md.certificate, - PrivateKey: h.md.privateKey, - NegotiatedProtocol: h.md.alpn, - CertPool: h.certPool, - MitmBypass: h.md.mitmBypass, - ReadTimeout: h.md.readTimeout, + Websocket: h.md.sniffingWebsocket, + WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, + Recorder: h.recorder.Recorder, + RecorderOptions: h.recorder.Options, + Certificate: h.md.certificate, + PrivateKey: h.md.privateKey, + NegotiatedProtocol: h.md.alpn, + CertPool: h.certPool, + MitmBypass: h.md.mitmBypass, + ReadTimeout: h.md.readTimeout, } conn = xnet.NewReadWriteConn(br, conn, conn) diff --git a/handler/socks/v5/metadata.go b/handler/socks/v5/metadata.go index b5128c73..898c846f 100644 --- a/handler/socks/v5/metadata.go +++ b/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" ) @@ -25,8 +25,10 @@ type metadata struct { muxCfg *mux.Config observePeriod time.Duration - sniffing bool - sniffingTimeout time.Duration + 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") diff --git a/handler/ss/handler.go b/handler/ss/handler.go index 3cf4e2e8..3ecf4067 100644 --- a/handler/ss/handler.go +++ b/handler/ss/handler.go @@ -192,14 +192,16 @@ func (h *ssHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.H return cc, nil } sniffer := &sniffing.Sniffer{ - Recorder: h.recorder.Recorder, - RecorderOptions: h.recorder.Options, - Certificate: h.md.certificate, - PrivateKey: h.md.privateKey, - NegotiatedProtocol: h.md.alpn, - CertPool: h.certPool, - MitmBypass: h.md.mitmBypass, - ReadTimeout: h.md.readTimeout, + Websocket: h.md.sniffingWebsocket, + WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, + Recorder: h.recorder.Recorder, + RecorderOptions: h.recorder.Options, + Certificate: h.md.certificate, + PrivateKey: h.md.privateKey, + NegotiatedProtocol: h.md.alpn, + CertPool: h.certPool, + MitmBypass: h.md.mitmBypass, + ReadTimeout: h.md.readTimeout, } conn = xnet.NewReadWriteConn(br, conn, conn) diff --git a/handler/ss/metadata.go b/handler/ss/metadata.go index e1944635..33c27a9c 100644 --- a/handler/ss/metadata.go +++ b/handler/ss/metadata.go @@ -17,8 +17,10 @@ type metadata struct { hash string readTimeout time.Duration - sniffing bool - sniffingTimeout time.Duration + 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") diff --git a/handler/sshd/handler.go b/handler/sshd/handler.go index 298206c7..87782a4f 100644 --- a/handler/sshd/handler.go +++ b/handler/sshd/handler.go @@ -180,14 +180,16 @@ func (h *forwardHandler) handleDirectForward(ctx context.Context, conn *sshd_uti return cc, nil } sniffer := &sniffing.Sniffer{ - Recorder: h.recorder.Recorder, - RecorderOptions: h.recorder.Options, - Certificate: h.md.certificate, - PrivateKey: h.md.privateKey, - NegotiatedProtocol: h.md.alpn, - CertPool: h.certPool, - MitmBypass: h.md.mitmBypass, - ReadTimeout: h.md.readTimeout, + Websocket: h.md.sniffingWebsocket, + WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, + Recorder: h.recorder.Recorder, + RecorderOptions: h.recorder.Options, + Certificate: h.md.certificate, + PrivateKey: h.md.privateKey, + NegotiatedProtocol: h.md.alpn, + CertPool: h.certPool, + MitmBypass: h.md.mitmBypass, + ReadTimeout: h.md.readTimeout, } switch proto { diff --git a/handler/sshd/metadata.go b/handler/sshd/metadata.go index ba2f56bc..01cbab57 100644 --- a/handler/sshd/metadata.go +++ b/handler/sshd/metadata.go @@ -15,8 +15,10 @@ import ( type metadata struct { readTimeout time.Duration - sniffing bool - sniffingTimeout time.Duration + 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") diff --git a/handler/tunnel/entrypoint.go b/handler/tunnel/entrypoint.go index 27d9f855..791ee37c 100644 --- a/handler/tunnel/entrypoint.go +++ b/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 - if maxSize <= 0 { - maxSize = defaultBodySize + bodySize := opts.MaxBodySize + if bodySize <= 0 { + bodySize = sniffing.DefaultBodySize } - reqBody = xhttp.NewBody(req.Body, maxSize) + if bodySize > sniffing.MaxBodySize { + bodySize = sniffing.MaxBodySize + } + reqBody = xhttp.NewBody(req.Body, bodySize) req.Body = reqBody } } ctx = ctxvalue.ContextWithClientAddr(ctx, ctxvalue.ClientAddr(clientAddr)) - ctx = contextWithRecorderObject(ctx, ro) + ctx = ctx_internal.ContextWithRecorderObject(ctx, ro) ctx = ctxvalue.ContextWithLogger(ctx, log) resp, err := ep.transport.RoundTrip(req.WithContext(ctx)) @@ -289,16 +349,19 @@ func (ep *entrypoint) httpRoundTrip(ctx context.Context, rw io.ReadWriter, req * } if resp.StatusCode == http.StatusSwitchingProtocols { - return ep.handleUpgradeResponse(rw, req, resp) + return ep.handleUpgradeResponse(ctx, rw, req, resp, ro, log) } var respBody *xhttp.Body if opts := ep.recorder.Options; opts != nil && opts.HTTPBody { - maxSize := opts.MaxBodySize - if maxSize <= 0 { - maxSize = defaultBodySize + bodySize := opts.MaxBodySize + if bodySize <= 0 { + bodySize = sniffing.DefaultBodySize } - respBody = xhttp.NewBody(resp.Body, maxSize) + if bodySize > sniffing.MaxBodySize { + bodySize = sniffing.MaxBodySize + } + respBody = xhttp.NewBody(resp.Body, bodySize) resp.Body = respBody } @@ -321,7 +384,7 @@ func upgradeType(h http.Header) string { return h.Get("Upgrade") } -func (ep *entrypoint) handleUpgradeResponse(rw io.ReadWriter, req *http.Request, res *http.Response) error { +func (ep *entrypoint) handleUpgradeResponse(ctx context.Context, rw io.ReadWriter, req *http.Request, res *http.Response, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error { reqUpType := upgradeType(req.Header) resUpType := upgradeType(res.Header) if !strings.EqualFold(reqUpType, resUpType) { @@ -332,81 +395,151 @@ func (ep *entrypoint) handleUpgradeResponse(rw io.ReadWriter, req *http.Request, if !ok { return fmt.Errorf("internal error: 101 switching protocols response with non-writable body") } + defer backConn.Close() res.Body = nil if err := res.Write(rw); err != nil { return fmt.Errorf("response write: %v", err) } + if reqUpType == "websocket" && ep.sniffingWebsocket { + return ep.sniffingWebsocketFrame(ctx, rw, backConn, ro, log) + } + return xnet.Transport(rw, backConn) } -func (h *entrypoint) dial(ctx context.Context, network, addr string) (conn net.Conn, err error) { - var tunnelID relay.TunnelID - if h.ingress != nil { - if rule := h.ingress.GetRule(ctx, addr); rule != nil { - tunnelID = parseTunnelID(rule.Endpoint) +func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWriter, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error { + errc := make(chan error, 1) + + sampleRate := ep.websocketSampleRate + if sampleRate <= 0 { + sampleRate = sniffing.DefaultSampleRate + } + if sampleRate > sniffing.MaxSampleRate { + sampleRate = sniffing.MaxSampleRate + } + d := time.Duration(1 / sampleRate * 1e9) + + go func() { + ro2 := &xrecorder.HandlerRecorderObject{} + *ro2 = *ro + ro = ro2 + + ticker := time.NewTicker(d) + defer ticker.Stop() + + buf := &bytes.Buffer{} + for { + start := time.Now() + + err := ep.copyWebsocketFrame(cc, rw, buf, "client", ro) + select { + case <-ticker.C: + if err != nil { + ro.Err = err.Error() + } + ro.Duration = time.Since(start) + ro.Time = time.Now() + if err := ro.Record(ctx, ep.recorder.Recorder); err != nil { + log.Errorf("record: %v", err) + } + default: + } + + if err != nil { + errc <- err + return + } } - } + }() - log := ctxvalue.LoggerFromContext(ctx) - if log == nil { - log = h.log - } - log.Debugf("dial: new connection to host %s", addr) + go func() { + ro2 := &xrecorder.HandlerRecorderObject{} + *ro2 = *ro + ro = ro2 - if tunnelID.IsZero() { - return nil, fmt.Errorf("%w %s", ErrTunnelRoute, addr) - } + ticker := time.NewTicker(d) + defer ticker.Stop() - if ro := recorderObjectFromContext(ctx); ro != nil { - ro.ClientID = tunnelID.String() - } + buf := &bytes.Buffer{} + for { + start := time.Now() - if tunnelID.IsPrivate() { - return nil, fmt.Errorf("%w: tunnel %s is private for host %s", ErrPrivateTunnel, tunnelID, addr) - } + err := ep.copyWebsocketFrame(rw, cc, buf, "server", ro) + select { + case <-ticker.C: + if err != nil { + ro.Err = err.Error() + } + ro.Duration = time.Since(start) + ro.Time = time.Now() + if err := ro.Record(ctx, ep.recorder.Recorder); err != nil { + log.Errorf("record: %v", err) + } + default: + } - log = log.WithFields(map[string]any{ - "tunnel": tunnelID.String(), - }) - - d := &Dialer{ - node: h.node, - pool: h.pool, - sd: h.sd, - retry: 3, - timeout: 15 * time.Second, - log: log, - } - conn, node, cid, err := d.Dial(ctx, "tcp", tunnelID.String()) - if err != nil { - return - } - log.Debugf("dial: connected to host %s, tunnel: %s, connector: %s", addr, tunnelID, cid) - - if node == h.node { - clientAddr := ctxvalue.ClientAddrFromContext(ctx) - var features []relay.Feature - af := &relay.AddrFeature{} - af.ParseFrom(string(clientAddr)) - features = append(features, af) // src address - - af = &relay.AddrFeature{} - af.ParseFrom(addr) - features = append(features, af) // dst address - - if _, err = (&relay.Response{ - Version: relay.Version1, - Status: relay.StatusOK, - Features: features, - }).WriteTo(conn); err != nil { - conn.Close() - return nil, err + if err != nil { + errc <- err + return + } } + }() + + <-errc + return nil +} + +func (ep *entrypoint) copyWebsocketFrame(w io.Writer, r io.Reader, buf *bytes.Buffer, from string, ro *xrecorder.HandlerRecorderObject) (err error) { + fr := ws_util.Frame{} + if _, err = fr.ReadFrom(r); err != nil { + return err } - return conn, nil + ws := &xrecorder.WebsocketRecorderObject{ + From: from, + Fin: fr.Header.Fin, + Rsv1: fr.Header.Rsv1, + Rsv2: fr.Header.Rsv2, + Rsv3: fr.Header.Rsv3, + OpCode: int(fr.Header.OpCode), + Masked: fr.Header.Masked, + MaskKey: fr.Header.MaskKey, + Length: fr.Header.PayloadLength, + } + if opts := ep.recorder.Options; opts != nil && opts.HTTPBody { + bodySize := opts.MaxBodySize + if bodySize <= 0 { + bodySize = sniffing.DefaultBodySize + } + if bodySize > sniffing.MaxBodySize { + bodySize = sniffing.MaxBodySize + } + + buf.Reset() + if _, err := io.Copy(buf, io.LimitReader(fr.Data, int64(bodySize))); err != nil { + return err + } + ws.Payload = buf.Bytes() + } + + ro.Websocket = ws + length := uint64(fr.Header.Length()) + uint64(fr.Header.PayloadLength) + if from == "client" { + ro.InputBytes = length + ro.OutputBytes = 0 + } else { + ro.InputBytes = 0 + ro.OutputBytes = length + } + + fr.Data = io.MultiReader(bytes.NewReader(buf.Bytes()), fr.Data) + if _, err := fr.WriteTo(w); err != nil { + return err + } + + return nil } func (ep *entrypoint) handleConnect(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) (err error) { diff --git a/handler/tunnel/handler.go b/handler/tunnel/handler.go index a4af23df..b47b7d2a 100644 --- a/handler/tunnel/handler.go +++ b/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, diff --git a/handler/tunnel/metadata.go b/handler/tunnel/metadata.go index 052f5aa5..d950e8c7 100644 --- a/handler/tunnel/metadata.go +++ b/handler/tunnel/metadata.go @@ -22,11 +22,13 @@ const ( type metadata struct { readTimeout time.Duration - entryPoint string - entryPointID relay.TunnelID - entryPointProxyProtocol int - entryPointKeepalive bool - entryPointReadTimeout time.Duration + entryPoint string + entryPointID relay.TunnelID + entryPointProxyProtocol int + entryPointKeepalive bool + entryPointReadTimeout time.Duration + sniffingWebsocket bool + sniffingWebsocketSampleRate float64 directTunnel bool tunnelTTL time.Duration @@ -53,6 +55,9 @@ func (h *tunnelHandler) parseMetadata(md mdata.Metadata) (err error) { h.md.entryPointReadTimeout = 15 * time.Second } + h.md.sniffingWebsocket = mdutil.GetBool(md, "sniffing.websocket") + h.md.sniffingWebsocketSampleRate = mdutil.GetFloat(md, "sniffing.websocket.sampleRate") + h.md.tunnelTTL = mdutil.GetDuration(md, "tunnel.ttl") if h.md.tunnelTTL <= 0 { h.md.tunnelTTL = defaultTTL diff --git a/internal/ctx/value.go b/internal/ctx/value.go new file mode 100644 index 00000000..1c4556c6 --- /dev/null +++ b/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 +} diff --git a/internal/util/forwarder/sniffer.go b/internal/util/forwarder/sniffer.go index 774b834e..f86f532d 100644 --- a/internal/util/forwarder/sniffer.go +++ b/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 - if maxSize <= 0 { - maxSize = DefaultBodySize + bodySize := opts.MaxBodySize + if bodySize <= 0 { + bodySize = sniffing.DefaultBodySize } - reqBody = xhttp.NewBody(req.Body, maxSize) + if bodySize > sniffing.MaxBodySize { + bodySize = sniffing.MaxBodySize + } + reqBody = xhttp.NewBody(req.Body, bodySize) req.Body = reqBody } } @@ -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 - if maxSize <= 0 { - maxSize = DefaultBodySize + bodySize := opts.MaxBodySize + if bodySize <= 0 { + bodySize = sniffing.DefaultBodySize } - respBody = xhttp.NewBody(resp.Body, maxSize) + if bodySize > sniffing.MaxBodySize { + bodySize = sniffing.MaxBodySize + } + respBody = xhttp.NewBody(resp.Body, bodySize) resp.Body = respBody } @@ -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 { - xnet.Transport(rw, cc) + return resp.Close, nil +} + +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 - if maxSize <= 0 { - maxSize = DefaultBodySize + bodySize := opts.MaxBodySize + if bodySize <= 0 { + bodySize = sniffing.DefaultBodySize } - reqBody = xhttp.NewBody(req.Body, maxSize) + if bodySize > sniffing.MaxBodySize { + bodySize = sniffing.MaxBodySize + } + reqBody = xhttp.NewBody(req.Body, bodySize) req.Body = reqBody } } @@ -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 - if maxSize <= 0 { - maxSize = DefaultBodySize + bodySize := opts.MaxBodySize + if bodySize <= 0 { + bodySize = sniffing.DefaultBodySize } - respBody = xhttp.NewBody(resp.Body, maxSize) + if bodySize > sniffing.MaxBodySize { + bodySize = sniffing.MaxBodySize + } + respBody = xhttp.NewBody(resp.Body, bodySize) resp.Body = respBody } diff --git a/internal/util/sniffing/sniffer.go b/internal/util/sniffing/sniffer.go index 47f498d6..c7791cbf 100644 --- a/internal/util/sniffing/sniffer.go +++ b/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 = 1024 * 1024 // 1MB + // 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 = 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 - if maxSize <= 0 { - maxSize = DefaultBodySize + bodySize := opts.MaxBodySize + if bodySize <= 0 { + 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 - if maxSize <= 0 { - maxSize = DefaultBodySize + bodySize := opts.MaxBodySize + if bodySize <= 0 { + 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 { - xnet.Transport(rw, cc) + return resp.Close, nil +} + +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 - if maxSize <= 0 { - maxSize = DefaultBodySize + bodySize := opts.MaxBodySize + if bodySize <= 0 { + 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 - if maxSize <= 0 { - maxSize = DefaultBodySize + bodySize := opts.MaxBodySize + if bodySize <= 0 { + bodySize = DefaultBodySize } - respBody = xhttp.NewBody(resp.Body, maxSize) + if bodySize > MaxBodySize { + bodySize = MaxBodySize + } + respBody = xhttp.NewBody(resp.Body, bodySize) resp.Body = respBody } diff --git a/internal/util/ws/frame.go b/internal/util/ws/frame.go new file mode 100644 index 00000000..6962083b --- /dev/null +++ b/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 +} diff --git a/recorder/recorder.go b/recorder/recorder.go index 1a3b70f2..f9262fb7 100644 --- a/recorder/recorder.go +++ b/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"` @@ -59,25 +72,27 @@ type DNSRecorderObject struct { } type HandlerRecorderObject struct { - Node string `json:"node,omitempty"` - Service string `json:"service"` - Network string `json:"network"` - RemoteAddr string `json:"remote"` - LocalAddr string `json:"local"` - Host string `json:"host"` - Proto string `json:"proto,omitempty"` - ClientIP string `json:"clientIP"` - ClientID string `json:"clientID,omitempty"` - HTTP *HTTPRecorderObject `json:"http,omitempty"` - TLS *TLSRecorderObject `json:"tls,omitempty"` - DNS *DNSRecorderObject `json:"dns,omitempty"` - Route string `json:"route,omitempty"` - InputBytes uint64 `json:"inputBytes"` - OutputBytes uint64 `json:"outputBytes"` - Err string `json:"err,omitempty"` - SID string `json:"sid"` - Duration time.Duration `json:"duration"` - Time time.Time `json:"time"` + Node string `json:"node,omitempty"` + Service string `json:"service"` + Network string `json:"network"` + 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"` + InputBytes uint64 `json:"inputBytes"` + OutputBytes uint64 `json:"outputBytes"` + Err string `json:"err,omitempty"` + SID string `json:"sid"` + Duration time.Duration `json:"duration"` + Time time.Time `json:"time"` } func (p *HandlerRecorderObject) Record(ctx context.Context, r recorder.Recorder) error {