wormhole-server/pkg/config/config.go

package config

import (
	"fmt"
	"os"
	"time"

	"github.com/sirupsen/logrus"
	"github.com/spf13/viper"
)

// ServerConfig 服务器配置
type ServerConfig struct {
	ServiceType string `mapstructure:"serviceType"`

	Proxy ProxyConfig `mapstructure:"proxy"`
	Auth  AuthConfig  `mapstructure:"auth"`

	Timeout  time.Duration `mapstructure:"timeout"`
	MaxConns int           `mapstructure:"maxConns"`
	LogLevel string        `mapstructure:"logLevel"`

	HealthCheck HealthCheckConfig `mapstructure:"healthCheck"`

	OptimizedServer OptimizedServerConfig `mapstructure:"optimizedServer"`
}

// ProxyConfig 代理配置
type ProxyConfig struct {
	Address string `mapstructure:"address"`
	Port    int    `mapstructure:"port"`
}

// AuthConfig 认证配置
type AuthConfig struct {
	Username string   `mapstructure:"username"`
	Password string   `mapstructure:"password"`
	Methods  []string `mapstructure:"methods"`
}

// HealthCheckConfig 健康检查配置
type HealthCheckConfig struct {
	Enabled bool   `mapstructure:"enabled"`
	Address string `mapstructure:"address"`
	Port    int    `mapstructure:"port"`
}

// OptimizedServerConfig 优化服务器配置
type OptimizedServerConfig struct {
	Enabled        bool          `mapstructure:"enabled"`
	MaxIdleTime    time.Duration `mapstructure:"maxIdleTime"`
	BufferSize     int           `mapstructure:"bufferSize"`
	LogConnections bool          `mapstructure:"logConnections"`

	DNSCache       DNSCacheConfig       `mapstructure:"dnsCache"`
	RateLimit      RateLimitConfig      `mapstructure:"rateLimit"`
	AccessControl  AccessControlConfig  `mapstructure:"accessControl"`
	Metrics        MetricsConfig        `mapstructure:"metrics"`
	ConnectionPool ConnectionPoolConfig `mapstructure:"connectionPool"`
	Memory         MemoryConfig         `mapstructure:"memory"`
	Transparent    TransparentConfig    `mapstructure:"transparent"`
}

// DNSCacheConfig DNS缓存配置
type DNSCacheConfig struct {
	Enabled bool          `mapstructure:"enabled"`
	MaxSize int           `mapstructure:"maxSize"`
	TTL     time.Duration `mapstructure:"ttl"`
}

// RateLimitConfig 速率限制配置
type RateLimitConfig struct {
	Enabled             bool          `mapstructure:"enabled"`
	RequestsPerSecond   int           `mapstructure:"requestsPerSecond"`
	BurstSize           int           `mapstructure:"burstSize"`
	PerIPRequestsPerSec int           `mapstructure:"perIPRequestsPerSec"`
	PerIPBurstSize      int           `mapstructure:"perIPBurstSize"`
	CleanupInterval     time.Duration `mapstructure:"cleanupInterval"`
}

// AccessControlConfig 访问控制配置
type AccessControlConfig struct {
	AllowedIPs []string `mapstructure:"allowedIPs"`
}

// MetricsConfig 指标配置
type MetricsConfig struct {
	Enabled  bool          `mapstructure:"enabled"`
	Interval time.Duration `mapstructure:"interval"`
}

// ConnectionPoolConfig 连接池配置
type ConnectionPoolConfig struct {
	Enabled     bool          `mapstructure:"enabled"`
	MaxSize     int           `mapstructure:"maxSize"`
	MaxLifetime time.Duration `mapstructure:"maxLifetime"`
	MaxIdle     time.Duration `mapstructure:"maxIdle"`
	InitialSize int           `mapstructure:"initialSize"`
}

// MemoryConfig 内存优化配置
type MemoryConfig struct {
	Enabled              bool          `mapstructure:"enabled"`
	BufferSizes          []int         `mapstructure:"bufferSizes"`
	MonitorInterval      time.Duration `mapstructure:"monitorInterval"`
	EnableAutoGC         bool          `mapstructure:"enableAutoGC"`
	HeapAllocThresholdMB int64         `mapstructure:"heapAllocThresholdMB"`
	HeapSysThresholdMB   int64         `mapstructure:"heapSysThresholdMB"`
	ForceGCThresholdMB   int64         `mapstructure:"forceGCThresholdMB"`
}

// TransparentConfig 透明代理配置
type TransparentConfig struct {
	Enabled         bool     `mapstructure:"enabled"`
	TransparentPort int      `mapstructure:"transparentPort"`
	DNSPort         int      `mapstructure:"dnsPort"`
	BypassIPs       []string `mapstructure:"bypassIPs"`
	BypassDomains   []string `mapstructure:"bypassDomains"`
}

// LoadConfig 加载配置文件
func LoadConfig(configPath string) (*ServerConfig, error) {
	// 检查配置文件是否存在
	if _, err := os.Stat(configPath); os.IsNotExist(err) {
		return nil, fmt.Errorf("config file not found: %s", configPath)
	}

	// 初始化viper
	viper.SetConfigFile(configPath)
	viper.SetConfigType("yaml")

	// 设置默认值
	setDefaults()

	// 读取配置文件
	if err := viper.ReadInConfig(); err != nil {
		return nil, fmt.Errorf("failed to read config file: %w", err)
	}

	// 解析配置
	var config ServerConfig
	if err := viper.Unmarshal(&config); err != nil {
		return nil, fmt.Errorf("failed to unmarshal config: %w", err)
	}

	// 验证配置
	if err := validateConfig(&config); err != nil {
		return nil, fmt.Errorf("invalid config: %w", err)
	}

	return &config, nil
}

