add handler recorder

This commit is contained in:
ginuerzh
2024-09-14 23:21:25 +08:00
parent e7581f0a9e
commit 69455ace9d
39 changed files with 1256 additions and 271 deletions
+51 -9
View File
@@ -12,8 +12,13 @@ import (
"github.com/go-gost/core/handler"
"github.com/go-gost/core/logger"
md "github.com/go-gost/core/metadata"
"github.com/go-gost/core/recorder"
xbypass "github.com/go-gost/x/bypass"
ctxvalue "github.com/go-gost/x/ctx"
netpkg "github.com/go-gost/x/internal/net"
sshd_util "github.com/go-gost/x/internal/util/sshd"
rate_limiter "github.com/go-gost/x/limiter/rate"
xrecorder "github.com/go-gost/x/recorder"
"github.com/go-gost/x/registry"
"golang.org/x/crypto/ssh"
)
@@ -28,8 +33,9 @@ func init() {
}
type forwardHandler struct {
md metadata
options handler.Options
md metadata
options handler.Options
recorder recorder.RecorderObject
}
func NewHandler(opts ...handler.Option) handler.Handler {
@@ -48,26 +54,60 @@ func (h *forwardHandler) Init(md md.Metadata) (err error) {
return
}
for _, ro := range h.options.Recorders {
if ro.Record == xrecorder.RecorderServiceHandler {
h.recorder = ro
break
}
}
return nil
}
func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.HandleOption) error {
func (h *forwardHandler) 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 !ro.Time.IsZero() {
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
}
switch cc := conn.(type) {
case *sshd_util.DirectForwardConn:
return h.handleDirectForward(ctx, cc, log)
return h.handleDirectForward(ctx, cc, ro, log)
case *sshd_util.RemoteForwardConn:
return h.handleRemoteForward(ctx, cc, log)
return h.handleRemoteForward(ctx, cc, ro, log)
default:
err := errors.New("sshd: wrong connection type")
log.Error(err)
@@ -75,9 +115,10 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
}
}
func (h *forwardHandler) handleDirectForward(ctx context.Context, conn *sshd_util.DirectForwardConn, log logger.Logger) error {
func (h *forwardHandler) handleDirectForward(ctx context.Context, conn *sshd_util.DirectForwardConn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
targetAddr := conn.DstAddr()
ro.Host = targetAddr
log = log.WithFields(map[string]any{
"dst": fmt.Sprintf("%s/%s", targetAddr, "tcp"),
"cmd": "connect",
@@ -87,7 +128,7 @@ func (h *forwardHandler) handleDirectForward(ctx context.Context, conn *sshd_uti
if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, "tcp", targetAddr) {
log.Debugf("bypass %s", targetAddr)
return nil
return xbypass.ErrBypass
}
cc, err := h.options.Router.Dial(ctx, "tcp", targetAddr)
@@ -106,7 +147,7 @@ func (h *forwardHandler) handleDirectForward(ctx context.Context, conn *sshd_uti
return nil
}
func (h *forwardHandler) handleRemoteForward(ctx context.Context, conn *sshd_util.RemoteForwardConn, log logger.Logger) error {
func (h *forwardHandler) handleRemoteForward(ctx context.Context, conn *sshd_util.RemoteForwardConn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
req := conn.Request()
t := tcpipForward{}
@@ -117,6 +158,7 @@ func (h *forwardHandler) handleRemoteForward(ctx context.Context, conn *sshd_uti
network := "tcp"
addr := net.JoinHostPort(t.Host, strconv.Itoa(int(t.Port)))
ro.Host = addr
log = log.WithFields(map[string]any{
"dst": fmt.Sprintf("%s/%s", addr, network),