package main import ( "fmt" "os" "strconv" "strings" "sync" "time" "github.com/pelletier/go-toml/v2" "github.com/spf13/viper" "github.com/fsnotify/fsnotify" ) // RegistryMapping Registry映射配置 type RegistryMapping struct { Upstream string `toml:"upstream"` // 上游Registry地址 AuthHost string `toml:"authHost"` // 认证服务器地址 AuthType string `toml:"authType"` // 认证类型: docker/github/google/basic Enabled bool `toml:"enabled"` // 是否启用 } // AppConfig 应用配置结构体 type AppConfig struct { Server struct { Host string `toml:"host"` // 监听地址 Port int `toml:"port"` // 监听端口 FileSize int64 `toml:"fileSize"` // 文件大小限制(字节) } `toml:"server"` RateLimit struct { RequestLimit int `toml:"requestLimit"` // 每小时请求限制 PeriodHours float64 `toml:"periodHours"` // 限制周期(小时) } `toml:"rateLimit"` Security struct { WhiteList []string `toml:"whiteList"` // 白名单IP/CIDR列表 BlackList []string `toml:"blackList"` // 黑名单IP/CIDR列表 } `toml:"security"` Proxy struct { WhiteList []string `toml:"whiteList"` // 代理白名单(仓库级别) BlackList []string `toml:"blackList"` // 代理黑名单(仓库级别) } `toml:"proxy"` Download struct { MaxImages int `toml:"maxImages"` // 单次下载最大镜像数量限制 } `toml:"download"` // 新增:Registry映射配置 Registries map[string]RegistryMapping `toml:"registries"` // Token缓存配置 TokenCache struct { Enabled bool `toml:"enabled"` // 是否启用token缓存 DefaultTTL string `toml:"defaultTTL"` // 默认缓存时间 } `toml:"tokenCache"` } var ( appConfig *AppConfig appConfigLock sync.RWMutex isViperEnabled bool viperInstance *viper.Viper ) // DefaultConfig 返回默认配置 func DefaultConfig() *AppConfig { return &AppConfig{ Server: struct { Host string `toml:"host"` Port int `toml:"port"` FileSize int64 `toml:"fileSize"` }{ Host: "0.0.0.0", Port: 5000, FileSize: 2 * 1024 * 1024 * 1024, // 2GB }, RateLimit: struct { RequestLimit int `toml:"requestLimit"` PeriodHours float64 `toml:"periodHours"` }{ RequestLimit: 20, PeriodHours: 1.0, }, Security: struct { WhiteList []string `toml:"whiteList"` BlackList []string `toml:"blackList"` }{ WhiteList: []string{}, BlackList: []string{}, }, Proxy: struct { WhiteList []string `toml:"whiteList"` BlackList []string `toml:"blackList"` }{ WhiteList: []string{}, BlackList: []string{}, }, Download: struct { MaxImages int `toml:"maxImages"` }{ MaxImages: 10, // 默认值:最多同时下载10个镜像 }, Registries: map[string]RegistryMapping{ "ghcr.io": { Upstream: "ghcr.io", AuthHost: "ghcr.io/token", AuthType: "github", Enabled: true, }, "gcr.io": { Upstream: "gcr.io", AuthHost: "gcr.io/v2/token", AuthType: "google", Enabled: true, }, "quay.io": { Upstream: "quay.io", AuthHost: "quay.io/v2/auth", AuthType: "quay", Enabled: true, }, "registry.k8s.io": { Upstream: "registry.k8s.io", AuthHost: "registry.k8s.io", AuthType: "anonymous", Enabled: true, }, }, TokenCache: struct { Enabled bool `toml:"enabled"` DefaultTTL string `toml:"defaultTTL"` }{ Enabled: true, // docker认证的匿名Token缓存配置,用于提升性能 DefaultTTL: "20m", }, } } // GetConfig 安全地获取配置副本 func GetConfig() *AppConfig { appConfigLock.RLock() defer appConfigLock.RUnlock() if appConfig == nil { return DefaultConfig() } // 返回配置的深拷贝 configCopy := *appConfig configCopy.Security.WhiteList = append([]string(nil), appConfig.Security.WhiteList...) configCopy.Security.BlackList = append([]string(nil), appConfig.Security.BlackList...) configCopy.Proxy.WhiteList = append([]string(nil), appConfig.Proxy.WhiteList...) configCopy.Proxy.BlackList = append([]string(nil), appConfig.Proxy.BlackList...) return &configCopy } // setConfig 安全地设置配置 func setConfig(cfg *AppConfig) { appConfigLock.Lock() defer appConfigLock.Unlock() appConfig = cfg } // LoadConfig 加载配置文件 func LoadConfig() error { // 首先使用默认配置 cfg := DefaultConfig() // 尝试加载TOML配置文件 if data, err := os.ReadFile("config.toml"); err == nil { if err := toml.Unmarshal(data, cfg); err != nil { return fmt.Errorf("解析配置文件失败: %v", err) } } else { fmt.Println("未找到config.toml,使用默认配置") } // 从环境变量覆盖配置 overrideFromEnv(cfg) // 设置配置 setConfig(cfg) // 🔥 首次加载后启用Viper热重载 if !isViperEnabled { go enableViperHotReload() } fmt.Printf("配置加载成功: 监听 %s:%d, 文件大小限制 %d MB, 限流 %d请求/%g小时, 离线镜像并发数 %d\n", cfg.Server.Host, cfg.Server.Port, cfg.Server.FileSize/(1024*1024), cfg.RateLimit.RequestLimit, cfg.RateLimit.PeriodHours, cfg.Download.MaxImages) return nil } // 🔥 启用Viper自动热重载 func enableViperHotReload() { if isViperEnabled { return } // 创建Viper实例 viperInstance = viper.New() // 配置Viper viperInstance.SetConfigName("config") viperInstance.SetConfigType("toml") viperInstance.AddConfigPath(".") // 读取配置文件 if err := viperInstance.ReadInConfig(); err != nil { fmt.Printf("读取配置失败,继续使用当前配置: %v\n", err) return } isViperEnabled = true fmt.Println("自动热重载已启用") // 🚀 启用文件监听 viperInstance.WatchConfig() viperInstance.OnConfigChange(func(e fsnotify.Event) { fmt.Printf("检测到配置文件变化: %s\n", e.Name) hotReloadWithViper() }) } // 🔥 使用Viper进行热重载 func hotReloadWithViper() { start := time.Now() fmt.Println("🔄 自动热重载...") // 创建新配置 cfg := DefaultConfig() // 使用Viper解析配置到结构体 if err := viperInstance.Unmarshal(cfg); err != nil { fmt.Printf("❌ 配置解析失败: %v\n", err) return } // 从环境变量覆盖(保持原有功能) overrideFromEnv(cfg) // 原子性更新配置 setConfig(cfg) // 异步更新受影响的组件 go func() { updateAffectedComponents() fmt.Printf("✅ Viper配置热重载完成,耗时: %v\n", time.Since(start)) }() } // 🔧 更新受配置影响的组件 func updateAffectedComponents() { // 重新初始化限流器 if globalLimiter != nil { fmt.Println("📡 重新初始化限流器...") initLimiter() } // 重新加载访问控制 fmt.Println("🔒 重新加载访问控制规则...") if GlobalAccessController != nil { GlobalAccessController.Reload() } // 🔥 刷新Registry配置映射 fmt.Println("🌐 更新Registry配置映射...") reloadRegistryConfig() // 其他需要重新初始化的组件可以在这里添加 fmt.Println("🔧 组件更新完成") } // 🔥 重新加载Registry配置 func reloadRegistryConfig() { cfg := GetConfig() enabledCount := 0 // 统计启用的Registry数量 for _, mapping := range cfg.Registries { if mapping.Enabled { enabledCount++ } } fmt.Printf("🌐 Registry配置已更新: %d个启用\n", enabledCount) // Registry配置是动态读取的,每次请求都会调用GetConfig() // 所以这里只需要简单通知,实际生效是自动的 } // overrideFromEnv 从环境变量覆盖配置 func overrideFromEnv(cfg *AppConfig) { // 服务器配置 if val := os.Getenv("SERVER_HOST"); val != "" { cfg.Server.Host = val } if val := os.Getenv("SERVER_PORT"); val != "" { if port, err := strconv.Atoi(val); err == nil && port > 0 { cfg.Server.Port = port } } if val := os.Getenv("MAX_FILE_SIZE"); val != "" { if size, err := strconv.ParseInt(val, 10, 64); err == nil && size > 0 { cfg.Server.FileSize = size } } // 限流配置 if val := os.Getenv("RATE_LIMIT"); val != "" { if limit, err := strconv.Atoi(val); err == nil && limit > 0 { cfg.RateLimit.RequestLimit = limit } } if val := os.Getenv("RATE_PERIOD_HOURS"); val != "" { if period, err := strconv.ParseFloat(val, 64); err == nil && period > 0 { cfg.RateLimit.PeriodHours = period } } // IP限制配置 if val := os.Getenv("IP_WHITELIST"); val != "" { cfg.Security.WhiteList = append(cfg.Security.WhiteList, strings.Split(val, ",")...) } if val := os.Getenv("IP_BLACKLIST"); val != "" { cfg.Security.BlackList = append(cfg.Security.BlackList, strings.Split(val, ",")...) } // 下载限制配置 if val := os.Getenv("MAX_IMAGES"); val != "" { if maxImages, err := strconv.Atoi(val); err == nil && maxImages > 0 { cfg.Download.MaxImages = maxImages } } } // CreateDefaultConfigFile 创建默认配置文件 func CreateDefaultConfigFile() error { cfg := DefaultConfig() data, err := toml.Marshal(cfg) if err != nil { return fmt.Errorf("序列化默认配置失败: %v", err) } return os.WriteFile("config.toml", data, 0644) }