From bc44baba55e4adafb54a71eb019077ebcc38fc41 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Sat, 13 Jun 2026 14:29:51 +0800
Subject: [PATCH] fix(handler/sshd): unwrap limiter wrappers to fix wrong
 connection type (issue #867)

The sshd handler performed a concrete type assertion on the accepted
connection, but the SSH listener's Accept() wraps each connection with
a traffic limiter (limitConn), obscuring the underlying *DirectForwardConn
type and causing 'sshd: wrong connection type' whenever a limiter was
configured.

Add an UnwrapConn() method to the traffic limiter wrapper and an
unwrapConn() helper in the handler that peels through wrapper layers
before the type switch. The handler uses the unwrapped conn only for
the type-specific DstAddr() call, while the original limiter-wrapped
conn drives data transfer so rate limiting remains effective.

Also track the last accept/bind error on service Status so callers can
retrieve the cause when a service enters the failed state.
---
 handler/sshd/handler.go         | 31 ++++++++++++++++++++++++++-----
 limiter/traffic/wrapper/conn.go |  6 ++++++
 service/service.go              |  2 ++
 service/status.go               | 19 +++++++++++++++++++
 4 files changed, 53 insertions(+), 5 deletions(-)

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 {