package main import ( "flag" "fmt" "github.com/go-redis/redis/v7" "github.com/spf13/viper" "meow.tf/joker/godns/cache" "meow.tf/joker/godns/hosts" "meow.tf/joker/godns/log" "meow.tf/joker/godns/resolver" "meow.tf/joker/godns/settings" "os" "os/signal" "runtime" "runtime/pprof" "time" ) const ( Version = "0.3.0" ) var ( cfgFile string ) func init() { flag.StringVar(&cfgFile, "config", "/etc/godns/godns.conf", "") } func main() { initLogger() viper.SetConfigFile(cfgFile) viper.AutomaticEnv() if err := viper.ReadInConfig(); err == nil { fmt.Println("Using config file:", viper.ConfigFileUsed()) } server := &Server{ host: viper.GetString("server.host"), networks: viper.GetStringSlice("server.nets"), rTimeout: 5 * time.Second, wTimeout: 5 * time.Second, } var resolverCache, negCache cache.Cache r := resolver.NewResolver(resolver.Settings{ Timeout: viper.GetInt("resolv.timeout"), Interval: viper.GetInt("resolv.interval"), SetEDNS0: viper.GetBool("resolv.edns0"), ServerListFile: viper.GetStringSlice("resolv.server-list"), ResolvFile: viper.GetString("resolv.file"), }) cacheDuration := viper.GetDuration("cache.expire") redisConfig := settings.RedisSettings{ Host: viper.GetString("cache.redis.host"), Port: viper.GetInt("cache.redis.port"), DB: viper.GetInt("cache.redis.db"), Password: viper.GetString("cache.redis.password"), } backend := viper.GetString("cache.backend") switch backend { case "memory": cacheMaxCount := viper.GetInt("cache.memory.maxCount") negCache = cache.NewMemoryCache(cacheDuration/2, cacheMaxCount) resolverCache = cache.NewMemoryCache(cacheDuration, cacheMaxCount) case "memcached": servers := viper.GetStringSlice("cache.memcached.servers") resolverCache = cache.NewMemcachedCache(servers, int32(cacheDuration.Seconds())) negCache = cache.NewMemcachedCache(servers, int32(cacheDuration.Seconds()/2)) case "redis": resolverCache = cache.NewRedisCache(redisConfig, cacheDuration) negCache = cache.NewRedisCache(redisConfig, cacheDuration/2) default: log.Error("Invalid cache backend %s", backend) panic("Invalid cache backend") } providers := make([]hosts.Provider, 0) if viper.GetBool("hosts.file.enable") { providers = append(providers, hosts.NewFileProvider(viper.GetString("hosts.file.file"), viper.GetDuration("hosts.file.ttl"))) } if viper.GetBool("hosts.redis.enable") { rc := redis.NewClient(&redis.Options{Addr: redisConfig.Addr(), DB: redisConfig.DB, Password: redisConfig.Password}) providers = append(providers, hosts.NewRedisProvider(rc, viper.GetString("hosts.redis.key"), viper.GetDuration("hosts.redis.ttl"))) } h := hosts.NewHosts(providers) server.Run(NewHandler(r, resolverCache, negCache, h)) log.Info("godns %s (%s)", Version, runtime.Version()) if viper.GetBool("debug") { go profileCPU() go profileMEM() } sig := make(chan os.Signal) signal.Notify(sig, os.Interrupt) <- sig log.Info("signal received, stopping") } func profileCPU() { f, err := os.Create("godns.cprof") if err != nil { log.Error("%s", err) return } pprof.StartCPUProfile(f) time.AfterFunc(6*time.Minute, func() { pprof.StopCPUProfile() f.Close() }) } func profileMEM() { f, err := os.Create("godns.mprof") if err != nil { log.Error("%s", err) return } time.AfterFunc(5*time.Minute, func() { pprof.WriteHeapProfile(f) f.Close() }) } func initLogger() { if viper.GetBool("log.stdout") { log.SetLogger("console", nil) } if file := viper.GetString("log.file"); file != "" { log.SetLogger("file", map[string]interface{}{"file": file}) } log.SetLevel(settings.LogLevelFor(viper.GetString("log.level"))) }