SimpleRemoter/server/go/server/server.go

package server

import (
	"context"
	"crypto/tls"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/yuanyuanxiang/SimpleRemoter/server/go/connection"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/logger"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/protocol"
)

// Config holds server configuration
type Config struct {
	Port            int
	MaxConnections  int
	ReadBufferSize  int
	WriteBufferSize int
	KeepAliveTime   time.Duration
	ReadTimeout     time.Duration
	WriteTimeout    time.Duration

	// MaxBufferedBytes is a hard cap on unparsed inbound bytes buffered per
	// connection. The protocol's own length field is validated to 50 MiB, but
	// a peer that sends bytes that never assemble into a complete packet (e.g.
	// an HTTP-mask envelope without its terminating CRLF CRLF, which Parse()
	// waits on indefinitely) could otherwise grow the buffer without bound.
	// Default leaves slack for the largest legitimate packet plus its mask
	// envelope and a read chunk.
	MaxBufferedBytes int

	// Per-IP connection limits. The TCP listener (port 6543) is exposed
	// directly to the public internet, so a single host must not be able to
	// exhaust the global connection budget or flood new-connection attempts.
	// Defaults are generous enough for a NAT office with dozens of devices
	// sharing one source IP; raise them if you run a larger NAT deployment.
	PerIPMaxConnections int           // max concurrent connections per source IP (0 = unlimited)
	PerIPConnRateLimit  int           // max new connections per PerIPConnRateWindow per IP (0 = unlimited)
	PerIPConnRateWindow time.Duration // window for PerIPConnRateLimit

	// TLSCertFile / TLSKeyFile enable optional one-way TLS on the device port.
	// When both are set the server loads the cert and auto-detects TLS clients
	// by their ClientHello; plaintext clients keep working unchanged. Both empty
	// (the default) disables TLS and behaves exactly as before.
	TLSCertFile string // server certificate (PEM)
	TLSKeyFile  string // server private key (PEM)
}

// DefaultConfig returns default configuration
func DefaultConfig() Config {
	return Config{
		Port:                6543,
		MaxConnections:      9999,
		ReadBufferSize:      8192,
		WriteBufferSize:     8192,
		KeepAliveTime:       time.Minute * 5,
		ReadTimeout:         time.Minute * 2,
		WriteTimeout:        time.Second * 30,
		MaxBufferedBytes:    50*1024*1024 + 64*1024,
		PerIPMaxConnections: 256,
		PerIPConnRateLimit:  120,
		PerIPConnRateWindow: time.Minute,
	}
}

// Handler defines the interface for handling client events
type Handler interface {
	OnConnect(ctx *connection.Context)
	OnDisconnect(ctx *connection.Context)
	OnReceive(ctx *connection.Context, data []byte)
}

// Server is the TCP server
type Server struct {
	config   Config
	listener net.Listener
	manager  *connection.Manager
	handler  Handler
	parser   *protocol.Parser
	ipLimit  *ipLimiter

	// tlsCfg is non-nil when one-way TLS is enabled (cert + key configured).
	// wrapTLS consults it to sniff and wrap incoming connections.
	tlsCfg *tls.Config

	running atomic.Bool
	wg      sync.WaitGroup
	ctx     context.Context
	cancel  context.CancelFunc

	// Logger
	logger *logger.Logger
}

// New creates a new server
func New(config Config) *Server {
	ctx, cancel := context.WithCancel(context.Background())

	s := &Server{
		config:  config,
		manager: connection.NewManager(config.MaxConnections),
		parser:  protocol.NewParser(),
		ipLimit: newIPLimiter(config.PerIPMaxConnections, config.PerIPConnRateLimit, config.PerIPConnRateWindow),
		ctx:     ctx,
		cancel:  cancel,
		logger:  logger.New(logger.DefaultConfig()).WithPrefix("Server"),
	}

	return s
}

// ipLimiter caps concurrent connections and the rate of new connections per
// source IP. Both limits are per-host, not global, so a single attacker (or a
// single misbehaving host) cannot starve legitimate devices of the global
// connection budget.
type ipLimiter struct {
	mu            sync.Mutex
	maxConcurrent int
	window        time.Duration
	maxPerWindow  int
	concurrent    map[string]int
	recent        map[string][]time.Time
}

func newIPLimiter(maxConcurrent, maxPerWindow int, window time.Duration) *ipLimiter {
	return &ipLimiter{
		maxConcurrent: maxConcurrent,
		maxPerWindow:  maxPerWindow,
		window:        window,
		concurrent:    make(map[string]int),
		recent:        make(map[string][]time.Time),
	}
}

