wormhole-server/internal/server/server.go

package server

import (
	"context"
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/azoic/wormhole-server/pkg/config"
	"github.com/azoic/wormhole-server/pkg/dns"
	"github.com/azoic/wormhole-server/pkg/health"
	"github.com/azoic/wormhole-server/pkg/logger"
	"github.com/azoic/wormhole-server/pkg/memory"
	"github.com/azoic/wormhole-server/pkg/metrics"
	"github.com/azoic/wormhole-server/pkg/pool"
	"github.com/azoic/wormhole-server/pkg/ratelimit"
	"github.com/azoic/wormhole-server/pkg/socks5"
	"github.com/azoic/wormhole-server/pkg/transparent"
	"github.com/sirupsen/logrus"
)

type Server struct {
	config       *config.ServerConfig
	logger       *logrus.Logger
	listener     net.Listener
	socks5Server *socks5.Server
	healthServer *health.HealthCheckServer
	metrics      *metrics.Metrics

	// 优化组件
	rateLimiter      *ratelimit.RateLimiter
	connectionPool   *pool.ConnectionPool
	memoryManager    *memory.Manager
	dnsProxy         *dns.DNSProxy
	transparentProxy *transparent.TransparentProxy

	ctx    context.Context
	cancel context.CancelFunc
	wg     sync.WaitGroup
}

// ConnectionFactory 连接工厂实现
type ConnectionFactory struct {
	network string
	timeout time.Duration
}

func (cf *ConnectionFactory) Create() (net.Conn, error) {
	// 连接池不应该预创建到未知目标的连接
	// 这里返回一个模拟连接或者直接返回错误
	return nil, fmt.Errorf("connection pool should not pre-create connections for SOCKS5 proxy")
}

func (cf *ConnectionFactory) Validate(conn net.Conn) bool {
	if conn == nil {
		return false
	}

	// 简单的连接验证
	conn.SetReadDeadline(time.Now().Add(1 * time.Second))
	defer conn.SetReadDeadline(time.Time{})

	// 尝试读取0字节来检查连接状态
	buf := make([]byte, 0)
	_, err := conn.Read(buf)
	return err == nil
}

func (cf *ConnectionFactory) Close(conn net.Conn) error {
	if conn != nil {
		return conn.Close()
	}
	return nil
}

func NewServer() *Server {
	return &Server{}
}

