From 64a16fafbc957d1350cdc7b3a42e31741a32bb7a Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Sun, 24 May 2026 23:58:28 +0800 Subject: [PATCH] fix(recorder): propagate context, handle short writes, add nil guards, idempotent Close, and doc comments - httpRecorder: use http.NewRequestWithContext instead of http.NewRequest - tcpRecorder: loop Write to avoid short-write data loss - fileRecorder: guard against nil out in Record and Close - plugin/grpc, plugin/http: handle json.Marshal errors instead of ignoring - fileRecorder, redis*Recorder: idempotent Close via sync.Once - All exported symbols: add doc comments --- recorder/file.go | 18 ++++++++++++++-- recorder/http.go | 6 +++++- recorder/plugin/grpc.go | 5 ++++- recorder/plugin/http.go | 5 ++++- recorder/recorder.go | 12 +++++++++++ recorder/redis.go | 46 ++++++++++++++++++++++++++++++----------- recorder/tcp.go | 14 +++++++++++-- 7 files changed, 87 insertions(+), 19 deletions(-) diff --git a/recorder/file.go b/recorder/file.go index 9b03a7ed..b4c58d5c 100644 --- a/recorder/file.go +++ b/recorder/file.go @@ -15,14 +15,17 @@ type fileRecorderOptions struct { sep string } +// FileRecorderOption configures FileRecorder options. type FileRecorderOption func(opts *fileRecorderOptions) +// RecorderFileRecorderOption sets the recorder name for metrics labeling. func RecorderFileRecorderOption(recorder string) FileRecorderOption { return func(opts *fileRecorderOptions) { opts.recorder = recorder } } +// SepFileRecorderOption sets the record separator written after each record. func SepFileRecorderOption(sep string) FileRecorderOption { return func(opts *fileRecorderOptions) { opts.sep = sep @@ -34,7 +37,8 @@ type fileRecorder struct { out io.WriteCloser sep string - mu sync.Mutex + mu sync.Mutex + closeOnce sync.Once } // FileRecorder records data to file. @@ -52,6 +56,10 @@ func FileRecorder(out io.WriteCloser, opts ...FileRecorderOption) recorder.Recor } func (r *fileRecorder) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error { + if r.out == nil { + return nil + } + xmetrics.GetCounter(xmetrics.MetricRecorderRecordsCounter, metrics.Labels{"recorder": r.recorder}).Inc() if r.sep != "" { @@ -72,5 +80,11 @@ func (r *fileRecorder) Record(ctx context.Context, b []byte, opts ...recorder.Re } func (r *fileRecorder) Close() error { - return r.out.Close() + var err error + r.closeOnce.Do(func() { + if r.out != nil { + err = r.out.Close() + } + }) + return err } diff --git a/recorder/http.go b/recorder/http.go index b6cf60fe..50719964 100644 --- a/recorder/http.go +++ b/recorder/http.go @@ -19,20 +19,24 @@ type httpRecorderOptions struct { header http.Header } +// HTTPRecorderOption configures HTTPRecorder options. type HTTPRecorderOption func(opts *httpRecorderOptions) +// RecorderHTTPRecorderOption sets the recorder name for metrics labeling. func RecorderHTTPRecorderOption(recorder string) HTTPRecorderOption { return func(opts *httpRecorderOptions) { opts.recorder = recorder } } +// TimeoutHTTPRecorderOption sets the HTTP client timeout. func TimeoutHTTPRecorderOption(timeout time.Duration) HTTPRecorderOption { return func(opts *httpRecorderOptions) { opts.timeout = timeout } } +// HeaderHTTPRecorderOption sets additional HTTP headers sent with each record. func HeaderHTTPRecorderOption(header http.Header) HTTPRecorderOption { return func(opts *httpRecorderOptions) { opts.header = header @@ -72,7 +76,7 @@ func HTTPRecorder(url string, opts ...HTTPRecorderOption) recorder.Recorder { func (r *httpRecorder) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error { xmetrics.GetCounter(xmetrics.MetricRecorderRecordsCounter, metrics.Labels{"recorder": r.recorder}).Inc() - req, err := http.NewRequest(http.MethodPost, r.url, bytes.NewReader(b)) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.url, bytes.NewReader(b)) if err != nil { return err } diff --git a/recorder/plugin/grpc.go b/recorder/plugin/grpc.go index cd6964b9..82b76bcd 100644 --- a/recorder/plugin/grpc.go +++ b/recorder/plugin/grpc.go @@ -55,7 +55,10 @@ func (p *grpcPlugin) Record(ctx context.Context, b []byte, opts ...recorder.Reco opt(&options) } - md, _ := json.Marshal(options.Metadata) + md, err := json.Marshal(options.Metadata) + if err != nil { + return err + } reply, err := p.client.Record(ctx, &proto.RecordRequest{ diff --git a/recorder/plugin/http.go b/recorder/plugin/http.go index 6547592f..99cc9b80 100644 --- a/recorder/plugin/http.go +++ b/recorder/plugin/http.go @@ -56,7 +56,10 @@ func (p *httpPlugin) Record(ctx context.Context, b []byte, opts ...recorder.Reco opt(&options) } - md, _ := json.Marshal(options.Metadata) + md, err := json.Marshal(options.Metadata) + if err != nil { + return err + } rb := httpPluginRequest{ Data: b, diff --git a/recorder/recorder.go b/recorder/recorder.go index 24869628..c2eb7916 100644 --- a/recorder/recorder.go +++ b/recorder/recorder.go @@ -15,18 +15,22 @@ const ( RecorderServiceHandlerTunnel = "recorder.service.handler.tunnel" ) +// HTTPRequestRecorderObject holds the recorded data of an HTTP request. type HTTPRequestRecorderObject struct { ContentLength int64 `json:"contentLength"` Header http.Header `json:"header"` Body []byte `json:"body"` } +// HTTPResponseRecorderObject holds the recorded data of an HTTP response. type HTTPResponseRecorderObject struct { ContentLength int64 `json:"contentLength"` Header http.Header `json:"header"` Body []byte `json:"body"` } +// HTTPRecorderObject holds the recorded HTTP request and response data +// for a single HTTP transaction. type HTTPRecorderObject struct { Host string `json:"host"` Method string `json:"method"` @@ -38,6 +42,7 @@ type HTTPRecorderObject struct { Response HTTPResponseRecorderObject `json:"response"` } +// WebsocketRecorderObject holds the recorded data of a WebSocket frame. type WebsocketRecorderObject struct { From string `json:"from"` Fin bool `json:"fin"` @@ -51,6 +56,7 @@ type WebsocketRecorderObject struct { Payload []byte `json:"payload"` } +// TLSRecorderObject holds the recorded data of a TLS handshake. type TLSRecorderObject struct { ServerName string `json:"serverName"` CipherSuite string `json:"cipherSuite"` @@ -61,6 +67,7 @@ type TLSRecorderObject struct { ServerHello string `json:"serverHello"` } +// DNSRecorderObject holds the recorded data of a DNS query and response. type DNSRecorderObject struct { ID int `json:"id"` Name string `json:"name"` @@ -71,6 +78,9 @@ type DNSRecorderObject struct { Cached bool `json:"cached"` } +// HandlerRecorderObject bundles traffic metadata recorded at the handler +// level — addresses, protocols, transferred bytes, and optional sub-records +// for HTTP, WebSocket, TLS, and DNS traffic. type HandlerRecorderObject struct { Node string `json:"node,omitempty"` Service string `json:"service"` @@ -98,6 +108,8 @@ type HandlerRecorderObject struct { Time time.Time `json:"time"` } +// Record serializes the HandlerRecorderObject as JSON and writes it to r. +// It returns nil if p or r is nil or if p.Time is the zero value. func (p *HandlerRecorderObject) Record(ctx context.Context, r recorder.Recorder) error { if p == nil || r == nil || p.Time.IsZero() { return nil diff --git a/recorder/redis.go b/recorder/redis.go index 78ba4d07..e7d743b2 100644 --- a/recorder/redis.go +++ b/recorder/redis.go @@ -2,6 +2,7 @@ package recorder import ( "context" + "sync" "github.com/go-gost/core/metrics" "github.com/go-gost/core/recorder" @@ -16,31 +17,37 @@ type redisRecorderOptions struct { password string key string } +// RedisRecorderOption configures Redis recorder options. type RedisRecorderOption func(opts *redisRecorderOptions) +// RecorderRedisRecorderOption sets the recorder name for metrics labeling. func RecorderRedisRecorderOption(recorder string) RedisRecorderOption { return func(opts *redisRecorderOptions) { opts.recorder = recorder } } +// DBRedisRecorderOption sets the Redis database number. func DBRedisRecorderOption(db int) RedisRecorderOption { return func(opts *redisRecorderOptions) { opts.db = db } } +// UsernameRedisRecorderOption sets the Redis username for authentication. func UsernameRedisRecorderOption(username string) RedisRecorderOption { return func(opts *redisRecorderOptions) { opts.username = username } } +// PasswordRedisRecorderOption sets the Redis password for authentication. func PasswordRedisRecorderOption(password string) RedisRecorderOption { return func(opts *redisRecorderOptions) { opts.password = password } } +// KeyRedisRecorderOption sets the Redis key for storing records. func KeyRedisRecorderOption(key string) RedisRecorderOption { return func(opts *redisRecorderOptions) { opts.key = key @@ -48,9 +55,10 @@ func KeyRedisRecorderOption(key string) RedisRecorderOption { } type redisSetRecorder struct { - recorder string - client *redis.Client - key string + recorder string + client *redis.Client + key string + closeOnce sync.Once } // RedisSetRecorder records data to a redis set. @@ -83,13 +91,18 @@ func (r *redisSetRecorder) Record(ctx context.Context, b []byte, opts ...recorde } func (r *redisSetRecorder) Close() error { - return r.client.Close() + var err error + r.closeOnce.Do(func() { + err = r.client.Close() + }) + return err } type redisListRecorder struct { - recorder string - client *redis.Client - key string + recorder string + client *redis.Client + key string + closeOnce sync.Once } // RedisListRecorder records data to a redis list. @@ -122,13 +135,18 @@ func (r *redisListRecorder) Record(ctx context.Context, b []byte, opts ...record } func (r *redisListRecorder) Close() error { - return r.client.Close() + var err error + r.closeOnce.Do(func() { + err = r.client.Close() + }) + return err } type redisSortedSetRecorder struct { - recorder string - client *redis.Client - key string + recorder string + client *redis.Client + key string + closeOnce sync.Once } // RedisSortedSetRecorder records data to a redis sorted set. @@ -164,5 +182,9 @@ func (r *redisSortedSetRecorder) Record(ctx context.Context, b []byte, opts ...r } func (r *redisSortedSetRecorder) Close() error { - return r.client.Close() + var err error + r.closeOnce.Do(func() { + err = r.client.Close() + }) + return err } diff --git a/recorder/tcp.go b/recorder/tcp.go index 08c9d033..f910821b 100644 --- a/recorder/tcp.go +++ b/recorder/tcp.go @@ -17,20 +17,24 @@ type tcpRecorderOptions struct { log logger.Logger } +// TCPRecorderOption configures TCPRecorder options. type TCPRecorderOption func(opts *tcpRecorderOptions) +// RecorderTCPRecorderOption sets the recorder name for metrics labeling. func RecorderTCPRecorderOption(recorder string) TCPRecorderOption { return func(opts *tcpRecorderOptions) { opts.recorder = recorder } } +// TimeoutTCPRecorderOption sets the dial timeout for TCP connections. func TimeoutTCPRecorderOption(timeout time.Duration) TCPRecorderOption { return func(opts *tcpRecorderOptions) { opts.timeout = timeout } } +// LogTCPRecorderOption sets the logger for the TCP recorder. func LogTCPRecorderOption(log logger.Logger) TCPRecorderOption { return func(opts *tcpRecorderOptions) { opts.log = log @@ -70,6 +74,12 @@ func (r *tcpRecorder) Record(ctx context.Context, b []byte, opts ...recorder.Rec } defer c.Close() - _, err = c.Write(b) - return err + for written := 0; written < len(b); { + n, err := c.Write(b[written:]) + if err != nil { + return err + } + written += n + } + return nil }