diff --git a/handler/http/connect.go b/handler/http/connect.go index 9d29886..2cf9123 100644 --- a/handler/http/connect.go +++ b/handler/http/connect.go @@ -78,13 +78,15 @@ func (h *httpHandler) handleConnect(ctx context.Context, conn net.Conn, ro *xrec } if h.md.sniffing { - snifferHandled, err = h.sniffAndHandle(ctx, conn, cc, ro, log) + snifferHandled, err = h.sniffAndHandle(ictx.ContextWithSession(ctx, nil), conn, cc, ro, log) if snifferHandled { ro.Time = time.Time{} return err } } + ictx.SessionFromContext(ctx).Start(*ro) + start := time.Now() log.Infof("%s <-> %s", conn.RemoteAddr(), addr) xnet.Pipe(ctx, conn, cc, xnet.WithReadTimeout(h.md.idleTimeout)) @@ -214,6 +216,7 @@ type SnifferBuilder struct { Websocket bool WebsocketSampleRate float64 Recorder recorder.Recorder + SessionRecorder *xrecorder.SessionRecorder RecorderOptions *recorder.Options Certificate *x509.Certificate PrivateKey crypto.PrivateKey @@ -234,6 +237,7 @@ func (b *SnifferBuilder) Build() *sniffing.Sniffer { Websocket: b.Websocket, WebsocketSampleRate: b.WebsocketSampleRate, Recorder: b.Recorder, + SessionRecorder: b.SessionRecorder, RecorderOptions: b.RecorderOptions, Certificate: b.Certificate, PrivateKey: b.PrivateKey, diff --git a/handler/http/handler.go b/handler/http/handler.go index bb52caf..09b9bd6 100644 --- a/handler/http/handler.go +++ b/handler/http/handler.go @@ -210,6 +210,7 @@ import ( "github.com/go-gost/core/recorder" xbypass "github.com/go-gost/x/bypass" xctx "github.com/go-gost/x/ctx" + ictx "github.com/go-gost/x/internal/ctx" xnet "github.com/go-gost/x/internal/net" xhttp "github.com/go-gost/x/internal/net/http" "github.com/go-gost/x/internal/util/httpcache" @@ -231,16 +232,17 @@ func init() { // parsed metadata, handler options, and runtime state such as the traffic // limiter, recorder, certificate pool, and the upstream HTTP transport. type httpHandler struct { - md metadata // parsed configuration - options handler.Options // handler options from the service config - auth *Authenticator // auth + probe resistance (constructed in Init) - sniffer *SnifferBuilder // builds sniffing.Sniffer per connection - stats *stats_util.HandlerStats // per-client stats, created when Observer is set - limiter traffic.TrafficLimiter // per-client traffic shaper (cached) - cancel context.CancelFunc // cancels the observeStats goroutine - recorder recorder.RecorderObject // first matching service-handler recorder - certPool tls_util.CertPool // in-memory cert pool for MITM TLS termination - transport http.RoundTripper // upstream HTTP transport (injectable for tests) + md metadata // parsed configuration + options handler.Options // handler options from the service config + auth *Authenticator // auth + probe resistance (constructed in Init) + sniffer *SnifferBuilder // builds sniffing.Sniffer per connection + stats *stats_util.HandlerStats // per-client stats, created when Observer is set + limiter traffic.TrafficLimiter // per-client traffic shaper (cached) + cancel context.CancelFunc // cancels the observeStats goroutine + recorder recorder.RecorderObject // first matching service-handler recorder + sessionRecorder *xrecorder.SessionRecorder // reports live sessions on an interval when recorder.period is set. + certPool tls_util.CertPool // in-memory cert pool for MITM TLS termination + transport http.RoundTripper // upstream HTTP transport (injectable for tests) } // NewHandler creates a new HTTP handler and applies the given options. @@ -295,10 +297,15 @@ func (h *httpHandler) Init(md md.Metadata) error { Log: h.options.Logger, } + h.sessionRecorder = xrecorder.NewSessionRecorder(h.recorder.Recorder, xrecorder.SessionRecorderOptions{ + Period: h.md.recorderPeriod, + }) + h.sniffer = &SnifferBuilder{ Websocket: h.md.sniffingWebsocket, WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, Recorder: h.recorder.Recorder, + SessionRecorder: h.sessionRecorder, RecorderOptions: h.recorder.Options, Certificate: h.md.certificate, PrivateKey: h.md.privateKey, @@ -364,6 +371,9 @@ func (h *httpHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler pStats := xstats.Stats{} conn = stats_wrapper.WrapConn(conn, &pStats) + session := h.sessionRecorder.NewSession(ctx, &pStats) + ctx = ictx.ContextWithSession(ctx, session) + defer func() { if err != nil { ro.Err = err.Error() @@ -371,7 +381,7 @@ func (h *httpHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler ro.InputBytes = pStats.Get(stats.KindInputBytes) ro.OutputBytes = pStats.Get(stats.KindOutputBytes) ro.Duration = time.Since(start) - if err := ro.Record(ctx, h.recorder.Recorder); err != nil { + if err := session.Finish(ctx, *ro); err != nil { log.Error("record: %v", err) } @@ -539,7 +549,7 @@ func (h *httpHandler) Close() error { if h.cancel != nil { h.cancel() } - return nil + return h.sessionRecorder.Close() } // buildHTTPRecorder creates an HTTPRecorderObject from the request metadata. diff --git a/handler/http/metadata.go b/handler/http/metadata.go index 41f0362..d6ab362 100644 --- a/handler/http/metadata.go +++ b/handler/http/metadata.go @@ -23,6 +23,8 @@ const ( // Fields are populated by parseMetadata and read throughout the handler's // request-processing methods. type metadata struct { + recorderPeriod time.Duration // reports a live session on this interval; 0 = one record per session, written when it ends. + readTimeout time.Duration // deadline for upstream response headers (http.Transport.ResponseHeaderTimeout); 0=15s default, negative=disabled idleTimeout time.Duration // idle read deadline per Pipe direction during CONNECT/forwarding; 0 or negative = disabled keepalive bool // enable HTTP keep-alive on the upstream transport @@ -59,6 +61,8 @@ type metadata struct { // observerPeriod (5s, min 1s), and proxyAgent ("gost/3.0"). MITM TLS is // enabled when both mitm.certFile and mitm.keyFile are provided. func (h *httpHandler) parseMetadata(md mdata.Metadata) error { + h.md.recorderPeriod = mdutil.GetDuration(md, "recorder.period", "recorder.reportPeriod") + h.md.readTimeout = mdutil.GetDuration(md, "readTimeout") if h.md.readTimeout == 0 { h.md.readTimeout = 15 * time.Second diff --git a/handler/http/proxy.go b/handler/http/proxy.go index f1339ea..3fdaa06 100644 --- a/handler/http/proxy.go +++ b/handler/http/proxy.go @@ -8,13 +8,14 @@ import ( "io" "net" "net/http" + "net/http/httptrace" "net/http/httputil" "strings" "time" "github.com/go-gost/core/bypass" - stats "github.com/go-gost/core/observer/stats" "github.com/go-gost/core/logger" + stats "github.com/go-gost/core/observer/stats" xbypass "github.com/go-gost/x/bypass" ictx "github.com/go-gost/x/internal/ctx" xio "github.com/go-gost/x/internal/io" @@ -118,6 +119,10 @@ func (h *httpHandler) proxyRoundTrip(ctx context.Context, rw io.ReadWriteCloser, "host": host, }) + // A keep-alive connection carries several exchanges; each is accounted for + // on its own, under the connection's shared counters. + session := h.sessionRecorder.NewSession(ctx, pStats) + log.Infof("%s <-> %s", ro.RemoteAddr, req.Host) defer func() { if err != nil { @@ -126,7 +131,7 @@ func (h *httpHandler) proxyRoundTrip(ctx context.Context, rw io.ReadWriteCloser, ro.InputBytes = pStats.Get(stats.KindInputBytes) ro.OutputBytes = pStats.Get(stats.KindOutputBytes) ro.Duration = time.Since(ro.Time) - if err := ro.Record(ctx, h.recorder.Recorder); err != nil { + if err := session.Finish(ctx, *ro); err != nil { log.Errorf("record: %v", err) } @@ -202,7 +207,18 @@ func (h *httpHandler) proxyRoundTrip(ctx context.Context, rw io.ReadWriteCloser, ctx = ictx.ContextWithRecorderObject(ctx, ro) ctx = ictx.ContextWithLogger(ctx, log) - resp, err := h.transport.RoundTrip(req.WithContext(ctx)) + // Report from the moment the upstream connection is in hand, and only for + // this request: a recorder of its own would otherwise inherit the callback + // and overwrite the addresses with its own. + traceCtx := httptrace.WithClientTrace(ctx, &httptrace.ClientTrace{ + GotConn: func(info httptrace.GotConnInfo) { + ro.SrcAddr = info.Conn.LocalAddr().String() + ro.DstAddr = info.Conn.RemoteAddr().String() + session.Start(*ro) + }, + }) + + resp, err := h.transport.RoundTrip(req.WithContext(traceCtx)) if reqBody != nil { ro.HTTP.Request.Body = reqBody.Content() diff --git a/handler/redirect/tcp/handler.go b/handler/redirect/tcp/handler.go index 74ac6d2..0f644d7 100644 --- a/handler/redirect/tcp/handler.go +++ b/handler/redirect/tcp/handler.go @@ -35,10 +35,11 @@ func init() { } type redirectHandler struct { - md metadata - options handler.Options - recorder recorder.RecorderObject - certPool tls_util.CertPool + md metadata + options handler.Options + recorder recorder.RecorderObject + sessionRecorder *xrecorder.SessionRecorder + certPool tls_util.CertPool } func NewHandler(opts ...handler.Option) handler.Handler { @@ -64,6 +65,10 @@ func (h *redirectHandler) Init(md md.Metadata) (err error) { } } + h.sessionRecorder = xrecorder.NewSessionRecorder(h.recorder.Recorder, xrecorder.SessionRecorderOptions{ + Period: h.md.recorderPeriod, + }) + if h.md.certificate != nil && h.md.privateKey != nil { h.certPool = tls_util.NewMemoryCertPool() } @@ -101,6 +106,9 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han pStats := xstats.Stats{} conn = stats_wrapper.WrapConn(conn, &pStats) + session := h.sessionRecorder.NewSession(ctx, &pStats) + ctx = ictx.ContextWithSession(ctx, session) + defer func() { if err != nil { ro.Err = err.Error() @@ -108,7 +116,7 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han ro.InputBytes = pStats.Get(stats.KindInputBytes) ro.OutputBytes = pStats.Get(stats.KindOutputBytes) ro.Duration = time.Since(start) - if err := ro.Record(ctx, h.recorder.Recorder); err != nil { + if err := session.Finish(ctx, *ro); err != nil { log.Errorf("record: %v", err) } @@ -220,6 +228,7 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han Websocket: h.md.sniffingWebsocket, WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, Recorder: h.recorder.Recorder, + SessionRecorder: h.sessionRecorder, RecorderOptions: h.recorder.Options, Certificate: h.md.certificate, PrivateKey: h.md.privateKey, @@ -298,6 +307,8 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han ro.SrcAddr = cc.LocalAddr().String() ro.DstAddr = cc.RemoteAddr().String() + session.Start(*ro) + t := time.Now() log.Infof("%s <-> %s", conn.RemoteAddr(), dstAddr) // xnet.Transport(conn, cc) @@ -320,3 +331,5 @@ func (h *redirectHandler) checkRateLimit(addr net.Addr) bool { return true } + +func (h *redirectHandler) Close() error { return h.sessionRecorder.Close() } diff --git a/handler/redirect/tcp/metadata.go b/handler/redirect/tcp/metadata.go index d34d166..eca7b40 100644 --- a/handler/redirect/tcp/metadata.go +++ b/handler/redirect/tcp/metadata.go @@ -13,6 +13,7 @@ import ( ) type metadata struct { + recorderPeriod time.Duration // reports a live session on this interval; 0 = one record per session, written when it ends. // readTimeout is passed to SnifferBuilder as the timeout for reading // upstream response headers during HTTP/TLS sniffing. It is NOT used // as a deadline on the initial client connection (unlike socks/ss @@ -35,6 +36,8 @@ type metadata struct { } func (h *redirectHandler) parseMetadata(md mdata.Metadata) (err error) { + h.md.recorderPeriod = mdutil.GetDuration(md, "recorder.period", "recorder.reportPeriod") + h.md.readTimeout = mdutil.GetDuration(md, "readTimeout") if h.md.readTimeout <= 0 { h.md.readTimeout = 15 * time.Second diff --git a/handler/redirect/udp/handler.go b/handler/redirect/udp/handler.go index 6f058c9..80dce78 100644 --- a/handler/redirect/udp/handler.go +++ b/handler/redirect/udp/handler.go @@ -31,9 +31,10 @@ func init() { } type redirectHandler struct { - md metadata - options handler.Options - recorder recorder.RecorderObject + md metadata + options handler.Options + recorder recorder.RecorderObject + sessionRecorder *xrecorder.SessionRecorder } func NewHandler(opts ...handler.Option) handler.Handler { @@ -59,6 +60,10 @@ func (h *redirectHandler) Init(md md.Metadata) (err error) { } } + h.sessionRecorder = xrecorder.NewSessionRecorder(h.recorder.Recorder, xrecorder.SessionRecorderOptions{ + Period: h.md.recorderPeriod, + }) + return } @@ -93,6 +98,8 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han pStats := xstats.Stats{} conn = stats_wrapper.WrapConn(conn, &pStats) + session := h.sessionRecorder.NewSession(ctx, &pStats) + defer func() { if err != nil { ro.Err = err.Error() @@ -100,7 +107,7 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han ro.Duration = time.Since(start) ro.InputBytes = pStats.Get(stats.KindInputBytes) ro.OutputBytes = pStats.Get(stats.KindOutputBytes) - if err := ro.Record(ctx, h.recorder.Recorder); err != nil { + if err := session.Finish(ctx, *ro); err != nil { log.Error("record: %v", err) } @@ -213,6 +220,8 @@ func (h *redirectHandler) Handle(ctx context.Context, conn net.Conn, opts ...han ro.SrcAddr = cc.LocalAddr().String() ro.DstAddr = cc.RemoteAddr().String() + session.Start(*ro) + t := time.Now() log.Infof("%s <-> %s", conn.RemoteAddr(), dstAddr) // xnet.Transport(conn, cc) @@ -235,3 +244,5 @@ func (h *redirectHandler) checkRateLimit(addr net.Addr) bool { return true } + +func (h *redirectHandler) Close() error { return h.sessionRecorder.Close() } diff --git a/handler/redirect/udp/metadata.go b/handler/redirect/udp/metadata.go index 0a1ed43..fe3d75e 100644 --- a/handler/redirect/udp/metadata.go +++ b/handler/redirect/udp/metadata.go @@ -8,11 +8,15 @@ import ( ) type metadata struct { + recorderPeriod time.Duration // reports a live session on this interval; 0 = one record per session, written when it ends. + sniffing bool sniffingTimeout time.Duration } func (h *redirectHandler) parseMetadata(md mdata.Metadata) (err error) { + h.md.recorderPeriod = mdutil.GetDuration(md, "recorder.period", "recorder.reportPeriod") + h.md.sniffing = mdutil.GetBool(md, "sniffing") h.md.sniffingTimeout = mdutil.GetDuration(md, "sniffing.timeout") diff --git a/handler/socks/v5/bind.go b/handler/socks/v5/bind.go index ec0894f..b15a732 100644 --- a/handler/socks/v5/bind.go +++ b/handler/socks/v5/bind.go @@ -10,6 +10,7 @@ import ( "github.com/go-gost/core/observer/stats" "github.com/go-gost/gosocks5" xctx "github.com/go-gost/x/ctx" + ictx "github.com/go-gost/x/internal/ctx" xnet "github.com/go-gost/x/internal/net" traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" stats_wrapper "github.com/go-gost/x/observer/stats/wrapper" @@ -79,6 +80,7 @@ func (h *socks5Handler) bindLocal(ctx context.Context, conn net.Conn, network, a "bind": ln.Addr().String(), }) ro.SrcAddr = ln.Addr().String() + ictx.SessionFromContext(ctx).Start(*ro) socksAddr := gosocks5.Addr{} if err := socksAddr.ParseFrom(ln.Addr().String()); err != nil { diff --git a/handler/socks/v5/connect.go b/handler/socks/v5/connect.go index 0f2160c..6410d80 100644 --- a/handler/socks/v5/connect.go +++ b/handler/socks/v5/connect.go @@ -115,6 +115,7 @@ func (h *socks5Handler) handleConnect(ctx context.Context, conn net.Conn, networ Websocket: h.md.sniffingWebsocket, WebsocketSampleRate: h.md.sniffingWebsocketSampleRate, Recorder: h.recorder.Recorder, + SessionRecorder: h.sessionRecorder, RecorderOptions: h.recorder.Options, Certificate: h.md.certificate, PrivateKey: h.md.privateKey, @@ -147,6 +148,8 @@ func (h *socks5Handler) handleConnect(ctx context.Context, conn net.Conn, networ } } + ictx.SessionFromContext(ctx).Start(*ro) + t := time.Now() log.Infof("%s <-> %s", conn.RemoteAddr(), address) // xnet.Transport(conn, cc) diff --git a/handler/socks/v5/handler.go b/handler/socks/v5/handler.go index c8137c6..ccff733 100644 --- a/handler/socks/v5/handler.go +++ b/handler/socks/v5/handler.go @@ -15,6 +15,7 @@ import ( "github.com/go-gost/core/recorder" "github.com/go-gost/gosocks5" xctx "github.com/go-gost/x/ctx" + ictx "github.com/go-gost/x/internal/ctx" "github.com/go-gost/x/internal/util/socks" stats_util "github.com/go-gost/x/internal/util/stats" tls_util "github.com/go-gost/x/internal/util/tls" @@ -36,14 +37,15 @@ func init() { } type socks5Handler struct { - selector gosocks5.Selector - md metadata - options handler.Options - stats *stats_util.HandlerStats - limiter traffic.TrafficLimiter - cancel context.CancelFunc - recorder recorder.RecorderObject - certPool tls_util.CertPool + selector gosocks5.Selector + md metadata + options handler.Options + stats *stats_util.HandlerStats + limiter traffic.TrafficLimiter + cancel context.CancelFunc + recorder recorder.RecorderObject + sessionRecorder *xrecorder.SessionRecorder + certPool tls_util.CertPool } func NewHandler(opts ...handler.Option) handler.Handler { @@ -93,6 +95,10 @@ func (h *socks5Handler) Init(md md.Metadata) (err error) { } } + h.sessionRecorder = xrecorder.NewSessionRecorder(h.recorder.Recorder, xrecorder.SessionRecorderOptions{ + Period: h.md.recorderPeriod, + }) + if h.md.certificate != nil && h.md.privateKey != nil { h.certPool = tls_util.NewMemoryCertPool() } @@ -150,6 +156,9 @@ func (h *socks5Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl pStats := xstats.Stats{} conn = stats_wrapper.WrapConn(conn, &pStats) + session := h.sessionRecorder.NewSession(ctx, &pStats) + ctx = ictx.ContextWithSession(ctx, session) + defer func() { if err != nil { ro.Err = err.Error() @@ -157,7 +166,7 @@ func (h *socks5Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl ro.InputBytes += pStats.Get(stats.KindInputBytes) ro.OutputBytes += pStats.Get(stats.KindOutputBytes) ro.Duration = time.Since(start) - if err := ro.Record(ctx, h.recorder.Recorder); err != nil { + if err := session.Finish(ctx, *ro); err != nil { log.Errorf("record: %v", err) } @@ -204,7 +213,7 @@ func (h *socks5Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl return h.handleMuxBind(ctx, conn, networkAddr("tcp", req.Addr), address, ro, log) case gosocks5.CmdUdp: ro.Network = "udp" - return h.handleUDP(ctx, conn, "udp", ro, log) + return h.handleUDP(ctx, conn, "udp", ro, &pStats, log) case socks.CmdUDPTun: ro.Network = "udp" return h.handleUDPTun(ctx, conn, networkAddr("udp", req.Addr), address, ro, log) @@ -226,7 +235,7 @@ func (h *socks5Handler) Close() error { if h.cancel != nil { h.cancel() } - return nil + return h.sessionRecorder.Close() } func (h *socks5Handler) checkRateLimit(addr net.Addr) bool { diff --git a/handler/socks/v5/mbind.go b/handler/socks/v5/mbind.go index 2dc0f12..8b6b594 100644 --- a/handler/socks/v5/mbind.go +++ b/handler/socks/v5/mbind.go @@ -10,6 +10,7 @@ import ( "github.com/go-gost/core/observer/stats" "github.com/go-gost/gosocks5" xctx "github.com/go-gost/x/ctx" + ictx "github.com/go-gost/x/internal/ctx" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/util/mux" traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" @@ -80,6 +81,7 @@ func (h *socks5Handler) muxBindLocal(ctx context.Context, conn net.Conn, network }) ro.SrcAddr = ln.Addr().String() + ictx.SessionFromContext(ctx).Start(*ro) socksAddr := gosocks5.Addr{} err = socksAddr.ParseFrom(ln.Addr().String()) if err != nil { diff --git a/handler/socks/v5/metadata.go b/handler/socks/v5/metadata.go index 23d166c..3d9373e 100644 --- a/handler/socks/v5/metadata.go +++ b/handler/socks/v5/metadata.go @@ -14,6 +14,8 @@ import ( ) type metadata struct { + recorderPeriod time.Duration // reports a live session on this interval; 0 = one record per session, written when it ends. + publicAddr string // readTimeout is the deadline for reading the initial SOCKS5 // handshake (auth + connect/associate/udp request) from the client @@ -52,6 +54,8 @@ type metadata struct { } func (h *socks5Handler) parseMetadata(md mdata.Metadata) (err error) { + h.md.recorderPeriod = mdutil.GetDuration(md, "recorder.period", "recorder.reportPeriod") + h.md.publicAddr = mdutil.GetString(md, "socks.publicAddr", "publicAddr") h.md.readTimeout = mdutil.GetDuration(md, "readTimeout") if h.md.readTimeout <= 0 { diff --git a/handler/socks/v5/udp.go b/handler/socks/v5/udp.go index 00f56f5..d36d270 100644 --- a/handler/socks/v5/udp.go +++ b/handler/socks/v5/udp.go @@ -25,12 +25,11 @@ import ( "github.com/go-gost/x/internal/util/socks" traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" - xstats "github.com/go-gost/x/observer/stats" stats_wrapper "github.com/go-gost/x/observer/stats/wrapper" xrecorder "github.com/go-gost/x/recorder" ) -func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, network string, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error { +func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, network string, ro *xrecorder.HandlerRecorderObject, pStats stats.Stats, log logger.Logger) error { log = log.WithFields(map[string]any{ "network": network, "cmd": network, @@ -129,13 +128,10 @@ func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, network st pc = metrics.WrapPacketConn(ro.Service, pc) { - pStats := xstats.Stats{} - cc = stats_wrapper.WrapPacketConn(cc, &pStats) - - defer func() { - ro.InputBytes += pStats.Get(stats.KindInputBytes) - ro.OutputBytes += pStats.Get(stats.KindOutputBytes) - }() + // The association's relay socket feeds the same counters as its control + // connection, so one session accounts for the whole association -- the + // two were already summed into the record. + cc = stats_wrapper.WrapPacketConn(cc, pStats) clientID := ctxvalue.ClientIDFromContext(ctx) cc = traffic_wrapper.WrapPacketConn( @@ -179,6 +175,8 @@ func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, network st } } + ictx.SessionFromContext(ctx).Start(*ro) + r := udp.NewRelay(pc1, pc). WithService(h.options.Service). WithBypass(h.options.Bypass). diff --git a/handler/socks/v5/udp_tun.go b/handler/socks/v5/udp_tun.go index 6218d77..a7ee600 100644 --- a/handler/socks/v5/udp_tun.go +++ b/handler/socks/v5/udp_tun.go @@ -135,6 +135,7 @@ func (h *socks5Handler) handleUDPTun(ctx context.Context, conn net.Conn, network "bind": pc.LocalAddr().String(), }) ro.SrcAddr = pc.LocalAddr().String() + ictx.SessionFromContext(ctx).Start(*ro) saddr := gosocks5.Addr{} saddr.ParseFrom(pc.LocalAddr().String()) diff --git a/internal/ctx/value.go b/internal/ctx/value.go index b5a7b3b..9df9056 100644 --- a/internal/ctx/value.go +++ b/internal/ctx/value.go @@ -64,3 +64,14 @@ func RecorderObjectFromContext(ctx context.Context) *xrecorder.HandlerRecorderOb v, _ := ctx.Value(recorderObjectCtxKey{}).(*xrecorder.HandlerRecorderObject) return v } + +type sessionCtxKey struct{} + +func ContextWithSession(ctx context.Context, s *xrecorder.Session) context.Context { + return context.WithValue(ctx, sessionCtxKey{}, s) +} + +func SessionFromContext(ctx context.Context) *xrecorder.Session { + v, _ := ctx.Value(sessionCtxKey{}).(*xrecorder.Session) + return v +} diff --git a/internal/util/sniffing/sniffer.go b/internal/util/sniffing/sniffer.go index c57ac5e..17c1789 100644 --- a/internal/util/sniffing/sniffer.go +++ b/internal/util/sniffing/sniffer.go @@ -127,6 +127,7 @@ type Sniffer struct { Recorder recorder.Recorder RecorderOptions *recorder.Options + SessionRecorder *xrecorder.SessionRecorder // MITM TLS termination Certificate *x509.Certificate @@ -187,3 +188,10 @@ func (h *Sniffer) effectiveReadTimeout() time.Duration { } return DefaultReadTimeout } + +func (h *Sniffer) sessionRecorder() *xrecorder.SessionRecorder { + if h.SessionRecorder != nil { + return h.SessionRecorder + } + return xrecorder.NewSessionRecorder(h.Recorder, xrecorder.SessionRecorderOptions{}) +} diff --git a/internal/util/sniffing/sniffer_h2.go b/internal/util/sniffing/sniffer_h2.go index ef8e93d..52ca96e 100644 --- a/internal/util/sniffing/sniffer_h2.go +++ b/internal/util/sniffing/sniffer_h2.go @@ -13,8 +13,12 @@ import ( "time" "github.com/go-gost/core/logger" + "github.com/go-gost/core/observer/stats" "github.com/go-gost/core/recorder" + xio "github.com/go-gost/x/internal/io" xhttp "github.com/go-gost/x/internal/net/http" + xstats "github.com/go-gost/x/observer/stats" + stats_wrapper "github.com/go-gost/x/observer/stats/wrapper" xrecorder "github.com/go-gost/x/recorder" "golang.org/x/net/http2" ) @@ -64,7 +68,7 @@ func (h *Sniffer) serveH2(ctx context.Context, network string, conn net.Conn, ho SawClientPreface: true, Handler: &h2Handler{ transport: tr, - recorder: h.Recorder, + sessionRecorder: h.sessionRecorder(), recorderOptions: h.RecorderOptions, recorderObject: ro, log: log, @@ -77,7 +81,7 @@ func (h *Sniffer) serveH2(ctx context.Context, network string, conn net.Conn, ho // http2.Transport while recording request and response metadata. type h2Handler struct { transport http.RoundTripper - recorder recorder.Recorder + sessionRecorder *xrecorder.SessionRecorder recorderOptions *recorder.Options recorderObject *xrecorder.HandlerRecorderObject log logger.Logger @@ -96,6 +100,9 @@ func (h *h2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { ro.Time = time.Now() + var counters xstats.Stats + session := h.sessionRecorder.NewSession(r.Context(), &counters) + var err error log.Infof("%s <-> %s", ro.RemoteAddr, r.Host) defer func() { @@ -103,7 +110,9 @@ func (h *h2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { if err != nil { ro.Err = err.Error() } - if rerr := ro.Record(r.Context(), h.recorder); rerr != nil { + ro.InputBytes = counters.Get(stats.KindInputBytes) + ro.OutputBytes = counters.Get(stats.KindOutputBytes) + if rerr := session.Finish(r.Context(), *ro); rerr != nil { log.Errorf("record: %v", rerr) } @@ -132,6 +141,10 @@ func (h *h2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { log.Trace(string(dump)) } + if r.Body != nil { + r.Body = countBody(r.Body, &counters) + } + url := r.URL url.Scheme = "https" url.Host = r.Host @@ -176,14 +189,15 @@ func (h *h2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.setHeader(w, resp.Header) w.WriteHeader(resp.StatusCode) + counted := stats_wrapper.WrapReadWriter(xio.NewReadWriter(nil, w), &counters) if bodySize := ClampBodySize(h.recorderOptions); bodySize > 0 && ro.RecordMode != "headers" && ro.RecordMode != "off" { respBody := xhttp.NewBody(resp.Body, bodySize) resp.Body = respBody - io.Copy(w, resp.Body) + io.Copy(counted, resp.Body) ro.HTTP.Response.Body = respBody.Content() ro.HTTP.Response.ContentLength = respBody.Length() } else { - io.Copy(w, resp.Body) + io.Copy(counted, resp.Body) } } @@ -194,3 +208,8 @@ func (h *h2Handler) setHeader(w http.ResponseWriter, header http.Header) { } } } + +func countBody(body io.ReadCloser, counters stats.Stats) io.ReadCloser { + counted := stats_wrapper.WrapReadWriter(xio.NewReadWriter(body, io.Discard), counters) + return xio.NewReadWriteCloser(counted, io.Discard, body) +} diff --git a/internal/util/sniffing/sniffer_http.go b/internal/util/sniffing/sniffer_http.go index f0da963..15ef722 100644 --- a/internal/util/sniffing/sniffer_http.go +++ b/internal/util/sniffing/sniffer_http.go @@ -239,6 +239,8 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, req.Header.Del("Gost-Record") ro.Time = time.Now() + session := h.sessionRecorder().NewSession(ctx, pStats) + session.Start(*ro) log.Infof("%s <-> %s", ro.RemoteAddr, req.Host) defer func() { if err != nil { @@ -247,7 +249,7 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, ro.InputBytes = pStats.Get(stats.KindInputBytes) ro.OutputBytes = pStats.Get(stats.KindOutputBytes) ro.Duration = time.Since(ro.Time) - if rerr := ro.Record(ctx, h.Recorder); rerr != nil { + if rerr := session.Finish(ctx, *ro); rerr != nil { log.Errorf("record: %v", rerr) } diff --git a/internal/util/sniffing/sniffer_test.go b/internal/util/sniffing/sniffer_test.go index c972605..1735009 100644 --- a/internal/util/sniffing/sniffer_test.go +++ b/internal/util/sniffing/sniffer_test.go @@ -787,7 +787,7 @@ func TestH2HandlerServeHTTP_Success(t *testing.T) { Body: io.NopCloser(strings.NewReader("h2 body")), }, }, - recorder: &noopRecorder{}, + sessionRecorder: xrecorder.NewSessionRecorder(&noopRecorder{}, xrecorder.SessionRecorderOptions{}), recorderOptions: &recorder.Options{HTTPBody: false}, recorderObject: &xrecorder.HandlerRecorderObject{}, log: xlogger.Nop(), @@ -812,7 +812,7 @@ func TestH2HandlerServeHTTP_RoundTripError(t *testing.T) { transport: &mockRoundTripper{ err: errors.New("connection refused"), }, - recorder: &noopRecorder{}, + sessionRecorder: xrecorder.NewSessionRecorder(&noopRecorder{}, xrecorder.SessionRecorderOptions{}), recorderOptions: &recorder.Options{HTTPBody: false}, recorderObject: &xrecorder.HandlerRecorderObject{}, log: xlogger.Nop(), @@ -838,7 +838,7 @@ func TestH2HandlerServeHTTP_WithBodyRecording(t *testing.T) { Body: io.NopCloser(strings.NewReader(`{"key":"value"}`)), }, }, - recorder: &noopRecorder{}, + sessionRecorder: xrecorder.NewSessionRecorder(&noopRecorder{}, xrecorder.SessionRecorderOptions{}), recorderOptions: &recorder.Options{HTTPBody: true, MaxBodySize: 1024}, recorderObject: &xrecorder.HandlerRecorderObject{}, log: xlogger.Nop(), diff --git a/internal/util/sniffing/sniffer_tls.go b/internal/util/sniffing/sniffer_tls.go index 9e6dbce..3c1c14a 100644 --- a/internal/util/sniffing/sniffer_tls.go +++ b/internal/util/sniffing/sniffer_tls.go @@ -13,6 +13,7 @@ import ( "github.com/go-gost/core/bypass" dissector "github.com/go-gost/tls-dissector" xbypass "github.com/go-gost/x/bypass" + ictx "github.com/go-gost/x/internal/ctx" xio "github.com/go-gost/x/internal/io" xnet "github.com/go-gost/x/internal/net" tls_util "github.com/go-gost/x/internal/util/tls" @@ -111,6 +112,7 @@ func (h *Sniffer) HandleTLS(ctx context.Context, network string, conn net.Conn, return err } + ictx.SessionFromContext(ctx).Start(*ro) log.Infof("%s <-> %s", ro.RemoteAddr, ro.Host) xnet.Pipe(ctx, conn, cc) log.WithFields(map[string]any{ diff --git a/recorder/session.go b/recorder/session.go new file mode 100644 index 0000000..5f36356 --- /dev/null +++ b/recorder/session.go @@ -0,0 +1,125 @@ +package recorder + +import ( + "cmp" + "context" + "sync" + "time" + + "github.com/go-gost/core/observer/stats" +) + +type Session struct { + recorder *SessionRecorder + counters stats.Stats + labels map[string]string + id string + + mu sync.Mutex + base HandlerRecorderObject + cursor cursor + finished bool +} + +func sessionMetadata(base HandlerRecorderObject, labels map[string]string) HandlerRecorderObject { + base.HTTP, base.TLS, base.DNS, base.Websocket, base.Redis = nil, nil, nil, nil, nil + if base.Labels == nil { + base.Labels = labels + } + + return base +} + +func (s *Session) Start(base HandlerRecorderObject) { + if s == nil || s.recorder == nil || !s.recorder.Enabled() { + return + } + + s.mu.Lock() + defer s.mu.Unlock() + if s.finished { + return + } + + base = sessionMetadata(base, s.labels) + if base.RecordMode == "off" { + return + } + + if !s.cursor.started() { + s.cursor.open(cmp.Or(base.Time, time.Now())) + s.recorder.track(s) + } + + base.Time = s.cursor.startedAt + s.base = base +} + +func (s *Session) Finish(ctx context.Context, final HandlerRecorderObject) error { + if s == nil || s.recorder == nil { + return nil + } + r := s.recorder + + s.mu.Lock() + defer s.mu.Unlock() + if s.finished { + return nil + } + s.finished = true + + r.untrack(s) + + if final.Time.IsZero() || final.RecordMode == "off" { + return nil + } + if final.Labels == nil { + final.Labels = s.labels + } + + if !r.Enabled() { + return r.write(ctx, final) + } + if r.isClosed() { + return ErrSessionRecorderClosed + } + if !s.cursor.started() { + s.cursor.open(final.Time) + } + + return s.emit(ctx, final, time.Now(), PhaseFinal) +} + +func (s *Session) Interim(ctx context.Context, now time.Time) error { + if s == nil || s.recorder == nil || !s.recorder.Enabled() { + return nil + } + + s.mu.Lock() + defer s.mu.Unlock() + if s.finished || s.counters == nil || !now.After(s.cursor.last) { + return nil + } + if s.base.RecordMode == "off" { + return nil + } + + ro := s.base + ro.InputBytes = s.counters.Get(stats.KindInputBytes) + ro.OutputBytes = s.counters.Get(stats.KindOutputBytes) + if s.cursor.idle(ro) { + return nil + } + + return s.emit(ctx, ro, now, PhaseInterim) +} + +func (s *Session) emit(ctx context.Context, ro HandlerRecorderObject, now time.Time, phase string) error { + record := s.cursor.record(ro, now, phase, s.id) + if err := s.recorder.write(ctx, record); err != nil { + return err + } + s.cursor.commit(record) + + return nil +} diff --git a/recorder/session_record.go b/recorder/session_record.go new file mode 100644 index 0000000..026ac95 --- /dev/null +++ b/recorder/session_record.go @@ -0,0 +1,59 @@ +package recorder + +import "time" + +const ( + PhaseInterim = "interim" + PhaseFinal = "final" +) + +type SessionRecord struct { + HandlerRecorderObject + SessionID string `json:"sessionID"` + RecordIndex uint64 `json:"recordIndex"` + Phase string `json:"phase"` + InputBytesDelta uint64 `json:"inputBytesDelta"` + OutputBytesDelta uint64 `json:"outputBytesDelta"` + DurationDelta time.Duration `json:"durationDelta"` +} + +type cursor struct { + startedAt time.Time + last time.Time + recordIndex uint64 + input uint64 + output uint64 +} + +func (c *cursor) open(at time.Time) { + c.startedAt, c.last = at, at +} + +func (c *cursor) started() bool { return !c.startedAt.IsZero() } + +func (c *cursor) idle(ro HandlerRecorderObject) bool { + return ro.InputBytes == c.input && ro.OutputBytes == c.output +} +func (c *cursor) record(ro HandlerRecorderObject, now time.Time, phase, id string) SessionRecord { + input := max(ro.InputBytes, c.input) + output := max(ro.OutputBytes, c.output) + ro.InputBytes, ro.OutputBytes = input, output + ro.Time = now + ro.Duration = now.Sub(c.startedAt) + + return SessionRecord{ + HandlerRecorderObject: ro, + SessionID: id, + RecordIndex: c.recordIndex + 1, + Phase: phase, + InputBytesDelta: input - c.input, + OutputBytesDelta: output - c.output, + DurationDelta: now.Sub(c.last), + } +} + +func (c *cursor) commit(record SessionRecord) { + c.input, c.output = record.InputBytes, record.OutputBytes + c.recordIndex = record.RecordIndex + c.last = record.Time +} diff --git a/recorder/session_recorder.go b/recorder/session_recorder.go new file mode 100644 index 0000000..8659938 --- /dev/null +++ b/recorder/session_recorder.go @@ -0,0 +1,158 @@ +package recorder + +import ( + "context" + "crypto/rand" + "encoding/json" + "errors" + "fmt" + "maps" + "sync" + "time" + + "github.com/go-gost/core/observer/stats" + "github.com/go-gost/core/recorder" + xctx "github.com/go-gost/x/ctx" +) + +const MinPeriod = time.Second + +var ErrSessionRecorderClosed = errors.New("session recorder is closed") + +type SessionRecorderOptions struct { + Period time.Duration + WriteTimeout time.Duration +} + +type SessionRecorder struct { + rec recorder.Recorder + opts SessionRecorderOptions + + stop chan struct{} + done chan struct{} + + mu sync.Mutex + sessions map[*Session]struct{} + closed bool + + closeOnce sync.Once + closeErr error +} + +func NewSessionRecorder(rec recorder.Recorder, opts SessionRecorderOptions) *SessionRecorder { + if opts.Period > 0 && opts.Period < MinPeriod { + opts.Period = MinPeriod + } + if opts.WriteTimeout <= 0 { + opts.WriteTimeout = 5 * time.Second + } + + r := &SessionRecorder{rec: rec, opts: opts, sessions: make(map[*Session]struct{})} + if !r.Enabled() { + return r + } + + r.stop = make(chan struct{}) + r.done = make(chan struct{}) + go r.sampleSessions() + + return r +} + +func (r *SessionRecorder) Enabled() bool { return r != nil && r.rec != nil && r.opts.Period > 0 } + +func (r *SessionRecorder) NewSession(ctx context.Context, counters stats.Stats) *Session { + s := &Session{recorder: r, counters: counters, labels: maps.Clone(xctx.LabelsFromContext(ctx))} + if r.Enabled() { + s.id = rand.Text() + } + + return s +} + +func (r *SessionRecorder) Close() error { + if !r.Enabled() { + return nil + } + + r.closeOnce.Do(func() { + r.mu.Lock() + r.closed = true + active := len(r.sessions) + r.mu.Unlock() + + close(r.stop) + <-r.done + + if active > 0 { + r.closeErr = fmt.Errorf("session recorder closed with %d active sessions", active) + } + }) + + return r.closeErr +} + +func (r *SessionRecorder) sampleSessions() { + defer close(r.done) + + ticker := time.NewTicker(r.opts.Period) + defer ticker.Stop() + + var batch []*Session + for { + select { + case <-r.stop: + return + case now := <-ticker.C: + r.mu.Lock() + batch = batch[:0] + for s := range r.sessions { + batch = append(batch, s) + } + r.mu.Unlock() + + for _, s := range batch { + _ = s.Interim(context.Background(), now) + } + } + } +} + +func (r *SessionRecorder) write(ctx context.Context, record any) error { + if r.rec == nil { + return nil + } + + data, err := json.Marshal(record) + if err != nil { + return err + } + + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), r.opts.WriteTimeout) + defer cancel() + + return r.rec.Record(ctx, data) +} + +func (r *SessionRecorder) track(s *Session) { + r.mu.Lock() + defer r.mu.Unlock() + + if !r.closed { + r.sessions[s] = struct{}{} + } +} + +func (r *SessionRecorder) untrack(s *Session) { + r.mu.Lock() + defer r.mu.Unlock() + + delete(r.sessions, s) +} + +func (r *SessionRecorder) isClosed() bool { + r.mu.Lock() + defer r.mu.Unlock() + + return r.closed +} diff --git a/recorder/session_test.go b/recorder/session_test.go new file mode 100644 index 0000000..e894d97 --- /dev/null +++ b/recorder/session_test.go @@ -0,0 +1,449 @@ +package recorder + +import ( + "context" + "encoding/json" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/go-gost/core/observer/stats" + "github.com/go-gost/core/recorder" + xstats "github.com/go-gost/x/observer/stats" +) + +// sessionSink collects what the reporter delivers. +type sessionSink struct { + mu sync.Mutex + records [][]byte +} + +func (s *sessionSink) Record(_ context.Context, b []byte, _ ...recorder.RecordOption) error { + s.mu.Lock() + defer s.mu.Unlock() + s.records = append(s.records, append([]byte(nil), b...)) + return nil +} + +func (s *sessionSink) snapshot() [][]byte { + s.mu.Lock() + defer s.mu.Unlock() + return append([][]byte(nil), s.records...) +} + +func (s *sessionSink) len() int { + s.mu.Lock() + defer s.mu.Unlock() + return len(s.records) +} + +// gatedSink holds every delivery until it is released: a recorder backend that +// has stopped answering looks like this from the reporter's side. +type gatedSink struct { + sessionSink + entered chan struct{} + release chan struct{} + arrived sync.Once + released sync.Once +} + +func newGatedSink() *gatedSink { + return &gatedSink{entered: make(chan struct{}), release: make(chan struct{})} +} + +func (s *gatedSink) resume() { s.released.Do(func() { close(s.release) }) } + +func (s *gatedSink) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error { + s.arrived.Do(func() { close(s.entered) }) + select { + case <-s.release: + return s.sessionSink.Record(ctx, b, opts...) + case <-ctx.Done(): + return ctx.Err() + } +} + +func decodeSessionRecord(t *testing.T, b []byte) SessionRecord { + t.Helper() + var r SessionRecord + if err := json.Unmarshal(b, &r); err != nil { + t.Fatalf("unmarshal record: %v", err) + } + return r +} + +func waitFor(t *testing.T, cond func() bool, msg string) { + t.Helper() + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if cond() { + return + } + time.Sleep(time.Millisecond) + } + t.Fatal(msg) +} + +type flakySink struct { + sessionSink + fail int32 + attempts *atomic.Int32 +} + +func (s *flakySink) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error { + if s.attempts.Add(1) <= s.fail { + return errors.New("sink is down") + } + return s.sessionSink.Record(ctx, b, opts...) +} + +func TestSessionRecorderEmitsDeltasAndFinalRemainder(t *testing.T) { + sink := new(sessionSink) + r := NewSessionRecorder(sink, SessionRecorderOptions{Period: MinPeriod}) + defer r.Close() + + if !r.Enabled() { + t.Fatal("recorder should be enabled with a period") + } + + counters := xstats.NewStats(false) + session := r.NewSession(context.Background(), counters) + start := time.Now() + session.Start(HandlerRecorderObject{Service: "test", Time: start}) + + counters.Add(stats.KindOutputBytes, 10) + waitFor(t, func() bool { return sink.len() > 0 }, "no record while the session ran") + + counters.Add(stats.KindOutputBytes, 7) + if err := session.Finish(context.Background(), HandlerRecorderObject{ + Service: "test", Time: start, OutputBytes: 17, + }); err != nil { + t.Fatalf("finish: %v", err) + } + waitFor(t, func() bool { return sink.len() >= 2 }, "no final record") + + raw := sink.snapshot() + var total uint64 + stream := decodeSessionRecord(t, raw[0]).SessionID + for i, b := range raw { + rec := decodeSessionRecord(t, b) + total += rec.OutputBytesDelta + if rec.SessionID != stream { + t.Errorf("record %d: sessionID = %q, want %q", i, rec.SessionID, stream) + } + if rec.RecordIndex != uint64(i+1) { + t.Errorf("record %d: recordIndex = %d", i, rec.RecordIndex) + } + } + if got := decodeSessionRecord(t, raw[len(raw)-1]); got.Phase != PhaseFinal || got.OutputBytes != 17 { + t.Errorf("last record = %+v, want the final 17", got) + } + // Deltas, so the records add up to the traffic instead of restating it. + if total != 17 { + t.Errorf("summed deltas = %d, want 17", total) + } +} + +func TestSessionFinishWithoutStart(t *testing.T) { + sink := new(sessionSink) + r := NewSessionRecorder(sink, SessionRecorderOptions{Period: MinPeriod}) + defer r.Close() + + session := r.NewSession(context.Background(), nil) + if err := session.Finish(context.Background(), HandlerRecorderObject{ + Time: time.Now(), Err: "dial failed", InputBytes: 3, + }); err != nil { + t.Fatalf("finish: %v", err) + } + + raw := sink.snapshot() + if len(raw) != 1 { + t.Fatalf("records = %d, want 1", len(raw)) + } + if rec := decodeSessionRecord(t, raw[0]); rec.Phase != PhaseFinal || rec.InputBytes != 3 { + t.Errorf("record = %+v, want a final with 3 bytes", rec) + } +} + +// With the period off the whole session is written when it ends, in the shape +// the recorder used before sessions existed. +func TestSessionRecorderDisabledKeepsLegacyRecord(t *testing.T) { + sink := new(sessionSink) + r := NewSessionRecorder(sink, SessionRecorderOptions{}) + defer r.Close() + + if r.Enabled() { + t.Fatal("recorder should be disabled without a period") + } + + counters := xstats.NewStats(false) + session := r.NewSession(context.Background(), counters) + session.Start(HandlerRecorderObject{Service: "test", Time: time.Now()}) + counters.Add(stats.KindOutputBytes, 21) + + time.Sleep(50 * time.Millisecond) + if n := sink.len(); n != 0 { + t.Fatalf("a disabled recorder wrote %d records mid-session", n) + } + + stamp := time.Now() + if err := session.Finish(context.Background(), HandlerRecorderObject{ + Service: "test", Time: stamp, OutputBytes: 21, + }); err != nil { + t.Fatalf("finish: %v", err) + } + + raw := sink.snapshot() + if len(raw) != 1 { + t.Fatalf("records = %d, want 1", len(raw)) + } + var legacy map[string]any + if err := json.Unmarshal(raw[0], &legacy); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if _, ok := legacy["phase"]; ok { + t.Error("the legacy record must keep its shape") + } + if legacy["outputBytes"] != float64(21) { + t.Errorf("outputBytes = %v, want 21", legacy["outputBytes"]) + } + if legacy["time"] != stamp.Format(time.RFC3339Nano) { + t.Errorf("time = %v, want the stamp its handler gave it", legacy["time"]) + } +} + +// A handler that zeroes the time of a recorder object means "do not write this": +// handleProxy does it to the outer object of a proxy connection, whose bytes its +// per-request objects carry. +func TestSessionRecorderDisabledDropsUnstampedObjects(t *testing.T) { + sink := new(sessionSink) + r := NewSessionRecorder(sink, SessionRecorderOptions{}) + defer r.Close() + + session := r.NewSession(context.Background(), xstats.NewStats(false)) + if err := session.Finish(context.Background(), HandlerRecorderObject{ + Service: "test", OutputBytes: 137, + }); err != nil { + t.Fatalf("finish: %v", err) + } + if n := sink.len(); n != 0 { + t.Fatalf("an object its handler never stamped was recorded %d time(s)", n) + } +} + +// "off" is the client's opt-out, and Record enforces it. Periodic reporting must +// not write around it: not the start record, not the final one. +func TestSessionRecorderHonoursRecordModeOff(t *testing.T) { + sink := new(sessionSink) + r := NewSessionRecorder(sink, SessionRecorderOptions{Period: MinPeriod}) + defer r.Close() + + counters := xstats.NewStats(false) + session := r.NewSession(context.Background(), counters) + session.Start(HandlerRecorderObject{ + Service: "test", Time: time.Now(), RecordMode: "off", + }) + counters.Add(stats.KindOutputBytes, 21) + + time.Sleep(MinPeriod + 500*time.Millisecond) + if err := session.Finish(context.Background(), HandlerRecorderObject{ + Service: "test", Time: time.Now(), RecordMode: "off", OutputBytes: 21, + }); err != nil { + t.Fatalf("finish: %v", err) + } + + if n := sink.len(); n != 0 { + t.Fatalf("a session recorded %d time(s) with recording off", n) + } +} + +// A sink that refuses a record must not cost the traffic it refused: the bytes +// stay in the delta and go out with the next one. +func TestSessionRecorderRefusedIntervalKeepsItsBytes(t *testing.T) { + var attempts atomic.Int32 + sink := &flakySink{fail: 2, attempts: &attempts} + r := NewSessionRecorder(sink, SessionRecorderOptions{Period: MinPeriod}) + defer r.Close() + + counters := xstats.NewStats(false) + session := r.NewSession(context.Background(), counters) + start := time.Now() + session.Start(HandlerRecorderObject{Service: "test", Time: start}) + counters.Add(stats.KindOutputBytes, 10) + + waitFor(t, func() bool { return attempts.Load() >= 2 }, "the sink was never tried again") + counters.Add(stats.KindOutputBytes, 7) + + if err := session.Finish(context.Background(), HandlerRecorderObject{ + Service: "test", Time: start, OutputBytes: 17, + }); err != nil { + t.Fatalf("finish: %v", err) + } + + var total uint64 + for _, b := range sink.snapshot() { + total += decodeSessionRecord(t, b).OutputBytesDelta + } + if total != 17 { + t.Errorf("summed deltas = %d, want 17: a refused record took its bytes with it", total) + } +} + +// The final record is written inline, so its handler learns what happened to it. +func TestSessionRecorderFinishReportsTheSinkError(t *testing.T) { + var attempts atomic.Int32 + sink := &flakySink{fail: 100, attempts: &attempts} + r := NewSessionRecorder(sink, SessionRecorderOptions{Period: MinPeriod}) + defer r.Close() + + session := r.NewSession(context.Background(), xstats.NewStats(false)) + err := session.Finish(context.Background(), HandlerRecorderObject{ + Service: "test", Time: time.Now(), OutputBytes: 21, + }) + if err == nil { + t.Fatal("a refused final record must reach the handler as an error") + } + if n := sink.len(); n != 0 { + t.Fatalf("sink took %d records, want none", n) + } +} + +// A sink that stopped answering costs the handler its write timeout, and no more. +func TestSessionRecorderFinishIsBoundedByWriteTimeout(t *testing.T) { + sink := newGatedSink() + r := NewSessionRecorder(sink, SessionRecorderOptions{ + Period: MinPeriod, WriteTimeout: 50 * time.Millisecond, + }) + t.Cleanup(sink.resume) + defer r.Close() + + session := r.NewSession(context.Background(), xstats.NewStats(false)) + done := make(chan error, 1) + go func() { + done <- session.Finish(context.Background(), HandlerRecorderObject{ + Service: "test", Time: time.Now(), OutputBytes: 21, + }) + }() + + select { + case err := <-done: + if err == nil { + t.Error("an unresponsive sink must not look like a successful write") + } + case <-time.After(2 * time.Second): + t.Fatal("Finish blocked on an unresponsive sink") + } +} + +// Close stops the sampling, says which sessions were still running, and refuses +// the ones that finish afterwards. +func TestSessionRecorderCloseStopsSampling(t *testing.T) { + sink := new(sessionSink) + r := NewSessionRecorder(sink, SessionRecorderOptions{Period: MinPeriod}) + + counters := xstats.NewStats(false) + session := r.NewSession(context.Background(), counters) + session.Start(HandlerRecorderObject{Service: "test", Time: time.Now()}) + counters.Add(stats.KindOutputBytes, 5) + waitFor(t, func() bool { return sink.len() > 0 }, "no interim record") + + if err := r.Close(); err == nil { + t.Error("Close must report the sessions still running") + } + + counters.Add(stats.KindOutputBytes, 16) + time.Sleep(2 * MinPeriod) + if n := sink.len(); n != 1 { + t.Errorf("records = %d, want 1: the recorder kept sampling after Close", n) + } + + if err := session.Finish(context.Background(), HandlerRecorderObject{ + Service: "test", Time: time.Now(), OutputBytes: 21, + }); !errors.Is(err, ErrSessionRecorderClosed) { + t.Errorf("finish after Close = %v, want %v", err, ErrSessionRecorderClosed) + } +} + +// A handler whose service configures no recorder still runs: every session is a +// no-op, and nothing dereferences the missing sink. +func TestSessionRecorderWithoutASink(t *testing.T) { + r := NewSessionRecorder(nil, SessionRecorderOptions{Period: MinPeriod}) + if r.Enabled() { + t.Fatal("a recorder with no sink cannot be enabled") + } + + session := r.NewSession(context.Background(), xstats.NewStats(false)) + session.Start(HandlerRecorderObject{Service: "test", Time: time.Now()}) + if err := session.Finish(context.Background(), HandlerRecorderObject{ + Service: "test", Time: time.Now(), OutputBytes: 21, + }); err != nil { + t.Errorf("finish: %v", err) + } + if err := r.Close(); err != nil { + t.Errorf("close: %v", err) + } +} + +// Traffic keeps moving while the recorder samples and the session finishes, so +// the records must still add up to exactly what crossed the connection. +func TestSessionRecorderAccountsConcurrentTraffic(t *testing.T) { + const sessions = 8 + + sink := new(sessionSink) + r := NewSessionRecorder(sink, SessionRecorderOptions{Period: MinPeriod}) + deadline := time.Now().Add(MinPeriod + 300*time.Millisecond) + + var ( + wg sync.WaitGroup + mu sync.Mutex + written uint64 + failed []error + ) + for range sessions { + wg.Go(func() { + counters := xstats.NewStats(false) + session := r.NewSession(context.Background(), counters) + session.Start(HandlerRecorderObject{Time: time.Now()}) + + var sent uint64 + for time.Now().Before(deadline) { + counters.Add(stats.KindInputBytes, 1) + sent++ + time.Sleep(time.Millisecond) + } + err := session.Finish(context.Background(), HandlerRecorderObject{ + Time: time.Now(), InputBytes: counters.Get(stats.KindInputBytes), + }) + + mu.Lock() + defer mu.Unlock() + written += sent + if err != nil { + failed = append(failed, err) + } + }) + } + wg.Wait() + if err := r.Close(); err != nil { + t.Errorf("close: %v", err) + } + if len(failed) > 0 { + t.Fatalf("finish failed: %v", failed) + } + + var total uint64 + seen := map[string]uint64{} + for _, raw := range sink.snapshot() { + rec := decodeSessionRecord(t, raw) + total += rec.InputBytesDelta + seen[rec.SessionID]++ + } + if total != written { + t.Errorf("summed deltas = %d, want %d: concurrent reporting lost or duplicated traffic", total, written) + } + if len(seen) != sessions { + t.Errorf("sessions reported = %d, want %d", len(seen), sessions) + } +}