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 } // 保存原始的读取截止时间 var originalDeadline time.Time if deadline, ok := conn.(interface{ RemoteAddr() net.Addr }); ok { // 尝试获取远程地址,如果连接已关闭会失败 if deadline.RemoteAddr() == nil { return false } } // 设置短暂的读取超时进行连接检查 conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond)) defer func() { // 恢复原始的读取截止时间 conn.SetReadDeadline(originalDeadline) }() // 使用非阻塞的方式检查连接状态 // 尝试读取1字节,但不期望有数据 buf := make([]byte, 1) _, err := conn.Read(buf) // 如果是超时错误,说明连接是活跃的(没有数据可读但连接正常) if netErr, ok := err.(net.Error); ok && netErr.Timeout() { return true } // 如果没有错误,说明有数据可读,连接正常 // 但是需要将读取的数据放回去(这在实际应用中比较困难) if err == nil { return true } // 其他错误表示连接有问题 return false } 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") } // 初始化连接池(注意:SOCKS5代理的连接池使用受限) if cfg.OptimizedServer.ConnectionPool.Enabled { s.logger.Warn("Connection pool for SOCKS5 proxy has limited effectiveness due to dynamic target addresses") 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, continuing without connection pooling") s.connectionPool = nil // 确保连接池为nil,避免后续使用 } else { s.logger.WithField("max_size", cfg.OptimizedServer.ConnectionPool.MaxSize).Info("Connection pool initialized (limited effectiveness 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") } }