diff --git a/hotreload.go b/hotreload.go
index 12e94be..231d2bd 100644
--- a/hotreload.go
+++ b/hotreload.go
@@ -334,6 +334,7 @@ func applyFullConfigReload(newCfg *Config) ConfigReloadReport {
// (picked up by new connections).
setDefaultLimits(newCfg.DefaultLimitMbpsUp, newCfg.DefaultLimitMbpsDown)
setSSHIdleTimeoutFromConfig(newCfg.SSHIdleTimeout)
+ setMaxTotalConnsFromConfig(newCfg.MaxTotalConnections)
// Quiet logging / user count display
if newCfg.Quiet {
diff --git a/main.go b/main.go
index 476cf8b..15982a0 100644
--- a/main.go
+++ b/main.go
@@ -23,6 +23,7 @@ import (
"strconv"
"strings"
"sync"
+ "sync/atomic"
"time"
_ "github.com/lib/pq"
@@ -96,6 +97,13 @@ type Config struct {
// released from the active user count. Empty = default 5m. Use "0s" to disable.
SSHIdleTimeout string `json:"ssh_idle_timeout,omitempty"`
+ // MaxTotalConnections caps the total number of concurrent SSH connections
+ // across all users. Once reached, new connections are rejected before the
+ // (expensive) SSH handshake, so a surge past normal load cannot exhaust
+ // memory/CPU and crash the server. 0 or empty uses the default (10000);
+ // set to -1 to disable the global cap entirely.
+ MaxTotalConnections int `json:"max_total_connections,omitempty"`
+
// NEW: Directory to serve the admin panel from
AdminDir string `json:"admin_dir"`
@@ -259,17 +267,21 @@ type UDPGWConfig struct {
// debug logs. A value of zero suppresses hex dumps. Default is 64.
HexdumpN int `json:"hexdump"`
// WriteChan sets the size of the buffered channel used for sending
- // reply frames back to the client. Larger values allow more queued
- // replies before blocking. The default is 4096.
+ // reply frames back to the client. Default and hard ceiling is 1024:
+ // values above the ceiling are clamped down, so an oversized value left in
+ // an old config cannot bloat per-client memory. Zero uses the default.
WriteChan int `json:"write_chan"`
// UDPBindIP, if non‑empty, causes each per‑client UDP socket to bind to
// the specified local IP address. The port is chosen automatically.
UDPBindIP string `json:"udp_bind"`
- // UDPRBuf sets the size of the UDP socket read buffer in bytes. The
- // default is 8 MiB. Setting this to zero uses the default.
+ // UDPRBuf sets the per-client UDP socket read buffer in bytes. This is
+ // KERNEL memory allocated per connected client. Default and hard ceiling is
+ // 512 KiB; larger values (including old 8 MiB configs) are clamped down so
+ // thousands of clients cannot exhaust kernel memory. Zero uses the default.
UDPRBuf int `json:"udp_rbuf"`
- // UDPWBuf sets the size of the UDP socket write buffer in bytes. The
- // default is 8 MiB. Setting this to zero uses the default.
+ // UDPWBuf sets the per-client UDP socket write buffer in bytes. Default and
+ // hard ceiling is 512 KiB; larger values are clamped down. Zero uses the
+ // default.
UDPWBuf int `json:"udp_wbuf"`
// MapTTL controls how long a destination->connID mapping remains
// valid after the last packet from that destination. Expressed as a
@@ -292,6 +304,12 @@ type UDPGWConfig struct {
// growth if a client sprays packets to many unique destinations.
// Default is 32768.
MaxMapEntries int `json:"max_map_entries"`
+ // MaxClients caps the total number of concurrent UDPGW client TCP
+ // connections. Once reached, new clients are rejected (their TCP socket is
+ // closed immediately) instead of being accepted, so a surge past normal
+ // load cannot exhaust memory and crash the server. Default is 10000.
+ // A value <= 0 uses the default.
+ MaxClients int `json:"max_clients"`
// AutoRestartInterval controls a watchdog that periodically hard-restarts
// the integrated UDPGW listener and closes all connected UDPGW clients. Empty,
@@ -494,25 +512,37 @@ func mbpsToBytesPerSec(mbps int) int64 {
return int64(mbps) * 1024 * 1024 / 8
}
+// copyBufSize is the per-direction relay buffer for tunnel traffic. SSH tunnel
+// clients (mobile injection apps) push many low-bandwidth flows, so a large
+// buffer wastes heap with no throughput benefit. 16 KiB keeps per-channel heap
+// low across thousands of concurrent direct-tcpip channels.
+const copyBufSize = 16 * 1024
+
var copyBufPool = sync.Pool{
- New: func() interface{} { b := make([]byte, 32*1024); return &b },
+ New: func() interface{} { b := make([]byte, copyBufSize); return &b },
}
+// copyWithRateLimit relays src->dst through a pooled buffer. Every copy path
+// (rate-limited or not) goes through the pool: previously the lim==nil path used
+// io.Copy, which allocates a fresh 32 KiB buffer per direction per channel and
+// never pools it — at thousands of channels that churn dominated GC pressure.
func copyWithRateLimit(dst io.Writer, src io.Reader, lim *rate.Limiter) (written int64, err error) {
- if lim == nil {
- return io.Copy(dst, src)
- }
-
bufp := copyBufPool.Get().(*[]byte)
buf := *bufp
defer copyBufPool.Put(bufp)
- ctx := context.Background()
+
+ var ctx context.Context
+ if lim != nil {
+ ctx = context.Background()
+ }
for {
nr, er := src.Read(buf)
if nr > 0 {
- if err := lim.WaitN(ctx, nr); err != nil {
- return written, err
+ if lim != nil {
+ if err := lim.WaitN(ctx, nr); err != nil {
+ return written, err
+ }
}
nw, ew := dst.Write(buf[:nr])
@@ -571,6 +601,35 @@ func getSSHIdleTimeout() time.Duration {
return d
}
+// defaultMaxTotalConns is the global concurrent SSH connection cap used when
+// max_total_connections is unset (0). Keeps headroom above normal load while
+// preventing an unbounded surge from exhausting memory/CPU.
+const defaultMaxTotalConns = 10000
+
+var (
+ // activeSSHConns is the current number of in-flight SSH connections.
+ activeSSHConns int64
+ // maxTotalConns is the global cap; 0 means unlimited. Set from config.
+ maxTotalConns int64 = defaultMaxTotalConns
+ // sshConnsRejected counts connections turned away at the global cap.
+ sshConnsRejected int64
+)
+
+// setMaxTotalConnsFromConfig applies the max_total_connections config value:
+// 0 -> default, <0 -> disabled (unlimited), >0 -> that exact cap.
+func setMaxTotalConnsFromConfig(v int) {
+ var lim int64
+ switch {
+ case v < 0:
+ lim = 0
+ case v == 0:
+ lim = defaultMaxTotalConns
+ default:
+ lim = int64(v)
+ }
+ atomic.StoreInt64(&maxTotalConns, lim)
+}
+
// activityConn tracks real SSH transport activity in both directions. The idle
// monitor uses this instead of a read deadline so download-only or upload-only
// tunnels are considered live and are not disconnected.
@@ -2018,6 +2077,19 @@ func handleConn(tcpConn net.Conn, config *ssh.ServerConfig) {
trackedConn := newActivityConn(tcpConn)
defer trackedConn.Close()
+ // Global connection cap: reject before the expensive SSH handshake so a
+ // surge past normal load cannot exhaust memory/CPU. Count first, then check,
+ // so the defer always balances the increment.
+ live := atomic.AddInt64(&activeSSHConns, 1)
+ defer atomic.AddInt64(&activeSSHConns, -1)
+ if lim := atomic.LoadInt64(&maxTotalConns); lim > 0 && live > lim {
+ r := atomic.AddInt64(&sshConnsRejected, 1)
+ if r == 1 || r%1000 == 0 {
+ log.Printf("ssh: global connection cap reached (%d); rejected %d connection(s) so far", lim, r)
+ }
+ return
+ }
+
// Prevent goroutine leaks from clients that connect but never complete the SSH handshake.
_ = trackedConn.SetReadDeadline(time.Now().Add(sshHandshakeTimeout))
@@ -2142,8 +2214,17 @@ func handleDirectTCPIP(newChan ssh.NewChannel, u *UserState, upLimiter, downLimi
return
}
+ // Shared closer: when either copy direction finishes (including a
+ // half-close that never completes), both sides are force-closed so the
+ // other direction unblocks. Close is idempotent, so calling it from both
+ // directions is safe and no separate waiter goroutine is needed.
+ closeAll := func() {
+ _ = backend.Close()
+ _ = ch.Close()
+ }
+
+ // Drain channel requests concurrently so the peer isn't left waiting.
go func() {
- defer ch.Close()
for req := range reqs {
if req.WantReply {
req.Reply(false, nil)
@@ -2151,22 +2232,8 @@ func handleDirectTCPIP(newChan ssh.NewChannel, u *UserState, upLimiter, downLimi
}
}()
- // Use a WaitGroup + shared closer so that when either copy direction
- // finishes (including a half-close that never completes), both sides
- // are force-closed and the other goroutine is unblocked. Without
- // this, a backend that issues CloseWrite but never closes the read
- // side would leave the downstream goroutine blocked indefinitely,
- // leaking a goroutine, a rate-limiter, and the SSH channel.
- var wg sync.WaitGroup
- closeAll := func() {
- _ = backend.Close()
- _ = ch.Close()
- }
-
- // upstream: SSH channel -> backend
- wg.Add(1)
+ // upstream: SSH channel -> backend, in its own goroutine.
go func() {
- defer wg.Done()
_, _ = copyWithRateLimit(backend, ch, upLimiter)
// Signal to the backend that we are done writing.
if cw, ok := backend.(interface{ CloseWrite() error }); ok {
@@ -2175,19 +2242,11 @@ func handleDirectTCPIP(newChan ssh.NewChannel, u *UserState, upLimiter, downLimi
closeAll()
}()
- // downstream: backend -> SSH channel
- wg.Add(1)
- go func() {
- defer wg.Done()
- _, _ = copyWithRateLimit(ch, backend, downLimiter)
- closeAll()
- }()
-
- // Wait for both goroutines to finish, then ensure everything is closed.
- go func() {
- wg.Wait()
- closeAll()
- }()
+ // downstream: backend -> SSH channel, run in this goroutine.
+ // handleDirectTCPIP already runs as its own goroutine (see handleConn),
+ // so reusing it here avoids spawning a third goroutine per channel.
+ _, _ = copyWithRateLimit(ch, backend, downLimiter)
+ closeAll()
}
func handleDummySession(newChan ssh.NewChannel) {
@@ -3006,6 +3065,7 @@ func main() {
// Initialise default per-connection bandwidth limits and SSH inactivity cleanup.
setDefaultLimits(cfg.DefaultLimitMbpsUp, cfg.DefaultLimitMbpsDown)
setSSHIdleTimeoutFromConfig(cfg.SSHIdleTimeout)
+ setMaxTotalConnsFromConfig(cfg.MaxTotalConnections)
// Initialise listener pools (used for initial startup and hot-reload alike).
publicPool = newListenerPool(serveHTTP80)
diff --git a/udpgw_integration.go b/udpgw_integration.go
index f5e0778..5cf557f 100644
--- a/udpgw_integration.go
+++ b/udpgw_integration.go
@@ -25,10 +25,24 @@ import (
"time"
)
+// Safe ceilings applied regardless of config, so oversized values left in an
+// existing config.json cannot bloat per-client memory at scale.
+const (
+ udpgwSocketBufferMax = 512 * 1024 // per-client UDP socket buffer (kernel memory)
+ udpgwWriteChanMax = 1024 // per-client reply queue slots
+ udpgwDefaultMaxClients = 10000 // total concurrent client cap
+)
+
var (
udpgwMu sync.Mutex
udpgwLn net.Listener
udpgwClients = make(map[net.Conn]struct{})
+ // udpgwClientLimit is the max concurrent clients (0 = unlimited). Set when
+ // the listener starts and read on every accept, all under udpgwMu.
+ udpgwClientLimit int
+ // udpgwClientsRejected counts clients turned away at the cap, for logging
+ // and future stats. Guarded by udpgwMu.
+ udpgwClientsRejected int64
udpgwAutoMu sync.Mutex
udpgwAutoCancel context.CancelFunc
@@ -71,6 +85,17 @@ func registerUDPGWClient(conn net.Conn) bool {
_ = conn.Close()
return false
}
+ // Reject past the hard client cap so a surge cannot exhaust memory.
+ if udpgwClientLimit > 0 && len(udpgwClients) >= udpgwClientLimit {
+ udpgwClientsRejected++
+ rejected := udpgwClientsRejected
+ _ = conn.Close()
+ // Log the first rejection and then every 1000th to avoid log spam.
+ if rejected == 1 || rejected%1000 == 0 {
+ log.Printf("udpgw: client cap reached (%d); rejected %d client(s) so far", udpgwClientLimit, rejected)
+ }
+ return false
+ }
udpgwClients[conn] = struct{}{}
return true
}
@@ -131,21 +156,29 @@ func startUDPGWInstance(cfg *UDPGWConfig) error {
} else {
c.hexdumpN = 64
}
- if cfg.WriteChan > 0 {
+ // Per-client outgoing frame queue. A large queue costs ~24 B/slot of heap
+ // per client even when empty; at thousands of clients that adds up. UDP is
+ // lossy by nature, so a smaller queue that drops under backpressure is fine.
+ // The value is clamped to udpgwWriteChanMax so an oversized value left in an
+ // old config.json is ignored; only smaller custom values are honored.
+ c.writeChan = udpgwWriteChanMax
+ if cfg.WriteChan > 0 && cfg.WriteChan < udpgwWriteChanMax {
c.writeChan = cfg.WriteChan
- } else {
- c.writeChan = 4096
}
c.udpBindIP = cfg.UDPBindIP
- if cfg.UDPRBuf > 0 {
+ // Per-client UDP socket buffers are KERNEL memory, allocated per connected
+ // client. 8 MB per socket is fine for a single process-wide listener, but
+ // here every tunnel user gets its own socket, so a large value multiplied by
+ // thousands of clients can exhaust kernel memory (especially if
+ // net.core.rmem_max was raised for DNSTT). Clamp to udpgwSocketBufferMax so
+ // old large config values are ignored; only smaller custom values are used.
+ c.udpRBuf = udpgwSocketBufferMax
+ if cfg.UDPRBuf > 0 && cfg.UDPRBuf < udpgwSocketBufferMax {
c.udpRBuf = cfg.UDPRBuf
- } else {
- c.udpRBuf = 8 * 1024 * 1024
}
- if cfg.UDPWBuf > 0 {
+ c.udpWBuf = udpgwSocketBufferMax
+ if cfg.UDPWBuf > 0 && cfg.UDPWBuf < udpgwSocketBufferMax {
c.udpWBuf = cfg.UDPWBuf
- } else {
- c.udpWBuf = 8 * 1024 * 1024
}
// Parse durations with fallback defaults.
if cfg.MapTTL != "" {
@@ -191,6 +224,13 @@ func startUDPGWInstance(cfg *UDPGWConfig) error {
} else {
c.maxMapEntries = 32768
}
+ // Total concurrent client cap. New clients past this are rejected so a
+ // surge (e.g. well past normal load) cannot exhaust memory and crash.
+ if cfg.MaxClients > 0 {
+ c.maxClients = cfg.MaxClients
+ } else {
+ c.maxClients = udpgwDefaultMaxClients
+ }
// Start listening.
ln, err := net.Listen("tcp", c.listen)
if err != nil {
@@ -198,12 +238,14 @@ func startUDPGWInstance(cfg *UDPGWConfig) error {
return fmt.Errorf("udpgw: listen failed on %s: %w", c.listen, err)
}
- // Register as the active listener so stopUDPGW can close it.
+ // Register as the active listener so stopUDPGW can close it, and publish
+ // the current client cap so registerUDPGWClient can enforce it.
udpgwMu.Lock()
if udpgwLn != nil {
_ = udpgwLn.Close()
}
udpgwLn = ln
+ udpgwClientLimit = c.maxClients
udpgwMu.Unlock()
if c.debug {
@@ -329,6 +371,7 @@ type internalUDPGWConfig struct {
idleTimeout time.Duration
maxClientConns int
maxMapEntries int
+ maxClients int
}
// udpDestKey identifies a destination IPv4:port for the UDP gateway. A
@@ -365,7 +408,11 @@ func handleUDPGWClient(conn net.Conn, c *internalUDPGWConfig) {
if tcp, ok := conn.(*net.TCPConn); ok {
_ = tcp.SetNoDelay(true)
}
- br := bufio.NewReaderSize(conn, 256*1024)
+ // 32 KiB is ample for reading length-prefixed UDPGW frames (which are
+ // MTU-sized in practice). bufio serves reads larger than its buffer by
+ // reading straight into the caller's slice, so max-frame reads still work.
+ // The old 256 KiB buffer wasted ~224 KiB of heap per connected client.
+ br := bufio.NewReaderSize(conn, 32*1024)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
// Bind a UDP socket for this client. Use cfg.udpBindIP if provided.