godns/main.go

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(settings.ResolvSettings{
		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")))
}