fix sshd connection type

This commit is contained in:
Jason Lyu
2026-01-11 13:24:27 -05:00
committed by ginuerzh
parent 180145189f
commit cd7bf9521f
+10 -9
View File
@@ -107,7 +107,6 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
log.Infof("%s <> %s", conn.RemoteAddr(), conn.LocalAddr()) log.Infof("%s <> %s", conn.RemoteAddr(), conn.LocalAddr())
pStats := xstats.Stats{} pStats := xstats.Stats{}
conn = stats_wrapper.WrapConn(conn, &pStats)
defer func() { defer func() {
if err != nil { if err != nil {
@@ -133,9 +132,9 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
switch cc := conn.(type) { switch cc := conn.(type) {
case *sshd_util.DirectForwardConn: case *sshd_util.DirectForwardConn:
return h.handleDirectForward(ctx, cc, ro, log) return h.handleDirectForward(ctx, cc, ro, log, &pStats)
case *sshd_util.RemoteForwardConn: case *sshd_util.RemoteForwardConn:
return h.handleRemoteForward(ctx, cc, ro, log) return h.handleRemoteForward(ctx, cc, ro, log, &pStats)
default: default:
err := errors.New("sshd: wrong connection type") err := errors.New("sshd: wrong connection type")
log.Error(err) log.Error(err)
@@ -143,8 +142,9 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
} }
} }
func (h *forwardHandler) handleDirectForward(ctx context.Context, conn *sshd_util.DirectForwardConn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error { func (h *forwardHandler) handleDirectForward(ctx context.Context, c *sshd_util.DirectForwardConn, ro *xrecorder.HandlerRecorderObject, log logger.Logger, stats *xstats.Stats) error {
targetAddr := conn.DstAddr() targetAddr := c.DstAddr()
conn := stats_wrapper.WrapConn(c, stats)
ro.Host = targetAddr ro.Host = targetAddr
log = log.WithFields(map[string]any{ log = log.WithFields(map[string]any{
@@ -235,8 +235,9 @@ func (h *forwardHandler) handleDirectForward(ctx context.Context, conn *sshd_uti
return nil return nil
} }
func (h *forwardHandler) handleRemoteForward(ctx context.Context, conn *sshd_util.RemoteForwardConn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error { func (h *forwardHandler) handleRemoteForward(ctx context.Context, c *sshd_util.RemoteForwardConn, ro *xrecorder.HandlerRecorderObject, log logger.Logger, stats *xstats.Stats) error {
req := conn.Request() req := c.Request()
conn := stats_wrapper.WrapConn(c, stats)
t := tcpipForward{} t := tcpipForward{}
if err := ssh.Unmarshal(req.Payload, &t); err != nil { if err := ssh.Unmarshal(req.Payload, &t); err != nil {
@@ -290,7 +291,7 @@ func (h *forwardHandler) handleRemoteForward(ctx context.Context, conn *sshd_uti
return err return err
} }
sshConn := conn.Conn() sshConn := c.Conn()
go func() { go func() {
for { for {
@@ -340,7 +341,7 @@ func (h *forwardHandler) handleRemoteForward(ctx context.Context, conn *sshd_uti
tm := time.Now() tm := time.Now()
log.Infof("%s <-> %s", conn.RemoteAddr(), addr) log.Infof("%s <-> %s", conn.RemoteAddr(), addr)
<-conn.Done() <-c.Done()
log.WithFields(map[string]any{ log.WithFields(map[string]any{
"duration": time.Since(tm), "duration": time.Since(tm),
}).Infof("%s >-< %s", conn.RemoteAddr(), addr) }).Infof("%s >-< %s", conn.RemoteAddr(), addr)