github/gost-x
4
0
Fork 0
mirror of https://github.com/go-gost/x.git synced 2026-10-19 08:15:33 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

recorder: emit periodic session snapshots and byte deltas

Browse Source 
Handlers write one recorder entry when a session ends, so long-lived
SOCKS5 tunnels, streaming HTTP responses and UDP associations stay
invisible while they run and report their traffic as one late record.

With recorder.period set, a session now emits interim records on the
interval while it is active and one final record when it ends. Each
record carries the cumulative counters alongside the increments since the
previous one, so a consumer sums records instead of replacing them, plus
sessionID, recordIndex and phase. A session that ends inside one period
still writes a single record, as it does today.

One goroutine per handler samples every live session on the tick and
writes its record straight to the recorder. There is no queue and no
retry loop, and none is needed: a session moves its cursor past a record
only once the recorder has taken it, so a refused write costs no bytes.
They stay in the delta and go out with the next record. Holding on to
what the sink cannot take right now is the sink's job.

Interim and Finish take a context and return an error, so a handler
learns what became of every record it asked for. Start writes nothing, so
it takes neither. Close stops the sampling and reports the sessions its
owner had not finished; nothing is buffered, so there is nothing to
drain.

With recorder.period unset or non-positive the previous behaviour is
kept: one record per session, written inline, in the old shape. An object
its handler never stamped, and one whose record mode is off, is dropped
in both modes.
This commit is contained in:
Leonov Vladislav
2026-09-15 15:05:02 +03:00
parent f70b114578
commit 1c42c638a6
25 changed files with 971 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
handler/http/connect.go
+5 -1
View File
@@ -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,
handler/http/handler.go
+22 -12
View File
@@ -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.
handler/http/metadata.go
+4
View File
@@ -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
handler/http/proxy.go
+19 -3
View File
@@ -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()
handler/redirect/tcp/handler.go
+18 -5
View File
@@ -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() }
handler/redirect/tcp/metadata.go
+3
View File
@@ -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
handler/redirect/udp/handler.go
+15 -4
View File
@@ -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() }
handler/redirect/udp/metadata.go
+4
View File
@@ -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")
handler/socks/v5/bind.go
+2
View File
@@ -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 {
handler/socks/v5/connect.go
+3
View File
@@ -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)
handler/socks/v5/handler.go
+20 -11
View File
@@ -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 {
handler/socks/v5/mbind.go
+2
View File
@@ -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 {
handler/socks/v5/metadata.go
+4
View File
@@ -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 {
handler/socks/v5/udp.go
+7 -9
View File
@@ -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).
handler/socks/v5/udp_tun.go
+1
View File
@@ -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())
internal/ctx/value.go
+11
View File
@@ -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
}
internal/util/sniffing/sniffer.go
+8
View File
@@ -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{})
}
internal/util/sniffing/sniffer_h2.go
+24 -5
View File
@@ -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)
}
internal/util/sniffing/sniffer_http.go
+3 -1
View File
@@ -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)
}
internal/util/sniffing/sniffer_test.go
+3 -3
View File
@@ -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(),
internal/util/sniffing/sniffer_tls.go
+2
View File
@@ -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{
recorder/session.go
+125
View File
@@ -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
}
recorder/session_record.go
+59
View File
@@ -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
}
recorder/session_recorder.go
+158
View File
@@ -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
}
recorder/session_test.go
+449
View File
@@ -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)
}
}