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

Merge pull request #126: emit periodic session snapshots and byte deltas

Browse Source
This commit is contained in:
ginuerzh
2026-10-07 02:08:41 +00:00
parent 89188536d3 1c42c638a6
commit 3e533c61c7
26 changed files with 977 additions and 60 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
@@ -53,6 +55,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,
@@ -133,13 +132,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(
@@ -183,6 +179,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
@@ -139,6 +139,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
@@ -79,3 +79,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
@@ -149,6 +149,7 @@ type Sniffer struct {
Recorder recorder.Recorder
RecorderOptions *recorder.Options
SessionRecorder *xrecorder.SessionRecorder
// MITM TLS termination
Certificate *x509.Certificate
@@ -209,3 +210,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_bypass_test.go
+6 -6
View File
@@ -310,12 +310,12 @@ func (t *countingTransport) count() int {
func TestServeH2_AuthorizesEveryStream(t *testing.T) {
newHandler := func(bp *policyBypass, tr http.RoundTripper) *h2Handler {
return &h2Handler{
transport: tr,
recorder: &noopRecorder{},
recorderObject: &xrecorder.HandlerRecorderObject{},
log: xlogger.Nop(),
handleOptions: &HandleOptions{Bypass: bp},
network: "tcp",
transport: tr,
sessionRecorder: xrecorder.NewSessionRecorder(&noopRecorder{}, xrecorder.SessionRecorderOptions{}),
recorderObject: &xrecorder.HandlerRecorderObject{},
log: xlogger.Nop(),
handleOptions: &HandleOptions{Bypass: bp},
network: "tcp",
}
}
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,
@@ -79,7 +83,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
@@ -100,6 +104,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() {
@@ -107,7 +114,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)
}
@@ -145,6 +154,10 @@ func (h *h2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
return
}
if r.Body != nil {
r.Body = countBody(r.Body, &counters)
}
url := r.URL
url.Scheme = "https"
url.Host = r.Host
@@ -189,14 +202,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)
}
}
@@ -207,3 +221,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
@@ -246,6 +246,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 {
@@ -254,7 +256,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)
}
}