diff --git a/handler/sshd/handler.go b/handler/sshd/handler.go index 463a861f..e36469c3 100644 --- a/handler/sshd/handler.go +++ b/handler/sshd/handler.go @@ -130,9 +130,9 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand return rate_limiter.ErrRateLimit } - switch cc := conn.(type) { + switch cc := unwrapConn(conn).(type) { case *sshd_util.DirectForwardConn: - return h.handleDirectForward(ctx, cc, ro, log, &pStats) + return h.handleDirectForward(ctx, cc.DstAddr(), conn, ro, log, &pStats) case *sshd_util.RemoteForwardConn: return h.handleRemoteForward(ctx, cc, ro, log, &pStats) default: @@ -142,9 +142,8 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand } } -func (h *forwardHandler) handleDirectForward(ctx context.Context, c *sshd_util.DirectForwardConn, ro *xrecorder.HandlerRecorderObject, log logger.Logger, stats *xstats.Stats) error { - targetAddr := c.DstAddr() - conn := stats_wrapper.WrapConn(c, stats) +func (h *forwardHandler) handleDirectForward(ctx context.Context, targetAddr string, origConn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger, stats *xstats.Stats) error { + conn := stats_wrapper.WrapConn(origConn, stats) ro.Host = targetAddr log = log.WithFields(map[string]any{ @@ -361,6 +360,28 @@ func (h *forwardHandler) checkRateLimit(addr net.Addr) bool { return true } +// connUnwrapper is an interface for unwrapping net.Conn wrappers to access +// the underlying connection type. +type connUnwrapper interface { + UnwrapConn() net.Conn +} + +// unwrapConn peels through net.Conn wrapper layers (e.g. traffic limiters) +// to find the underlying concrete connection type. +func unwrapConn(c net.Conn) net.Conn { + for { + uw, ok := c.(connUnwrapper) + if !ok { + return c + } + if inner := uw.UnwrapConn(); inner != nil { + c = inner + } else { + return c + } + } +} + func getHostPortFromAddr(addr net.Addr) (host string, port int, err error) { host, portString, err := net.SplitHostPort(addr.String()) if err != nil { diff --git a/limiter/traffic/wrapper/conn.go b/limiter/traffic/wrapper/conn.go index ee670787..a96b2555 100644 --- a/limiter/traffic/wrapper/conn.go +++ b/limiter/traffic/wrapper/conn.go @@ -111,6 +111,12 @@ func (c *limitConn) Write(b []byte) (n int, err error) { return } +// UnwrapConn returns the underlying connection, allowing type assertions +// through wrapper layers. +func (c *limitConn) UnwrapConn() net.Conn { + return c.Conn +} + func (c *limitConn) SyscallConn() (rc syscall.RawConn, err error) { if sc, ok := c.Conn.(syscall.Conn); ok { rc, err = sc.SyscallConn() diff --git a/service/service.go b/service/service.go index 6d1b79e0..20be9aff 100644 --- a/service/service.go +++ b/service/service.go @@ -223,6 +223,7 @@ func (s *defaultService) Serve() error { } s.setState(StateFailed) + s.status.setLastError(e) log.Warnf("accept: %v, retrying in %v", e, tempDelay) time.Sleep(tempDelay) @@ -242,6 +243,7 @@ func (s *defaultService) Serve() error { if tempDelay > 0 { tempDelay = 0 s.setState(StateReady) + s.status.setLastError(nil) } ctx := gctx diff --git a/service/status.go b/service/status.go index ce8b933b..f753cbd3 100644 --- a/service/status.go +++ b/service/status.go @@ -37,6 +37,7 @@ type Status struct { state State events []Event stats stats.Stats + lastError error mu sync.RWMutex } @@ -81,6 +82,24 @@ func (p *Status) addEvent(event Event) { p.events = append(p.events, event) } +// LastError returns the last accept/bind error that caused the service to +// enter StateFailed. Returns nil if the service never entered the failed state. +func (p *Status) LastError() error { + if p == nil { + return nil + } + p.mu.RLock() + defer p.mu.RUnlock() + return p.lastError +} + +// setLastError stores the last accept/bind error for retrieval via LastError(). +func (p *Status) setLastError(err error) { + p.mu.Lock() + defer p.mu.Unlock() + p.lastError = err +} + // Stats returns the traffic statistics for the service. It safely handles nil // receivers by returning nil. func (p *Status) Stats() stats.Stats {