parent
77128910b3
commit
52f0fba78d
@ -0,0 +1,289 @@ |
||||
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"` |
||||
} |
||||
@ -0,0 +1,379 @@ |
||||
package memory |
||||
|
||||
import ( |
||||
"runtime" |
||||
"sync" |
||||
"time" |
||||
|
||||
"github.com/sirupsen/logrus" |
||||
) |
||||
|
||||
// BufferPool 缓冲区池
|
||||
type BufferPool struct { |
||||
pools map[int]*sync.Pool |
||||
sizes []int |
||||
logger *logrus.Logger |
||||
stats BufferStats |
||||
} |
||||
|
||||
// BufferStats 缓冲区统计
|
||||
type BufferStats struct { |
||||
mu sync.RWMutex |
||||
Gets int64 `json:"gets"` |
||||
Puts int64 `json:"puts"` |
||||
News int64 `json:"news"` |
||||
Reuses int64 `json:"reuses"` |
||||
TotalAlloc int64 `json:"totalAlloc"` |
||||
TotalReused int64 `json:"totalReused"` |
||||
} |
||||
|
||||
// MemoryMonitor 内存监控器
|
||||
type MemoryMonitor struct { |
||||
logger *logrus.Logger |
||||
ticker *time.Ticker |
||||
stopCh chan struct{} |
||||
thresholds Thresholds |
||||
lastStats runtime.MemStats |
||||
callbacks []MemoryCallback |
||||
mu sync.RWMutex |
||||
} |
||||
|
||||
// Thresholds 内存阈值配置
|
||||
type Thresholds struct { |
||||
HeapAllocMB int64 `json:"heapAllocMB"` |
||||
HeapSysMB int64 `json:"heapSysMB"` |
||||
GCPercent int `json:"gcPercent"` |
||||
ForceGCThreshMB int64 `json:"forceGCThreshMB"` |
||||
} |
||||
|
||||
// MemoryCallback 内存回调函数
|
||||
type MemoryCallback func(stats runtime.MemStats) |
||||
|
||||
// Config 内存管理配置
|
||||
type Config struct { |
||||
BufferSizes []int `json:"bufferSizes"` |
||||
MonitorInterval time.Duration `json:"monitorInterval"` |
||||
Thresholds Thresholds `json:"thresholds"` |
||||
EnableAutoGC bool `json:"enableAutoGC"` |
||||
EnableOptimization bool `json:"enableOptimization"` |
||||
} |
||||
|
||||
// Manager 内存管理器
|
||||
type Manager struct { |
||||
bufferPool *BufferPool |
||||
monitor *MemoryMonitor |
||||
config Config |
||||
logger *logrus.Logger |
||||
} |
||||
|
||||
// NewManager 创建内存管理器
|
||||
func NewManager(config Config, logger *logrus.Logger) *Manager { |
||||
// 设置默认值
|
||||
if len(config.BufferSizes) == 0 { |
||||
config.BufferSizes = []int{512, 1024, 2048, 4096, 8192, 16384, 32768, 65536} |
||||
} |
||||
if config.MonitorInterval == 0 { |
||||
config.MonitorInterval = 30 * time.Second |
||||
} |
||||
if config.Thresholds.HeapAllocMB == 0 { |
||||
config.Thresholds.HeapAllocMB = 100 |
||||
} |
||||
if config.Thresholds.HeapSysMB == 0 { |
||||
config.Thresholds.HeapSysMB = 200 |
||||
} |
||||
if config.Thresholds.GCPercent == 0 { |
||||
config.Thresholds.GCPercent = 100 |
||||
} |
||||
if config.Thresholds.ForceGCThreshMB == 0 { |
||||
config.Thresholds.ForceGCThreshMB = 500 |
||||
} |
||||
|
||||
manager := &Manager{ |
||||
config: config, |
||||
logger: logger, |
||||
} |
||||
|
||||
// 创建缓冲区池
|
||||
if config.EnableOptimization { |
||||
manager.bufferPool = NewBufferPool(config.BufferSizes, logger) |
||||
} |
||||
|
||||
// 创建内存监控器
|
||||
manager.monitor = NewMemoryMonitor(config.MonitorInterval, config.Thresholds, logger) |
||||
|
||||
// 启用自动GC
|
||||
if config.EnableAutoGC { |
||||
manager.monitor.AddCallback(manager.autoGCCallback) |
||||
} |
||||
|
||||
return manager |
||||
} |
||||
|
||||
// NewBufferPool 创建缓冲区池
|
||||
func NewBufferPool(sizes []int, logger *logrus.Logger) *BufferPool { |
||||
bp := &BufferPool{ |
||||
pools: make(map[int]*sync.Pool), |
||||
sizes: make([]int, len(sizes)), |
||||
logger: logger, |
||||
} |
||||
|
||||
copy(bp.sizes, sizes) |
||||
|
||||
// 为每个大小创建池
|
||||
for _, size := range sizes { |
||||
sz := size // 捕获循环变量
|
||||
bp.pools[sz] = &sync.Pool{ |
||||
New: func() interface{} { |
||||
bp.stats.incNews() |
||||
bp.stats.addTotalAlloc(int64(sz)) |
||||
return make([]byte, sz) |
||||
}, |
||||
} |
||||
} |
||||
|
||||
return bp |
||||
} |
||||
|
||||
// Get 获取缓冲区
|
||||
func (bp *BufferPool) Get(size int) []byte { |
||||
bp.stats.incGets() |
||||
|
||||
// 找到最合适的大小
|
||||
for _, poolSize := range bp.sizes { |
||||
if size <= poolSize { |
||||
buf := bp.pools[poolSize].Get().([]byte) |
||||
bp.stats.incReuses() |
||||
bp.stats.addTotalReused(int64(poolSize)) |
||||
return buf[:size] |
||||
} |
||||
} |
||||
|
||||
// 没有合适的池，创建新缓冲区
|
||||
bp.stats.incNews() |
||||
bp.stats.addTotalAlloc(int64(size)) |
||||
return make([]byte, size) |
||||
} |
||||
|
||||
// Put 归还缓冲区
|
||||
func (bp *BufferPool) Put(buf []byte) { |
||||
if buf == nil { |
||||
return |
||||
} |
||||
|
||||
bp.stats.incPuts() |
||||
capacity := cap(buf) |
||||
|
||||
// 找到对应的池
|
||||
for _, poolSize := range bp.sizes { |
||||
if capacity == poolSize { |
||||
bp.pools[poolSize].Put(buf[:poolSize]) |
||||
return |
||||
} |
||||
} |
||||
} |
||||
|
||||
// GetStats 获取缓冲区统计
|
||||
func (bp *BufferPool) GetStats() BufferStats { |
||||
bp.stats.mu.RLock() |
||||
defer bp.stats.mu.RUnlock() |
||||
return bp.stats |
||||
} |
||||
|
||||
// NewMemoryMonitor 创建内存监控器
|
||||
func NewMemoryMonitor(interval time.Duration, thresholds Thresholds, logger *logrus.Logger) *MemoryMonitor { |
||||
monitor := &MemoryMonitor{ |
||||
logger: logger, |
||||
ticker: time.NewTicker(interval), |
||||
stopCh: make(chan struct{}), |
||||
thresholds: thresholds, |
||||
callbacks: make([]MemoryCallback, 0), |
||||
} |
||||
|
||||
// 启动监控
|
||||
go monitor.run() |
||||
|
||||
return monitor |
||||
} |
||||
|
||||
// AddCallback 添加内存回调
|
||||
func (mm *MemoryMonitor) AddCallback(callback MemoryCallback) { |
||||
mm.mu.Lock() |
||||
mm.callbacks = append(mm.callbacks, callback) |
||||
mm.mu.Unlock() |
||||
} |
||||
|
||||
// Stop 停止监控
|
||||
func (mm *MemoryMonitor) Stop() { |
||||
close(mm.stopCh) |
||||
mm.ticker.Stop() |
||||
} |
||||
|
||||
// run 运行监控
|
||||
func (mm *MemoryMonitor) run() { |
||||
for { |
||||
select { |
||||
case <-mm.stopCh: |
||||
return |
||||
case <-mm.ticker.C: |
||||
mm.check() |
||||
} |
||||
} |
||||
} |
||||
|
||||
// check 检查内存状态
|
||||
func (mm *MemoryMonitor) check() { |
||||
var stats runtime.MemStats |
||||
runtime.ReadMemStats(&stats) |
||||
|
||||
// 记录内存统计
|
||||
heapAllocMB := stats.HeapAlloc / 1024 / 1024 |
||||
heapSysMB := stats.HeapSys / 1024 / 1024 |
||||
|
||||
mm.logger.WithFields(logrus.Fields{ |
||||
"heap_alloc_mb": heapAllocMB, |
||||
"heap_sys_mb": heapSysMB, |
||||
"gc_num": stats.NumGC, |
||||
"goroutines": runtime.NumGoroutine(), |
||||
}).Debug("Memory stats") |
||||
|
||||
// 检查阈值
|
||||
if int64(heapAllocMB) > mm.thresholds.HeapAllocMB { |
||||
mm.logger.WithField("heap_alloc_mb", heapAllocMB).Warn("Heap allocation threshold exceeded") |
||||
} |
||||
|
||||
if int64(heapSysMB) > mm.thresholds.HeapSysMB { |
||||
mm.logger.WithField("heap_sys_mb", heapSysMB).Warn("Heap system threshold exceeded") |
||||
} |
||||
|
||||
// 强制GC阈值
|
||||
if int64(heapAllocMB) > mm.thresholds.ForceGCThreshMB { |
||||
mm.logger.WithField("heap_alloc_mb", heapAllocMB).Info("Force GC triggered") |
||||
runtime.GC() |
||||
} |
||||
|
||||
// 调用回调函数
|
||||
mm.mu.RLock() |
||||
for _, callback := range mm.callbacks { |
||||
go callback(stats) |
||||
} |
||||
mm.mu.RUnlock() |
||||
|
||||
mm.lastStats = stats |
||||
} |
||||
|
||||
// GetStats 获取最新内存统计
|
||||
func (mm *MemoryMonitor) GetStats() runtime.MemStats { |
||||
return mm.lastStats |
||||
} |
||||
|
||||
// Manager 方法
|
||||
|
||||
// GetBuffer 获取缓冲区
|
||||
func (m *Manager) GetBuffer(size int) []byte { |
||||
if m.bufferPool != nil { |
||||
return m.bufferPool.Get(size) |
||||
} |
||||
return make([]byte, size) |
||||
} |
||||
|
||||
// PutBuffer 归还缓冲区
|
||||
func (m *Manager) PutBuffer(buf []byte) { |
||||
if m.bufferPool != nil { |
||||
m.bufferPool.Put(buf) |
||||
} |
||||
} |
||||
|
||||
// GetBufferStats 获取缓冲区统计
|
||||
func (m *Manager) GetBufferStats() BufferStats { |
||||
if m.bufferPool != nil { |
||||
return m.bufferPool.GetStats() |
||||
} |
||||
return BufferStats{} |
||||
} |
||||
|
||||
// GetMemoryStats 获取内存统计
|
||||
func (m *Manager) GetMemoryStats() runtime.MemStats { |
||||
return m.monitor.GetStats() |
||||
} |
||||
|
||||
// Stop 停止内存管理器
|
||||
func (m *Manager) Stop() { |
||||
if m.monitor != nil { |
||||
m.monitor.Stop() |
||||
} |
||||
} |
||||
|
||||
// ForceGC 强制垃圾回收
|
||||
func (m *Manager) ForceGC() { |
||||
m.logger.Info("Manual GC triggered") |
||||
runtime.GC() |
||||
} |
||||
|
||||
// autoGCCallback 自动GC回调
|
||||
func (m *Manager) autoGCCallback(stats runtime.MemStats) { |
||||
heapAllocMB := stats.HeapAlloc / 1024 / 1024 |
||||
|
||||
// 当堆分配超过阈值时触发GC
|
||||
if int64(heapAllocMB) > m.config.Thresholds.ForceGCThreshMB { |
||||
m.ForceGC() |
||||
} |
||||
} |
||||
|
||||
// GetOverallStats 获取总体统计信息
|
||||
func (m *Manager) GetOverallStats() map[string]interface{} { |
||||
memStats := m.GetMemoryStats() |
||||
bufferStats := m.GetBufferStats() |
||||
|
||||
return map[string]interface{}{ |
||||
"memory": map[string]interface{}{ |
||||
"heap_alloc_mb": memStats.HeapAlloc / 1024 / 1024, |
||||
"heap_sys_mb": memStats.HeapSys / 1024 / 1024, |
||||
"gc_num": memStats.NumGC, |
||||
"goroutines": runtime.NumGoroutine(), |
||||
}, |
||||
"buffers": bufferStats, |
||||
"config": map[string]interface{}{ |
||||
"optimization_enabled": m.config.EnableOptimization, |
||||
"auto_gc_enabled": m.config.EnableAutoGC, |
||||
"buffer_sizes": m.config.BufferSizes, |
||||
}, |
||||
} |
||||
} |
||||
|
||||
// BufferStats 方法
|
||||
|
||||
func (bs *BufferStats) incGets() { |
||||
bs.mu.Lock() |
||||
bs.Gets++ |
||||
bs.mu.Unlock() |
||||
} |
||||
|
||||
func (bs *BufferStats) incPuts() { |
||||
bs.mu.Lock() |
||||
bs.Puts++ |
||||
bs.mu.Unlock() |
||||
} |
||||
|
||||
func (bs *BufferStats) incNews() { |
||||
bs.mu.Lock() |
||||
bs.News++ |
||||
bs.mu.Unlock() |
||||
} |
||||
|
||||
func (bs *BufferStats) incReuses() { |
||||
bs.mu.Lock() |
||||
bs.Reuses++ |
||||
bs.mu.Unlock() |
||||
} |
||||
|
||||
func (bs *BufferStats) addTotalAlloc(size int64) { |
||||
bs.mu.Lock() |
||||
bs.TotalAlloc += size |
||||
bs.mu.Unlock() |
||||
} |
||||
|
||||
func (bs *BufferStats) addTotalReused(size int64) { |
||||
bs.mu.Lock() |
||||
bs.TotalReused += size |
||||
bs.mu.Unlock() |
||||
} |
||||
@ -0,0 +1,413 @@ |
||||
package pool |
||||
|
||||
import ( |
||||
"errors" |
||||
"net" |
||||
"sync" |
||||
"time" |
||||
|
||||
"github.com/sirupsen/logrus" |
||||
) |
||||
|
||||
var ( |
||||
ErrPoolClosed = errors.New("connection pool is closed") |
||||
ErrPoolFull = errors.New("connection pool is full") |
||||
ErrConnExpired = errors.New("connection has expired") |
||||
ErrConnInvalid = errors.New("connection is invalid") |
||||
) |
||||
|
||||
// ConnectionPool 连接池
|
||||
type ConnectionPool struct { |
||||
logger *logrus.Logger |
||||
factory Factory |
||||
pool chan *PooledConnection |
||||
mu sync.RWMutex |
||||
closed bool |
||||
maxSize int |
||||
maxLifetime time.Duration |
||||
maxIdle time.Duration |
||||
|
||||
// 统计信息
|
||||
stats Stats |
||||
} |
||||
|
||||
// PooledConnection 池化连接
|
||||
type PooledConnection struct { |
||||
conn net.Conn |
||||
createdAt time.Time |
||||
lastUsed time.Time |
||||
pool *ConnectionPool |
||||
} |
||||
|
||||
// Factory 连接工厂
|
||||
type Factory interface { |
||||
Create() (net.Conn, error) |
||||
Validate(net.Conn) bool |
||||
Close(net.Conn) error |
||||
} |
||||
|
||||
// Config 连接池配置
|
||||
type Config struct { |
||||
MaxSize int `json:"maxSize"` |
||||
MaxLifetime time.Duration `json:"maxLifetime"` |
||||
MaxIdle time.Duration `json:"maxIdle"` |
||||
InitialSize int `json:"initialSize"` |
||||
} |
||||
|
||||
// Stats 统计信息
|
||||
type Stats struct { |
||||
mu sync.RWMutex |
||||
Created int64 `json:"created"` |
||||
Reused int64 `json:"reused"` |
||||
Closed int64 `json:"closed"` |
||||
Active int64 `json:"active"` |
||||
Idle int64 `json:"idle"` |
||||
Failures int64 `json:"failures"` |
||||
} |
||||
|
||||
// NewConnectionPool 创建新的连接池
|
||||
func NewConnectionPool(config Config, factory Factory, logger *logrus.Logger) (*ConnectionPool, error) { |
||||
if config.MaxSize <= 0 { |
||||
config.MaxSize = 100 |
||||
} |
||||
if config.MaxLifetime == 0 { |
||||
config.MaxLifetime = 30 * time.Minute |
||||
} |
||||
if config.MaxIdle == 0 { |
||||
config.MaxIdle = 5 * time.Minute |
||||
} |
||||
|
||||
pool := &ConnectionPool{ |
||||
logger: logger, |
||||
factory: factory, |
||||
pool: make(chan *PooledConnection, config.MaxSize), |
||||
maxSize: config.MaxSize, |
||||
maxLifetime: config.MaxLifetime, |
||||
maxIdle: config.MaxIdle, |
||||
} |
||||
|
||||
// 预创建连接
|
||||
for i := 0; i < config.InitialSize && i < config.MaxSize; i++ { |
||||
conn, err := pool.factory.Create() |
||||
if err != nil { |
||||
pool.logger.WithError(err).Warn("Failed to create initial connection") |
||||
continue |
||||
} |
||||
|
||||
pooledConn := &PooledConnection{ |
||||
conn: conn, |
||||
createdAt: time.Now(), |
||||
lastUsed: time.Now(), |
||||
pool: pool, |
||||
} |
||||
|
||||
select { |
||||
case pool.pool <- pooledConn: |
||||
pool.stats.incCreated() |
||||
pool.stats.incIdle() |
||||
default: |
||||
conn.Close() |
||||
} |
||||
} |
||||
|
||||
// 启动清理goroutine
|
||||
go pool.cleaner() |
||||
|
||||
return pool, nil |
||||
} |
||||
|
||||
// Get 获取连接
|
||||
func (p *ConnectionPool) Get() (*PooledConnection, error) { |
||||
p.mu.RLock() |
||||
if p.closed { |
||||
p.mu.RUnlock() |
||||
return nil, ErrPoolClosed |
||||
} |
||||
p.mu.RUnlock() |
||||
|
||||
// 尝试从池中获取连接
|
||||
for { |
||||
select { |
||||
case conn := <-p.pool: |
||||
p.stats.decIdle() |
||||
|
||||
// 检查连接是否有效
|
||||
if p.isConnValid(conn) { |
||||
conn.lastUsed = time.Now() |
||||
p.stats.incReused() |
||||
p.stats.incActive() |
||||
return conn, nil |
||||
} |
||||
|
||||
// 连接无效，关闭并继续尝试
|
||||
p.closeConn(conn) |
||||
continue |
||||
|
||||
default: |
||||
// 池中没有可用连接，创建新连接
|
||||
return p.createConnection() |
||||
} |
||||
} |
||||
} |
||||
|
||||
// Put 归还连接到池
|
||||
func (p *ConnectionPool) Put(conn *PooledConnection) error { |
||||
if conn == nil { |
||||
return nil |
||||
} |
||||
|
||||
p.mu.RLock() |
||||
if p.closed { |
||||
p.mu.RUnlock() |
||||
p.closeConn(conn) |
||||
return ErrPoolClosed |
||||
} |
||||
p.mu.RUnlock() |
||||
|
||||
p.stats.decActive() |
||||
|
||||
// 检查连接是否有效
|
||||
if !p.isConnValid(conn) { |
||||
p.closeConn(conn) |
||||
return ErrConnInvalid |
||||
} |
||||
|
||||
// 尝试归还到池
|
||||
select { |
||||
case p.pool <- conn: |
||||
p.stats.incIdle() |
||||
return nil |
||||
default: |
||||
// 池已满，关闭连接
|
||||
p.closeConn(conn) |
||||
return ErrPoolFull |
||||
} |
||||
} |
||||
|
||||
// Close 关闭连接池
|
||||
func (p *ConnectionPool) Close() error { |
||||
p.mu.Lock() |
||||
if p.closed { |
||||
p.mu.Unlock() |
||||
return nil |
||||
} |
||||
p.closed = true |
||||
p.mu.Unlock() |
||||
|
||||
// 关闭所有池中的连接
|
||||
close(p.pool) |
||||
for conn := range p.pool { |
||||
p.closeConn(conn) |
||||
} |
||||
|
||||
p.logger.Info("Connection pool closed") |
||||
return nil |
||||
} |
||||
|
||||
// GetStats 获取统计信息
|
||||
func (p *ConnectionPool) GetStats() Stats { |
||||
p.stats.mu.RLock() |
||||
defer p.stats.mu.RUnlock() |
||||
|
||||
stats := p.stats |
||||
stats.Idle = int64(len(p.pool)) |
||||
return stats |
||||
} |
||||
|
||||
// createConnection 创建新连接
|
||||
func (p *ConnectionPool) createConnection() (*PooledConnection, error) { |
||||
conn, err := p.factory.Create() |
||||
if err != nil { |
||||
p.stats.incFailures() |
||||
return nil, err |
||||
} |
||||
|
||||
pooledConn := &PooledConnection{ |
||||
conn: conn, |
||||
createdAt: time.Now(), |
||||
lastUsed: time.Now(), |
||||
pool: p, |
||||
} |
||||
|
||||
p.stats.incCreated() |
||||
p.stats.incActive() |
||||
return pooledConn, nil |
||||
} |
||||
|
||||
// isConnValid 检查连接是否有效
|
||||
func (p *ConnectionPool) isConnValid(conn *PooledConnection) bool { |
||||
now := time.Now() |
||||
|
||||
// 检查连接生命周期
|
||||
if now.Sub(conn.createdAt) > p.maxLifetime { |
||||
return false |
||||
} |
||||
|
||||
// 检查空闲时间
|
||||
if now.Sub(conn.lastUsed) > p.maxIdle { |
||||
return false |
||||
} |
||||
|
||||
// 使用工厂验证连接
|
||||
return p.factory.Validate(conn.conn) |
||||
} |
||||
|
||||
// closeConn 关闭连接
|
||||
func (p *ConnectionPool) closeConn(conn *PooledConnection) { |
||||
if conn != nil && conn.conn != nil { |
||||
p.factory.Close(conn.conn) |
||||
p.stats.incClosed() |
||||
} |
||||
} |
||||
|
||||
// cleaner 清理过期连接
|
||||
func (p *ConnectionPool) cleaner() { |
||||
ticker := time.NewTicker(1 * time.Minute) |
||||
defer ticker.Stop() |
||||
|
||||
for { |
||||
select { |
||||
case <-ticker.C: |
||||
p.cleanExpiredConnections() |
||||
} |
||||
} |
||||
} |
||||
|
||||
// cleanExpiredConnections 清理过期连接
|
||||
func (p *ConnectionPool) cleanExpiredConnections() { |
||||
p.mu.RLock() |
||||
if p.closed { |
||||
p.mu.RUnlock() |
||||
return |
||||
} |
||||
p.mu.RUnlock() |
||||
|
||||
// 检查池中的连接
|
||||
poolSize := len(p.pool) |
||||
cleaned := 0 |
||||
|
||||
for i := 0; i < poolSize; i++ { |
||||
select { |
||||
case conn := <-p.pool: |
||||
if p.isConnValid(conn) { |
||||
// 连接有效，放回池中
|
||||
select { |
||||
case p.pool <- conn: |
||||
default: |
||||
// 池满了，关闭连接
|
||||
p.closeConn(conn) |
||||
cleaned++ |
||||
} |
||||
} else { |
||||
// 连接无效，关闭
|
||||
p.closeConn(conn) |
||||
p.stats.decIdle() |
||||
cleaned++ |
||||
} |
||||
default: |
||||
break |
||||
} |
||||
} |
||||
|
||||
if cleaned > 0 { |
||||
p.logger.WithField("count", cleaned).Debug("Cleaned expired connections") |
||||
} |
||||
} |
||||
|
||||
// PooledConnection 方法
|
||||
|
||||
// Read 读取数据
|
||||
func (pc *PooledConnection) Read(b []byte) (n int, err error) { |
||||
return pc.conn.Read(b) |
||||
} |
||||
|
||||
// Write 写入数据
|
||||
func (pc *PooledConnection) Write(b []byte) (n int, err error) { |
||||
return pc.conn.Write(b) |
||||
} |
||||
|
||||
// Close 关闭连接（归还到池）
|
||||
func (pc *PooledConnection) Close() error { |
||||
return pc.pool.Put(pc) |
||||
} |
||||
|
||||
// ForceClose 强制关闭连接
|
||||
func (pc *PooledConnection) ForceClose() error { |
||||
pc.pool.closeConn(pc) |
||||
return nil |
||||
} |
||||
|
||||
// LocalAddr 获取本地地址
|
||||
func (pc *PooledConnection) LocalAddr() net.Addr { |
||||
return pc.conn.LocalAddr() |
||||
} |
||||
|
||||
// RemoteAddr 获取远程地址
|
||||
func (pc *PooledConnection) RemoteAddr() net.Addr { |
||||
return pc.conn.RemoteAddr() |
||||
} |
||||
|
||||
// SetDeadline 设置截止时间
|
||||
func (pc *PooledConnection) SetDeadline(t time.Time) error { |
||||
return pc.conn.SetDeadline(t) |
||||
} |
||||
|
||||
// SetReadDeadline 设置读截止时间
|
||||
func (pc *PooledConnection) SetReadDeadline(t time.Time) error { |
||||
return pc.conn.SetReadDeadline(t) |
||||
} |
||||
|
||||
// SetWriteDeadline 设置写截止时间
|
||||
func (pc *PooledConnection) SetWriteDeadline(t time.Time) error { |
||||
return pc.conn.SetWriteDeadline(t) |
||||
} |
||||
|
||||
// Stats 方法
|
||||
|
||||
func (s *Stats) incCreated() { |
||||
s.mu.Lock() |
||||
s.Created++ |
||||
s.mu.Unlock() |
||||
} |
||||
|
||||
func (s *Stats) incReused() { |
||||
s.mu.Lock() |
||||
s.Reused++ |
||||
s.mu.Unlock() |
||||
} |
||||
|
||||
func (s *Stats) incClosed() { |
||||
s.mu.Lock() |
||||
s.Closed++ |
||||
s.mu.Unlock() |
||||
} |
||||
|
||||
func (s *Stats) incActive() { |
||||
s.mu.Lock() |
||||
s.Active++ |
||||
s.mu.Unlock() |
||||
} |
||||
|
||||
func (s *Stats) decActive() { |
||||
s.mu.Lock() |
||||
s.Active-- |
||||
s.mu.Unlock() |
||||
} |
||||
|
||||
func (s *Stats) incIdle() { |
||||
s.mu.Lock() |
||||
s.Idle++ |
||||
s.mu.Unlock() |
||||
} |
||||
|
||||
func (s *Stats) decIdle() { |
||||
s.mu.Lock() |
||||
s.Idle-- |
||||
s.mu.Unlock() |
||||
} |
||||
|
||||
func (s *Stats) incFailures() { |
||||
s.mu.Lock() |
||||
s.Failures++ |
||||
s.mu.Unlock() |
||||
} |
||||
@ -0,0 +1,231 @@ |
||||
package ratelimit |
||||
|
||||
import ( |
||||
"net" |
||||
"sync" |
||||
"time" |
||||
|
||||
"github.com/sirupsen/logrus" |
||||
) |
||||
|
||||
// RateLimiter 速率限制器
|
||||
type RateLimiter struct { |
||||
logger *logrus.Logger |
||||
globalLimit *TokenBucket |
||||
perIPLimiters map[string]*TokenBucket |
||||
cleanupInterval time.Duration |
||||
mu sync.RWMutex |
||||
stopCh chan struct{} |
||||
} |
||||
|
||||
// TokenBucket Token桶算法实现
|
||||
type TokenBucket struct { |
||||
capacity int64 |
||||
tokens int64 |
||||
refillRate int64 // tokens per second
|
||||
lastRefill time.Time |
||||
mu sync.Mutex |
||||
} |
||||
|
||||
// Config 速率限制配置
|
||||
type Config struct { |
||||
Enabled bool `json:"enabled"` |
||||
RequestsPerSecond int `json:"requestsPerSecond"` |
||||
BurstSize int `json:"burstSize"` |
||||
PerIPRequestsPerSec int `json:"perIPRequestsPerSec"` |
||||
PerIPBurstSize int `json:"perIPBurstSize"` |
||||
CleanupInterval time.Duration `json:"cleanupInterval"` |
||||
} |
||||
|
||||
// NewRateLimiter 创建新的速率限制器
|
||||
func NewRateLimiter(config Config, logger *logrus.Logger) *RateLimiter { |
||||
if !config.Enabled { |
||||
return &RateLimiter{ |
||||
logger: logger, |
||||
} |
||||
} |
||||
|
||||
// 默认值
|
||||
if config.BurstSize == 0 { |
||||
config.BurstSize = config.RequestsPerSecond * 2 |
||||
} |
||||
if config.PerIPBurstSize == 0 { |
||||
config.PerIPBurstSize = config.PerIPRequestsPerSec * 2 |
||||
} |
||||
if config.CleanupInterval == 0 { |
||||
config.CleanupInterval = 5 * time.Minute |
||||
} |
||||
|
||||
rl := &RateLimiter{ |
||||
logger: logger, |
||||
globalLimit: NewTokenBucket(int64(config.BurstSize), int64(config.RequestsPerSecond)), |
||||
perIPLimiters: make(map[string]*TokenBucket), |
||||
cleanupInterval: config.CleanupInterval, |
||||
stopCh: make(chan struct{}), |
||||
} |
||||
|
||||
// 启动清理goroutine
|
||||
go rl.cleanup() |
||||
|
||||
return rl |
||||
} |
||||
|
||||
// Allow 检查是否允许请求
|
||||
func (rl *RateLimiter) Allow(remoteAddr string) bool { |
||||
if rl.globalLimit == nil { |
||||
return true // 未启用速率限制
|
||||
} |
||||
|
||||
// 全局速率限制
|
||||
if !rl.globalLimit.Allow() { |
||||
rl.logger.WithField("type", "global").Debug("Rate limit exceeded") |
||||
return false |
||||
} |
||||
|
||||
// 获取客户端IP
|
||||
host, _, err := net.SplitHostPort(remoteAddr) |
||||
if err != nil { |
||||
host = remoteAddr |
||||
} |
||||
|
||||
// 单IP速率限制
|
||||
rl.mu.RLock() |
||||
bucket, exists := rl.perIPLimiters[host] |
||||
rl.mu.RUnlock() |
||||
|
||||
if !exists { |
||||
rl.mu.Lock() |
||||
bucket, exists = rl.perIPLimiters[host] |
||||
if !exists { |
||||
bucket = NewTokenBucket(20, 10) // 默认每IP 10 rps, burst 20
|
||||
rl.perIPLimiters[host] = bucket |
||||
} |
||||
rl.mu.Unlock() |
||||
} |
||||
|
||||
if !bucket.Allow() { |
||||
rl.logger.WithFields(logrus.Fields{ |
||||
"type": "per_ip", |
||||
"ip": host, |
||||
}).Debug("Per-IP rate limit exceeded") |
||||
return false |
||||
} |
||||
|
||||
return true |
||||
} |
||||
|
||||
// Stop 停止速率限制器
|
||||
func (rl *RateLimiter) Stop() { |
||||
if rl.stopCh != nil { |
||||
close(rl.stopCh) |
||||
} |
||||
} |
||||
|
||||
// GetStats 获取统计信息
|
||||
func (rl *RateLimiter) GetStats() map[string]interface{} { |
||||
if rl.globalLimit == nil { |
||||
return map[string]interface{}{ |
||||
"enabled": false, |
||||
} |
||||
} |
||||
|
||||
rl.mu.RLock() |
||||
perIPCount := len(rl.perIPLimiters) |
||||
rl.mu.RUnlock() |
||||
|
||||
return map[string]interface{}{ |
||||
"enabled": true, |
||||
"global_tokens": rl.globalLimit.GetTokens(), |
||||
"per_ip_buckets": perIPCount, |
||||
} |
||||
} |
||||
|
||||
// cleanup 清理过期的IP限制器
|
||||
func (rl *RateLimiter) cleanup() { |
||||
ticker := time.NewTicker(rl.cleanupInterval) |
||||
defer ticker.Stop() |
||||
|
||||
for { |
||||
select { |
||||
case <-rl.stopCh: |
||||
return |
||||
case <-ticker.C: |
||||
rl.cleanupExpiredBuckets() |
||||
} |
||||
} |
||||
} |
||||
|
||||
// cleanupExpiredBuckets 清理过期的IP桶
|
||||
func (rl *RateLimiter) cleanupExpiredBuckets() { |
||||
rl.mu.Lock() |
||||
defer rl.mu.Unlock() |
||||
|
||||
now := time.Now() |
||||
expiredIPs := make([]string, 0) |
||||
|
||||
for ip, bucket := range rl.perIPLimiters { |
||||
// 如果桶超过10分钟没有活动，清理掉
|
||||
if now.Sub(bucket.lastRefill) > 10*time.Minute { |
||||
expiredIPs = append(expiredIPs, ip) |
||||
} |
||||
} |
||||
|
||||
for _, ip := range expiredIPs { |
||||
delete(rl.perIPLimiters, ip) |
||||
} |
||||
|
||||
if len(expiredIPs) > 0 { |
||||
rl.logger.WithField("count", len(expiredIPs)).Debug("Cleaned up expired rate limit buckets") |
||||
} |
||||
} |
||||
|
||||
// NewTokenBucket 创建新的Token桶
|
||||
func NewTokenBucket(capacity, refillRate int64) *TokenBucket { |
||||
return &TokenBucket{ |
||||
capacity: capacity, |
||||
tokens: capacity, |
||||
refillRate: refillRate, |
||||
lastRefill: time.Now(), |
||||
} |
||||
} |
||||
|
||||
// Allow 检查是否允许请求（消耗一个token）
|
||||
func (tb *TokenBucket) Allow() bool { |
||||
tb.mu.Lock() |
||||
defer tb.mu.Unlock() |
||||
|
||||
tb.refill() |
||||
|
||||
if tb.tokens > 0 { |
||||
tb.tokens-- |
||||
return true |
||||
} |
||||
|
||||
return false |
||||
} |
||||
|
||||
// refill 补充tokens
|
||||
func (tb *TokenBucket) refill() { |
||||
now := time.Now() |
||||
elapsed := now.Sub(tb.lastRefill) |
||||
|
||||
// 计算应该补充的tokens
|
||||
tokensToAdd := int64(elapsed.Seconds()) * tb.refillRate |
||||
|
||||
if tokensToAdd > 0 { |
||||
tb.tokens += tokensToAdd |
||||
if tb.tokens > tb.capacity { |
||||
tb.tokens = tb.capacity |
||||
} |
||||
tb.lastRefill = now |
||||
} |
||||
} |
||||
|
||||
// GetTokens 获取当前token数量
|
||||
func (tb *TokenBucket) GetTokens() int64 { |
||||
tb.mu.Lock() |
||||
defer tb.mu.Unlock() |
||||
|
||||
tb.refill() |
||||
return tb.tokens |
||||
} |
||||
@ -0,0 +1,58 @@ |
||||
package socks5 |
||||
|
||||
import "strings" |
||||
|
||||
// SimpleAuthHandler 简单认证处理器
|
||||
type SimpleAuthHandler struct { |
||||
username string |
||||
password string |
||||
methods []byte |
||||
} |
||||
|
||||
// NewAuthHandler 创建认证处理器
|
||||
func NewAuthHandler(config AuthConfig) AuthHandler { |
||||
handler := &SimpleAuthHandler{ |
||||
username: config.Username, |
||||
password: config.Password, |
||||
} |
||||
|
||||
// 设置支持的认证方法
|
||||
if config.Username != "" && config.Password != "" { |
||||
handler.methods = []byte{AuthPassword} |
||||
} else { |
||||
handler.methods = []byte{AuthNone} |
||||
} |
||||
|
||||
// 如果配置中指定了方法，使用配置的方法
|
||||
if len(config.Methods) > 0 { |
||||
handler.methods = []byte{} |
||||
for _, method := range config.Methods { |
||||
switch strings.ToLower(method) { |
||||
case "none": |
||||
handler.methods = append(handler.methods, AuthNone) |
||||
case "password": |
||||
handler.methods = append(handler.methods, AuthPassword) |
||||
} |
||||
} |
||||
} |
||||
|
||||
return handler |
||||
} |
||||
|
||||
// Authenticate 验证用户名和密码
|
||||
func (h *SimpleAuthHandler) Authenticate(username, password string) bool { |
||||
// 如果支持无认证，直接返回true
|
||||
for _, method := range h.methods { |
||||
if method == AuthNone { |
||||
return true |
||||
} |
||||
} |
||||
|
||||
// 检查用户名密码
|
||||
return h.username == username && h.password == password |
||||
} |
||||
|
||||
// Methods 返回支持的认证方法
|
||||
func (h *SimpleAuthHandler) Methods() []byte { |
||||
return h.methods |
||||
} |
||||
@ -0,0 +1,34 @@ |
||||
package socks5 |
||||
|
||||
import ( |
||||
"net" |
||||
"time" |
||||
) |
||||
|
||||
// DirectDialer 直接连接拨号器
|
||||
type DirectDialer struct { |
||||
timeout time.Duration |
||||
} |
||||
|
||||
// Dial 建立连接
|
||||
func (d *DirectDialer) Dial(network, address string) (net.Conn, error) { |
||||
if d.timeout > 0 { |
||||
return net.DialTimeout(network, address, d.timeout) |
||||
} |
||||
return net.Dial(network, address) |
||||
} |
||||
|
||||
// DirectResolver 直接DNS解析器
|
||||
type DirectResolver struct{} |
||||
|
||||
// Resolve 解析域名
|
||||
func (r *DirectResolver) Resolve(domain string) (net.IP, error) { |
||||
ips, err := net.LookupIP(domain) |
||||
if err != nil { |
||||
return nil, err |
||||
} |
||||
if len(ips) == 0 { |
||||
return nil, net.ErrClosed |
||||
} |
||||
return ips[0], nil |
||||
} |
||||
@ -0,0 +1,114 @@ |
||||
package socks5 |
||||
|
||||
import ( |
||||
"net" |
||||
"strconv" |
||||
"strings" |
||||
) |
||||
|
||||
// SimpleRule 简单访问规则
|
||||
type SimpleRule struct { |
||||
action string // allow, deny
|
||||
networks []*net.IPNet |
||||
ports []int |
||||
} |
||||
|
||||
// NewRule 创建访问规则
|
||||
func NewRule(config RuleConfig) Rule { |
||||
rule := &SimpleRule{ |
||||
action: strings.ToLower(config.Action), |
||||
ports: config.Ports, |
||||
} |
||||
|
||||
// 解析IP网段
|
||||
for _, ipStr := range config.IPs { |
||||
if !strings.Contains(ipStr, "/") { |
||||
// 单个IP，添加适当的子网掩码
|
||||
if strings.Contains(ipStr, ":") { |
||||
ipStr += "/128" // IPv6
|
||||
} else { |
||||
ipStr += "/32" // IPv4
|
||||
} |
||||
} |
||||
|
||||
_, network, err := net.ParseCIDR(ipStr) |
||||
if err != nil { |
||||
// 如果解析失败，尝试作为单个IP处理
|
||||
ip := net.ParseIP(ipStr) |
||||
if ip != nil { |
||||
if ip.To4() != nil { |
||||
_, network, _ = net.ParseCIDR(ip.String() + "/32") |
||||
} else { |
||||
_, network, _ = net.ParseCIDR(ip.String() + "/128") |
||||
} |
||||
} |
||||
} |
||||
|
||||
if network != nil { |
||||
rule.networks = append(rule.networks, network) |
||||
} |
||||
} |
||||
|
||||
return rule |
||||
} |
||||
|
||||
// Allow 检查是否允许访问
|
||||
func (r *SimpleRule) Allow(addr net.Addr) bool { |
||||
// 解析地址
|
||||
var ip net.IP |
||||
var port int |
||||
|
||||
switch a := addr.(type) { |
||||
case *net.IPAddr: |
||||
ip = a.IP |
||||
case *net.TCPAddr: |
||||
ip = a.IP |
||||
port = a.Port |
||||
case *net.UDPAddr: |
||||
ip = a.IP |
||||
port = a.Port |
||||
default: |
||||
// 尝试从字符串解析
|
||||
host, portStr, err := net.SplitHostPort(addr.String()) |
||||
if err != nil { |
||||
return r.action == "allow" // 默认策略
|
||||
} |
||||
ip = net.ParseIP(host) |
||||
if p, err := strconv.Atoi(portStr); err == nil { |
||||
port = p |
||||
} |
||||
} |
||||
|
||||
if ip == nil { |
||||
return r.action == "allow" // 默认策略
|
||||
} |
||||
|
||||
// 检查IP是否匹配
|
||||
ipMatches := len(r.networks) == 0 // 如果没有指定网段，默认匹配
|
||||
for _, network := range r.networks { |
||||
if network.Contains(ip) { |
||||
ipMatches = true |
||||
break |
||||
} |
||||
} |
||||
|
||||
// 检查端口是否匹配
|
||||
portMatches := len(r.ports) == 0 // 如果没有指定端口，默认匹配
|
||||
for _, p := range r.ports { |
||||
if p == port { |
||||
portMatches = true |
||||
break |
||||
} |
||||
} |
||||
|
||||
// 根据规则类型返回结果
|
||||
matches := ipMatches && portMatches |
||||
switch r.action { |
||||
case "allow": |
||||
return matches |
||||
case "deny": |
||||
return !matches |
||||
default: |
||||
return true // 默认允许
|
||||
} |
||||
} |
||||
@ -0,0 +1,432 @@ |
||||
package socks5 |
||||
|
||||
import ( |
||||
"encoding/binary" |
||||
"fmt" |
||||
"io" |
||||
"net" |
||||
"time" |
||||
|
||||
"github.com/sirupsen/logrus" |
||||
) |
||||
|
||||
// SOCKS5 协议常量
|
||||
const ( |
||||
// SOCKS版本
|
||||
Version5 = 0x05 |
||||
|
||||
// 认证方法
|
||||
AuthNone = 0x00 |
||||
AuthPassword = 0x02 |
||||
AuthNoSupported = 0xFF |
||||
|
||||
// 命令类型
|
||||
CmdConnect = 0x01 |
||||
CmdBind = 0x02 |
||||
CmdUDP = 0x03 |
||||
|
||||
// 地址类型
|
||||
AddrIPv4 = 0x01 |
||||
AddrDomain = 0x03 |
||||
AddrIPv6 = 0x04 |
||||
|
||||
// 响应状态
|
||||
StatusSuccess = 0x00 |
||||
StatusServerFailure = 0x01 |
||||
StatusConnectionNotAllowed = 0x02 |
||||
StatusNetworkUnreachable = 0x03 |
||||
StatusHostUnreachable = 0x04 |
||||
StatusConnectionRefused = 0x05 |
||||
StatusTTLExpired = 0x06 |
||||
StatusCommandNotSupported = 0x07 |
||||
StatusAddressNotSupported = 0x08 |
||||
) |
||||
|
||||
type Server struct { |
||||
logger *logrus.Logger |
||||
auth AuthHandler |
||||
dialer Dialer |
||||
resolver Resolver |
||||
rules []Rule |
||||
} |
||||
|
||||
type Config struct { |
||||
Auth AuthConfig |
||||
Timeout time.Duration |
||||
Rules []RuleConfig |
||||
} |
||||
|
||||
type AuthConfig struct { |
||||
Methods []string |
||||
Username string |
||||
Password string |
||||
} |
||||
|
||||
type RuleConfig struct { |
||||
Action string // allow, deny
|
||||
IPs []string |
||||
Ports []int |
||||
} |
||||
|
||||
type AuthHandler interface { |
||||
Authenticate(username, password string) bool |
||||
Methods() []byte |
||||
} |
||||
|
||||
type Dialer interface { |
||||
Dial(network, address string) (net.Conn, error) |
||||
} |
||||
|
||||
type Resolver interface { |
||||
Resolve(domain string) (net.IP, error) |
||||
} |
||||
|
||||
type Rule interface { |
||||
Allow(addr net.Addr) bool |
||||
} |
||||
|
||||
// NewServer 创建新的SOCKS5服务器
|
||||
func NewServer(config Config, logger *logrus.Logger) *Server { |
||||
server := &Server{ |
||||
logger: logger, |
||||
dialer: &DirectDialer{timeout: config.Timeout}, |
||||
resolver: &DirectResolver{}, |
||||
} |
||||
|
||||
// 设置认证处理器
|
||||
server.auth = NewAuthHandler(config.Auth) |
||||
|
||||
// 设置访问规则
|
||||
for _, ruleConfig := range config.Rules { |
||||
rule := NewRule(ruleConfig) |
||||
server.rules = append(server.rules, rule) |
||||
} |
||||
|
||||
return server |
||||
} |
||||
|
||||
// HandleConnection 处理SOCKS5连接
|
||||
func (s *Server) HandleConnection(conn net.Conn) error { |
||||
defer conn.Close() |
||||
|
||||
s.logger.WithField("remote_addr", conn.RemoteAddr()).Debug("New SOCKS5 connection") |
||||
|
||||
// 设置连接超时
|
||||
conn.SetDeadline(time.Now().Add(30 * time.Second)) |
||||
|
||||
// 1. 协议版本协商
|
||||
if err := s.handleVersionNegotiation(conn); err != nil { |
||||
s.logger.WithError(err).Error("Version negotiation failed") |
||||
return err |
||||
} |
||||
|
||||
// 2. 认证
|
||||
if err := s.handleAuthentication(conn); err != nil { |
||||
s.logger.WithError(err).Error("Authentication failed") |
||||
return err |
||||
} |
||||
|
||||
// 3. 处理请求
|
||||
return s.handleRequest(conn) |
||||
} |
||||
|
||||
// handleVersionNegotiation 处理版本协商
|
||||
func (s *Server) handleVersionNegotiation(conn net.Conn) error { |
||||
// 读取客户端版本协商请求
|
||||
buf := make([]byte, 2) |
||||
if _, err := io.ReadFull(conn, buf); err != nil { |
||||
return fmt.Errorf("failed to read version: %w", err) |
||||
} |
||||
|
||||
version := buf[0] |
||||
nmethods := buf[1] |
||||
|
||||
if version != Version5 { |
||||
return fmt.Errorf("unsupported SOCKS version: %d", version) |
||||
} |
||||
|
||||
// 读取支持的认证方法
|
||||
methods := make([]byte, nmethods) |
||||
if _, err := io.ReadFull(conn, methods); err != nil { |
||||
return fmt.Errorf("failed to read methods: %w", err) |
||||
} |
||||
|
||||
// 选择认证方法
|
||||
authMethods := s.auth.Methods() |
||||
selectedMethod := byte(AuthNoSupported) |
||||
|
||||
for _, method := range methods { |
||||
for _, supported := range authMethods { |
||||
if method == supported { |
||||
selectedMethod = method |
||||
break |
||||
} |
||||
} |
||||
if selectedMethod != AuthNoSupported { |
||||
break |
||||
} |
||||
} |
||||
|
||||
// 发送选择的认证方法
|
||||
response := []byte{Version5, selectedMethod} |
||||
if _, err := conn.Write(response); err != nil { |
||||
return fmt.Errorf("failed to write method selection: %w", err) |
||||
} |
||||
|
||||
if selectedMethod == AuthNoSupported { |
||||
return fmt.Errorf("no supported authentication method") |
||||
} |
||||
|
||||
return nil |
||||
} |
||||
|
||||
// handleAuthentication 处理认证
|
||||
func (s *Server) handleAuthentication(conn net.Conn) error { |
||||
// 读取认证请求
|
||||
buf := make([]byte, 2) |
||||
if _, err := io.ReadFull(conn, buf); err != nil { |
||||
return fmt.Errorf("failed to read auth version: %w", err) |
||||
} |
||||
|
||||
version := buf[0] |
||||
usernameLen := buf[1] |
||||
|
||||
if version != 0x01 { |
||||
return fmt.Errorf("unsupported auth version: %d", version) |
||||
} |
||||
|
||||
// 读取用户名
|
||||
username := make([]byte, usernameLen) |
||||
if _, err := io.ReadFull(conn, username); err != nil { |
||||
return fmt.Errorf("failed to read username: %w", err) |
||||
} |
||||
|
||||
// 读取密码长度
|
||||
if _, err := io.ReadFull(conn, buf[:1]); err != nil { |
||||
return fmt.Errorf("failed to read password length: %w", err) |
||||
} |
||||
passwordLen := buf[0] |
||||
|
||||
// 读取密码
|
||||
password := make([]byte, passwordLen) |
||||
if _, err := io.ReadFull(conn, password); err != nil { |
||||
return fmt.Errorf("failed to read password: %w", err) |
||||
} |
||||
|
||||
// 验证认证
|
||||
success := s.auth.Authenticate(string(username), string(password)) |
||||
|
||||
// 发送认证结果
|
||||
status := byte(0x01) // 失败
|
||||
if success { |
||||
status = 0x00 // 成功
|
||||
} |
||||
|
||||
response := []byte{0x01, status} |
||||
if _, err := conn.Write(response); err != nil { |
||||
return fmt.Errorf("failed to write auth response: %w", err) |
||||
} |
||||
|
||||
if !success { |
||||
return fmt.Errorf("authentication failed") |
||||
} |
||||
|
||||
return nil |
||||
} |
||||
|
||||
// handleRequest 处理SOCKS5请求
|
||||
func (s *Server) handleRequest(conn net.Conn) error { |
||||
// 读取请求头
|
||||
buf := make([]byte, 4) |
||||
if _, err := io.ReadFull(conn, buf); err != nil { |
||||
return fmt.Errorf("failed to read request header: %w", err) |
||||
} |
||||
|
||||
version := buf[0] |
||||
cmd := buf[1] |
||||
// rsv := buf[2] // 保留字段
|
||||
addrType := buf[3] |
||||
|
||||
if version != Version5 { |
||||
return fmt.Errorf("invalid SOCKS version: %d", version) |
||||
} |
||||
|
||||
// 读取目标地址
|
||||
addr, err := s.readAddress(conn, addrType) |
||||
if err != nil { |
||||
s.sendResponse(conn, StatusServerFailure, "0.0.0.0", 0) |
||||
return fmt.Errorf("failed to read address: %w", err) |
||||
} |
||||
|
||||
// 检查访问规则
|
||||
if !s.checkRules(addr) { |
||||
s.sendResponse(conn, StatusConnectionNotAllowed, "0.0.0.0", 0) |
||||
return fmt.Errorf("connection not allowed: %s", addr) |
||||
} |
||||
|
||||
// 处理不同的命令
|
||||
switch cmd { |
||||
case CmdConnect: |
||||
return s.handleConnect(conn, addr) |
||||
case CmdBind: |
||||
return s.handleBind(conn, addr) |
||||
case CmdUDP: |
||||
return s.handleUDP(conn, addr) |
||||
default: |
||||
s.sendResponse(conn, StatusCommandNotSupported, "0.0.0.0", 0) |
||||
return fmt.Errorf("unsupported command: %d", cmd) |
||||
} |
||||
} |
||||
|
||||
// readAddress 读取地址信息
|
||||
func (s *Server) readAddress(conn net.Conn, addrType byte) (string, error) { |
||||
switch addrType { |
||||
case AddrIPv4: |
||||
buf := make([]byte, 6) // 4字节IP + 2字节端口
|
||||
if _, err := io.ReadFull(conn, buf); err != nil { |
||||
return "", err |
||||
} |
||||
ip := net.IP(buf[:4]) |
||||
port := binary.BigEndian.Uint16(buf[4:6]) |
||||
return fmt.Sprintf("%s:%d", ip.String(), port), nil |
||||
|
||||
case AddrDomain: |
||||
buf := make([]byte, 1) |
||||
if _, err := io.ReadFull(conn, buf); err != nil { |
||||
return "", err |
||||
} |
||||
domainLen := buf[0] |
||||
|
||||
domain := make([]byte, domainLen+2) // 域名 + 2字节端口
|
||||
if _, err := io.ReadFull(conn, domain); err != nil { |
||||
return "", err |
||||
} |
||||
port := binary.BigEndian.Uint16(domain[domainLen:]) |
||||
return fmt.Sprintf("%s:%d", string(domain[:domainLen]), port), nil |
||||
|
||||
case AddrIPv6: |
||||
buf := make([]byte, 18) // 16字节IP + 2字节端口
|
||||
if _, err := io.ReadFull(conn, buf); err != nil { |
||||
return "", err |
||||
} |
||||
ip := net.IP(buf[:16]) |
||||
port := binary.BigEndian.Uint16(buf[16:18]) |
||||
return fmt.Sprintf("[%s]:%d", ip.String(), port), nil |
||||
|
||||
default: |
||||
return "", fmt.Errorf("unsupported address type: %d", addrType) |
||||
} |
||||
} |
||||
|
||||
// handleConnect 处理CONNECT命令
|
||||
func (s *Server) handleConnect(conn net.Conn, addr string) error { |
||||
s.logger.WithField("target", addr).Debug("Handling CONNECT request") |
||||
|
||||
// 连接到目标服务器
|
||||
target, err := s.dialer.Dial("tcp", addr) |
||||
if err != nil { |
||||
s.logger.WithError(err).WithField("target", addr).Error("Failed to connect to target") |
||||
s.sendResponse(conn, StatusConnectionRefused, "0.0.0.0", 0) |
||||
return err |
||||
} |
||||
defer target.Close() |
||||
|
||||
// 发送成功响应
|
||||
localAddr := target.LocalAddr().(*net.TCPAddr) |
||||
s.sendResponse(conn, StatusSuccess, localAddr.IP.String(), uint16(localAddr.Port)) |
||||
|
||||
// 开始数据转发
|
||||
s.logger.WithField("target", addr).Info("Starting data relay") |
||||
return s.relay(conn, target) |
||||
} |
||||
|
||||
// handleBind 处理BIND命令
|
||||
func (s *Server) handleBind(conn net.Conn, addr string) error { |
||||
// BIND命令实现（简化版本）
|
||||
s.sendResponse(conn, StatusCommandNotSupported, "0.0.0.0", 0) |
||||
return fmt.Errorf("BIND command not implemented") |
||||
} |
||||
|
||||
// handleUDP 处理UDP命令
|
||||
func (s *Server) handleUDP(conn net.Conn, addr string) error { |
||||
// UDP关联实现（简化版本）
|
||||
s.sendResponse(conn, StatusCommandNotSupported, "0.0.0.0", 0) |
||||
return fmt.Errorf("UDP command not implemented") |
||||
} |
||||
|
||||
// sendResponse 发送SOCKS5响应
|
||||
func (s *Server) sendResponse(conn net.Conn, status byte, ip string, port uint16) error { |
||||
response := []byte{ |
||||
Version5, |
||||
status, |
||||
0x00, // 保留字段
|
||||
AddrIPv4, |
||||
} |
||||
|
||||
// 添加IP地址
|
||||
ipAddr := net.ParseIP(ip) |
||||
if ipAddr == nil { |
||||
ipAddr = net.ParseIP("0.0.0.0") |
||||
} |
||||
if ipv4 := ipAddr.To4(); ipv4 != nil { |
||||
response = append(response, ipv4...) |
||||
} else { |
||||
response[3] = AddrIPv6 |
||||
response = append(response, ipAddr.To16()...) |
||||
} |
||||
|
||||
// 添加端口
|
||||
portBytes := make([]byte, 2) |
||||
binary.BigEndian.PutUint16(portBytes, port) |
||||
response = append(response, portBytes...) |
||||
|
||||
_, err := conn.Write(response) |
||||
return err |
||||
} |
||||
|
||||
// checkRules 检查访问规则
|
||||
func (s *Server) checkRules(addr string) bool { |
||||
if len(s.rules) == 0 { |
||||
return true // 没有规则时允许所有连接
|
||||
} |
||||
|
||||
host, _, err := net.SplitHostPort(addr) |
||||
if err != nil { |
||||
return false |
||||
} |
||||
|
||||
targetAddr, err := net.ResolveIPAddr("ip", host) |
||||
if err != nil { |
||||
return false |
||||
} |
||||
|
||||
for _, rule := range s.rules { |
||||
if rule.Allow(targetAddr) { |
||||
return true |
||||
} |
||||
} |
||||
|
||||
return false |
||||
} |
||||
|
||||
// relay 数据转发
|
||||
func (s *Server) relay(client, target net.Conn) error { |
||||
// 创建双向数据转发
|
||||
errChan := make(chan error, 2) |
||||
|
||||
// 客户端到目标服务器
|
||||
go func() { |
||||
_, err := io.Copy(target, client) |
||||
errChan <- err |
||||
}() |
||||
|
||||
// 目标服务器到客户端
|
||||
go func() { |
||||
_, err := io.Copy(client, target) |
||||
errChan <- err |
||||
}() |
||||
|
||||
// 等待任一方向的连接关闭
|
||||
err := <-errChan |
||||
return err |
||||
} |
||||
@ -0,0 +1,54 @@ |
||||
package socks5 |
||||
|
||||
import ( |
||||
"testing" |
||||
"time" |
||||
|
||||
"github.com/azoic/wormhole-server/pkg/memory" |
||||
"github.com/sirupsen/logrus" |
||||
) |
||||
|
||||
func BenchmarkBufferPool(b *testing.B) { |
||||
logger := logrus.New() |
||||
logger.SetLevel(logrus.ErrorLevel) |
||||
|
||||
memConfig := memory.Config{ |
||||
BufferSizes: []int{512, 1024, 2048, 4096, 8192, 16384, 32768, 65536}, |
||||
EnableOptimization: true, |
||||
EnableAutoGC: false, |
||||
MonitorInterval: time.Minute, |
||||
} |
||||
memManager := memory.NewManager(memConfig, logger) |
||||
defer memManager.Stop() |
||||
|
||||
b.ResetTimer() |
||||
b.RunParallel(func(pb *testing.PB) { |
||||
for pb.Next() { |
||||
// 测试不同大小的缓冲区
|
||||
sizes := []int{512, 1024, 2048, 4096, 8192} |
||||
for _, size := range sizes { |
||||
buf := memManager.GetBuffer(size) |
||||
memManager.PutBuffer(buf) |
||||
} |
||||
} |
||||
}) |
||||
} |
||||
|
||||
func BenchmarkMemoryManagerStats(b *testing.B) { |
||||
logger := logrus.New() |
||||
logger.SetLevel(logrus.ErrorLevel) |
||||
|
||||
memConfig := memory.Config{ |
||||
BufferSizes: []int{512, 1024, 2048, 4096}, |
||||
EnableOptimization: true, |
||||
EnableAutoGC: false, |
||||
MonitorInterval: time.Minute, |
||||
} |
||||
memManager := memory.NewManager(memConfig, logger) |
||||
defer memManager.Stop() |
||||
|
||||
b.ResetTimer() |
||||
for i := 0; i < b.N; i++ { |
||||
_ = memManager.GetOverallStats() |
||||
} |
||||
} |
||||
@ -0,0 +1,71 @@ |
||||
package socks5 |
||||
|
||||
import ( |
||||
"testing" |
||||
"time" |
||||
|
||||
"github.com/sirupsen/logrus" |
||||
) |
||||
|
||||
func TestNewServer(t *testing.T) { |
||||
logger := logrus.New() |
||||
logger.SetLevel(logrus.ErrorLevel) // 减少测试输出
|
||||
|
||||
config := Config{ |
||||
Auth: AuthConfig{ |
||||
Username: "test", |
||||
Password: "pass", |
||||
}, |
||||
Timeout: 30 * time.Second, |
||||
} |
||||
|
||||
server := NewServer(config, logger) |
||||
if server == nil { |
||||
t.Fatal("Expected server to be created, got nil") |
||||
} |
||||
|
||||
if server.logger != logger { |
||||
t.Error("Logger not set correctly") |
||||
} |
||||
} |
||||
|
||||
func TestAuthHandler(t *testing.T) { |
||||
config := AuthConfig{ |
||||
Username: "testuser", |
||||
Password: "testpass", |
||||
Methods: []string{"password"}, |
||||
} |
||||
|
||||
handler := NewAuthHandler(config) |
||||
|
||||
// 测试正确的认证
|
||||
if !handler.Authenticate("testuser", "testpass") { |
||||
t.Error("Valid authentication failed") |
||||
} |
||||
|
||||
// 测试错误的认证
|
||||
if handler.Authenticate("wronguser", "wrongpass") { |
||||
t.Error("Invalid authentication succeeded") |
||||
} |
||||
|
||||
// 检查支持的方法
|
||||
methods := handler.Methods() |
||||
if len(methods) != 1 || methods[0] != AuthPassword { |
||||
t.Error("Expected password authentication method") |
||||
} |
||||
} |
||||
|
||||
func TestRule(t *testing.T) { |
||||
config := RuleConfig{ |
||||
Action: "allow", |
||||
IPs: []string{"192.168.1.0/24", "127.0.0.1"}, |
||||
Ports: []int{80, 443}, |
||||
} |
||||
|
||||
rule := NewRule(config) |
||||
if rule == nil { |
||||
t.Fatal("Expected rule to be created, got nil") |
||||
} |
||||
|
||||
// 这里可以添加更多的规则测试
|
||||
} |
||||
@ -1,368 +0,0 @@ |
||||
package system |
||||
|
||||
import ( |
||||
"fmt" |
||||
"os" |
||||
"os/exec" |
||||
"runtime" |
||||
"strings" |
||||
|
||||
"github.com/sirupsen/logrus" |
||||
) |
||||
|
||||
type SystemProxy struct { |
||||
logger *logrus.Logger |
||||
proxyAddr string |
||||
backupConfig map[string]string |
||||
} |
||||
|
||||
type Config struct { |
||||
HTTPProxy string |
||||
HTTPSProxy string |
||||
SOCKSProxy string |
||||
NoProxy []string |
||||
} |
||||
|
||||
func NewSystemProxy(proxyAddr string, logger *logrus.Logger) *SystemProxy { |
||||
return &SystemProxy{ |
||||
logger: logger, |
||||
proxyAddr: proxyAddr, |
||||
backupConfig: make(map[string]string), |
||||
} |
||||
} |
||||
|
||||
// SetSystemProxy 设置系统代理
|
||||
func (sp *SystemProxy) SetSystemProxy(config Config) error { |
||||
sp.logger.Info("Setting system proxy configuration") |
||||
|
||||
// 备份当前配置
|
||||
if err := sp.backupCurrentConfig(); err != nil { |
||||
sp.logger.WithError(err).Warn("Failed to backup current proxy config") |
||||
} |
||||
|
||||
switch runtime.GOOS { |
||||
case "darwin": |
||||
return sp.setMacOSProxy(config) |
||||
case "linux": |
||||
return sp.setLinuxProxy(config) |
||||
case "windows": |
||||
return sp.setWindowsProxy(config) |
||||
default: |
||||
return fmt.Errorf("unsupported operating system: %s", runtime.GOOS) |
||||
} |
||||
} |
||||
|
||||
// RestoreSystemProxy 恢复系统代理设置
|
||||
func (sp *SystemProxy) RestoreSystemProxy() error { |
||||
sp.logger.Info("Restoring system proxy configuration") |
||||
|
||||
switch runtime.GOOS { |
||||
case "darwin": |
||||
return sp.restoreMacOSProxy() |
||||
case "linux": |
||||
return sp.restoreLinuxProxy() |
||||
case "windows": |
||||
return sp.restoreWindowsProxy() |
||||
default: |
||||
return fmt.Errorf("unsupported operating system: %s", runtime.GOOS) |
||||
} |
||||
} |
||||
|
||||
// setMacOSProxy 设置macOS系统代理
|
||||
func (sp *SystemProxy) setMacOSProxy(config Config) error { |
||||
// 获取所有网络服务
|
||||
services, err := sp.getMacOSNetworkServices() |
||||
if err != nil { |
||||
return fmt.Errorf("failed to get network services: %w", err) |
||||
} |
||||
|
||||
for _, service := range services { |
||||
// 设置HTTP代理
|
||||
if config.HTTPProxy != "" { |
||||
parts := strings.Split(config.HTTPProxy, ":") |
||||
if len(parts) == 2 { |
||||
if err := sp.runCommand("networksetup", "-setwebproxy", service, parts[0], parts[1]); err != nil { |
||||
sp.logger.WithError(err).WithField("service", service).Warn("Failed to set HTTP proxy") |
||||
} |
||||
if err := sp.runCommand("networksetup", "-setwebproxystate", service, "on"); err != nil { |
||||
sp.logger.WithError(err).WithField("service", service).Warn("Failed to enable HTTP proxy") |
||||
} |
||||
} |
||||
} |
||||
|
||||
// 设置HTTPS代理
|
||||
if config.HTTPSProxy != "" { |
||||
parts := strings.Split(config.HTTPSProxy, ":") |
||||
if len(parts) == 2 { |
||||
if err := sp.runCommand("networksetup", "-setsecurewebproxy", service, parts[0], parts[1]); err != nil { |
||||
sp.logger.WithError(err).WithField("service", service).Warn("Failed to set HTTPS proxy") |
||||
} |
||||
if err := sp.runCommand("networksetup", "-setsecurewebproxystate", service, "on"); err != nil { |
||||
sp.logger.WithError(err).WithField("service", service).Warn("Failed to enable HTTPS proxy") |
||||
} |
||||
} |
||||
} |
||||
|
||||
// 设置SOCKS代理
|
||||
if config.SOCKSProxy != "" { |
||||
parts := strings.Split(config.SOCKSProxy, ":") |
||||
if len(parts) == 2 { |
||||
if err := sp.runCommand("networksetup", "-setsocksfirewallproxy", service, parts[0], parts[1]); err != nil { |
||||
sp.logger.WithError(err).WithField("service", service).Warn("Failed to set SOCKS proxy") |
||||
} |
||||
if err := sp.runCommand("networksetup", "-setsocksfirewallproxystate", service, "on"); err != nil { |
||||
sp.logger.WithError(err).WithField("service", service).Warn("Failed to enable SOCKS proxy") |
||||
} |
||||
} |
||||
} |
||||
|
||||
// 设置代理绕过列表
|
||||
if len(config.NoProxy) > 0 { |
||||
bypassList := strings.Join(config.NoProxy, " ") |
||||
if err := sp.runCommand("networksetup", "-setproxybypassdomains", service, bypassList); err != nil { |
||||
sp.logger.WithError(err).WithField("service", service).Warn("Failed to set proxy bypass list") |
||||
} |
||||
} |
||||
} |
||||
|
||||
return nil |
||||
} |
||||
|
||||
// setLinuxProxy 设置Linux系统代理（通过环境变量）
|
||||
func (sp *SystemProxy) setLinuxProxy(config Config) error { |
||||
envVars := make(map[string]string) |
||||
|
||||
if config.HTTPProxy != "" { |
||||
envVars["http_proxy"] = "http://" + config.HTTPProxy |
||||
envVars["HTTP_PROXY"] = "http://" + config.HTTPProxy |
||||
} |
||||
|
||||
if config.HTTPSProxy != "" { |
||||
envVars["https_proxy"] = "https://" + config.HTTPSProxy |
||||
envVars["HTTPS_PROXY"] = "https://" + config.HTTPSProxy |
||||
} |
||||
|
||||
if len(config.NoProxy) > 0 { |
||||
noProxy := strings.Join(config.NoProxy, ",") |
||||
envVars["no_proxy"] = noProxy |
||||
envVars["NO_PROXY"] = noProxy |
||||
} |
||||
|
||||
// 写入到用户的shell配置文件
|
||||
return sp.writeLinuxProxyConfig(envVars) |
||||
} |
||||
|
||||
// setWindowsProxy 设置Windows系统代理
|
||||
func (sp *SystemProxy) setWindowsProxy(config Config) error { |
||||
// Windows代理设置通过注册表
|
||||
if config.HTTPProxy != "" { |
||||
// 启用代理
|
||||
if err := sp.runCommand("reg", "add", `HKCU\Software\Microsoft\Windows\CurrentVersion\Internet Settings`, "/v", "ProxyEnable", "/t", "REG_DWORD", "/d", "1", "/f"); err != nil { |
||||
return fmt.Errorf("failed to enable proxy: %w", err) |
||||
} |
||||
|
||||
// 设置代理服务器
|
||||
if err := sp.runCommand("reg", "add", `HKCU\Software\Microsoft\Windows\CurrentVersion\Internet Settings`, "/v", "ProxyServer", "/t", "REG_SZ", "/d", config.HTTPProxy, "/f"); err != nil { |
||||
return fmt.Errorf("failed to set proxy server: %w", err) |
||||
} |
||||
} |
||||
|
||||
// 设置代理绕过列表
|
||||
if len(config.NoProxy) > 0 { |
||||
bypassList := strings.Join(config.NoProxy, ";") |
||||
if err := sp.runCommand("reg", "add", `HKCU\Software\Microsoft\Windows\CurrentVersion\Internet Settings`, "/v", "ProxyOverride", "/t", "REG_SZ", "/d", bypassList, "/f"); err != nil { |
||||
return fmt.Errorf("failed to set proxy bypass list: %w", err) |
||||
} |
||||
} |
||||
|
||||
return nil |
||||
} |
||||
|
||||
func (sp *SystemProxy) getMacOSNetworkServices() ([]string, error) { |
||||
output, err := exec.Command("networksetup", "-listallnetworkservices").Output() |
||||
if err != nil { |
||||
return nil, err |
||||
} |
||||
|
||||
lines := strings.Split(string(output), "\n") |
||||
var services []string |
||||
|
||||
for _, line := range lines { |
||||
line = strings.TrimSpace(line) |
||||
if line != "" && !strings.HasPrefix(line, "*") && line != "An asterisk (*) denotes that a network service is disabled." { |
||||
services = append(services, line) |
||||
} |
||||
} |
||||
|
||||
return services, nil |
||||
} |
||||
|
||||
func (sp *SystemProxy) backupCurrentConfig() error { |
||||
switch runtime.GOOS { |
||||
case "darwin": |
||||
return sp.backupMacOSConfig() |
||||
case "linux": |
||||
return sp.backupLinuxConfig() |
||||
case "windows": |
||||
return sp.backupWindowsConfig() |
||||
} |
||||
return nil |
||||
} |
||||
|
||||
func (sp *SystemProxy) backupMacOSConfig() error { |
||||
services, err := sp.getMacOSNetworkServices() |
||||
if err != nil { |
||||
return err |
||||
} |
||||
|
||||
for _, service := range services { |
||||
// 备份HTTP代理设置
|
||||
output, _ := exec.Command("networksetup", "-getwebproxy", service).Output() |
||||
sp.backupConfig[fmt.Sprintf("http_%s", service)] = string(output) |
||||
|
||||
// 备份HTTPS代理设置
|
||||
output, _ = exec.Command("networksetup", "-getsecurewebproxy", service).Output() |
||||
sp.backupConfig[fmt.Sprintf("https_%s", service)] = string(output) |
||||
|
||||
// 备份SOCKS代理设置
|
||||
output, _ = exec.Command("networksetup", "-getsocksfirewallproxy", service).Output() |
||||
sp.backupConfig[fmt.Sprintf("socks_%s", service)] = string(output) |
||||
} |
||||
|
||||
return nil |
||||
} |
||||
|
||||
func (sp *SystemProxy) backupLinuxConfig() error { |
||||
// 备份环境变量
|
||||
sp.backupConfig["http_proxy"] = os.Getenv("http_proxy") |
||||
sp.backupConfig["https_proxy"] = os.Getenv("https_proxy") |
||||
sp.backupConfig["no_proxy"] = os.Getenv("no_proxy") |
||||
return nil |
||||
} |
||||
|
||||
func (sp *SystemProxy) backupWindowsConfig() error { |
||||
// 备份Windows注册表设置
|
||||
output, _ := exec.Command("reg", "query", `HKCU\Software\Microsoft\Windows\CurrentVersion\Internet Settings`, "/v", "ProxyEnable").Output() |
||||
sp.backupConfig["ProxyEnable"] = string(output) |
||||
|
||||
output, _ = exec.Command("reg", "query", `HKCU\Software\Microsoft\Windows\CurrentVersion\Internet Settings`, "/v", "ProxyServer").Output() |
||||
sp.backupConfig["ProxyServer"] = string(output) |
||||
|
||||
return nil |
||||
} |
||||
|
||||
func (sp *SystemProxy) restoreMacOSProxy() error { |
||||
services, err := sp.getMacOSNetworkServices() |
||||
if err != nil { |
||||
return err |
||||
} |
||||
|
||||
for _, service := range services { |
||||
// 禁用所有代理
|
||||
sp.runCommand("networksetup", "-setwebproxystate", service, "off") |
||||
sp.runCommand("networksetup", "-setsecurewebproxystate", service, "off") |
||||
sp.runCommand("networksetup", "-setsocksfirewallproxystate", service, "off") |
||||
} |
||||
|
||||
return nil |
||||
} |
||||
|
||||
func (sp *SystemProxy) restoreLinuxProxy() error { |
||||
// 清除环境变量（这里简化处理）
|
||||
os.Unsetenv("http_proxy") |
||||
os.Unsetenv("https_proxy") |
||||
os.Unsetenv("no_proxy") |
||||
os.Unsetenv("HTTP_PROXY") |
||||
os.Unsetenv("HTTPS_PROXY") |
||||
os.Unsetenv("NO_PROXY") |
||||
return nil |
||||
} |
||||
|
||||
func (sp *SystemProxy) restoreWindowsProxy() error { |
||||
// 禁用代理
|
||||
return sp.runCommand("reg", "add", `HKCU\Software\Microsoft\Windows\CurrentVersion\Internet Settings`, "/v", "ProxyEnable", "/t", "REG_DWORD", "/d", "0", "/f") |
||||
} |
||||
|
||||
func (sp *SystemProxy) writeLinuxProxyConfig(envVars map[string]string) error { |
||||
// 写入到 ~/.bashrc 和 ~/.profile
|
||||
homeDir, err := os.UserHomeDir() |
||||
if err != nil { |
||||
return err |
||||
} |
||||
|
||||
configFiles := []string{ |
||||
homeDir + "/.bashrc", |
||||
homeDir + "/.profile", |
||||
} |
||||
|
||||
proxyLines := []string{"\n# Wormhole SOCKS5 Proxy Configuration"} |
||||
for key, value := range envVars { |
||||
proxyLines = append(proxyLines, fmt.Sprintf("export %s=%s", key, value)) |
||||
} |
||||
proxyLines = append(proxyLines, "# End Wormhole Configuration\n") |
||||
|
||||
for _, configFile := range configFiles { |
||||
file, err := os.OpenFile(configFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) |
||||
if err != nil { |
||||
sp.logger.WithError(err).WithField("file", configFile).Warn("Failed to open config file") |
||||
continue |
||||
} |
||||
|
||||
for _, line := range proxyLines { |
||||
if _, err := file.WriteString(line + "\n"); err != nil { |
||||
sp.logger.WithError(err).WithField("file", configFile).Warn("Failed to write to config file") |
||||
} |
||||
} |
||||
|
||||
file.Close() |
||||
} |
||||
|
||||
return nil |
||||
} |
||||
|
||||
func (sp *SystemProxy) runCommand(name string, args ...string) error { |
||||
cmd := exec.Command(name, args...) |
||||
output, err := cmd.CombinedOutput() |
||||
if err != nil { |
||||
sp.logger.WithFields(logrus.Fields{ |
||||
"command": fmt.Sprintf("%s %s", name, strings.Join(args, " ")), |
||||
"output": string(output), |
||||
}).WithError(err).Debug("Command execution failed") |
||||
return err |
||||
} |
||||
return nil |
||||
} |
||||
|
||||
// GetCurrentConfig 获取当前系统代理配置
|
||||
func (sp *SystemProxy) GetCurrentConfig() (Config, error) { |
||||
config := Config{ |
||||
NoProxy: []string{}, |
||||
} |
||||
|
||||
switch runtime.GOOS { |
||||
case "linux": |
||||
config.HTTPProxy = os.Getenv("http_proxy") |
||||
config.HTTPSProxy = os.Getenv("https_proxy") |
||||
noProxy := os.Getenv("no_proxy") |
||||
if noProxy != "" { |
||||
config.NoProxy = strings.Split(noProxy, ",") |
||||
} |
||||
} |
||||
|
||||
return config, nil |
||||
} |
||||
|
||||
// IsProxySet 检查是否已设置代理
|
||||
func (sp *SystemProxy) IsProxySet() bool { |
||||
switch runtime.GOOS { |
||||
case "linux": |
||||
return os.Getenv("http_proxy") != "" || os.Getenv("https_proxy") != "" |
||||
case "darwin": |
||||
// 检查macOS代理设置（简化）
|
||||
return false |
||||
case "windows": |
||||
// 检查Windows代理设置（简化）
|
||||
return false |
||||
} |
||||
return false |
||||
} |
||||
Loading…
Reference in new issue