func (s *Server) Start(configPath string) error {
	// 加载配置
	cfg, err := config.LoadConfig(configPath)
	if err != nil {
		return fmt.Errorf("failed to load config: %w", err)
	}
	s.config = cfg

	// 初始化日志
	s.logger = logger.NewLogger(config.GetLogLevel(cfg.LogLevel))
	s.logger.WithField("config", configPath).Info("Server starting")

	// 创建上下文
	s.ctx, s.cancel = context.WithCancel(context.Background())

	// 初始化指标
	s.metrics = metrics.NewMetrics(s.logger)

	// 初始化内存管理器
	if cfg.OptimizedServer.Memory.Enabled {
		memConfig := memory.Config{
			BufferSizes:        cfg.OptimizedServer.Memory.BufferSizes,
			MonitorInterval:    cfg.OptimizedServer.Memory.MonitorInterval,
			EnableAutoGC:       cfg.OptimizedServer.Memory.EnableAutoGC,
			EnableOptimization: true,
			Thresholds: memory.Thresholds{
				HeapAllocMB:     cfg.OptimizedServer.Memory.HeapAllocThresholdMB,
				HeapSysMB:       cfg.OptimizedServer.Memory.HeapSysThresholdMB,
				ForceGCThreshMB: cfg.OptimizedServer.Memory.ForceGCThresholdMB,
			},
		}
		s.memoryManager = memory.NewManager(memConfig, s.logger)
		s.logger.Info("Memory optimization enabled")
	}

	// 初始化速率限制器
	if cfg.OptimizedServer.RateLimit.Enabled {
		rateLimitConfig := ratelimit.Config{
			Enabled:             true,
			RequestsPerSecond:   cfg.OptimizedServer.RateLimit.RequestsPerSecond,
			BurstSize:           cfg.OptimizedServer.RateLimit.BurstSize,
			PerIPRequestsPerSec: cfg.OptimizedServer.RateLimit.PerIPRequestsPerSec,
			PerIPBurstSize:      cfg.OptimizedServer.RateLimit.PerIPBurstSize,
			CleanupInterval:     cfg.OptimizedServer.RateLimit.CleanupInterval,
		}
		s.rateLimiter = ratelimit.NewRateLimiter(rateLimitConfig, s.logger)
		s.logger.WithField("rps", cfg.OptimizedServer.RateLimit.RequestsPerSecond).Info("Rate limiting enabled")
	}

	// 初始化连接池
	if cfg.OptimizedServer.ConnectionPool.Enabled {
		poolConfig := pool.Config{
			MaxSize:     cfg.OptimizedServer.ConnectionPool.MaxSize,
			MaxLifetime: cfg.OptimizedServer.ConnectionPool.MaxLifetime,
			MaxIdle:     cfg.OptimizedServer.ConnectionPool.MaxIdle,
			InitialSize: 0, // 禁用预创建，SOCKS5代理无法预知目标
		}
		factory := &ConnectionFactory{
			network: "tcp",
			timeout: cfg.Timeout,
		}
		s.connectionPool, err = pool.NewConnectionPool(poolConfig, factory, s.logger)
		if err != nil {
			s.logger.WithError(err).Warn("Failed to create connection pool")
		} else {
			s.logger.WithField("max_size", cfg.OptimizedServer.ConnectionPool.MaxSize).Info("Connection pool enabled (no pre-connections for SOCKS5)")
		}
	}

	// 初始化DNS代理
	if cfg.OptimizedServer.DNSCache.Enabled {
		dnsConfig := dns.Config{
			ListenPort:  cfg.OptimizedServer.Transparent.DNSPort,
			UpstreamDNS: []string{"8.8.8.8:53", "8.8.4.4:53"},
			CacheTTL:    cfg.OptimizedServer.DNSCache.TTL,
		}
		s.dnsProxy = dns.NewDNSProxy(dnsConfig, s.logger)

		// 启动DNS代理
		go func() {
			if err := s.dnsProxy.Start(s.ctx); err != nil {
				s.logger.WithError(err).Error("DNS proxy failed")
			}
		}()
		s.logger.WithField("ttl", cfg.OptimizedServer.DNSCache.TTL).Info("DNS caching enabled")
	}

	// 初始化透明代理
	if cfg.OptimizedServer.Transparent.Enabled {
		transparentConfig := transparent.Config{
			ProxyPort:       cfg.Proxy.Port,
			TransparentPort: cfg.OptimizedServer.Transparent.TransparentPort,
			DNSPort:         cfg.OptimizedServer.Transparent.DNSPort,
			BypassIPs:       cfg.OptimizedServer.Transparent.BypassIPs,
			BypassDomains:   cfg.OptimizedServer.Transparent.BypassDomains,
		}
		s.transparentProxy = transparent.NewTransparentProxy(transparentConfig, s.logger)

		// 设置透明代理规则
		if err := s.transparentProxy.SetupTransparentProxy(s.ctx); err != nil {
			s.logger.WithError(err).Warn("Failed to setup transparent proxy")
		} else {
			s.logger.Info("Transparent proxy enabled")
		}
	}

	// 启动指标定期记录
	if cfg.OptimizedServer.Metrics.Enabled {
		s.metrics.StartPeriodicLogging(cfg.OptimizedServer.Metrics.Interval)
	}

	// 初始化SOCKS5服务器
	if err := s.initSOCKS5Server(); err != nil {
		return fmt.Errorf("failed to initialize SOCKS5 server: %w", err)
	}

	// 启动健康检查服务器
	if cfg.HealthCheck.Enabled {
		if err := s.startHealthServer(); err != nil {
			s.logger.WithError(err).Warn("Failed to start health check server")
		}
	}

	// 启动主SOCKS5服务器
	return s.startSOCKS5Server()
}

func (s *Server) Stop() error {
	s.logger.Info("Stopping server...")

	if s.cancel != nil {
		s.cancel()
	}

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

	if s.healthServer != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		s.healthServer.Stop(ctx)
	}

	// 停止优化组件
	if s.rateLimiter != nil {
		s.rateLimiter.Stop()
	}

	if s.connectionPool != nil {
		s.connectionPool.Close()
	}

	if s.memoryManager != nil {
		s.memoryManager.Stop()
	}

	if s.dnsProxy != nil {
		s.dnsProxy.Stop()
	}

	if s.transparentProxy != nil {
		s.transparentProxy.CleanupTransparentProxy()
	}

	// 等待所有goroutine完成
	s.wg.Wait()

	s.logger.Info("Server stopped")
	return nil
}

func (s *Server) initSOCKS5Server() error {
	// 转换配置
	socks5Config := socks5.Config{
		Auth: socks5.AuthConfig{
			Methods:  s.config.Auth.Methods,
			Username: s.config.Auth.Username,
			Password: s.config.Auth.Password,
		},
		Timeout: s.config.Timeout,
		Rules: []socks5.RuleConfig{
			{
				Action: "allow",
				IPs:    s.config.OptimizedServer.AccessControl.AllowedIPs,
			},
		},
	}

	s.socks5Server = socks5.NewServer(socks5Config, s.logger)

	// TODO: 集成优化的转发器
	// 需要更新SOCKS5服务器以支持可配置的转发器
	// if s.config.OptimizedServer.Enabled && s.memoryManager != nil {
	//     // 设置优化的转发器
	// }

	return nil
}