// allow records a new connection from ip and reports whether it is permitted.
// It enforces both a concurrent cap and a rolling accept-rate cap. An empty ip
// (which should not occur for a real TCP peer) is not throttled.
func (l *ipLimiter) allow(ip string) bool {
	if ip == "" {
		return true
	}
	now := time.Now()
	l.mu.Lock()
	defer l.mu.Unlock()

	if l.maxConcurrent > 0 && l.concurrent[ip] >= l.maxConcurrent {
		return false
	}

	if l.maxPerWindow > 0 && l.window > 0 {
		cutoff := now.Add(-l.window)
		times := l.recent[ip]
		keep := times[:0]
		for _, t := range times {
			if t.After(cutoff) {
				keep = append(keep, t)
			}
		}
		if len(keep) >= l.maxPerWindow {
			l.recent[ip] = keep
			return false
		}
		l.recent[ip] = append(keep, now)
	}

	l.concurrent[ip]++
	return true
}

// release decrements the concurrent count for ip. Safe to call on an ip that
// was never allowed.
func (l *ipLimiter) release(ip string) {
	if ip == "" {
		return
	}
	l.mu.Lock()
	if n := l.concurrent[ip]; n > 1 {
		l.concurrent[ip] = n - 1
	} else {
		delete(l.concurrent, ip)
	}
	l.mu.Unlock()
}

// sweep drops per-IP rate-limit entries whose timestamps have all aged out,
// keeping the map bounded over a long-running process.
func (l *ipLimiter) sweep() {
	l.mu.Lock()
	defer l.mu.Unlock()
	cutoff := time.Now().Add(-l.window)
	for ip, times := range l.recent {
		keep := times[:0]
		for _, t := range times {
			if t.After(cutoff) {
				keep = append(keep, t)
			}
		}
		if len(keep) == 0 {
			delete(l.recent, ip)
		} else {
			l.recent[ip] = keep
		}
	}
}

// peerIP returns the host part of a net.Addr's string form, mirroring
// connection.Context.GetPeerIP so accept-time and release-time keys agree.
func peerIP(addr net.Addr) string {
	host, _, err := net.SplitHostPort(addr.String())
	if err != nil {
		return addr.String()
	}
	return host
}

// SetHandler sets the event handler
func (s *Server) SetHandler(h Handler) {
	s.handler = h
	s.manager.SetCallbacks(
		func(ctx *connection.Context) {
			if s.handler != nil {
				s.handler.OnConnect(ctx)
			}
		},
		func(ctx *connection.Context) {
			if s.handler != nil {
				s.handler.OnDisconnect(ctx)
			}
		},
		func(ctx *connection.Context, data []byte) {
			if s.handler != nil {
				s.handler.OnReceive(ctx, data)
			}
		},
	)
}

// SetLogger sets the logger
func (s *Server) SetLogger(l *logger.Logger) {
	s.logger = l
}

// Start starts the server
func (s *Server) Start() error {
	// Load optional one-way TLS cert up front so a bad path/key fails startup
	// loudly rather than silently serving plaintext.
	if err := s.loadTLSConfig(); err != nil {
		return err
	}

	addr := net.JoinHostPort("0.0.0.0", itoa(s.config.Port))
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	s.listener = listener
	s.running.Store(true)

	s.logger.Info("Server started on port %d", s.config.Port)

	s.wg.Add(1)
	go s.acceptLoop()

	// Periodically prune aged-out per-IP rate-limit entries so the map stays
	// bounded over a long-running process. Stopped by s.ctx.Done() in Stop().
	s.wg.Add(1)
	go s.ipSweepLoop()

	return nil
}

func (s *Server) ipSweepLoop() {
	defer s.wg.Done()
	interval := s.config.PerIPConnRateWindow
	if interval <= 0 {
		interval = time.Minute
	}
	t := time.NewTicker(interval)
	defer t.Stop()
	for {
		select {
		case <-s.ctx.Done():
			return
		case <-t.C:
			s.ipLimit.sweep()
		}
	}
}

func itoa(i int) string {
	if i == 0 {
		return "0"
	}
	var b [20]byte
	pos := len(b)
	neg := i < 0
	if neg {
		// Handle math.MinInt safely by using unsigned
		u := uint(-i)
		for u > 0 {
			pos--
			b[pos] = byte('0' + u%10)
			u /= 10
		}
		pos--
		b[pos] = '-'
	} else {
		for i > 0 {
			pos--
			b[pos] = byte('0' + i%10)
			i /= 10
		}
	}
	return string(b[pos:])
}

// Stop stops the server gracefully
func (s *Server) Stop() {
	if !s.running.Swap(false) {
		return
	}

	s.cancel()

	if s.listener != nil {
		_ = s.listener.Close()
	}

	// Close all connections
	s.manager.CloseAll()

	// Close parser resources
	if s.parser != nil {
		s.parser.Close()
	}

	s.wg.Wait()
	s.logger.Info("Server stopped")
}

// acceptLoop accepts incoming connections
func (s *Server) acceptLoop() {
	defer s.wg.Done()

	for s.running.Load() {
		conn, err := s.listener.Accept()
		if err != nil {
			if s.running.Load() {
				s.logger.Error("Accept error: %v", err)
			}
			continue
		}

		// Check connection limit before spawning goroutine
		if s.manager.Count() >= s.config.MaxConnections {
			s.logger.Warn("Max connections reached, rejecting new connection from %s", conn.RemoteAddr())
			_ = conn.Close()
			continue
		}

		// Per-IP connection cap + accept rate. Rejecting here (before any
		// buffer allocation or goroutine) is cheap, and a rejected device
		// simply retries on its reconnect loop.
		if !s.ipLimit.allow(peerIP(conn.RemoteAddr())) {
			s.logger.Warn("Per-IP connection limit reached, rejecting %s", conn.RemoteAddr())
			_ = conn.Close()
			continue
		}

		// Handle each connection in its own goroutine
		go s.handleConnection(conn)
	}
}

