From 4a0383dce1bae5972aa54a8def4944c42e180a46 Mon Sep 17 00:00:00 2001 From: penguinehis Date: Thu, 9 Jul 2026 15:58:22 -0300 Subject: [PATCH] Optimization --- admin/assets/js/08-server-config.js | 4 + admin/index.html | 8 ++ hotreload.go | 1 + main.go | 146 ++++++++++++++++++++-------- udpgw_integration.go | 69 ++++++++++--- 5 files changed, 174 insertions(+), 54 deletions(-) diff --git a/admin/assets/js/08-server-config.js b/admin/assets/js/08-server-config.js index 111cbe3..799458c 100644 --- a/admin/assets/js/08-server-config.js +++ b/admin/assets/js/08-server-config.js @@ -95,6 +95,7 @@ async function loadServerConfig() { // SSH / general document.getElementById("cfgLimitUp").value = c.default_limit_mbps_up || 0; document.getElementById("cfgLimitDown").value = c.default_limit_mbps_down || 0; + document.getElementById("cfgMaxTotalConns").value = c.max_total_connections || 0; document.getElementById("cfgQuiet").checked = !!c.quiet; document.getElementById("cfgUserCount").checked = !!c.user_count; @@ -133,6 +134,7 @@ async function loadServerConfig() { const u = c.udpgw || {}; document.getElementById("cfgUdpgwListen").value = u.listen || ""; document.getElementById("cfgUdpgwMaxConns").value = u.max_client_conns || 0; + document.getElementById("cfgUdpgwMaxClients").value = u.max_clients || 0; document.getElementById("cfgUdpgwIdle").value = u.idle_timeout || ""; document.getElementById("cfgUdpgwMapTTL").value = u.map_ttl || ""; document.getElementById("cfgUdpgwAutoRestart").value = u.auto_restart_interval || ""; @@ -175,6 +177,7 @@ async function saveServerConfig() { admin_dir: "/opt/sshpanel/admin", default_limit_mbps_up: parseInt(document.getElementById("cfgLimitUp").value || "0", 10), default_limit_mbps_down: parseInt(document.getElementById("cfgLimitDown").value || "0", 10), + max_total_connections: parseInt(document.getElementById("cfgMaxTotalConns").value || "0", 10), quiet: document.getElementById("cfgQuiet").checked, user_count: document.getElementById("cfgUserCount").checked, banner: document.getElementById("cfgBanner").value, @@ -204,6 +207,7 @@ async function saveServerConfig() { udpgw: document.getElementById("cfgUdpgwEnabled").checked ? { listen: document.getElementById("cfgUdpgwListen").value.trim(), max_client_conns: parseInt(document.getElementById("cfgUdpgwMaxConns").value || "0", 10), + max_clients: parseInt(document.getElementById("cfgUdpgwMaxClients").value || "0", 10), idle_timeout: document.getElementById("cfgUdpgwIdle").value.trim(), map_ttl: document.getElementById("cfgUdpgwMapTTL").value.trim(), auto_restart_interval: document.getElementById("cfgUdpgwAutoRestart").value.trim(), diff --git a/admin/index.html b/admin/index.html index eca7eb6..3766414 100644 --- a/admin/index.html +++ b/admin/index.html @@ -976,6 +976,10 @@ +
+ + +
+
+ + +
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.