func (s *Server) startSOCKS5Server() error {
	address := fmt.Sprintf("%s:%d", s.config.Proxy.Address, s.config.Proxy.Port)

	listener, err := net.Listen("tcp", address)
	if err != nil {
		return fmt.Errorf("failed to listen on %s: %w", address, err)
	}
	s.listener = listener

	s.logger.WithFields(logrus.Fields{
		"address": s.config.Proxy.Address,
		"port":    s.config.Proxy.Port,
	}).Info("SOCKS5 server started")

	// 打印启动信息
	s.printStartupInfo()

	// 接受连接
	for {
		conn, err := listener.Accept()
		if err != nil {
			select {
			case <-s.ctx.Done():
				return nil // 正常关闭
			default:
				s.logger.WithError(err).Error("Failed to accept connection")
				continue
			}
		}

		// 速率限制检查
		if s.rateLimiter != nil {
			if !s.rateLimiter.Allow(conn.RemoteAddr().String()) {
				s.logger.WithField("remote_addr", conn.RemoteAddr()).Warn("Rate limit exceeded")
				conn.Close()
				continue
			}
		}

		// 检查连接数限制
		if s.metrics != nil && s.config.MaxConns > 0 {
			if s.metrics.ActiveConnections >= int64(s.config.MaxConns) {
				s.logger.Warn("Max connections reached, rejecting connection")
				conn.Close()
				continue
			}
		}

		// 处理连接
		s.wg.Add(1)
		go s.handleConnection(conn)
	}
}

func (s *Server) handleConnection(conn net.Conn) {
	defer s.wg.Done()
	defer conn.Close()

	if s.metrics != nil {
		s.metrics.IncActiveConnections()
		s.metrics.IncTotalRequests()
		defer s.metrics.DecActiveConnections()
	}

	remoteAddr := conn.RemoteAddr()

	if s.config.OptimizedServer.LogConnections {
		s.logger.WithField("remote_addr", remoteAddr).Info("New connection")
	}

	// 设置连接超时
	if s.config.Timeout > 0 {
		conn.SetDeadline(time.Now().Add(s.config.Timeout))
	}

	// 处理SOCKS5连接
	if err := s.socks5Server.HandleConnection(conn); err != nil {
		if s.metrics != nil {
			s.metrics.IncFailedRequests()
		}
		s.logger.WithError(err).WithField("remote_addr", remoteAddr).Debug("Connection handling failed")
	}
}

func (s *Server) startHealthServer() error {
	s.healthServer = health.NewHealthCheckServer(
		s.config.HealthCheck.Address,
		fmt.Sprintf("%d", s.config.HealthCheck.Port),
		s.logger,
		s.metrics,
	)

	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		if err := s.healthServer.Start(); err != nil {
			s.logger.WithError(err).Error("Health check server failed")
		}
	}()

	s.logger.WithFields(logrus.Fields{
		"address": s.config.HealthCheck.Address,
		"port":    s.config.HealthCheck.Port,
	}).Info("Health check server started")

	return nil
}

func (s *Server) printStartupInfo() {
	s.logger.Info("🚀 Wormhole SOCKS5 Server started successfully!")

	features := []string{}
	if s.config.OptimizedServer.Enabled {
		features = append(features, "optimized")
	}
	if s.rateLimiter != nil {
		features = append(features, "rate-limiting")
	}
	if s.connectionPool != nil {
		features = append(features, "connection-pooling")
	}
	if s.memoryManager != nil {
		features = append(features, "memory-optimization")
	}
	if s.dnsProxy != nil {
		features = append(features, "dns-caching")
	}
	if s.transparentProxy != nil {
		features = append(features, "transparent-proxy")
	}

	s.logger.WithFields(logrus.Fields{
		"proxy_address": fmt.Sprintf("%s:%d", s.config.Proxy.Address, s.config.Proxy.Port),
		"auth_enabled":  s.config.Auth.Username != "",
		"health_check":  s.config.HealthCheck.Enabled,
		"features":      features,
	}).Info("Server configuration")

	if s.config.HealthCheck.Enabled {
		s.logger.WithField("url", fmt.Sprintf("http://%s:%d",
			s.config.HealthCheck.Address, s.config.HealthCheck.Port)).Info("Health check available")
	}

	// 打印使用说明
	s.logger.Info("📖 Usage:")
	s.logger.Infof("   SOCKS5 Proxy: %s:%d", s.config.Proxy.Address, s.config.Proxy.Port)
	if s.config.Auth.Username != "" {
		s.logger.Info("   Authentication: Required")
		s.logger.Infof("   Username: %s", s.config.Auth.Username)
	} else {
		s.logger.Info("   Authentication: None")
	}

	if len(features) > 0 {
		s.logger.WithField("features", features).Info("✨ Optimization features enabled")
	}
}