// setDefaults 设置默认配置值
func setDefaults() {
	// 基本配置默认值
	viper.SetDefault("serviceType", "server")
	viper.SetDefault("proxy.address", "0.0.0.0")
	viper.SetDefault("proxy.port", 1080)
	viper.SetDefault("timeout", "30s")
	viper.SetDefault("maxConns", 5000)
	viper.SetDefault("logLevel", "info")

	// 健康检查默认值
	viper.SetDefault("healthCheck.enabled", true)
	viper.SetDefault("healthCheck.address", "127.0.0.1")
	viper.SetDefault("healthCheck.port", 8090)

	// 优化服务器默认值
	viper.SetDefault("optimizedServer.enabled", true)
	viper.SetDefault("optimizedServer.maxIdleTime", "5m")
	viper.SetDefault("optimizedServer.bufferSize", 65536)
	viper.SetDefault("optimizedServer.logConnections", true)

	// DNS缓存默认值
	viper.SetDefault("optimizedServer.dnsCache.enabled", true)
	viper.SetDefault("optimizedServer.dnsCache.maxSize", 10000)
	viper.SetDefault("optimizedServer.dnsCache.ttl", "10m")

	// 速率限制默认值
	viper.SetDefault("optimizedServer.rateLimit.enabled", true)
	viper.SetDefault("optimizedServer.rateLimit.requestsPerSecond", 100)

	// 指标默认值
	viper.SetDefault("optimizedServer.metrics.enabled", true)
	viper.SetDefault("optimizedServer.metrics.interval", "5m")

	// 连接池默认值
	viper.SetDefault("optimizedServer.connectionPool.enabled", true)
	viper.SetDefault("optimizedServer.connectionPool.maxSize", 1000)
	viper.SetDefault("optimizedServer.connectionPool.maxLifetime", "30m")
	viper.SetDefault("optimizedServer.connectionPool.maxIdle", "5m")
	viper.SetDefault("optimizedServer.connectionPool.initialSize", 100)

	// 内存优化默认值
	viper.SetDefault("optimizedServer.memory.enabled", true)
	viper.SetDefault("optimizedServer.memory.bufferSizes", []int{64, 128, 256, 512, 1024})
	viper.SetDefault("optimizedServer.memory.monitorInterval", "5m")
	viper.SetDefault("optimizedServer.memory.enableAutoGC", true)
	viper.SetDefault("optimizedServer.memory.heapAllocThresholdMB", 1024)
	viper.SetDefault("optimizedServer.memory.heapSysThresholdMB", 2048)
	viper.SetDefault("optimizedServer.memory.forceGCThresholdMB", 512)

	// 透明代理默认值
	viper.SetDefault("optimizedServer.transparent.enabled", false)
	viper.SetDefault("optimizedServer.transparent.transparentPort", 8080)
	viper.SetDefault("optimizedServer.transparent.dnsPort", 53)
	viper.SetDefault("optimizedServer.transparent.bypassIPs", []string{})
	viper.SetDefault("optimizedServer.transparent.bypassDomains", []string{})
}

// validateConfig 验证配置
func validateConfig(config *ServerConfig) error {
	// 验证端口范围
	if config.Proxy.Port < 1 || config.Proxy.Port > 65535 {
		return fmt.Errorf("invalid proxy port: %d", config.Proxy.Port)
	}

	if config.HealthCheck.Enabled {
		if config.HealthCheck.Port < 1 || config.HealthCheck.Port > 65535 {
			return fmt.Errorf("invalid health check port: %d", config.HealthCheck.Port)
		}
	}

	// 验证认证配置
	if config.Auth.Username != "" && config.Auth.Password == "" {
		return fmt.Errorf("password is required when username is set")
	}

	// 验证日志级别
	switch config.LogLevel {
	case "debug", "info", "warn", "error":
		// 有效的日志级别
	default:
		return fmt.Errorf("invalid log level: %s", config.LogLevel)
	}

	return nil
}

// GetLogLevel 获取logrus日志级别
func GetLogLevel(level string) logrus.Level {
	switch level {
	case "debug":
		return logrus.DebugLevel
	case "info":
		return logrus.InfoLevel
	case "warn":
		return logrus.WarnLevel
	case "error":
		return logrus.ErrorLevel
	default:
		return logrus.InfoLevel
	}
}

// ToSOCKS5Config 转换为SOCKS5配置
func (c *ServerConfig) ToSOCKS5Config() SOCKS5Config {
	return SOCKS5Config{
		Auth: SOCKS5AuthConfig{
			Methods:  c.Auth.Methods,
			Username: c.Auth.Username,
			Password: c.Auth.Password,
		},
		Timeout: c.Timeout,
		Rules:   []SOCKS5RuleConfig{}, // 从访问控制配置转换
	}
}

// SOCKS5Config SOCKS5特定配置
type SOCKS5Config struct {
	Auth    SOCKS5AuthConfig   `json:"auth"`
	Timeout time.Duration      `json:"timeout"`
	Rules   []SOCKS5RuleConfig `json:"rules"`
}

// SOCKS5AuthConfig SOCKS5认证配置
type SOCKS5AuthConfig struct {
	Methods  []string `json:"methods"`
	Username string   `json:"username"`
	Password string   `json:"password"`
}

// SOCKS5RuleConfig SOCKS5规则配置
type SOCKS5RuleConfig struct {
	Action string   `json:"action"`
	IPs    []string `json:"ips"`
	Ports  []int    `json:"ports"`
}