From 41f527aa467544c76d0d886db85339cb00e95e67 Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Sat, 19 Oct 2024 20:40:07 +0800 Subject: [PATCH] sniffing: full websocket frame recording --- handler/http/handler.go | 27 +++++++++++++++++++-------- handler/tunnel/entrypoint.go | 27 +++++++++++++++++++-------- internal/util/forwarder/sniffer.go | 27 +++++++++++++++++++-------- internal/util/sniffing/sniffer.go | 29 +++++++++++++++++++---------- 4 files changed, 76 insertions(+), 34 deletions(-) diff --git a/handler/http/handler.go b/handler/http/handler.go index 7f2e8318..fa49bfe5 100644 --- a/handler/http/handler.go +++ b/handler/http/handler.go @@ -673,12 +673,9 @@ func (h *httpHandler) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read errc := make(chan error, 1) sampleRate := h.md.sniffingWebsocketSampleRate - if sampleRate <= 0 { + if sampleRate == 0 { sampleRate = sniffing.DefaultSampleRate } - if sampleRate > sniffing.MaxSampleRate { - sampleRate = sniffing.MaxSampleRate - } d := time.Duration(1 / sampleRate * 1e9) go func() { @@ -686,8 +683,15 @@ func (h *httpHandler) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read *ro2 = *ro ro = ro2 - ticker := time.NewTicker(d) - defer ticker.Stop() + ticker := &time.Ticker{} + if d > 0 { + ticker = time.NewTicker(d) + defer ticker.Stop() + } else { + tc := make(chan time.Time) + close(tc) + ticker.C = tc + } buf := &bytes.Buffer{} for { @@ -719,8 +723,15 @@ func (h *httpHandler) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read *ro2 = *ro ro = ro2 - ticker := time.NewTicker(d) - defer ticker.Stop() + ticker := &time.Ticker{} + if d > 0 { + ticker = time.NewTicker(d) + defer ticker.Stop() + } else { + tc := make(chan time.Time) + close(tc) + ticker.C = tc + } buf := &bytes.Buffer{} for { diff --git a/handler/tunnel/entrypoint.go b/handler/tunnel/entrypoint.go index 791ee37c..8ceb63c4 100644 --- a/handler/tunnel/entrypoint.go +++ b/handler/tunnel/entrypoint.go @@ -413,12 +413,9 @@ func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read errc := make(chan error, 1) sampleRate := ep.websocketSampleRate - if sampleRate <= 0 { + if sampleRate == 0 { sampleRate = sniffing.DefaultSampleRate } - if sampleRate > sniffing.MaxSampleRate { - sampleRate = sniffing.MaxSampleRate - } d := time.Duration(1 / sampleRate * 1e9) go func() { @@ -426,8 +423,15 @@ func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read *ro2 = *ro ro = ro2 - ticker := time.NewTicker(d) - defer ticker.Stop() + ticker := &time.Ticker{} + if d > 0 { + ticker = time.NewTicker(d) + defer ticker.Stop() + } else { + tc := make(chan time.Time) + close(tc) + ticker.C = tc + } buf := &bytes.Buffer{} for { @@ -459,8 +463,15 @@ func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read *ro2 = *ro ro = ro2 - ticker := time.NewTicker(d) - defer ticker.Stop() + ticker := &time.Ticker{} + if d > 0 { + ticker = time.NewTicker(d) + defer ticker.Stop() + } else { + tc := make(chan time.Time) + close(tc) + ticker.C = tc + } buf := &bytes.Buffer{} for { diff --git a/internal/util/forwarder/sniffer.go b/internal/util/forwarder/sniffer.go index f86f532d..7a3ebb9e 100644 --- a/internal/util/forwarder/sniffer.go +++ b/internal/util/forwarder/sniffer.go @@ -557,12 +557,9 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit errc := make(chan error, 1) sampleRate := h.WebsocketSampleRate - if sampleRate <= 0 { + if sampleRate == 0 { sampleRate = sniffing.DefaultSampleRate } - if sampleRate > sniffing.MaxSampleRate { - sampleRate = sniffing.MaxSampleRate - } d := time.Duration(1 / sampleRate * 1e9) go func() { @@ -570,8 +567,15 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit *ro2 = *ro ro = ro2 - ticker := time.NewTicker(d) - defer ticker.Stop() + ticker := &time.Ticker{} + if d > 0 { + ticker = time.NewTicker(d) + defer ticker.Stop() + } else { + tc := make(chan time.Time) + close(tc) + ticker.C = tc + } buf := &bytes.Buffer{} for { @@ -603,8 +607,15 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit *ro2 = *ro ro = ro2 - ticker := time.NewTicker(d) - defer ticker.Stop() + ticker := &time.Ticker{} + if d > 0 { + ticker = time.NewTicker(d) + defer ticker.Stop() + } else { + tc := make(chan time.Time) + close(tc) + ticker.C = tc + } buf := &bytes.Buffer{} for { diff --git a/internal/util/sniffing/sniffer.go b/internal/util/sniffing/sniffer.go index c7791cbf..7c0aab16 100644 --- a/internal/util/sniffing/sniffer.go +++ b/internal/util/sniffing/sniffer.go @@ -42,8 +42,6 @@ const ( 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 ( @@ -422,12 +420,9 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit errc := make(chan error, 1) sampleRate := h.WebsocketSampleRate - if sampleRate <= 0 { + if sampleRate == 0 { sampleRate = DefaultSampleRate } - if sampleRate > MaxSampleRate { - sampleRate = MaxSampleRate - } d := time.Duration(1 / sampleRate * 1e9) go func() { @@ -435,8 +430,15 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit *ro2 = *ro ro = ro2 - ticker := time.NewTicker(d) - defer ticker.Stop() + ticker := &time.Ticker{} + if d > 0 { + ticker = time.NewTicker(d) + defer ticker.Stop() + } else { + tc := make(chan time.Time) + close(tc) + ticker.C = tc + } buf := &bytes.Buffer{} for { @@ -468,8 +470,15 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit *ro2 = *ro ro = ro2 - ticker := time.NewTicker(d) - defer ticker.Stop() + ticker := &time.Ticker{} + if d > 0 { + ticker = time.NewTicker(d) + defer ticker.Stop() + } else { + tc := make(chan time.Time) + close(tc) + ticker.C = tc + } buf := &bytes.Buffer{} for {