diff --git a/handler/http/handler.go b/handler/http/handler.go index 520c4a5f..3234286f 100644 --- a/handler/http/handler.go +++ b/handler/http/handler.go @@ -11,6 +11,7 @@ import ( "fmt" "hash/crc32" "io" + "math" "net" "net/http" "net/http/httputil" @@ -44,6 +45,7 @@ import ( xrecorder "github.com/go-gost/x/recorder" "github.com/go-gost/x/registry" "golang.org/x/net/http/httpguts" + "golang.org/x/time/rate" ) func init() { @@ -679,44 +681,33 @@ func (h *httpHandler) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read if sampleRate == 0 { sampleRate = sniffing.DefaultSampleRate } - d := time.Duration(1 / sampleRate * 1e9) + if sampleRate < 0 { + sampleRate = math.MaxFloat64 + } go func() { ro2 := &xrecorder.HandlerRecorderObject{} *ro2 = *ro ro := ro2 - ticker := &time.Ticker{} - if d > 0 { - ticker = time.NewTicker(d) - defer ticker.Stop() - } else { - tc := make(chan time.Time) - close(tc) - ticker.C = tc - } + limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate)) buf := &bytes.Buffer{} for { start := time.Now() - err := h.copyWebsocketFrame(cc, rw, buf, "client", ro) - select { - case <-ticker.C: - if err == nil { - 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 { + if err := h.copyWebsocketFrame(cc, rw, buf, "client", ro); err != nil { errc <- err return } + + if limiter.Allow() { + ro.Duration = time.Since(start) + ro.Time = time.Now() + if err := ro.Record(ctx, h.recorder.Recorder); err != nil { + log.Errorf("record: %v", err) + } + } } }() @@ -725,37 +716,24 @@ func (h *httpHandler) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read *ro2 = *ro ro := ro2 - ticker := &time.Ticker{} - if d > 0 { - ticker = time.NewTicker(d) - defer ticker.Stop() - } else { - tc := make(chan time.Time) - close(tc) - ticker.C = tc - } + limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate)) buf := &bytes.Buffer{} for { start := time.Now() - err := h.copyWebsocketFrame(rw, cc, buf, "server", ro) - select { - case <-ticker.C: - if err == nil { - 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 { + if err := h.copyWebsocketFrame(rw, cc, buf, "server", ro); err != nil { errc <- err return } + + if limiter.Allow() { + ro.Duration = time.Since(start) + ro.Time = time.Now() + if err := ro.Record(ctx, h.recorder.Recorder); err != nil { + log.Errorf("record: %v", err) + } + } } }() diff --git a/handler/tunnel/entrypoint.go b/handler/tunnel/entrypoint.go index 57660bfd..ddd5397e 100644 --- a/handler/tunnel/entrypoint.go +++ b/handler/tunnel/entrypoint.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "io" + "math" "net" "net/http" "net/http/httputil" @@ -37,6 +38,12 @@ import ( 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/time/rate" +) + +const ( + httpHeaderSID = "Gost-Sid" + httpHeaderForwardedNode = "Gost-Forwarded-Node" ) type entrypoint struct { @@ -108,17 +115,17 @@ 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) { +func (ep *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 { + if ep.ingress != nil { + if rule := ep.ingress.GetRule(ctx, addr); rule != nil { tunnelID = parseTunnelID(rule.Endpoint) } } log := ctxvalue.LoggerFromContext(ctx) if log == nil { - log = h.log + log = ep.log } log.Debugf("dial: new connection to host %s", addr) @@ -126,9 +133,8 @@ func (h *entrypoint) dial(ctx context.Context, network, addr string) (conn net.C return nil, fmt.Errorf("%w %s", ErrTunnelRoute, addr) } - if ro := ctx_internal.RecorderObjectFromContext(ctx); ro != nil { - ro.ClientID = tunnelID.String() - } + ro := ctx_internal.RecorderObjectFromContext(ctx) + ro.ClientID = tunnelID.String() if tunnelID.IsPrivate() { return nil, fmt.Errorf("%w: tunnel %s is private for host %s", ErrPrivateTunnel, tunnelID, addr) @@ -139,9 +145,9 @@ func (h *entrypoint) dial(ctx context.Context, network, addr string) (conn net.C }) d := &Dialer{ - node: h.node, - pool: h.pool, - sd: h.sd, + node: ep.node, + pool: ep.pool, + sd: ep.sd, retry: 3, timeout: 15 * time.Second, log: log, @@ -152,7 +158,9 @@ func (h *entrypoint) dial(ctx context.Context, network, addr string) (conn net.C } log.Debugf("dial: connected to host %s, tunnel: %s, connector: %s", addr, tunnelID, cid) - if node == h.node { + if node == ep.node { + ro.Redirect = "" + clientAddr := ctxvalue.ClientAddrFromContext(ctx) var features []relay.Feature af := &relay.AddrFeature{} @@ -171,6 +179,8 @@ func (h *entrypoint) dial(ctx context.Context, network, addr string) (conn net.C conn.Close() return nil, err } + } else { + ro.Redirect = node } return conn, nil @@ -237,6 +247,13 @@ func (ep *entrypoint) httpRoundTrip(ctx context.Context, rw io.ReadWriter, req * *ro2 = *ro ro = ro2 + if sid := req.Header.Get(httpHeaderSID); sid != "" { + ro.SID = sid + } else { + req.Header.Set(httpHeaderSID, ro.SID) + } + req.Header.Set(httpHeaderForwardedNode, ro.Node) + host := req.Host if _, port, _ := net.SplitHostPort(host); port == "" { host = net.JoinHostPort(strings.Trim(host, "[]"), "80") @@ -419,44 +436,33 @@ func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read if sampleRate == 0 { sampleRate = sniffing.DefaultSampleRate } - d := time.Duration(1 / sampleRate * 1e9) + if sampleRate < 0 { + sampleRate = math.MaxFloat64 + } go func() { ro2 := &xrecorder.HandlerRecorderObject{} *ro2 = *ro ro := ro2 - ticker := &time.Ticker{} - if d > 0 { - ticker = time.NewTicker(d) - defer ticker.Stop() - } else { - tc := make(chan time.Time) - close(tc) - ticker.C = tc - } + limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate)) buf := &bytes.Buffer{} for { start := time.Now() - err := ep.copyWebsocketFrame(cc, rw, buf, "client", ro) - select { - case <-ticker.C: - if err == nil { - 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 { + if err := ep.copyWebsocketFrame(cc, rw, buf, "client", ro); err != nil { errc <- err return } + + if limiter.Allow() { + ro.Duration = time.Since(start) + ro.Time = time.Now() + if err := ro.Record(ctx, ep.recorder.Recorder); err != nil { + log.Errorf("record: %v", err) + } + } } }() @@ -465,37 +471,24 @@ func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read *ro2 = *ro ro := ro2 - ticker := &time.Ticker{} - if d > 0 { - ticker = time.NewTicker(d) - defer ticker.Stop() - } else { - tc := make(chan time.Time) - close(tc) - ticker.C = tc - } + limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate)) buf := &bytes.Buffer{} for { start := time.Now() - err := ep.copyWebsocketFrame(rw, cc, buf, "server", ro) - select { - case <-ticker.C: - if err == nil { - 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 { + if err := ep.copyWebsocketFrame(rw, cc, buf, "server", ro); err != nil { errc <- err return } + + if limiter.Allow() { + ro.Duration = time.Since(start) + ro.Time = time.Now() + if err := ro.Record(ctx, ep.recorder.Recorder); err != nil { + log.Errorf("record: %v", err) + } + } } }() diff --git a/internal/util/forwarder/sniffer.go b/internal/util/forwarder/sniffer.go index 28dc6d6f..f0ec5c72 100644 --- a/internal/util/forwarder/sniffer.go +++ b/internal/util/forwarder/sniffer.go @@ -11,6 +11,7 @@ import ( "errors" "fmt" "io" + "math" "net" "net/http" "net/http/httputil" @@ -37,6 +38,7 @@ import ( xrecorder "github.com/go-gost/x/recorder" "golang.org/x/net/http/httpguts" "golang.org/x/net/http2" + "golang.org/x/time/rate" ) var ( @@ -564,44 +566,33 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit if sampleRate == 0 { sampleRate = sniffing.DefaultSampleRate } - d := time.Duration(1 / sampleRate * 1e9) + if sampleRate < 0 { + sampleRate = math.MaxFloat64 + } go func() { ro2 := &xrecorder.HandlerRecorderObject{} *ro2 = *ro ro := ro2 - ticker := &time.Ticker{} - if d > 0 { - ticker = time.NewTicker(d) - defer ticker.Stop() - } else { - tc := make(chan time.Time) - close(tc) - ticker.C = tc - } + limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate)) buf := &bytes.Buffer{} for { start := time.Now() - err := h.copyWebsocketFrame(cc, rw, buf, "client", ro) - select { - case <-ticker.C: - if err == nil { - 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 { + if err := h.copyWebsocketFrame(cc, rw, buf, "client", ro); err != nil { errc <- err return } + + if limiter.Allow() { + ro.Duration = time.Since(start) + ro.Time = time.Now() + if err := ro.Record(ctx, h.Recorder); err != nil { + log.Errorf("record: %v", err) + } + } } }() @@ -610,37 +601,24 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit *ro2 = *ro ro := ro2 - ticker := &time.Ticker{} - if d > 0 { - ticker = time.NewTicker(d) - defer ticker.Stop() - } else { - tc := make(chan time.Time) - close(tc) - ticker.C = tc - } + limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate)) buf := &bytes.Buffer{} for { start := time.Now() - err := h.copyWebsocketFrame(rw, cc, buf, "server", ro) - select { - case <-ticker.C: - if err == nil { - 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 { + if err := h.copyWebsocketFrame(rw, cc, buf, "server", ro); err != nil { errc <- err return } + + if limiter.Allow() { + ro.Duration = time.Since(start) + ro.Time = time.Now() + if err := ro.Record(ctx, h.Recorder); err != nil { + log.Errorf("record: %v", err) + } + } } }() diff --git a/internal/util/sniffing/sniffer.go b/internal/util/sniffing/sniffer.go index ef9ab3b3..bd4e5b4e 100644 --- a/internal/util/sniffing/sniffer.go +++ b/internal/util/sniffing/sniffer.go @@ -11,6 +11,7 @@ import ( "errors" "fmt" "io" + "math" "net" "net/http" "net/http/httputil" @@ -33,6 +34,7 @@ import ( xrecorder "github.com/go-gost/x/recorder" "golang.org/x/net/http/httpguts" "golang.org/x/net/http2" + "golang.org/x/time/rate" ) const ( @@ -427,44 +429,33 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit if sampleRate == 0 { sampleRate = DefaultSampleRate } - d := time.Duration(1 / sampleRate * 1e9) + if sampleRate < 0 { + sampleRate = math.MaxFloat64 + } go func() { ro2 := &xrecorder.HandlerRecorderObject{} *ro2 = *ro ro := ro2 - ticker := &time.Ticker{} - if d > 0 { - ticker = time.NewTicker(d) - defer ticker.Stop() - } else { - tc := make(chan time.Time) - close(tc) - ticker.C = tc - } + limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate)) buf := &bytes.Buffer{} for { start := time.Now() - err := h.copyWebsocketFrame(cc, rw, buf, "client", ro) - select { - case <-ticker.C: - if err == nil { - 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 { + if err := h.copyWebsocketFrame(cc, rw, buf, "client", ro); err != nil { errc <- err return } + + if limiter.Allow() { + ro.Duration = time.Since(start) + ro.Time = time.Now() + if err := ro.Record(ctx, h.Recorder); err != nil { + log.Errorf("record: %v", err) + } + } } }() @@ -473,37 +464,24 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit *ro2 = *ro ro := ro2 - ticker := &time.Ticker{} - if d > 0 { - ticker = time.NewTicker(d) - defer ticker.Stop() - } else { - tc := make(chan time.Time) - close(tc) - ticker.C = tc - } + limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate)) buf := &bytes.Buffer{} for { start := time.Now() - err := h.copyWebsocketFrame(rw, cc, buf, "server", ro) - select { - case <-ticker.C: - if err == nil { - 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 { + if err := h.copyWebsocketFrame(rw, cc, buf, "server", ro); err != nil { errc <- err return } + + if limiter.Allow() { + ro.Duration = time.Since(start) + ro.Time = time.Now() + if err := ro.Record(ctx, h.Recorder); err != nil { + log.Errorf("record: %v", err) + } + } } }() diff --git a/recorder/recorder.go b/recorder/recorder.go index f9262fb7..2190678e 100644 --- a/recorder/recorder.go +++ b/recorder/recorder.go @@ -89,6 +89,7 @@ type HandlerRecorderObject struct { Route string `json:"route,omitempty"` InputBytes uint64 `json:"inputBytes"` OutputBytes uint64 `json:"outputBytes"` + Redirect string `json:"redirect,omitempty"` Err string `json:"err,omitempty"` SID string `json:"sid"` Duration time.Duration `json:"duration"`