// handleConnection handles a single connection
func (s *Server) handleConnection(rawConn net.Conn) {
	// Sniff and (when TLS is enabled) wrap TLS peers; plaintext passes through.
	conn := s.wrapTLS(rawConn)

	// Create context
	ctx := connection.NewContext(conn, nil)

	// Always release the per-IP slot on every exit path, including the
	// manager.Add failure below (a rare accept/add race).
	defer s.ipLimit.release(ctx.GetPeerIP())

	// Add to manager
	if err := s.manager.Add(ctx); err != nil {
		s.logger.Warn("Failed to add connection: %v", err)
		_ = conn.Close()
		return
	}

	defer func() {
		s.manager.Remove(ctx)
		_ = ctx.Close()
	}()

	// Read loop
	buf := make([]byte, s.config.ReadBufferSize)
	for !ctx.IsClosed() && s.running.Load() {
		// Set read deadline
		if s.config.ReadTimeout > 0 {
			_ = conn.SetReadDeadline(time.Now().Add(s.config.ReadTimeout))
		}

		n, err := conn.Read(buf)
		if err != nil {
			if err != io.EOF && s.running.Load() {
				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
					// Timeout - check keepalive
					if s.config.KeepAliveTime > 0 && ctx.TimeSinceLastActive() > s.config.KeepAliveTime {
						s.logger.Info("Connection %d timed out", ctx.ID)
						break
					}
					continue
				}
			}
			break
		}

		if n > 0 {
			ctx.UpdateLastActive()

			// Write to input buffer
			_, _ = ctx.InBuffer.Write(buf[:n])

			// Hard cap on unparsed inbound bytes (memory-DoS guard). A peer
			// that sends bytes that never assemble into a complete packet can
			// otherwise grow this buffer without bound.
			if s.config.MaxBufferedBytes > 0 && ctx.InBuffer.Len() > s.config.MaxBufferedBytes {
				s.logger.Warn("Inbound buffer exceeded %d bytes from ip=%s conn=%d — closing",
					s.config.MaxBufferedBytes, ctx.GetPeerIP(), ctx.ID)
				break
			}

			// Process received data
			s.processData(ctx)
		}
	}
}

// processData processes received data and calls handler
func (s *Server) processData(ctx *connection.Context) {
	for ctx.InBuffer.Len() > 0 {
		// Try to parse a complete packet
		data, err := s.parser.Parse(ctx)
		if err != nil {
			if err == protocol.ErrNeedMore {
				return
			}
			// Log more details for unsupported protocol errors
			if err == protocol.ErrUnsupported {
				// Get first 32 bytes for debugging
				peekLen := 32
				if ctx.InBuffer.Len() < peekLen {
					peekLen = ctx.InBuffer.Len()
				}
				rawData := ctx.InBuffer.Peek(peekLen)
				// Sanitize for safe logging (replace non-printable chars)
				ascii := make([]byte, len(rawData))
				for i, b := range rawData {
					if b >= 32 && b < 127 {
						ascii[i] = b
					} else {
						ascii[i] = '.'
					}
				}
				s.logger.Warn("Unsupported protocol from ip=%s conn=%d raw=%x ascii=%s",
					ctx.GetPeerIP(), ctx.ID, rawData, string(ascii))
			} else {
				s.logger.Error("Parse error for connection %d: %v", ctx.ID, err)
			}
			_ = ctx.Close()
			return
		}

		if data != nil {
			// Call handler
			s.manager.OnReceive(ctx, data)
		}
	}
}

// Send sends data to a specific connection
func (s *Server) Send(ctx *connection.Context, data []byte) error {
	if ctx == nil || ctx.IsClosed() {
		return connection.ErrConnectionClosed
	}

	// Encode and compress data
	encoded, err := s.parser.Encode(ctx, data)
	if err != nil {
		return err
	}

	return ctx.Send(encoded)
}

// Broadcast sends data to all connections
func (s *Server) Broadcast(data []byte) {
	s.manager.Range(func(ctx *connection.Context) bool {
		_ = s.Send(ctx, data)
		return true
	})
}

// GetConnection returns a connection by ID
func (s *Server) GetConnection(id uint64) *connection.Context {
	return s.manager.Get(id)
}

// ConnectionCount returns the current connection count
func (s *Server) ConnectionCount() int {
	return s.manager.Count()
}

// Port returns the server port
func (s *Server) Port() int {
	return s.config.Port
}

// UpdateMaxConnections updates the max connections limit
func (s *Server) UpdateMaxConnections(max int) {
	s.manager.UpdateMaxConnections(max)
}