feat: add proxy health monitoring utilities (monitoredCopy, watchdog, TCP tuning)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: acmore

internal/cli/proxy_health.go

@@ -0,0 +1,93 @@
+package cli
+
+import (
+	"context"
+	"errors"
+	"io"
+	"log/slog"
+	"net"
+	"sync/atomic"
+	"syscall"
+	"time"
+)
+
+// ErrProxyHealthDisconnect is a sentinel error returned when the proxy
+// exits due to a health-check-triggered disconnection. It should be
+// suppressed from stderr in main.go since the ssh client will report
+// the disconnect naturally.
+var ErrProxyHealthDisconnect = errors.New("proxy health disconnect")
+
+// monitoredReader wraps an io.Reader and updates an atomic timestamp
+// on every successful Read.
+type monitoredReader struct {
+	r        io.Reader
+	lastData *atomic.Int64
+}
+
+func (m *monitoredReader) Read(p []byte) (int, error) {
+	n, err := m.r.Read(p)
+	if n > 0 {
+		m.lastData.Store(time.Now().UnixNano())
+	}
+	return n, err
+}
+
+// monitoredCopy copies from src to dst, updating lastData on every
+// successful read. Returns bytes copied and any error.
+func monitoredCopy(dst io.Writer, src io.Reader, lastData *atomic.Int64) (int64, error) {
+	return io.Copy(dst, &monitoredReader{r: src, lastData: lastData})
+}
+
+// proxyDataFlowWatchdog checks lastData every checkInterval. If no data
+// has flowed for idleThreshold, it closes conn and returns.
+// It also returns when ctx is cancelled.
+func proxyDataFlowWatchdog(ctx context.Context, conn net.Conn, lastData *atomic.Int64, checkInterval, idleThreshold time.Duration) {
+	ticker := time.NewTicker(checkInterval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			last := lastData.Load()
+			if last == 0 {
+				continue // no data yet, skip
+			}
+			idle := time.Since(time.Unix(0, last))
+			if idle >= idleThreshold {
+				slog.Info("ssh-proxy data flow watchdog: idle timeout",
+					"idle", idle.Round(time.Millisecond),
+					"threshold", idleThreshold,
+				)
+				_ = conn.Close()
+				return
+			}
+		}
+	}
+}
+
+// setTCPKeepAliveProxyTuning reduces the TCP keepalive period and
+// optionally sets TCP_KEEPCNT where supported.
+func setTCPKeepAliveProxyTuning(conn net.Conn) {
+	tc, ok := conn.(*net.TCPConn)
+	if !ok {
+		return
+	}
+	_ = tc.SetKeepAlive(true)
+	_ = tc.SetKeepAlivePeriod(5 * time.Second)
+
+	// Try to set TCP_KEEPCNT=2 via syscall. This is best-effort.
+	raw, err := tc.SyscallConn()
+	if err != nil {
+		slog.Debug("ssh-proxy: failed to get syscall conn for TCP_KEEPCNT", "error", err)
+		return
+	}
+	var sysErr error
+	err = raw.Control(func(fd uintptr) {
+		sysErr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_KEEPCNT, 2)
+	})
+	if err != nil || sysErr != nil {
+		slog.Debug("ssh-proxy: TCP_KEEPCNT not set (unsupported or failed)",
+			"controlErr", err, "sysErr", sysErr)
+	}
+}

internal/cli/proxy_health_test.go

@@ -152,7 +152,12 @@ func TestSetTCPKeepAliveProxyTuningSmoke(t *testing.T) {
 		t.Fatal(err)
 	}
 	defer ln.Close()
-	go func() { c, _ := ln.Accept(); _ = c.Close() }()
+	go func() {
+		c, err := ln.Accept()
+		if err == nil {
+			_ = c.Close()
+		}
+	}()
 
 	conn, err := net.Dial("tcp", ln.Addr().String())
 	if err != nil {