add handler recorder
This commit is contained in:
@@ -12,13 +12,16 @@ import (
|
||||
"github.com/go-gost/core/logger"
|
||||
md "github.com/go-gost/core/metadata"
|
||||
"github.com/go-gost/core/observer/stats"
|
||||
"github.com/go-gost/core/recorder"
|
||||
"github.com/go-gost/gosocks4"
|
||||
ctxvalue "github.com/go-gost/x/ctx"
|
||||
netpkg "github.com/go-gost/x/internal/net"
|
||||
limiter_util "github.com/go-gost/x/internal/util/limiter"
|
||||
stats_util "github.com/go-gost/x/internal/util/stats"
|
||||
rate_limiter "github.com/go-gost/x/limiter/rate"
|
||||
traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper"
|
||||
stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
|
||||
xrecorder "github.com/go-gost/x/recorder"
|
||||
"github.com/go-gost/x/registry"
|
||||
)
|
||||
|
||||
@@ -33,11 +36,12 @@ func init() {
|
||||
}
|
||||
|
||||
type socks4Handler struct {
|
||||
md metadata
|
||||
options handler.Options
|
||||
stats *stats_util.HandlerStats
|
||||
limiter traffic.TrafficLimiter
|
||||
cancel context.CancelFunc
|
||||
md metadata
|
||||
options handler.Options
|
||||
stats *stats_util.HandlerStats
|
||||
limiter traffic.TrafficLimiter
|
||||
cancel context.CancelFunc
|
||||
recorder recorder.RecorderObject
|
||||
}
|
||||
|
||||
func NewHandler(opts ...handler.Option) handler.Handler {
|
||||
@@ -68,28 +72,51 @@ func (h *socks4Handler) Init(md md.Metadata) (err error) {
|
||||
h.limiter = limiter_util.NewCachedTrafficLimiter(limiter, 30*time.Second, 60*time.Second)
|
||||
}
|
||||
|
||||
for _, ro := range h.options.Recorders {
|
||||
if ro.Record == xrecorder.RecorderServiceHandler {
|
||||
h.recorder = ro
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *socks4Handler) Handle(ctx context.Context, conn net.Conn, opts ...handler.HandleOption) error {
|
||||
func (h *socks4Handler) Handle(ctx context.Context, conn net.Conn, opts ...handler.HandleOption) (err error) {
|
||||
defer conn.Close()
|
||||
|
||||
start := time.Now()
|
||||
|
||||
ro := &xrecorder.HandlerRecorderObject{
|
||||
Service: h.options.Service,
|
||||
Network: "tcp",
|
||||
RemoteAddr: conn.RemoteAddr().String(),
|
||||
LocalAddr: conn.LocalAddr().String(),
|
||||
Time: start,
|
||||
SID: string(ctxvalue.SidFromContext(ctx)),
|
||||
}
|
||||
|
||||
log := h.options.Logger.WithFields(map[string]any{
|
||||
"remote": conn.RemoteAddr().String(),
|
||||
"local": conn.LocalAddr().String(),
|
||||
"sid": ctxvalue.SidFromContext(ctx),
|
||||
})
|
||||
|
||||
log.Infof("%s <> %s", conn.RemoteAddr(), conn.LocalAddr())
|
||||
defer func() {
|
||||
if err != nil {
|
||||
ro.Err = err.Error()
|
||||
}
|
||||
ro.Duration = time.Since(start)
|
||||
ro.Record(ctx, h.recorder.Recorder)
|
||||
|
||||
log.WithFields(map[string]any{
|
||||
"duration": time.Since(start),
|
||||
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
|
||||
}()
|
||||
|
||||
if !h.checkRateLimit(conn.RemoteAddr()) {
|
||||
return nil
|
||||
return rate_limiter.ErrRateLimit
|
||||
}
|
||||
|
||||
if h.md.readTimeout > 0 {
|
||||
@@ -101,6 +128,8 @@ func (h *socks4Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl
|
||||
log.Error(err)
|
||||
return err
|
||||
}
|
||||
|
||||
ro.Host = req.Addr.String()
|
||||
log.Trace(req)
|
||||
|
||||
conn.SetReadDeadline(time.Time{})
|
||||
|
||||
Reference in New Issue
Block a user